1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function, division
8import os
9import base64
11import numpy as num
13from pyrocko import automap
14from pyrocko.guts import String, Float, StringChoice
15from pyrocko.plot import AutoScaler, AutoScaleMode
16from pyrocko.dataset import topo
18from pyrocko.gui.talkie import TalkieRoot
19from pyrocko.gui.qt_compat import qc, qw
20from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable
22from .. import common
23from ..state import \
24 state_bind_combobox, state_bind
27def random_id():
28 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
31class ElementState(TalkieRoot):
33 element_id = String.T()
35 def __init__(self, **kwargs):
36 if 'element_id' not in kwargs:
37 kwargs['element_id'] = random_id()
39 TalkieRoot.__init__(self, **kwargs)
42class Element(object):
43 def __init__(self):
44 self._listeners = []
45 self._parent = None
46 self._state = None
48 def register_state_listener(self, listener):
49 self._listeners.append(listener) # keep listeners alive
51 def register_state_listener3(self, listener, state, path):
52 self.register_state_listener(state.add_listener(listener, path))
54 def remove(self):
55 if self._parent and self._state:
56 self._parent.state.elements.remove(self._state)
58 def set_parent(self, parent):
59 self._parent = parent
61 def unset_parent(self):
62 print(self)
63 raise NotImplementedError
65 def bind_state(self, state):
66 self._state = state
68 def unbind_state(self):
69 for listener in self._listeners:
70 try:
71 listener.release()
72 except Exception:
73 pass
75 self._listeners = []
76 self._state = None
78 def update_visibility(self, visible):
79 self._state.visible = visible
81 def get_title_control_remove(self):
82 button = common.MyDockWidgetTitleBarButton('\u00d7')
83 button.setStatusTip('Remove Element')
84 button.clicked.connect(self.remove)
85 return button
87 def get_title_control_visible(self):
88 assert hasattr(self._state, 'visible')
90 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
91 button.setStatusTip('Toggle Element Visibility')
92 button.toggled.connect(self.update_visibility)
94 def set_button_checked(*args):
95 button.blockSignals(True)
96 button.set_checked(self._state.visible)
97 button.blockSignals(False)
99 set_button_checked()
101 self.register_state_listener3(
102 set_button_checked, self._state, 'visible')
104 return button
107class CPTChoice(StringChoice):
108 choices = [
109 'slip_colors', 'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
112class CPTState(ElementState):
113 cpt_name = String.T(default=CPTChoice.choices[0])
114 cpt_mode = String.T(default=AutoScaleMode.choices[1])
115 cpt_scale_min = Float.T(optional=True)
116 cpt_scale_max = Float.T(optional=True)
119class CPTHandler(Element):
121 def __init__(self):
123 Element.__init__(self)
124 self._cpts = {}
125 self._autoscaler = None
126 self._lookuptable = None
127 self._cpt_combobox = None
128 self._values = None
129 self._state = None
130 self._cpt_scale_lineedit = None
132 def bind_state(self, cpt_state, update_function):
133 for state_attr in [
134 'cpt_name', 'cpt_mode', 'cpt_scale_min', 'cpt_scale_max']:
136 self.register_state_listener3(
137 update_function, cpt_state, state_attr)
139 self._state = cpt_state
141 def unbind_state(self):
142 self._cpts = {}
143 self._lookuptable = None
144 self._values = None
145 self._autoscaler = None
147 def open_cpt_load_dialog(self):
148 caption = 'Select one *.cpt file to open'
150 fns, _ = qw.QFileDialog.getOpenFileNames(
151 self._parent, caption, options=common.qfiledialog_options)
153 if fns:
154 self.load_cpt_file(fns[0])
156 def load_cpt_file(self, path):
157 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
158 self._cpts.update([(cpt_name, automap.read_cpt(path))])
160 self._state.cpt_name = cpt_name
162 self._update_cpt_combobox()
163 self.update_cpt()
165 def _update_cpt_combobox(self):
166 from pyrocko import config
167 conf = config.config()
169 if self._cpt_combobox is None:
170 raise ValueError('CPT combobox needs init before updating!')
172 cb = self._cpt_combobox
174 if cb is not None:
175 cb.clear()
177 for s in CPTChoice.choices:
178 if s not in self._cpts:
179 try:
180 cpt = automap.read_cpt(topo.cpt(s))
181 except Exception:
182 from matplotlib import pyplot as plt
183 cmap = plt.cm.get_cmap(s)
184 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
186 self._cpts.update([(s, cpt)])
188 cpt_dir = conf.colortables_dir
189 if os.path.isdir(cpt_dir):
190 for f in [
191 f for f in os.listdir(cpt_dir)
192 if f.lower().endswith('.cpt')]:
194 s = 'USR' + os.path.basename(f).split('.')[0]
195 self._cpts.update(
196 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
198 for i, (s, cpt) in enumerate(self._cpts.items()):
199 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
200 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
202 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
204 def _update_cptscale_lineedit(self):
205 le = self._cpt_scale_lineedit
206 if le is not None:
207 le.clear()
209 self._cptscale_to_lineedit(self._state, le)
211 def _cptscale_to_lineedit(self, state, widget):
212 # sel = widget.selectedText() == widget.text()
214 crange = (None, None)
215 if self._lookuptable is not None:
216 crange = self._lookuptable.GetRange()
218 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
219 crange = state.cpt_scale_min, state.cpt_scale_max
221 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
223 widget.setText(fmt % crange)
225 # if sel:
226 # widget.selectAll()
228 def update_cpt(self):
229 state = self._state
231 if self._autoscaler is None:
232 self._autoscaler = AutoScaler()
234 if self._cpt_scale_lineedit:
235 if state.cpt_mode == 'off':
236 self._cpt_scale_lineedit.setEnabled(True)
237 else:
238 self._cpt_scale_lineedit.setEnabled(False)
240 if state.cpt_scale_min is not None:
241 state.cpt_scale_min = None
243 if state.cpt_scale_max is not None:
244 state.cpt_scale_max = None
246 if state.cpt_name is not None and self._values is not None:
247 vscale = (num.nanmin(self._values), num.nanmax(self._values))
249 vmin, vmax = None, None
250 if None not in (state.cpt_scale_min, state.cpt_scale_max):
251 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
252 else:
253 vmin, vmax, _ = self._autoscaler.make_scale(
254 vscale, override_mode=state.cpt_mode)
256 self._cpts[state.cpt_name].scale(vmin, vmax)
257 cpt = self._cpts[state.cpt_name]
259 vtk_lut = cpt_to_vtk_lookuptable(cpt)
260 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
262 self._lookuptable = vtk_lut
263 self._update_cptscale_lineedit()
265 elif state.cpt_name and self._values is None:
266 raise ValueError('No values passed to colormapper!')
268 def cpt_controls(self, parent, state, layout):
269 self._parent = parent
271 iy = layout.rowCount() + 1
273 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
275 cb = common.CPTComboBox()
276 layout.addWidget(cb, iy, 1)
277 state_bind_combobox(
278 self, state, 'cpt_name', cb)
280 self._cpt_combobox = cb
282 pb = qw.QPushButton('Load CPT')
283 layout.addWidget(pb, iy, 2)
284 pb.clicked.connect(self.open_cpt_load_dialog)
286 iy += 1
287 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
289 cb = common.string_choices_to_combobox(AutoScaleMode)
290 layout.addWidget(cb, iy, 1)
291 state_bind_combobox(
292 self, state, 'cpt_mode', cb)
294 le = qw.QLineEdit()
295 le.setEnabled(False)
296 layout.addWidget(le, iy, 2)
297 state_bind(
298 self, state,
299 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
300 le, [le.editingFinished, le.returnPressed],
301 self._cptscale_to_lineedit)
303 self._cpt_scale_lineedit = le
306def _lineedit_to_cptscale(widget, cpt_state):
307 s = str(widget.text())
308 s = s.replace(',', ' ')
310 crange = tuple((float(i) for i in s.split()))
311 crange = tuple((
312 crange[0],
313 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
315 try:
316 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
317 except Exception:
318 raise ValueError(
319 'need two numerical values: <vmin>, <vmax>')
322__all__ = [
323 'Element',
324 'ElementState',
325 'random_id',
326]