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 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
17from pyrocko.gui.qt_compat import qc, qw
18from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable
20from .. import common
21from ..state import \
22 state_bind_combobox, state_bind
25def random_id():
26 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
29class ElementState(TalkieRoot):
31 element_id = String.T()
33 def __init__(self, **kwargs):
34 if 'element_id' not in kwargs:
35 kwargs['element_id'] = random_id()
37 TalkieRoot.__init__(self, **kwargs)
40class Element(object):
41 def __init__(self):
42 self._listeners = []
43 self._parent = None
44 self._state = None
46 def register_state_listener(self, listener):
47 self._listeners.append(listener) # keep listeners alive
49 def register_state_listener3(self, listener, state, path):
50 self.register_state_listener(state.add_listener(listener, path))
52 def remove(self):
53 if self._parent and self._state:
54 self._parent.state.elements.remove(self._state)
56 def set_parent(self, parent):
57 self._parent = parent
59 def unset_parent(self):
60 print(self)
61 raise NotImplementedError
63 def bind_state(self, state):
64 self._state = state
66 def unbind_state(self):
67 for listener in self._listeners:
68 try:
69 listener.release()
70 except Exception:
71 pass
73 self._listeners = []
74 self._state = None
76 def update_visibility(self, visible):
77 self._state.visible = visible
79 def get_title_control_remove(self):
80 button = common.MyDockWidgetTitleBarButton('\u00d7')
81 button.setStatusTip('Remove Element')
82 button.clicked.connect(self.remove)
83 return button
85 def get_title_control_visible(self):
86 assert hasattr(self._state, 'visible')
88 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
89 button.setStatusTip('Toggle Element Visibility')
90 button.toggled.connect(self.update_visibility)
92 def set_button_checked(*args):
93 button.blockSignals(True)
94 button.set_checked(self._state.visible)
95 button.blockSignals(False)
97 set_button_checked()
99 self.register_state_listener3(
100 set_button_checked, self._state, 'visible')
102 return button
105class CPTChoice(StringChoice):
106 choices = [
107 'slip_colors', 'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
110class CPTState(ElementState):
111 cpt_name = String.T(default=CPTChoice.choices[0])
112 cpt_mode = String.T(default=AutoScaleMode.choices[1])
113 cpt_scale_min = Float.T(optional=True)
114 cpt_scale_max = Float.T(optional=True)
117class CPTHandler(Element):
119 def __init__(self):
121 Element.__init__(self)
122 self._cpts = {}
123 self._autoscaler = None
124 self._lookuptable = None
125 self._cpt_combobox = None
126 self._values = None
127 self._state = None
128 self._cpt_scale_lineedit = None
130 def bind_state(self, cpt_state, update_function):
131 for state_attr in [
132 'cpt_name', 'cpt_mode', 'cpt_scale_min', 'cpt_scale_max']:
134 self.register_state_listener3(
135 update_function, cpt_state, state_attr)
137 self._state = cpt_state
139 def unbind_state(self):
140 self._cpts = {}
141 self._lookuptable = None
142 self._values = None
143 self._autoscaler = None
145 def open_cpt_load_dialog(self):
146 caption = 'Select one *.cpt file to open'
148 fns, _ = qw.QFileDialog.getOpenFileNames(
149 self._parent, caption, options=common.qfiledialog_options)
151 if fns:
152 self.load_cpt_file(fns[0])
154 def load_cpt_file(self, path):
155 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
156 self._cpts.update([(cpt_name, automap.read_cpt(path))])
158 self._state.cpt_name = cpt_name
160 self._update_cpt_combobox()
161 self.update_cpt()
163 def _update_cpt_combobox(self):
164 from pyrocko import config
165 conf = config.config()
167 if self._cpt_combobox is None:
168 raise ValueError('CPT combobox needs init before updating!')
170 cb = self._cpt_combobox
172 if cb is not None:
173 cb.clear()
175 for s in CPTChoice.choices:
176 if s not in self._cpts:
177 try:
178 cpt = automap.read_cpt(topo.cpt(s))
179 except Exception:
180 from matplotlib import pyplot as plt
181 cmap = plt.cm.get_cmap(s)
182 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
184 self._cpts.update([(s, cpt)])
186 cpt_dir = conf.colortables_dir
187 if os.path.isdir(cpt_dir):
188 for f in [
189 f for f in os.listdir(cpt_dir)
190 if f.lower().endswith('.cpt')]:
192 s = 'USR' + os.path.basename(f).split('.')[0]
193 self._cpts.update(
194 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
196 for i, (s, cpt) in enumerate(self._cpts.items()):
197 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
198 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
200 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
202 def _update_cptscale_lineedit(self):
203 le = self._cpt_scale_lineedit
204 if le is not None:
205 le.clear()
207 self._cptscale_to_lineedit(self._state, le)
209 def _cptscale_to_lineedit(self, state, widget):
210 # sel = widget.selectedText() == widget.text()
212 crange = (None, None)
213 if self._lookuptable is not None:
214 crange = self._lookuptable.GetRange()
216 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
217 crange = state.cpt_scale_min, state.cpt_scale_max
219 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
221 widget.setText(fmt % crange)
223 # if sel:
224 # widget.selectAll()
226 def update_cpt(self):
227 state = self._state
229 if self._autoscaler is None:
230 self._autoscaler = AutoScaler()
232 if self._cpt_scale_lineedit:
233 if state.cpt_mode == 'off':
234 self._cpt_scale_lineedit.setEnabled(True)
235 else:
236 self._cpt_scale_lineedit.setEnabled(False)
238 if state.cpt_scale_min is not None:
239 state.cpt_scale_min = None
241 if state.cpt_scale_max is not None:
242 state.cpt_scale_max = None
244 if state.cpt_name is not None and self._values is not None:
245 vscale = (num.nanmin(self._values), num.nanmax(self._values))
247 vmin, vmax = None, None
248 if None not in (state.cpt_scale_min, state.cpt_scale_max):
249 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
250 else:
251 vmin, vmax, _ = self._autoscaler.make_scale(
252 vscale, override_mode=state.cpt_mode)
254 self._cpts[state.cpt_name].scale(vmin, vmax)
255 cpt = self._cpts[state.cpt_name]
257 vtk_lut = cpt_to_vtk_lookuptable(cpt)
258 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
260 self._lookuptable = vtk_lut
261 self._update_cptscale_lineedit()
263 elif state.cpt_name and self._values is None:
264 raise ValueError('No values passed to colormapper!')
266 def cpt_controls(self, parent, state, layout):
267 self._parent = parent
269 iy = layout.rowCount() + 1
271 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
273 cb = common.CPTComboBox()
274 layout.addWidget(cb, iy, 1)
275 state_bind_combobox(
276 self, state, 'cpt_name', cb)
278 self._cpt_combobox = cb
280 pb = qw.QPushButton('Load CPT')
281 layout.addWidget(pb, iy, 2)
282 pb.clicked.connect(self.open_cpt_load_dialog)
284 iy += 1
285 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
287 cb = common.string_choices_to_combobox(AutoScaleMode)
288 layout.addWidget(cb, iy, 1)
289 state_bind_combobox(
290 self, state, 'cpt_mode', cb)
292 le = qw.QLineEdit()
293 le.setEnabled(False)
294 layout.addWidget(le, iy, 2)
295 state_bind(
296 self, state,
297 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
298 le, [le.editingFinished, le.returnPressed],
299 self._cptscale_to_lineedit)
301 self._cpt_scale_lineedit = le
304def _lineedit_to_cptscale(widget, cpt_state):
305 s = str(widget.text())
306 s = s.replace(',', ' ')
308 crange = tuple((float(i) for i in s.split()))
309 crange = tuple((
310 crange[0],
311 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
313 try:
314 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
315 except Exception:
316 raise ValueError(
317 'need two numerical values: <vmin>, <vmax>')
320__all__ = [
321 'Element',
322 'ElementState',
323 'random_id',
324]