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
21from .. import common
22from ..state import \
23 state_bind_combobox, state_bind
26mpl_cmap_blacklist = [
27 "prism", "flag",
28 "Accent", "Dark2",
29 "Paired", "Pastel1", "Pastel2",
30 "Set1", "Set2", "Set3",
31 "tab10", "tab20", "tab20b", "tab20c"
32]
35def get_mpl_cmap_choices():
36 try:
37 from matplotlib import colormaps as mpl_cmaps
38 mpl_cmap_choices = list(mpl_cmaps.keys())
40 for cmap_name in mpl_cmap_blacklist:
41 try:
42 mpl_cmap_choices.remove(cmap_name)
43 mpl_cmap_choices.remove("%s_r" % cmap_name)
44 except ValueError:
45 pass
47 except ImportError:
48 mpl_cmap_choices = [
49 'seismic', 'seismic_r', 'jet', 'hot_r', 'gist_earth_r']
51 mpl_cmap_choices.sort()
52 return mpl_cmap_choices
55def random_id():
56 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')
59class ElementState(TalkieRoot):
61 element_id = String.T()
63 def __init__(self, **kwargs):
64 if 'element_id' not in kwargs:
65 kwargs['element_id'] = random_id()
67 TalkieRoot.__init__(self, **kwargs)
70class Element(TalkieConnectionOwner):
71 def __init__(self):
72 TalkieConnectionOwner.__init__(self)
73 self._parent = None
74 self._state = None
76 def remove(self):
77 if self._parent and self._state:
78 self._parent.state.elements.remove(self._state)
80 def set_parent(self, parent):
81 self._parent = parent
83 def unset_parent(self):
84 print(self)
85 raise NotImplementedError
87 def bind_state(self, state):
88 self._state = state
90 def unbind_state(self):
91 self.talkie_disconnect_all()
92 self._state = None
94 def update_visibility(self, visible):
95 self._state.visible = visible
97 def get_title_label(self):
98 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())
100 def update_label(*args):
101 title_label.set_slug(self._state.element_id)
103 self.talkie_connect(
104 self._state, 'element_id', update_label)
106 update_label()
107 return title_label
109 def get_title_control_remove(self):
110 button = common.MyDockWidgetTitleBarButton('\u2716')
111 button.setStatusTip('Remove Element')
112 button.clicked.connect(self.remove)
113 return button
115 def get_title_control_visible(self):
116 assert hasattr(self._state, 'visible')
118 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')
119 button.setStatusTip('Toggle Element Visibility')
120 button.toggled.connect(self.update_visibility)
122 def set_button_checked(*args):
123 button.blockSignals(True)
124 button.set_checked(self._state.visible)
125 button.blockSignals(False)
127 set_button_checked()
129 self.talkie_connect(
130 self._state, 'visible', set_button_checked)
132 return button
135class CPTChoice(StringChoice):
137 choices = ['slip_colors'] + get_mpl_cmap_choices()
140class CPTState(ElementState):
141 cpt_name = String.T(default=CPTChoice.choices[0])
142 cpt_mode = String.T(default=AutoScaleMode.choices[1])
143 cpt_scale_min = Float.T(optional=True)
144 cpt_scale_max = Float.T(optional=True)
147class CPTHandler(Element):
149 def __init__(self):
151 Element.__init__(self)
152 self._cpts = {}
153 self._autoscaler = None
154 self._lookuptable = None
155 self._cpt_combobox = None
156 self._values = None
157 self._state = None
158 self._cpt_scale_lineedit = None
160 def bind_state(self, cpt_state, update_function):
161 for state_attr in [
162 'cpt_name', 'cpt_mode', 'cpt_scale_min', 'cpt_scale_max']:
164 self.talkie_connect(
165 cpt_state, state_attr, update_function)
167 self._state = cpt_state
169 def unbind_state(self):
170 Element.unbind_state(self)
171 self._cpts = {}
172 self._lookuptable = None
173 self._values = None
174 self._autoscaler = None
176 def open_cpt_load_dialog(self):
177 caption = 'Select one *.cpt file to open'
179 fns, _ = qw.QFileDialog.getOpenFileNames(
180 self._parent, caption, options=common.qfiledialog_options)
182 if fns:
183 self.load_cpt_file(fns[0])
185 def load_cpt_file(self, path):
186 cpt_name = 'USR' + os.path.basename(path).split('.')[0]
187 self._cpts.update([(cpt_name, automap.read_cpt(path))])
189 self._state.cpt_name = cpt_name
191 self._update_cpt_combobox()
192 self.update_cpt()
194 def _update_cpt_combobox(self):
195 from pyrocko import config
196 conf = config.config()
198 if self._cpt_combobox is None:
199 raise ValueError('CPT combobox needs init before updating!')
201 cb = self._cpt_combobox
203 if cb is not None:
204 cb.clear()
206 for s in CPTChoice.choices:
207 if s not in self._cpts:
208 try:
209 cpt = automap.read_cpt(topo.cpt(s))
210 except Exception:
211 from matplotlib import pyplot as plt
212 cmap = plt.cm.get_cmap(s)
213 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])
215 self._cpts.update([(s, cpt)])
217 cpt_dir = conf.colortables_dir
218 if os.path.isdir(cpt_dir):
219 for f in [
220 f for f in os.listdir(cpt_dir)
221 if f.lower().endswith('.cpt')]:
223 s = 'USR' + os.path.basename(f).split('.')[0]
224 self._cpts.update(
225 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])
227 for i, (s, cpt) in enumerate(self._cpts.items()):
228 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))
229 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)
231 cb.setCurrentIndex(cb.findText(self._state.cpt_name))
233 def _update_cptscale_lineedit(self):
234 le = self._cpt_scale_lineedit
235 if le is not None:
236 le.clear()
238 self._cptscale_to_lineedit(self._state, le)
240 def _cptscale_to_lineedit(self, state, widget):
241 # sel = widget.selectedText() == widget.text()
243 crange = (None, None)
244 if self._lookuptable is not None:
245 crange = self._lookuptable.GetRange()
247 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:
248 crange = state.cpt_scale_min, state.cpt_scale_max
250 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])
252 widget.setText(fmt % crange)
254 # if sel:
255 # widget.selectAll()
257 def update_cpt(self):
258 state = self._state
260 if self._autoscaler is None:
261 self._autoscaler = AutoScaler()
263 if self._cpt_scale_lineedit:
264 if state.cpt_mode == 'off':
265 self._cpt_scale_lineedit.setEnabled(True)
266 else:
267 self._cpt_scale_lineedit.setEnabled(False)
269 if state.cpt_scale_min is not None:
270 state.cpt_scale_min = None
272 if state.cpt_scale_max is not None:
273 state.cpt_scale_max = None
275 if state.cpt_name is not None and self._values is not None:
276 if self._values.size == 0:
277 vscale = (0., 1.)
278 else:
279 vscale = (num.nanmin(self._values), num.nanmax(self._values))
281 vmin, vmax = None, None
282 if None not in (state.cpt_scale_min, state.cpt_scale_max):
283 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max
284 else:
285 vmin, vmax, _ = self._autoscaler.make_scale(
286 vscale, override_mode=state.cpt_mode)
288 self._cpts[state.cpt_name].scale(vmin, vmax)
289 cpt = self._cpts[state.cpt_name]
291 vtk_lut = cpt_to_vtk_lookuptable(cpt)
292 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)
294 self._lookuptable = vtk_lut
295 self._update_cptscale_lineedit()
297 elif state.cpt_name and self._values is None:
298 raise ValueError('No values passed to colormapper!')
300 def cpt_controls(self, parent, state, layout):
301 self._parent = parent
303 iy = layout.rowCount() + 1
305 layout.addWidget(qw.QLabel('Color Map'), iy, 0)
307 cb = common.CPTComboBox()
308 layout.addWidget(cb, iy, 1)
309 state_bind_combobox(
310 self, state, 'cpt_name', cb)
312 self._cpt_combobox = cb
314 pb = qw.QPushButton('Load CPT')
315 layout.addWidget(pb, iy, 2)
316 pb.clicked.connect(self.open_cpt_load_dialog)
318 iy += 1
319 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)
321 cb = common.string_choices_to_combobox(AutoScaleMode)
322 layout.addWidget(cb, iy, 1)
323 state_bind_combobox(
324 self, state, 'cpt_mode', cb)
326 le = qw.QLineEdit()
327 le.setEnabled(False)
328 layout.addWidget(le, iy, 2)
329 state_bind(
330 self, state,
331 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,
332 le, [le.editingFinished, le.returnPressed],
333 self._cptscale_to_lineedit)
335 self._cpt_scale_lineedit = le
338def _lineedit_to_cptscale(widget, cpt_state):
339 s = str(widget.text())
340 s = s.replace(',', ' ')
342 crange = tuple((float(i) for i in s.split()))
343 crange = tuple((
344 crange[0],
345 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))
347 try:
348 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange
349 except Exception:
350 raise ValueError(
351 'need two numerical values: <vmin>, <vmax>')
354__all__ = [
355 'Element',
356 'ElementState',
357 'random_id',
358]