Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/sparrow/elements/base.py: 85%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6import os 

7import base64 

8 

9import numpy as num 

10 

11from pyrocko.plot import automap, mpl_get_cmap_names, mpl_get_cmap 

12from pyrocko.guts import String, Float, StringChoice, Bool 

13from pyrocko.plot import AutoScaler, AutoScaleMode 

14from pyrocko.dataset import topo 

15 

16from pyrocko.gui.talkie import (TalkieRoot, TalkieConnectionOwner, 

17 has_computed, computed) 

18 

19from pyrocko.gui.qt_compat import qc, qw 

20from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable 

21 

22 

23from .. import common 

24from ..state import \ 

25 state_bind_combobox, state_bind, state_bind_checkbox 

26 

27 

28mpl_cmap_blacklist = [ 

29 "prism", "flag", 

30 "Accent", "Dark2", 

31 "Paired", "Pastel1", "Pastel2", 

32 "Set1", "Set2", "Set3", 

33 "tab10", "tab20", "tab20b", "tab20c" 

34] 

35 

36 

37def get_mpl_cmap_choices(): 

38 names = mpl_get_cmap_names() 

39 for cmap_name in mpl_cmap_blacklist: 

40 try: 

41 names.remove(cmap_name) 

42 names.remove("%s_r" % cmap_name) 

43 except ValueError: 

44 pass 

45 

46 return names 

47 

48 

49def random_id(): 

50 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii') 

51 

52 

53class ElementState(TalkieRoot): 

54 

55 element_id = String.T() 

56 

57 def __init__(self, **kwargs): 

58 if 'element_id' not in kwargs: 

59 kwargs['element_id'] = random_id() 

60 

61 TalkieRoot.__init__(self, **kwargs) 

62 

63 

64class Element(TalkieConnectionOwner): 

65 def __init__(self): 

66 TalkieConnectionOwner.__init__(self) 

67 self._parent = None 

68 self._state = None 

69 

70 def remove(self): 

71 if self._parent and self._state: 

72 self._parent.state.elements.remove(self._state) 

73 

74 def set_parent(self, parent): 

75 self._parent = parent 

76 

77 def unset_parent(self): 

78 print(self) 

79 raise NotImplementedError 

80 

81 def bind_state(self, state): 

82 self._state = state 

83 

84 def unbind_state(self): 

85 self.talkie_disconnect_all() 

86 self._state = None 

87 

88 def update_visibility(self, visible): 

89 self._state.visible = visible 

90 

91 def get_title_label(self): 

92 title_label = common.MyDockWidgetTitleBarLabel(self.get_name()) 

93 

94 def update_label(*args): 

95 title_label.set_slug(self._state.element_id) 

96 

97 self.talkie_connect( 

98 self._state, 'element_id', update_label) 

99 

100 update_label() 

101 return title_label 

102 

103 def get_title_control_remove(self): 

104 button = common.MyDockWidgetTitleBarButton('\u2716') 

105 button.setStatusTip('Remove Element') 

106 button.clicked.connect(self.remove) 

107 return button 

108 

109 def get_title_control_visible(self): 

110 assert hasattr(self._state, 'visible') 

111 

112 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54') 

113 button.setStatusTip('Toggle Element Visibility') 

114 button.toggled.connect(self.update_visibility) 

115 

116 def set_button_checked(*args): 

117 button.blockSignals(True) 

118 button.set_checked(self._state.visible) 

119 button.blockSignals(False) 

120 

121 set_button_checked() 

122 

123 self.talkie_connect( 

124 self._state, 'visible', set_button_checked) 

125 

126 return button 

127 

128 

129class CPTChoice(StringChoice): 

130 

131 choices = ['slip_colors'] + get_mpl_cmap_choices() 

132 

133 

134@has_computed 

135class CPTState(ElementState): 

136 cpt_name = String.T(default=CPTChoice.choices[0]) 

137 cpt_mode = String.T(default=AutoScaleMode.choices[1]) 

138 cpt_scale_min = Float.T(optional=True) 

139 cpt_scale_max = Float.T(optional=True) 

140 cpt_revert = Bool.T(default=False) 

141 

142 @computed(['cpt_name', 'cpt_revert']) 

143 def effective_cpt_name(self): 

144 if self.cpt_revert: 

145 return '%s_r' % self.cpt_name 

146 else: 

147 return self.cpt_name 

148 

149 

150class CPTHandler(Element): 

151 

152 def __init__(self): 

153 

154 Element.__init__(self) 

155 self._cpts = {} 

156 self._autoscaler = None 

157 self._lookuptable = None 

158 self._cpt_combobox = None 

159 self._values = None 

160 self._state = None 

161 self._cpt_scale_lineedit = None 

162 

163 def bind_state(self, cpt_state, update_function): 

164 for state_attr in [ 

165 'effective_cpt_name', 'cpt_mode', 

166 'cpt_scale_min', 'cpt_scale_max']: 

167 

168 self.talkie_connect( 

169 cpt_state, state_attr, update_function) 

170 

171 self._state = cpt_state 

172 

173 def unbind_state(self): 

174 Element.unbind_state(self) 

175 self._cpts = {} 

176 self._lookuptable = None 

177 self._values = None 

178 self._autoscaler = None 

179 

180 def open_cpt_load_dialog(self): 

181 caption = 'Select one *.cpt file to open' 

182 

183 fns, _ = qw.QFileDialog.getOpenFileNames( 

184 self._parent, caption, options=common.qfiledialog_options) 

185 

186 if fns: 

187 self.load_cpt_file(fns[0]) 

188 

189 def load_cpt_file(self, path): 

190 cpt_name = 'USR' + os.path.basename(path).split('.')[0] 

191 self._cpts.update([(cpt_name, automap.read_cpt(path))]) 

192 

193 self._state.cpt_name = cpt_name 

194 

195 self._update_cpt_combobox() 

196 self.update_cpt() 

197 

198 def _update_cpt_combobox(self): 

199 from pyrocko import config 

200 conf = config.config() 

201 

202 if self._cpt_combobox is None: 

203 raise ValueError('CPT combobox needs init before updating!') 

204 

205 cb = self._cpt_combobox 

206 

207 if cb is not None: 

208 cb.clear() 

209 

210 for s in CPTChoice.choices: 

211 if s not in self._cpts: 

212 try: 

213 cpt = automap.read_cpt(topo.cpt(s)) 

214 except Exception: 

215 cmap = mpl_get_cmap(s) 

216 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1]) 

217 

218 self._cpts.update([(s, cpt)]) 

219 

220 cpt_dir = conf.colortables_dir 

221 if os.path.isdir(cpt_dir): 

222 for f in [ 

223 f for f in os.listdir(cpt_dir) 

224 if f.lower().endswith('.cpt')]: 

225 

226 s = 'USR' + os.path.basename(f).split('.')[0] 

227 self._cpts.update( 

228 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))]) 

229 

230 for i, (s, cpt) in enumerate(self._cpts.items()): 

231 if s[-2::] != "_r": 

232 cb.insertItem(i, s, qc.QVariant(self._cpts[s])) 

233 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole) 

234 

235 cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name)) 

236 

237 def _update_cptscale_lineedit(self): 

238 le = self._cpt_scale_lineedit 

239 if le is not None: 

240 le.clear() 

241 

242 self._cptscale_to_lineedit(self._state, le) 

243 

244 def _cptscale_to_lineedit(self, state, widget): 

245 # sel = widget.selectedText() == widget.text() 

246 

247 crange = (None, None) 

248 if self._lookuptable is not None: 

249 crange = self._lookuptable.GetRange() 

250 

251 if state.cpt_scale_min is not None and state.cpt_scale_max is not None: 

252 crange = state.cpt_scale_min, state.cpt_scale_max 

253 

254 fmt = ', '.join(['%s' if item is None else '%g' for item in crange]) 

255 

256 widget.setText(fmt % crange) 

257 

258 # if sel: 

259 # widget.selectAll() 

260 

261 def update_cpt(self): 

262 state = self._state 

263 

264 if self._autoscaler is None: 

265 self._autoscaler = AutoScaler() 

266 

267 if self._cpt_scale_lineedit: 

268 if state.cpt_mode == 'off': 

269 self._cpt_scale_lineedit.setEnabled(True) 

270 else: 

271 self._cpt_scale_lineedit.setEnabled(False) 

272 

273 if state.cpt_scale_min is not None: 

274 state.cpt_scale_min = None 

275 

276 if state.cpt_scale_max is not None: 

277 state.cpt_scale_max = None 

278 

279 if state.effective_cpt_name is not None and self._values is not None: 

280 if self._values.size == 0: 

281 vscale = (0., 1.) 

282 else: 

283 vscale = (num.nanmin(self._values), num.nanmax(self._values)) 

284 

285 vmin, vmax = None, None 

286 if None not in (state.cpt_scale_min, state.cpt_scale_max): 

287 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max 

288 else: 

289 vmin, vmax, _ = self._autoscaler.make_scale( 

290 vscale, override_mode=state.cpt_mode) 

291 

292 self._cpts[state.effective_cpt_name].scale(vmin, vmax) 

293 cpt = self._cpts[state.effective_cpt_name] 

294 vtk_lut = cpt_to_vtk_lookuptable(cpt) 

295 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0) 

296 

297 self._lookuptable = vtk_lut 

298 self._update_cptscale_lineedit() 

299 

300 elif state.effective_cpt_name and self._values is None: 

301 raise ValueError('No values passed to colormapper!') 

302 

303 def cpt_controls(self, parent, state, layout): 

304 self._parent = parent 

305 

306 iy = layout.rowCount() + 1 

307 

308 layout.addWidget(qw.QLabel('Color Map'), iy, 0) 

309 

310 cb = common.CPTComboBox() 

311 layout.addWidget(cb, iy, 1) 

312 state_bind_combobox( 

313 self, state, 'cpt_name', cb) 

314 

315 self._cpt_combobox = cb 

316 

317 pb = qw.QPushButton('Load CPT') 

318 layout.addWidget(pb, iy, 2) 

319 pb.clicked.connect(self.open_cpt_load_dialog) 

320 

321 iy += 1 

322 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0) 

323 

324 cb = common.string_choices_to_combobox(AutoScaleMode) 

325 layout.addWidget(cb, iy, 1) 

326 state_bind_combobox( 

327 self, state, 'cpt_mode', cb) 

328 

329 le = qw.QLineEdit() 

330 le.setEnabled(False) 

331 layout.addWidget(le, iy, 2) 

332 state_bind( 

333 self, state, 

334 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale, 

335 le, [le.editingFinished, le.returnPressed], 

336 self._cptscale_to_lineedit) 

337 

338 self._cpt_scale_lineedit = le 

339 

340 iy += 1 

341 cb = qw.QCheckBox('Revert') 

342 layout.addWidget(cb, iy, 1) 

343 state_bind_checkbox(self, state, 'cpt_revert', cb) 

344 

345 

346def _lineedit_to_cptscale(widget, cpt_state): 

347 s = str(widget.text()) 

348 s = s.replace(',', ' ') 

349 

350 crange = tuple((float(i) for i in s.split())) 

351 crange = tuple(( 

352 crange[0], 

353 crange[0]+0.01 if crange[0] >= crange[1] else crange[1])) 

354 

355 try: 

356 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange 

357 except Exception: 

358 raise ValueError( 

359 'need two numerical values: <vmin>, <vmax>') 

360 

361 

362__all__ = [ 

363 'Element', 

364 'ElementState', 

365 'random_id', 

366]