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
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(TalkieConnectionOwner):
41 def __init__(self):
42 TalkieConnectionOwner.__init__(self)
43 self._parent = None
44 self._state = None
46 def remove(self):
47 if self._parent and self._state:
48 self._parent.state.elements.remove(self._state)
50 def set_parent(self, parent):
51 self._parent = parent
53 def unset_parent(self):
54 print(self)
55 raise NotImplementedError
57 def bind_state(self, state):
58 self._state = state
60 def unbind_state(self):
61 self.talkie_disconnect_all()
62 self._state = None
64 def update_visibility(self, visible):
65 self._state.visible = visible
67 def get_title_label(self):
68 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())
70 def update_label(*args):
71 title_label.set_slug(self._state.element_id)
73 self.talkie_connect(
74 self._state, 'element_id', update_label)
76 update_label()
77 return title_label
79 def get_title_control_remove(self):
80 button = common.MyDockWidgetTitleBarButton('\u2716')
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.talkie_connect(
100 self._state, 'visible', set_button_checked)
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.talkie_connect(
135 cpt_state, state_attr, update_function)
137 self._state = cpt_state
139 def unbind_state(self):
140 Element.unbind_state(self)
141 self._cpts = {}
142 self._lookuptable = None
143 self._values = None
144 self._autoscaler = None
146 def open_cpt_load_dialog(self):
147 caption = 'Select one *.cpt file to open'
149 fns, _ = qw.QFileDialog.getOpenFileNames(
150 self._parent, caption, options=common.qfiledialog_options)
152 if fns:
153 self.load_cpt_file(fns[0])
155 def load_cpt_file(self, path):
156 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
157 self._cpts.update([(cpt_name, automap.read_cpt(path))])
159 self._state.cpt_name = cpt_name
161 self._update_cpt_combobox()
162 self.update_cpt()
164 def _update_cpt_combobox(self):
165 from pyrocko import config
166 conf = config.config()
168 if self._cpt_combobox is None:
169 raise ValueError('CPT combobox needs init before updating!')
171 cb = self._cpt_combobox
173 if cb is not None:
174 cb.clear()
176 for s in CPTChoice.choices:
177 if s not in self._cpts:
178 try:
179 cpt = automap.read_cpt(topo.cpt(s))
180 except Exception:
181 from matplotlib import pyplot as plt
182 cmap = plt.cm.get_cmap(s)
183 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
185 self._cpts.update([(s, cpt)])
187 cpt_dir = conf.colortables_dir
188 if os.path.isdir(cpt_dir):
189 for f in [
190 f for f in os.listdir(cpt_dir)
191 if f.lower().endswith('.cpt')]:
193 s = 'USR' + os.path.basename(f).split('.')[0]
194 self._cpts.update(
195 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
197 for i, (s, cpt) in enumerate(self._cpts.items()):
198 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
199 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
201 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
203 def _update_cptscale_lineedit(self):
204 le = self._cpt_scale_lineedit
205 if le is not None:
206 le.clear()
208 self._cptscale_to_lineedit(self._state, le)
210 def _cptscale_to_lineedit(self, state, widget):
211 # sel = widget.selectedText() == widget.text()
213 crange = (None, None)
214 if self._lookuptable is not None:
215 crange = self._lookuptable.GetRange()
217 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
218 crange = state.cpt_scale_min, state.cpt_scale_max
220 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
222 widget.setText(fmt % crange)
224 # if sel:
225 # widget.selectAll()
227 def update_cpt(self):
228 state = self._state
230 if self._autoscaler is None:
231 self._autoscaler = AutoScaler()
233 if self._cpt_scale_lineedit:
234 if state.cpt_mode == 'off':
235 self._cpt_scale_lineedit.setEnabled(True)
236 else:
237 self._cpt_scale_lineedit.setEnabled(False)
239 if state.cpt_scale_min is not None:
240 state.cpt_scale_min = None
242 if state.cpt_scale_max is not None:
243 state.cpt_scale_max = None
245 if state.cpt_name is not None and self._values is not None:
246 if self._values.size == 0:
247 vscale = (0., 1.)
248 else:
249 vscale = (num.nanmin(self._values), num.nanmax(self._values))
251 vmin, vmax = None, None
252 if None not in (state.cpt_scale_min, state.cpt_scale_max):
253 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
254 else:
255 vmin, vmax, _ = self._autoscaler.make_scale(
256 vscale, override_mode=state.cpt_mode)
258 self._cpts[state.cpt_name].scale(vmin, vmax)
259 cpt = self._cpts[state.cpt_name]
261 vtk_lut = cpt_to_vtk_lookuptable(cpt)
262 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
264 self._lookuptable = vtk_lut
265 self._update_cptscale_lineedit()
267 elif state.cpt_name and self._values is None:
268 raise ValueError('No values passed to colormapper!')
270 def cpt_controls(self, parent, state, layout):
271 self._parent = parent
273 iy = layout.rowCount() + 1
275 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
277 cb = common.CPTComboBox()
278 layout.addWidget(cb, iy, 1)
279 state_bind_combobox(
280 self, state, 'cpt_name', cb)
282 self._cpt_combobox = cb
284 pb = qw.QPushButton('Load CPT')
285 layout.addWidget(pb, iy, 2)
286 pb.clicked.connect(self.open_cpt_load_dialog)
288 iy += 1
289 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
291 cb = common.string_choices_to_combobox(AutoScaleMode)
292 layout.addWidget(cb, iy, 1)
293 state_bind_combobox(
294 self, state, 'cpt_mode', cb)
296 le = qw.QLineEdit()
297 le.setEnabled(False)
298 layout.addWidget(le, iy, 2)
299 state_bind(
300 self, state,
301 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
302 le, [le.editingFinished, le.returnPressed],
303 self._cptscale_to_lineedit)
305 self._cpt_scale_lineedit = le
308def _lineedit_to_cptscale(widget, cpt_state):
309 s = str(widget.text())
310 s = s.replace(',', ' ')
312 crange = tuple((float(i) for i in s.split()))
313 crange = tuple((
314 crange[0],
315 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
317 try:
318 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
319 except Exception:
320 raise ValueError(
321 'need two numerical values: <vmin>, <vmax>')
324__all__ = [
325 'Element',
326 'ElementState',
327 'random_id',
328]