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
23from .. import common
24from ..state import \
25 state_bind_combobox, state_bind, state_bind_checkbox
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()
142@has_computed
143class CPTState(ElementState):
144 cpt_name = String.T(default=CPTChoice.choices[0])
145 cpt_mode = String.T(default=AutoScaleMode.choices[1])
146 cpt_scale_min = Float.T(optional=True)
147 cpt_scale_max = Float.T(optional=True)
148 cpt_revert = Bool.T(default=False)
150 @computed(['cpt_name', 'cpt_revert'])
151 def effective_cpt_name(self):
152 if self.cpt_revert:
153 return '%s_r' % self.cpt_name
154 else:
155 return self.cpt_name
158class CPTHandler(Element):
160 def __init__(self):
162 Element.__init__(self)
163 self._cpts = {}
164 self._autoscaler = None
165 self._lookuptable = None
166 self._cpt_combobox = None
167 self._values = None
168 self._state = None
169 self._cpt_scale_lineedit = None
171 def bind_state(self, cpt_state, update_function):
172 for state_attr in [
173 'effective_cpt_name', 'cpt_mode',
174 'cpt_scale_min', 'cpt_scale_max']:
176 self.talkie_connect(
177 cpt_state, state_attr, update_function)
179 self._state = cpt_state
181 def unbind_state(self):
182 Element.unbind_state(self)
183 self._cpts = {}
184 self._lookuptable = None
185 self._values = None
186 self._autoscaler = None
188 def open_cpt_load_dialog(self):
189 caption = 'Select one *.cpt file to open'
191 fns, _ = qw.QFileDialog.getOpenFileNames(
192 self._parent, caption, options=common.qfiledialog_options)
194 if fns:
195 self.load_cpt_file(fns[0])
197 def load_cpt_file(self, path):
198 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
199 self._cpts.update([(cpt_name, automap.read_cpt(path))])
201 self._state.cpt_name = cpt_name
203 self._update_cpt_combobox()
204 self.update_cpt()
206 def _update_cpt_combobox(self):
207 from pyrocko import config
208 conf = config.config()
210 if self._cpt_combobox is None:
211 raise ValueError('CPT combobox needs init before updating!')
213 cb = self._cpt_combobox
215 if cb is not None:
216 cb.clear()
218 for s in CPTChoice.choices:
219 if s not in self._cpts:
220 try:
221 cpt = automap.read_cpt(topo.cpt(s))
222 except Exception:
223 from matplotlib import pyplot as plt
224 cmap = plt.cm.get_cmap(s)
225 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
227 self._cpts.update([(s, cpt)])
229 cpt_dir = conf.colortables_dir
230 if os.path.isdir(cpt_dir):
231 for f in [
232 f for f in os.listdir(cpt_dir)
233 if f.lower().endswith('.cpt')]:
235 s = 'USR' + os.path.basename(f).split('.')[0]
236 self._cpts.update(
237 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
239 for i, (s, cpt) in enumerate(self._cpts.items()):
240 if s[-2::] != "_r":
241 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
242 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
244 cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name))
246 def _update_cptscale_lineedit(self):
247 le = self._cpt_scale_lineedit
248 if le is not None:
249 le.clear()
251 self._cptscale_to_lineedit(self._state, le)
253 def _cptscale_to_lineedit(self, state, widget):
254 # sel = widget.selectedText() == widget.text()
256 crange = (None, None)
257 if self._lookuptable is not None:
258 crange = self._lookuptable.GetRange()
260 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
261 crange = state.cpt_scale_min, state.cpt_scale_max
263 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
265 widget.setText(fmt % crange)
267 # if sel:
268 # widget.selectAll()
270 def update_cpt(self):
271 state = self._state
273 if self._autoscaler is None:
274 self._autoscaler = AutoScaler()
276 if self._cpt_scale_lineedit:
277 if state.cpt_mode == 'off':
278 self._cpt_scale_lineedit.setEnabled(True)
279 else:
280 self._cpt_scale_lineedit.setEnabled(False)
282 if state.cpt_scale_min is not None:
283 state.cpt_scale_min = None
285 if state.cpt_scale_max is not None:
286 state.cpt_scale_max = None
288 if state.effective_cpt_name is not None and self._values is not None:
289 if self._values.size == 0:
290 vscale = (0., 1.)
291 else:
292 vscale = (num.nanmin(self._values), num.nanmax(self._values))
294 vmin, vmax = None, None
295 if None not in (state.cpt_scale_min, state.cpt_scale_max):
296 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
297 else:
298 vmin, vmax, _ = self._autoscaler.make_scale(
299 vscale, override_mode=state.cpt_mode)
301 self._cpts[state.effective_cpt_name].scale(vmin, vmax)
302 cpt = self._cpts[state.effective_cpt_name]
303 vtk_lut = cpt_to_vtk_lookuptable(cpt)
304 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
306 self._lookuptable = vtk_lut
307 self._update_cptscale_lineedit()
309 elif state.effective_cpt_name and self._values is None:
310 raise ValueError('No values passed to colormapper!')
312 def cpt_controls(self, parent, state, layout):
313 self._parent = parent
315 iy = layout.rowCount() + 1
317 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
319 cb = common.CPTComboBox()
320 layout.addWidget(cb, iy, 1)
321 state_bind_combobox(
322 self, state, 'cpt_name', cb)
324 self._cpt_combobox = cb
326 pb = qw.QPushButton('Load CPT')
327 layout.addWidget(pb, iy, 2)
328 pb.clicked.connect(self.open_cpt_load_dialog)
330 iy += 1
331 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
333 cb = common.string_choices_to_combobox(AutoScaleMode)
334 layout.addWidget(cb, iy, 1)
335 state_bind_combobox(
336 self, state, 'cpt_mode', cb)
338 le = qw.QLineEdit()
339 le.setEnabled(False)
340 layout.addWidget(le, iy, 2)
341 state_bind(
342 self, state,
343 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
344 le, [le.editingFinished, le.returnPressed],
345 self._cptscale_to_lineedit)
347 self._cpt_scale_lineedit = le
349 iy += 1
350 cb = qw.QCheckBox('Revert')
351 layout.addWidget(cb, iy, 1)
352 state_bind_checkbox(self, state, 'cpt_revert', cb)
355def _lineedit_to_cptscale(widget, cpt_state):
356 s = str(widget.text())
357 s = s.replace(',', ' ')
359 crange = tuple((float(i) for i in s.split()))
360 crange = tuple((
361 crange[0],
362 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
364 try:
365 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
366 except Exception:
367 raise ValueError(
368 'need two numerical values: <vmin>, <vmax>')
371__all__ = [
372 'Element',
373 'ElementState',
374 'random_id',
375]