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, 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_control_remove(self):
68 button = common.MyDockWidgetTitleBarButton('\u00d7')
69 button.setStatusTip('Remove Element')
70 button.clicked.connect(self.remove)
71 return button
73 def get_title_control_visible(self):
74 assert hasattr(self._state, 'visible')
76 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
77 button.setStatusTip('Toggle Element Visibility')
78 button.toggled.connect(self.update_visibility)
80 def set_button_checked(*args):
81 button.blockSignals(True)
82 button.set_checked(self._state.visible)
83 button.blockSignals(False)
85 set_button_checked()
87 self.talkie_connect(
88 self._state, 'visible', set_button_checked)
90 return button
93class CPTChoice(StringChoice):
94 choices = [
95 'slip_colors', 'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
98class CPTState(ElementState):
99 cpt_name = String.T(default=CPTChoice.choices[0])
100 cpt_mode = String.T(default=AutoScaleMode.choices[1])
101 cpt_scale_min = Float.T(optional=True)
102 cpt_scale_max = Float.T(optional=True)
105class CPTHandler(Element):
107 def __init__(self):
109 Element.__init__(self)
110 self._cpts = {}
111 self._autoscaler = None
112 self._lookuptable = None
113 self._cpt_combobox = None
114 self._values = None
115 self._state = None
116 self._cpt_scale_lineedit = None
118 def bind_state(self, cpt_state, update_function):
119 for state_attr in [
120 'cpt_name', 'cpt_mode', 'cpt_scale_min', 'cpt_scale_max']:
122 self.talkie_connect(
123 cpt_state, state_attr, update_function)
125 self._state = cpt_state
127 def unbind_state(self):
128 Element.unbind_state(self)
129 self._cpts = {}
130 self._lookuptable = None
131 self._values = None
132 self._autoscaler = None
134 def open_cpt_load_dialog(self):
135 caption = 'Select one *.cpt file to open'
137 fns, _ = qw.QFileDialog.getOpenFileNames(
138 self._parent, caption, options=common.qfiledialog_options)
140 if fns:
141 self.load_cpt_file(fns[0])
143 def load_cpt_file(self, path):
144 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
145 self._cpts.update([(cpt_name, automap.read_cpt(path))])
147 self._state.cpt_name = cpt_name
149 self._update_cpt_combobox()
150 self.update_cpt()
152 def _update_cpt_combobox(self):
153 from pyrocko import config
154 conf = config.config()
156 if self._cpt_combobox is None:
157 raise ValueError('CPT combobox needs init before updating!')
159 cb = self._cpt_combobox
161 if cb is not None:
162 cb.clear()
164 for s in CPTChoice.choices:
165 if s not in self._cpts:
166 try:
167 cpt = automap.read_cpt(topo.cpt(s))
168 except Exception:
169 from matplotlib import pyplot as plt
170 cmap = plt.cm.get_cmap(s)
171 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
173 self._cpts.update([(s, cpt)])
175 cpt_dir = conf.colortables_dir
176 if os.path.isdir(cpt_dir):
177 for f in [
178 f for f in os.listdir(cpt_dir)
179 if f.lower().endswith('.cpt')]:
181 s = 'USR' + os.path.basename(f).split('.')[0]
182 self._cpts.update(
183 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
185 for i, (s, cpt) in enumerate(self._cpts.items()):
186 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
187 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
189 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
191 def _update_cptscale_lineedit(self):
192 le = self._cpt_scale_lineedit
193 if le is not None:
194 le.clear()
196 self._cptscale_to_lineedit(self._state, le)
198 def _cptscale_to_lineedit(self, state, widget):
199 # sel = widget.selectedText() == widget.text()
201 crange = (None, None)
202 if self._lookuptable is not None:
203 crange = self._lookuptable.GetRange()
205 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
206 crange = state.cpt_scale_min, state.cpt_scale_max
208 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
210 widget.setText(fmt % crange)
212 # if sel:
213 # widget.selectAll()
215 def update_cpt(self):
216 state = self._state
218 if self._autoscaler is None:
219 self._autoscaler = AutoScaler()
221 if self._cpt_scale_lineedit:
222 if state.cpt_mode == 'off':
223 self._cpt_scale_lineedit.setEnabled(True)
224 else:
225 self._cpt_scale_lineedit.setEnabled(False)
227 if state.cpt_scale_min is not None:
228 state.cpt_scale_min = None
230 if state.cpt_scale_max is not None:
231 state.cpt_scale_max = None
233 if state.cpt_name is not None and self._values is not None:
234 vscale = (num.nanmin(self._values), num.nanmax(self._values))
236 vmin, vmax = None, None
237 if None not in (state.cpt_scale_min, state.cpt_scale_max):
238 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
239 else:
240 vmin, vmax, _ = self._autoscaler.make_scale(
241 vscale, override_mode=state.cpt_mode)
243 self._cpts[state.cpt_name].scale(vmin, vmax)
244 cpt = self._cpts[state.cpt_name]
246 vtk_lut = cpt_to_vtk_lookuptable(cpt)
247 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
249 self._lookuptable = vtk_lut
250 self._update_cptscale_lineedit()
252 elif state.cpt_name and self._values is None:
253 raise ValueError('No values passed to colormapper!')
255 def cpt_controls(self, parent, state, layout):
256 self._parent = parent
258 iy = layout.rowCount() + 1
260 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
262 cb = common.CPTComboBox()
263 layout.addWidget(cb, iy, 1)
264 state_bind_combobox(
265 self, state, 'cpt_name', cb)
267 self._cpt_combobox = cb
269 pb = qw.QPushButton('Load CPT')
270 layout.addWidget(pb, iy, 2)
271 pb.clicked.connect(self.open_cpt_load_dialog)
273 iy += 1
274 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
276 cb = common.string_choices_to_combobox(AutoScaleMode)
277 layout.addWidget(cb, iy, 1)
278 state_bind_combobox(
279 self, state, 'cpt_mode', cb)
281 le = qw.QLineEdit()
282 le.setEnabled(False)
283 layout.addWidget(le, iy, 2)
284 state_bind(
285 self, state,
286 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
287 le, [le.editingFinished, le.returnPressed],
288 self._cptscale_to_lineedit)
290 self._cpt_scale_lineedit = le
293def _lineedit_to_cptscale(widget, cpt_state):
294 s = str(widget.text())
295 s = s.replace(',', ' ')
297 crange = tuple((float(i) for i in s.split()))
298 crange = tuple((
299 crange[0],
300 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
302 try:
303 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
304 except Exception:
305 raise ValueError(
306 'need two numerical values: <vmin>, <vmax>')
309__all__ = [
310 'Element',
311 'ElementState',
312 'random_id',
313]