1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
7import numpy as num
9from pyrocko.guts import Bool, String, load, Float, StringChoice
10from pyrocko.geometry import arr_vertices, arr_faces
11from pyrocko.gui.qt_compat import qw, qc
12from pyrocko.gui.vtk_util import TrimeshPipe, ColorbarPipe, OutlinesPipe, Color
13from pyrocko.orthodrome import geographic_midpoint
15from pyrocko.model import Geometry
17from . import base
18from .. import common
19from ..state import state_bind_combobox, state_bind_slider, \
20 state_bind_combobox_color, state_bind_checkbox
23logger = logging.getLogger('geometry')
25guts_prefix = 'sparrow'
27km = 1e3
30class ColorBarPositionChoice(StringChoice):
31 choices = ['bottom-left', 'bottom-right', 'top-left', 'top-right']
34class GeometryState(base.ElementState):
35 opacity = Float.T(default=1.0)
36 visible = Bool.T(default=True)
37 geometry = Geometry.T(default=None, optional=True)
38 display_parameter = String.T(default="")
39 time = Float.T(default=0., optional=True)
40 cpt = base.CPTState.T(default=base.CPTState.D())
41 color = Color.T(default=Color.D('white'))
42 line_width = Float.T(default=1.0)
43 show_color_bar = Bool.T(default=True)
44 position_color_bar = ColorBarPositionChoice.T(default='bottom-right')
46 def create(self):
47 element = GeometryElement()
48 return element
51class GeometryElement(base.Element):
53 def __init__(self):
54 base.Element.__init__(self)
55 self._parent = None
56 self._state = None
57 self._controls = None
59 self._pipe = None
60 self._cbar_pipe = None
61 self._outlines_pipe = []
63 self.cpt_handler = base.CPTHandler()
65 def remove(self):
66 if self._parent and self._state:
67 self._parent.state.elements.remove(self._state)
69 def init_pipeslots(self):
70 if not self._pipe:
71 self._pipe.append([])
73 def remove_cbar_pipe(self):
74 if self._cbar_pipe is not None:
75 self._parent.remove_actor(self._cbar_pipe.actor)
77 self._cbar_pipe = None
79 def remove_pipes(self):
80 if self._pipe is not None:
81 self._parent.remove_actor(self._pipe.actor)
83 self.remove_cbar_pipe()
85 if len(self._outlines_pipe) > 0:
86 for pipe in self._outlines_pipe:
87 self._parent.remove_actor(pipe.actor)
89 self._pipe = None
90 self._outlines_pipe = []
92 def set_parent(self, parent):
93 self._parent = parent
94 self._parent.add_panel(
95 self.get_title_label(),
96 self._get_controls(),
97 visible=True,
98 title_controls=[
99 self.get_title_control_remove(),
100 self.get_title_control_visible()])
102 self.talkie_connect(
103 self._parent.state,
104 ['tmin', 'tmax', 'lat', 'lon', 'tmax_effective'],
105 self.update)
107 self.update()
109 def unset_parent(self):
110 self.unbind_state()
111 if self._parent:
112 if self._pipe or self._cbar_pipe or self._outlines_pipe:
113 self.remove_pipes()
115 if self._controls:
116 self._parent.remove_panel(self._controls)
117 self._controls = None
119 self._parent.update_view()
120 self._parent = None
122 def bind_state(self, state):
123 base.Element.bind_state(self, state)
125 self.talkie_connect(
126 state,
127 ['visible', 'geometry', 'display_parameter', 'time',
128 'opacity', 'color', 'line_width',
129 'show_color_bar', 'position_color_bar'],
130 self.update)
132 self.cpt_handler.bind_state(state.cpt, self.update)
134 def unbind_state(self):
135 self.cpt_handler.unbind_state()
136 base.Element.unbind_state(self)
138 def update_cpt(self, state):
140 if len(state.display_parameter) != 0:
141 values = state.geometry.get_property(state.display_parameter)
142 # TODO Check
143 # if values.ndim == 2:
144 # values = values.sum(1)
146 self.cpt_handler._values = values
147 self.cpt_handler.update_cpt(mask_zeros=True)
149 def get_name(self):
150 return 'Geometry'
152 def open_file_load_dialog(self):
153 caption = 'Select one file containing a geometry to open'
154 fns, _ = qw.QFileDialog.getOpenFileNames(
155 self._parent, caption, options=common.qfiledialog_options)
157 if fns:
158 self.load_file(str(fns[0]))
159 else:
160 return
162 def load_file(self, path):
164 loaded_geometry = load(filename=path)
165 props = loaded_geometry.properties.get_col_names(sub_headers=False)
167 if props:
168 if self._state.display_parameter not in props:
169 self._state.display_parameter = props[0]
171 self._parent.remove_panel(self._controls)
172 self._controls = None
173 self._state.geometry = loaded_geometry
175 self._parent.add_panel(
176 self.get_title_label(),
177 self._get_controls(),
178 visible=True,
179 title_controls=[
180 self.get_title_control_remove(),
181 self.get_title_control_visible()])
183 self.update()
185 def get_values(self, geom):
186 values = geom.get_property(self._state.display_parameter)
188 if geom.event is not None:
189 ref_time = geom.event.time
190 else:
191 ref_time = 0.
193 if len(values.shape) == 2:
194 tmin = self._parent.state.tmin
195 tmax = self._parent.state.tmax_effective
197 if tmin is not None:
198 ref_tmin = tmin - ref_time
199 ref_idx_min = geom.time2idx(ref_tmin)
200 else:
201 ref_idx_min = geom.time2idx(self._state.time)
203 if tmax is not None:
204 ref_tmax = tmax - ref_time
205 ref_idx_max = geom.time2idx(ref_tmax)
206 else:
207 ref_idx_max = geom.time2idx(self._state.time)
209 if ref_idx_min == ref_idx_max:
210 out = values[:, ref_idx_min]
211 elif ref_idx_min > ref_idx_max:
212 out = values[:, ref_idx_min]
213 elif ref_idx_max < ref_idx_min:
214 out = values[:, ref_idx_max]
215 else:
216 # for cumulative display
217 out = values[:, ref_idx_min:ref_idx_max].sum(1)
218 # out = values[:, ref_idx_max]
219 else:
220 out = values.ravel()
221 return out
223 def update_view(self, *args):
224 pstate = self._parent.state
225 geom = self._state.geometry
227 if geom.no_faces() > 0:
228 latlon = geom.get_vertices('latlon')
229 pstate.lat, pstate.lon = geographic_midpoint(
230 latlon[:, 0],
231 latlon[:, 1])
232 elif geom.outlines:
233 latlon = num.concatenate([
234 outline.get_col('latlon') for outline in geom.outlines
235 ])
236 pstate.lat, pstate.lon = geographic_midpoint(
237 latlon[:, 0],
238 latlon[:, 1])
239 elif geom.event:
240 pstate.lat = geom.event.lat
241 pstate.lon = geom.event.lon
242 else:
243 raise ValueError('Geometry Element has no location information.')
245 self.update()
247 def clear(self):
248 self._parent.remove_panel(self._controls)
249 self._controls = None
250 self._state.geometry = None
252 self._parent.add_panel(
253 self.get_title_label(),
254 self._get_controls(),
255 visible=True,
256 title_controls=[
257 self.get_title_control_remove(),
258 self.get_title_control_visible()])
260 self.update()
262 def update_outlines(self, geo):
263 state = self._state
264 if len(self._outlines_pipe) == 0:
265 for cs in ['latlondepth']:
266 outline_pipe = OutlinesPipe(
267 geo, color=state.color, cs=cs)
268 outline_pipe.set_line_width(state.line_width)
269 self._outlines_pipe.append(outline_pipe)
270 self._parent.add_actor(
271 self._outlines_pipe[-1].actor)
273 else:
274 for outline_pipe in self._outlines_pipe:
275 outline_pipe.set_color(state.color)
276 outline_pipe.set_line_width(state.line_width)
278 def update(self, *args):
280 state = self._state
282 if state.geometry and self._controls:
283 self._update_controls()
284 self.update_cpt(state)
286 if state.visible:
287 geo = state.geometry
288 lut = self.cpt_handler._lookuptable
289 no_faces = geo.no_faces()
290 if no_faces:
291 values = self.get_values(geo)
292 if not isinstance(self._pipe, TrimeshPipe):
293 vertices = arr_vertices(geo.get_vertices('xyz'))
294 faces = arr_faces(geo.get_faces())
295 self._pipe = TrimeshPipe(
296 vertices, faces,
297 values=values,
298 lut=lut,
299 backface_culling=False)
300 else:
301 self._pipe.set_values(values)
302 self._pipe.set_lookuptable(lut)
303 self._pipe.set_opacity(self._state.opacity)
305 if state.show_color_bar:
306 sx, sy = 1, 1
307 off = 0.08 * sy
308 pos = {
309 'top-left': (off, sy/2 + off, 0, 2),
310 'top-right': (sx - off, sy/2 + off, 2, 2),
311 'bottom-left': (off, off, 0, 0),
312 'bottom-right': (sx - off, off, 2, 0)}
313 x, y, _, _ = pos[state.position_color_bar]
315 if not isinstance(self._cbar_pipe, ColorbarPipe):
316 self._cbar_pipe = ColorbarPipe(
317 lut=lut,
318 cbar_title=state.display_parameter,
319 position=(x, y))
320 self._parent.add_actor(self._pipe.actor)
321 self._parent.add_actor(self._cbar_pipe.actor)
322 else:
323 self._cbar_pipe.set_lookuptable(lut)
324 self._cbar_pipe.set_title(state.display_parameter)
325 self._cbar_pipe._set_position(x, y)
326 else:
327 self.remove_cbar_pipe()
329 if geo.outlines:
330 self.update_outlines(geo)
331 else:
332 self.remove_pipes()
334 else:
335 self.remove_pipes()
337 self._parent.update_view()
339 def _get_controls(self):
340 state = self._state
341 if not self._controls:
343 frame = qw.QFrame()
344 layout = qw.QGridLayout()
345 layout.setAlignment(qc.Qt.AlignTop)
346 frame.setLayout(layout)
348 # load geometry
349 il = 0
350 if not state.geometry:
351 pb = qw.QPushButton('Load')
352 layout.addWidget(pb, il, 0)
353 pb.clicked.connect(self.open_file_load_dialog)
355 # property choice
356 else:
357 props = []
358 for prop in state.geometry.properties.get_col_names(
359 sub_headers=False):
360 props.append(prop)
362 layout.addWidget(qw.QLabel('Display Parameter'), il, 0)
363 cb = qw.QComboBox()
365 unique_props = list(set(props))
366 for i, s in enumerate(unique_props):
367 cb.insertItem(i, s)
369 layout.addWidget(cb, il, 1)
370 state_bind_combobox(self, state, 'display_parameter', cb)
372 if state.geometry.no_faces != 0:
373 # color maps
374 self.cpt_handler.cpt_controls(
375 self._parent, self._state.cpt, layout)
377 il += 1
378 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
380 self.cpt_handler._update_cpt_combobox()
381 self.cpt_handler._update_cptscale_lineedit()
383 # color scale checkbox
384 il = layout.rowCount() + 1
385 layout.addWidget(qw.QLabel('Color Bar'), il, 0)
387 chb = qw.QCheckBox('show')
388 layout.addWidget(chb, il, 1)
389 state_bind_checkbox(self, state, 'show_color_bar', chb)
391 cb = common.string_choices_to_combobox(
392 ColorBarPositionChoice)
393 layout.addWidget(cb, il, 2)
394 state_bind_combobox(
395 self, self._state, 'position_color_bar', cb)
397 # times slider
398 if state.geometry.times is not None:
399 il = layout.rowCount() + 1
400 slider = qw.QSlider(qc.Qt.Horizontal)
401 slider.setSizePolicy(
402 qw.QSizePolicy(
403 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
405 def iround(x):
406 return int(round(x))
408 slider.setMinimum(iround(state.geometry.times.min()))
409 slider.setMaximum(iround(state.geometry.times.max()))
410 slider.setSingleStep(iround(state.geometry.deltat))
411 slider.setPageStep(iround(state.geometry.deltat))
413 time_label = qw.QLabel('Time')
414 layout.addWidget(time_label, il, 0)
415 layout.addWidget(slider, il, 1)
417 state_bind_slider(
418 self, state, 'time', slider, dtype=int)
420 self._time_label = time_label
421 self._time_slider = slider
423 il = layout.rowCount() + 1
424 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
425 slider_opacity.setSizePolicy(
426 qw.QSizePolicy(
427 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
428 slider_opacity.setMinimum(0)
429 slider_opacity.setMaximum(1000)
431 opacity_label = qw.QLabel('Opacity')
432 layout.addWidget(opacity_label, il, 0)
433 layout.addWidget(slider_opacity, il, 1)
435 state_bind_slider(
436 self, state, 'opacity', slider_opacity, factor=0.001)
438 self._opacity_label = opacity_label
439 self._opacity_slider = slider_opacity
441 # color
442 il += 1
443 layout.addWidget(qw.QLabel('Color'), il, 0)
445 cb = common.strings_to_combobox(
446 ['black', 'white', 'blue', 'red'])
448 layout.addWidget(cb, il, 1)
449 state_bind_combobox_color(self, state, 'color', cb)
451 # linewidth outline
452 il += 1
453 layout.addWidget(qw.QLabel('Line Width'), il, 0)
455 slider = qw.QSlider(qc.Qt.Horizontal)
456 slider.setSizePolicy(
457 qw.QSizePolicy(
458 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
459 slider.setMinimum(0)
460 slider.setMaximum(100)
461 layout.addWidget(slider, il, 1)
462 state_bind_slider(
463 self, state, 'line_width', slider, factor=0.1)
465 # Clear scene
466 il += 1
467 pb = qw.QPushButton('Clear')
468 layout.addWidget(pb, il, 1)
469 pb.clicked.connect(self.clear)
471 # Change view to source
472 pb = qw.QPushButton('Move To')
473 layout.addWidget(pb, il, 2)
474 pb.clicked.connect(self.update_view)
476 self._controls = frame
478 self._update_controls()
480 return self._controls
482 def _update_controls(self):
483 state = self._state
484 if state.geometry:
485 if len(state.display_parameter) != 0:
486 values = state.geometry.get_property(state.display_parameter)
488 if values.ndim == 2:
489 self._time_label.setVisible(True)
490 self._time_slider.setVisible(True)
491 self._opacity_label.setVisible(True)
492 self._opacity_slider.setVisible(True)
493 else:
494 self._time_label.setVisible(False)
495 self._time_slider.setVisible(False)
496 self._opacity_label.setVisible(False)
497 self._opacity_slider.setVisible(False)
500__all__ = [
501 'GeometryElement',
502 'GeometryState'
503]