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
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 try:
39 from matplotlib import colormaps as mpl_cmaps
40 mpl_cmap_choices = list(mpl_cmaps.keys())
42 for cmap_name in mpl_cmap_blacklist:
43 try:
44 mpl_cmap_choices.remove(cmap_name)
45 mpl_cmap_choices.remove("%s_r" % cmap_name)
46 except ValueError:
47 pass
49 except ImportError:
50 mpl_cmap_choices = [
51 'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
53 mpl_cmap_choices.sort()
54 return mpl_cmap_choices
57def random_id():
58 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
61class ElementState(TalkieRoot):
63 element_id = String.T()
65 def __init__(self, **kwargs):
66 if 'element_id' not in kwargs:
67 kwargs['element_id'] = random_id()
69 TalkieRoot.__init__(self, **kwargs)
72class Element(TalkieConnectionOwner):
73 def __init__(self):
74 TalkieConnectionOwner.__init__(self)
75 self._parent = None
76 self._state = None
78 def remove(self):
79 if self._parent and self._state:
80 self._parent.state.elements.remove(self._state)
82 def set_parent(self, parent):
83 self._parent = parent
85 def unset_parent(self):
86 print(self)
87 raise NotImplementedError
89 def bind_state(self, state):
90 self._state = state
92 def unbind_state(self):
93 self.talkie_disconnect_all()
94 self._state = None
96 def update_visibility(self, visible):
97 self._state.visible = visible
99 def get_title_label(self):
100 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())
102 def update_label(*args):
103 title_label.set_slug(self._state.element_id)
105 self.talkie_connect(
106 self._state, 'element_id', update_label)
108 update_label()
109 return title_label
111 def get_title_control_remove(self):
112 button = common.MyDockWidgetTitleBarButton('\u2716')
113 button.setStatusTip('Remove Element')
114 button.clicked.connect(self.remove)
115 return button
117 def get_title_control_visible(self):
118 assert hasattr(self._state, 'visible')
120 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
121 button.setStatusTip('Toggle Element Visibility')
122 button.toggled.connect(self.update_visibility)
124 def set_button_checked(*args):
125 button.blockSignals(True)
126 button.set_checked(self._state.visible)
127 button.blockSignals(False)
129 set_button_checked()
131 self.talkie_connect(
132 self._state, 'visible', set_button_checked)
134 return button
137class CPTChoice(StringChoice):
139 choices = ['slip_colors'] + get_mpl_cmap_choices()
142class ColorBarPositionChoice(StringChoice):
143 choices = ['bottom-left', 'bottom-right', 'top-left', 'top-right']
146@has_computed
147class CPTState(ElementState):
148 cpt_name = String.T(default=CPTChoice.choices[0])
149 cpt_mode = String.T(default=AutoScaleMode.choices[1])
150 cpt_scale_min = Float.T(optional=True)
151 cpt_scale_max = Float.T(optional=True)
152 cpt_revert = Bool.T(default=False)
153 cbar_show = Bool.T(default=True)
154 cbar_position = ColorBarPositionChoice.T(default='bottom-right')
155 cbar_annotation_lightness = Float.T(default=1.0)
156 cbar_annotation_fontsize = Float.T(default=0.03)
157 cbar_height = Float.T(default=1.)
158 cbar_width = Float.T(default=1.)
160 @computed(['cpt_name', 'cpt_revert'])
161 def effective_cpt_name(self):
162 if self.cpt_revert:
163 return '%s_r' % self.cpt_name
164 else:
165 return self.cpt_name
168class CPTHandler(Element):
170 def __init__(self):
172 Element.__init__(self)
173 self._cpts = {}
174 self._autoscaler = None
175 self._lookuptable = None
176 self._cpt_combobox = None
177 self._values = None
178 self._state = None
179 self._cpt_scale_lineedit = None
180 self._cbar_pipe = None
182 def bind_state(self, cpt_state, update_function):
183 for state_attr in [
184 'effective_cpt_name', 'cpt_mode',
185 'cpt_scale_min', 'cpt_scale_max',
186 'cbar_show', 'cbar_position',
187 'cbar_annotation_lightness',
188 'cbar_annotation_fontsize',
189 'cbar_height', 'cbar_width']:
191 self.talkie_connect(
192 cpt_state, state_attr, update_function)
194 self._state = cpt_state
196 def unbind_state(self):
197 Element.unbind_state(self)
198 self._cpts = {}
199 self._lookuptable = None
200 self._values = None
201 self._autoscaler = None
203 def open_cpt_load_dialog(self):
204 caption = 'Select one *.cpt file to open'
206 fns, _ = qw.QFileDialog.getOpenFileNames(
207 self._parent, caption, options=common.qfiledialog_options)
209 if fns:
210 self.load_cpt_file(fns[0])
212 def load_cpt_file(self, path):
213 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
214 self._cpts.update([(cpt_name, automap.read_cpt(path))])
216 self._state.cpt_name = cpt_name
218 self._update_cpt_combobox()
219 self.update_cpt()
221 def _update_cpt_combobox(self):
222 from pyrocko import config
223 conf = config.config()
225 if self._cpt_combobox is None:
226 raise ValueError('CPT combobox needs init before updating!')
228 cb = self._cpt_combobox
230 if cb is not None:
231 cb.clear()
233 for s in CPTChoice.choices:
234 if s not in self._cpts:
235 try:
236 cpt = automap.read_cpt(topo.cpt(s))
237 except Exception:
238 from matplotlib import pyplot as plt
239 cmap = plt.cm.get_cmap(s)
240 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
242 self._cpts.update([(s, cpt)])
244 cpt_dir = conf.colortables_dir
245 if os.path.isdir(cpt_dir):
246 for f in [
247 f for f in os.listdir(cpt_dir)
248 if f.lower().endswith('.cpt')]:
250 s = 'USR' + os.path.basename(f).split('.')[0]
251 self._cpts.update(
252 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
254 for i, (s, cpt) in enumerate(self._cpts.items()):
255 if s[-2::] != "_r":
256 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
257 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
259 cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name))
261 def _update_cptscale_lineedit(self):
262 le = self._cpt_scale_lineedit
263 if le is not None:
264 le.clear()
266 self._cptscale_to_lineedit(self._state, le)
268 def _cptscale_to_lineedit(self, state, widget):
269 # sel = widget.selectedText() == widget.text()
271 crange = (None, None)
272 if self._lookuptable is not None:
273 crange = self._lookuptable.GetRange()
275 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
276 crange = state.cpt_scale_min, state.cpt_scale_max
278 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
280 widget.setText(fmt % crange)
282 # if sel:
283 # widget.selectAll()
285 def update_cpt(self, mask_zeros=False):
286 state = self._state
288 if self._autoscaler is None:
289 self._autoscaler = AutoScaler()
291 if self._cpt_scale_lineedit:
292 if state.cpt_mode == 'off':
293 self._cpt_scale_lineedit.setEnabled(True)
294 else:
295 self._cpt_scale_lineedit.setEnabled(False)
297 if state.cpt_scale_min is not None:
298 state.cpt_scale_min = None
300 if state.cpt_scale_max is not None:
301 state.cpt_scale_max = None
303 if state.effective_cpt_name is not None and self._values is not None:
304 if self._values.size == 0:
305 vscale = (0., 1.)
306 else:
307 vscale = (num.nanmin(self._values), num.nanmax(self._values))
309 vmin, vmax = None, None
310 if None not in (state.cpt_scale_min, state.cpt_scale_max):
311 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
312 else:
313 vmin, vmax, _ = self._autoscaler.make_scale(
314 vscale, override_mode=state.cpt_mode)
316 self._cpts[state.effective_cpt_name].scale(vmin, vmax)
317 cpt = self._cpts[state.effective_cpt_name]
318 vtk_lut = cpt_to_vtk_lookuptable(cpt, mask_zeros=mask_zeros)
319 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
321 self._lookuptable = vtk_lut
322 self._update_cptscale_lineedit()
324 elif state.effective_cpt_name and self._values is None:
325 raise ValueError('No values passed to colormapper!')
327 def update_cbar(self, display_parameter):
329 state = self._state
330 lut = self._lookuptable
332 if state.cbar_show and lut:
333 sx, sy = 1, 1
334 off = 0.08 * sy
335 pos = {
336 'top-left': (off, sy/2 + off, 0, 2),
337 'top-right': (sx - off, sy/2 + off, 2, 2),
338 'bottom-left': (off, off, 0, 0),
339 'bottom-right': (sx - off, off, 2, 0)}
340 x, y, _, _ = pos[state.cbar_position]
342 if not isinstance(self._cbar_pipe, ColorbarPipe):
343 self._cbar_pipe = ColorbarPipe(
344 parent_pipe=self._parent,
345 lut=lut,
346 cbar_title=display_parameter,
347 position=(x, y))
348 self._parent.add_actor(self._cbar_pipe.actor)
349 else:
350 self._cbar_pipe.set_lookuptable(lut)
351 self._cbar_pipe.set_title(display_parameter)
352 self._cbar_pipe._set_position(x, y)
354 sx, sy = self._parent.gui_state.size
355 fontsize = round(state.cbar_annotation_fontsize*sy)
356 lightness = 0.9 * state.cbar_annotation_lightness
357 self._cbar_pipe._format_text(
358 lightness=lightness, fontsize=fontsize)
360 height_px = int(round(sy / 3 * state.cbar_height))
361 width_px = int(round(50 * state.cbar_width))
362 self._cbar_pipe._format_size(height_px, width_px)
364 else:
365 self.remove_cbar_pipe()
367 def remove_cbar_pipe(self):
368 if self._cbar_pipe is not None:
369 self._parent.remove_actor(self._cbar_pipe.actor)
371 self._cbar_pipe = None
373 def cpt_controls(self, parent, state, layout):
374 self._parent = parent
376 iy = layout.rowCount() + 1
378 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
380 cb = common.CPTComboBox()
381 layout.addWidget(cb, iy, 1)
382 state_bind_combobox(
383 self, state, 'cpt_name', cb)
385 self._cpt_combobox = cb
387 pb = qw.QPushButton('Load CPT')
388 layout.addWidget(pb, iy, 2)
389 pb.clicked.connect(self.open_cpt_load_dialog)
391 iy += 1
392 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
394 cb = common.string_choices_to_combobox(AutoScaleMode)
395 layout.addWidget(cb, iy, 1)
396 state_bind_combobox(
397 self, state, 'cpt_mode', cb)
399 le = qw.QLineEdit()
400 le.setEnabled(False)
401 layout.addWidget(le, iy, 2)
402 state_bind(
403 self, state,
404 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
405 le, [le.editingFinished, le.returnPressed],
406 self._cptscale_to_lineedit)
408 self._cpt_scale_lineedit = le
410 iy += 1
411 cb = qw.QCheckBox('Revert')
412 layout.addWidget(cb, iy, 1)
413 state_bind_checkbox(self, state, 'cpt_revert', cb)
415 # color bar
416 iy += 1
417 layout.addWidget(qw.QLabel('Color Bar'), iy, 0)
419 chb = qw.QCheckBox('show')
420 layout.addWidget(chb, iy, 1)
421 state_bind_checkbox(self, state, 'cbar_show', chb)
423 cb = common.string_choices_to_combobox(
424 ColorBarPositionChoice)
425 layout.addWidget(cb, iy, 2)
426 state_bind_combobox(
427 self, self._state, 'cbar_position', cb)
429 # cbar text
430 iy += 1
431 layout.addWidget(qw.QLabel('Lightness'), iy, 1)
433 slider = qw.QSlider(qc.Qt.Horizontal)
434 slider.setSizePolicy(
435 qw.QSizePolicy(
436 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
437 slider.setMinimum(0)
438 slider.setMaximum(1000)
439 layout.addWidget(slider, iy, 2)
441 state_bind_slider(
442 self,
443 self._state,
444 'cbar_annotation_lightness',
445 slider,
446 factor=0.001)
448 iy += 1
449 layout.addWidget(qw.QLabel('Fontsize'), iy, 1)
451 slider = qw.QSlider(qc.Qt.Horizontal)
452 slider.setSizePolicy(
453 qw.QSizePolicy(
454 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
455 slider.setMinimum(0)
456 slider.setMaximum(100)
457 layout.addWidget(slider, iy, 2)
459 state_bind_slider(
460 self,
461 self._state,
462 'cbar_annotation_fontsize',
463 slider,
464 factor=0.001)
466 # cbar size
467 iy += 1
468 layout.addWidget(qw.QLabel('Height'), iy, 1)
470 slider = qw.QSlider(qc.Qt.Horizontal)
471 slider.setSizePolicy(
472 qw.QSizePolicy(
473 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
474 slider.setMinimum(1)
475 slider.setMaximum(200)
476 layout.addWidget(slider, iy, 2)
478 state_bind_slider(
479 self,
480 self._state,
481 'cbar_height',
482 slider,
483 factor=0.01)
485 iy += 1
486 layout.addWidget(qw.QLabel('Width'), iy, 1)
488 slider = qw.QSlider(qc.Qt.Horizontal)
489 slider.setSizePolicy(
490 qw.QSizePolicy(
491 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
492 slider.setMinimum(1)
493 slider.setMaximum(200)
494 layout.addWidget(slider, iy, 2)
496 state_bind_slider(
497 self,
498 self._state,
499 'cbar_width',
500 slider,
501 factor=0.01)
504def _lineedit_to_cptscale(widget, cpt_state):
505 s = str(widget.text())
506 s = s.replace(',', ' ')
508 crange = tuple((float(i) for i in s.split()))
509 crange = tuple((
510 crange[0],
511 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
513 try:
514 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
515 except Exception:
516 raise ValueError(
517 'need two numerical values: <vmin>, <vmax>')
520__all__ = [
521 'Element',
522 'ElementState',
523 'random_id',
524]