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
13from pyrocko.plot import AutoScaler, AutoScaleMode
14from pyrocko.dataset import topo
16from pyrocko.gui.talkie import TalkieRoot, TalkieConnectionOwner
17from pyrocko.gui.qt_compat import qc, qw
18from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable
20try:
21 from matplotlib import colormaps as mpl_cmaps
22 mpl_cmap_choices = list(mpl_cmaps.keys())
23except ImportError:
24 mpl_cmap_choices = ['seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
27from .. import common
28from ..state import \
29 state_bind_combobox, state_bind
32def random_id():
33 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
36class ElementState(TalkieRoot):
38 element_id = String.T()
40 def __init__(self, **kwargs):
41 if 'element_id' not in kwargs:
42 kwargs['element_id'] = random_id()
44 TalkieRoot.__init__(self, **kwargs)
47class Element(TalkieConnectionOwner):
48 def __init__(self):
49 TalkieConnectionOwner.__init__(self)
50 self._parent = None
51 self._state = None
53 def remove(self):
54 if self._parent and self._state:
55 self._parent.state.elements.remove(self._state)
57 def set_parent(self, parent):
58 self._parent = parent
60 def unset_parent(self):
61 print(self)
62 raise NotImplementedError
64 def bind_state(self, state):
65 self._state = state
67 def unbind_state(self):
68 self.talkie_disconnect_all()
69 self._state = None
71 def update_visibility(self, visible):
72 self._state.visible = visible
74 def get_title_label(self):
75 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())
77 def update_label(*args):
78 title_label.set_slug(self._state.element_id)
80 self.talkie_connect(
81 self._state, 'element_id', update_label)
83 update_label()
84 return title_label
86 def get_title_control_remove(self):
87 button = common.MyDockWidgetTitleBarButton('\u2716')
88 button.setStatusTip('Remove Element')
89 button.clicked.connect(self.remove)
90 return button
92 def get_title_control_visible(self):
93 assert hasattr(self._state, 'visible')
95 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
96 button.setStatusTip('Toggle Element Visibility')
97 button.toggled.connect(self.update_visibility)
99 def set_button_checked(*args):
100 button.blockSignals(True)
101 button.set_checked(self._state.visible)
102 button.blockSignals(False)
104 set_button_checked()
106 self.talkie_connect(
107 self._state, 'visible', set_button_checked)
109 return button
112class CPTChoice(StringChoice):
114 choices = ['slip_colors'] + mpl_cmap_choices
117class CPTState(ElementState):
118 cpt_name = String.T(default=CPTChoice.choices[0])
119 cpt_mode = String.T(default=AutoScaleMode.choices[1])
120 cpt_scale_min = Float.T(optional=True)
121 cpt_scale_max = Float.T(optional=True)
124class CPTHandler(Element):
126 def __init__(self):
128 Element.__init__(self)
129 self._cpts = {}
130 self._autoscaler = None
131 self._lookuptable = None
132 self._cpt_combobox = None
133 self._values = None
134 self._state = None
135 self._cpt_scale_lineedit = None
137 def bind_state(self, cpt_state, update_function):
138 for state_attr in [
139 'cpt_name', 'cpt_mode', 'cpt_scale_min', 'cpt_scale_max']:
141 self.talkie_connect(
142 cpt_state, state_attr, update_function)
144 self._state = cpt_state
146 def unbind_state(self):
147 Element.unbind_state(self)
148 self._cpts = {}
149 self._lookuptable = None
150 self._values = None
151 self._autoscaler = None
153 def open_cpt_load_dialog(self):
154 caption = 'Select one *.cpt file to open'
156 fns, _ = qw.QFileDialog.getOpenFileNames(
157 self._parent, caption, options=common.qfiledialog_options)
159 if fns:
160 self.load_cpt_file(fns[0])
162 def load_cpt_file(self, path):
163 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
164 self._cpts.update([(cpt_name, automap.read_cpt(path))])
166 self._state.cpt_name = cpt_name
168 self._update_cpt_combobox()
169 self.update_cpt()
171 def _update_cpt_combobox(self):
172 from pyrocko import config
173 conf = config.config()
175 if self._cpt_combobox is None:
176 raise ValueError('CPT combobox needs init before updating!')
178 cb = self._cpt_combobox
180 if cb is not None:
181 cb.clear()
183 for s in CPTChoice.choices:
184 if s not in self._cpts:
185 try:
186 cpt = automap.read_cpt(topo.cpt(s))
187 except Exception:
188 from matplotlib import pyplot as plt
189 cmap = plt.cm.get_cmap(s)
190 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
192 self._cpts.update([(s, cpt)])
194 cpt_dir = conf.colortables_dir
195 if os.path.isdir(cpt_dir):
196 for f in [
197 f for f in os.listdir(cpt_dir)
198 if f.lower().endswith('.cpt')]:
200 s = 'USR' + os.path.basename(f).split('.')[0]
201 self._cpts.update(
202 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
204 for i, (s, cpt) in enumerate(self._cpts.items()):
205 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
206 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
208 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
210 def _update_cptscale_lineedit(self):
211 le = self._cpt_scale_lineedit
212 if le is not None:
213 le.clear()
215 self._cptscale_to_lineedit(self._state, le)
217 def _cptscale_to_lineedit(self, state, widget):
218 # sel = widget.selectedText() == widget.text()
220 crange = (None, None)
221 if self._lookuptable is not None:
222 crange = self._lookuptable.GetRange()
224 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
225 crange = state.cpt_scale_min, state.cpt_scale_max
227 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
229 widget.setText(fmt % crange)
231 # if sel:
232 # widget.selectAll()
234 def update_cpt(self):
235 state = self._state
237 if self._autoscaler is None:
238 self._autoscaler = AutoScaler()
240 if self._cpt_scale_lineedit:
241 if state.cpt_mode == 'off':
242 self._cpt_scale_lineedit.setEnabled(True)
243 else:
244 self._cpt_scale_lineedit.setEnabled(False)
246 if state.cpt_scale_min is not None:
247 state.cpt_scale_min = None
249 if state.cpt_scale_max is not None:
250 state.cpt_scale_max = None
252 if state.cpt_name is not None and self._values is not None:
253 if self._values.size == 0:
254 vscale = (0., 1.)
255 else:
256 vscale = (num.nanmin(self._values), num.nanmax(self._values))
258 vmin, vmax = None, None
259 if None not in (state.cpt_scale_min, state.cpt_scale_max):
260 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
261 else:
262 vmin, vmax, _ = self._autoscaler.make_scale(
263 vscale, override_mode=state.cpt_mode)
265 self._cpts[state.cpt_name].scale(vmin, vmax)
266 cpt = self._cpts[state.cpt_name]
268 vtk_lut = cpt_to_vtk_lookuptable(cpt)
269 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
271 self._lookuptable = vtk_lut
272 self._update_cptscale_lineedit()
274 elif state.cpt_name and self._values is None:
275 raise ValueError('No values passed to colormapper!')
277 def cpt_controls(self, parent, state, layout):
278 self._parent = parent
280 iy = layout.rowCount() + 1
282 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
284 cb = common.CPTComboBox()
285 layout.addWidget(cb, iy, 1)
286 state_bind_combobox(
287 self, state, 'cpt_name', cb)
289 self._cpt_combobox = cb
291 pb = qw.QPushButton('Load CPT')
292 layout.addWidget(pb, iy, 2)
293 pb.clicked.connect(self.open_cpt_load_dialog)
295 iy += 1
296 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
298 cb = common.string_choices_to_combobox(AutoScaleMode)
299 layout.addWidget(cb, iy, 1)
300 state_bind_combobox(
301 self, state, 'cpt_mode', cb)
303 le = qw.QLineEdit()
304 le.setEnabled(False)
305 layout.addWidget(le, iy, 2)
306 state_bind(
307 self, state,
308 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
309 le, [le.editingFinished, le.returnPressed],
310 self._cptscale_to_lineedit)
312 self._cpt_scale_lineedit = le
315def _lineedit_to_cptscale(widget, cpt_state):
316 s = str(widget.text())
317 s = s.replace(',', ' ')
319 crange = tuple((float(i) for i in s.split()))
320 crange = tuple((
321 crange[0],
322 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
324 try:
325 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
326 except Exception:
327 raise ValueError(
328 'need two numerical values: <vmin>, <vmax>')
331__all__ = [
332 'Element',
333 'ElementState',
334 'random_id',
335]