Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/sparrow/elements/base.py: 87%
292 statements
« prev ^ index » next coverage.py v6.5.0, created at 2024-01-15 12:05 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2024-01-15 12:05 +0000
1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import os
7import base64
9import numpy as num
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
16from pyrocko.gui.talkie import (TalkieRoot, TalkieConnectionOwner,
17 has_computed, computed)
19from pyrocko.gui.qt_compat import qc, qw
20from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable, ColorbarPipe
23from .. import common
24from ..state import \
25 state_bind_combobox, state_bind, state_bind_checkbox, state_bind_slider
28mpl_cmap_blacklist = [
29 "prism", "flag",
30 "Accent", "Dark2",
31 "Paired", "Pastel1", "Pastel2",
32 "Set1", "Set2", "Set3",
33 "tab10", "tab20", "tab20b", "tab20c"
34]
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
46 return names
49def random_id():
50 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
53class ElementState(TalkieRoot):
55 element_id = String.T()
57 def __init__(self, **kwargs):
58 if 'element_id' not in kwargs:
59 kwargs['element_id'] = random_id()
61 TalkieRoot.__init__(self, **kwargs)
64class Element(TalkieConnectionOwner):
65 def __init__(self):
66 TalkieConnectionOwner.__init__(self)
67 self._parent = None
68 self._state = None
70 def remove(self):
71 if self._parent and self._state:
72 self._parent.state.elements.remove(self._state)
74 def set_parent(self, parent):
75 self._parent = parent
77 def unset_parent(self):
78 print(self)
79 raise NotImplementedError
81 def bind_state(self, state):
82 self._state = state
84 def unbind_state(self):
85 self.talkie_disconnect_all()
86 self._state = None
88 def update_visibility(self, visible):
89 self._state.visible = visible
91 def get_title_label(self):
92 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())
94 def update_label(*args):
95 title_label.set_slug(self._state.element_id)
97 self.talkie_connect(
98 self._state, 'element_id', update_label)
100 update_label()
101 return title_label
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
109 def get_title_control_visible(self):
110 assert hasattr(self._state, 'visible')
112 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
113 button.setStatusTip('Toggle Element Visibility')
114 button.toggled.connect(self.update_visibility)
116 def set_button_checked(*args):
117 button.blockSignals(True)
118 button.set_checked(self._state.visible)
119 button.blockSignals(False)
121 set_button_checked()
123 self.talkie_connect(
124 self._state, 'visible', set_button_checked)
126 return button
129class CPTChoice(StringChoice):
131 choices = ['slip_colors'] + get_mpl_cmap_choices()
134class ColorBarPositionChoice(StringChoice):
135 choices = ['bottom-left', 'bottom-right', 'top-left', 'top-right']
138@has_computed
139class CPTState(ElementState):
140 cpt_name = String.T(default=CPTChoice.choices[0])
141 cpt_mode = String.T(default=AutoScaleMode.choices[1])
142 cpt_scale_min = Float.T(optional=True)
143 cpt_scale_max = Float.T(optional=True)
144 cpt_revert = Bool.T(default=False)
145 cbar_show = Bool.T(default=True)
146 cbar_position = ColorBarPositionChoice.T(default='bottom-right')
147 cbar_annotation_lightness = Float.T(default=1.0)
148 cbar_annotation_fontsize = Float.T(default=0.03)
149 cbar_height = Float.T(default=1.)
150 cbar_width = Float.T(default=1.)
152 @computed(['cpt_name', 'cpt_revert'])
153 def effective_cpt_name(self):
154 if self.cpt_revert:
155 return '%s_r' % self.cpt_name
156 else:
157 return self.cpt_name
160class CPTHandler(Element):
162 def __init__(self):
164 Element.__init__(self)
165 self._cpts = {}
166 self._autoscaler = None
167 self._lookuptable = None
168 self._cpt_combobox = None
169 self._values = None
170 self._state = None
171 self._cpt_scale_lineedit = None
172 self._cbar_pipe = None
174 def bind_state(self, cpt_state, update_function):
175 for state_attr in [
176 'effective_cpt_name', 'cpt_mode',
177 'cpt_scale_min', 'cpt_scale_max',
178 'cbar_show', 'cbar_position',
179 'cbar_annotation_lightness',
180 'cbar_annotation_fontsize',
181 'cbar_height', 'cbar_width']:
183 self.talkie_connect(
184 cpt_state, state_attr, update_function)
186 self._state = cpt_state
188 def unbind_state(self):
189 Element.unbind_state(self)
190 self._cpts = {}
191 self._lookuptable = None
192 self._values = None
193 self._autoscaler = None
195 def open_cpt_load_dialog(self):
196 caption = 'Select one *.cpt file to open'
198 fns, _ = qw.QFileDialog.getOpenFileNames(
199 self._parent, caption, options=common.qfiledialog_options)
201 if fns:
202 self.load_cpt_file(fns[0])
204 def load_cpt_file(self, path):
205 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
206 self._cpts.update([(cpt_name, automap.read_cpt(path))])
208 self._state.cpt_name = cpt_name
210 self._update_cpt_combobox()
211 self.update_cpt()
213 def _update_cpt_combobox(self):
214 from pyrocko import config
215 conf = config.config()
217 if self._cpt_combobox is None:
218 raise ValueError('CPT combobox needs init before updating!')
220 cb = self._cpt_combobox
222 if cb is not None:
223 cb.clear()
225 for s in CPTChoice.choices:
226 if s not in self._cpts:
227 try:
228 cpt = automap.read_cpt(topo.cpt(s))
229 except Exception:
230 cmap = mpl_get_cmap(s)
231 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
233 self._cpts.update([(s, cpt)])
235 cpt_dir = conf.colortables_dir
236 if os.path.isdir(cpt_dir):
237 for f in [
238 f for f in os.listdir(cpt_dir)
239 if f.lower().endswith('.cpt')]:
241 s = 'USR' + os.path.basename(f).split('.')[0]
242 self._cpts.update(
243 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
245 for i, (s, cpt) in enumerate(self._cpts.items()):
246 if s[-2::] != "_r":
247 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
248 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
250 cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name))
252 def _update_cptscale_lineedit(self):
253 le = self._cpt_scale_lineedit
254 if le is not None:
255 le.clear()
257 self._cptscale_to_lineedit(self._state, le)
259 def _cptscale_to_lineedit(self, state, widget):
260 # sel = widget.selectedText() == widget.text()
262 crange = (None, None)
263 if self._lookuptable is not None:
264 crange = self._lookuptable.GetRange()
266 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
267 crange = state.cpt_scale_min, state.cpt_scale_max
269 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
271 widget.setText(fmt % crange)
273 # if sel:
274 # widget.selectAll()
276 def update_cpt(self, mask_zeros=False):
277 state = self._state
279 if self._autoscaler is None:
280 self._autoscaler = AutoScaler()
282 if self._cpt_scale_lineedit:
283 if state.cpt_mode == 'off':
284 self._cpt_scale_lineedit.setEnabled(True)
285 else:
286 self._cpt_scale_lineedit.setEnabled(False)
288 if state.cpt_scale_min is not None:
289 state.cpt_scale_min = None
291 if state.cpt_scale_max is not None:
292 state.cpt_scale_max = None
294 if state.effective_cpt_name is not None and self._values is not None:
295 if self._values.size == 0:
296 vscale = (0., 1.)
297 else:
298 vscale = (num.nanmin(self._values), num.nanmax(self._values))
300 vmin, vmax = None, None
301 if None not in (state.cpt_scale_min, state.cpt_scale_max):
302 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
303 else:
304 vmin, vmax, _ = self._autoscaler.make_scale(
305 vscale, override_mode=state.cpt_mode)
307 self._cpts[state.effective_cpt_name].scale(vmin, vmax)
308 cpt = self._cpts[state.effective_cpt_name]
309 vtk_lut = cpt_to_vtk_lookuptable(cpt, mask_zeros=mask_zeros)
310 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
312 self._lookuptable = vtk_lut
313 self._update_cptscale_lineedit()
315 elif state.effective_cpt_name and self._values is None:
316 raise ValueError('No values passed to colormapper!')
318 def update_cbar(self, display_parameter):
320 state = self._state
321 lut = self._lookuptable
323 if state.cbar_show and lut:
324 sx, sy = 1, 1
325 off = 0.08 * sy
326 pos = {
327 'top-left': (off, sy/2 + off, 0, 2),
328 'top-right': (sx - off, sy/2 + off, 2, 2),
329 'bottom-left': (off, off, 0, 0),
330 'bottom-right': (sx - off, off, 2, 0)}
331 x, y, _, _ = pos[state.cbar_position]
333 if not isinstance(self._cbar_pipe, ColorbarPipe):
334 self._cbar_pipe = ColorbarPipe(
335 parent_pipe=self._parent,
336 lut=lut,
337 cbar_title=display_parameter,
338 position=(x, y))
339 self._parent.add_actor(self._cbar_pipe.actor)
340 else:
341 self._cbar_pipe.set_lookuptable(lut)
342 self._cbar_pipe.set_title(display_parameter)
343 self._cbar_pipe._set_position(x, y)
345 sx, sy = self._parent.gui_state.size
346 fontsize = round(state.cbar_annotation_fontsize*sy)
347 lightness = 0.9 * state.cbar_annotation_lightness
348 self._cbar_pipe._format_text(
349 lightness=lightness, fontsize=fontsize)
351 height_px = int(round(sy / 3 * state.cbar_height))
352 width_px = int(round(50 * state.cbar_width))
353 self._cbar_pipe._format_size(height_px, width_px)
355 else:
356 self.remove_cbar_pipe()
358 def remove_cbar_pipe(self):
359 if self._cbar_pipe is not None:
360 self._parent.remove_actor(self._cbar_pipe.actor)
362 self._cbar_pipe = None
364 def cpt_controls(self, parent, state, layout):
365 self._parent = parent
367 iy = layout.rowCount() + 1
369 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
371 cb = common.CPTComboBox()
372 layout.addWidget(cb, iy, 1)
373 state_bind_combobox(
374 self, state, 'cpt_name', cb)
376 self._cpt_combobox = cb
378 pb = qw.QPushButton('Load CPT')
379 layout.addWidget(pb, iy, 2)
380 pb.clicked.connect(self.open_cpt_load_dialog)
382 iy += 1
383 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
385 cb = common.string_choices_to_combobox(AutoScaleMode)
386 layout.addWidget(cb, iy, 1)
387 state_bind_combobox(
388 self, state, 'cpt_mode', cb)
390 le = qw.QLineEdit()
391 le.setEnabled(False)
392 layout.addWidget(le, iy, 2)
393 state_bind(
394 self, state,
395 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
396 le, [le.editingFinished, le.returnPressed],
397 self._cptscale_to_lineedit)
399 self._cpt_scale_lineedit = le
401 iy += 1
402 cb = qw.QCheckBox('Revert')
403 layout.addWidget(cb, iy, 1)
404 state_bind_checkbox(self, state, 'cpt_revert', cb)
406 # color bar
407 iy += 1
408 layout.addWidget(qw.QLabel('Color Bar'), iy, 0)
410 chb = qw.QCheckBox('show')
411 layout.addWidget(chb, iy, 1)
412 state_bind_checkbox(self, state, 'cbar_show', chb)
414 cb = common.string_choices_to_combobox(
415 ColorBarPositionChoice)
416 layout.addWidget(cb, iy, 2)
417 state_bind_combobox(
418 self, self._state, 'cbar_position', cb)
420 # cbar text
421 iy += 1
422 layout.addWidget(qw.QLabel('Lightness'), iy, 1)
424 slider = qw.QSlider(qc.Qt.Horizontal)
425 slider.setSizePolicy(
426 qw.QSizePolicy(
427 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
428 slider.setMinimum(0)
429 slider.setMaximum(1000)
430 layout.addWidget(slider, iy, 2)
432 state_bind_slider(
433 self,
434 self._state,
435 'cbar_annotation_lightness',
436 slider,
437 factor=0.001)
439 iy += 1
440 layout.addWidget(qw.QLabel('Fontsize'), iy, 1)
442 slider = qw.QSlider(qc.Qt.Horizontal)
443 slider.setSizePolicy(
444 qw.QSizePolicy(
445 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
446 slider.setMinimum(0)
447 slider.setMaximum(100)
448 layout.addWidget(slider, iy, 2)
450 state_bind_slider(
451 self,
452 self._state,
453 'cbar_annotation_fontsize',
454 slider,
455 factor=0.001)
457 # cbar size
458 iy += 1
459 layout.addWidget(qw.QLabel('Height'), iy, 1)
461 slider = qw.QSlider(qc.Qt.Horizontal)
462 slider.setSizePolicy(
463 qw.QSizePolicy(
464 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
465 slider.setMinimum(1)
466 slider.setMaximum(200)
467 layout.addWidget(slider, iy, 2)
469 state_bind_slider(
470 self,
471 self._state,
472 'cbar_height',
473 slider,
474 factor=0.01)
476 iy += 1
477 layout.addWidget(qw.QLabel('Width'), iy, 1)
479 slider = qw.QSlider(qc.Qt.Horizontal)
480 slider.setSizePolicy(
481 qw.QSizePolicy(
482 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
483 slider.setMinimum(1)
484 slider.setMaximum(200)
485 layout.addWidget(slider, iy, 2)
487 state_bind_slider(
488 self,
489 self._state,
490 'cbar_width',
491 slider,
492 factor=0.01)
495def _lineedit_to_cptscale(widget, cpt_state):
496 s = str(widget.text())
497 s = s.replace(',', ' ')
499 crange = tuple((float(i) for i in s.split()))
500 crange = tuple((
501 crange[0],
502 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
504 try:
505 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
506 except Exception:
507 raise ValueError(
508 'need two numerical values: <vmin>, <vmax>')
511__all__ = [
512 'Element',
513 'ElementState',
514 'random_id',
515]