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
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
21logger = logging.getLogger('geometry')
23guts_prefix = 'sparrow'
25km = 1e3
28class GeometryState(base.ElementState):
29 opacity = Float.T(default=1.0)
30 visible = Bool.T(default=True)
31 geometry = Geometry.T(default=None, optional=True)
32 display_parameter = String.T(default="")
33 time = Float.T(default=0., optional=True)
34 cpt = base.CPTState.T(default=base.CPTState.D())
35 color = Color.T(default=Color.D('white'))
36 line_width = Float.T(default=1.0)
38 def create(self):
39 element = GeometryElement()
40 return element
43class GeometryElement(base.Element):
45 def __init__(self):
46 base.Element.__init__(self)
47 self._parent = None
48 self._state = None
49 self._controls = None
51 self._pipe = None
52 self._cbar_pipe = None
53 self._outlines_pipe = []
55 self.cpt_handler = base.CPTHandler()
57 def remove(self):
58 if self._parent and self._state:
59 self._parent.state.elements.remove(self._state)
61 def init_pipeslots(self):
62 if not self._pipe:
63 self._pipe.append([])
65 def remove_pipes(self):
66 if self._pipe is not None:
67 self._parent.remove_actor(self._pipe.actor)
69 if self._cbar_pipe is not None:
70 self._parent.remove_actor(self._cbar_pipe.actor)
72 if len(self._outlines_pipe) > 0:
73 for pipe in self._outlines_pipe:
74 self._parent.remove_actor(pipe.actor)
76 self._pipe = None
77 self._cbar_pipe = None
78 self._outlines_pipe = []
80 def set_parent(self, parent):
81 self._parent = parent
82 self._parent.add_panel(
83 self.get_title_label(),
84 self._get_controls(),
85 visible=True,
86 title_controls=[
87 self.get_title_control_remove(),
88 self.get_title_control_visible()])
90 self.talkie_connect(
91 self._parent.state,
92 ['tmin', 'tmax', 'lat', 'lon'],
93 self.update)
95 self.update()
97 def unset_parent(self):
98 self.unbind_state()
99 if self._parent:
100 if self._pipe or self._cbar_pipe or self._outlines_pipe:
101 self.remove_pipes()
103 if self._controls:
104 self._parent.remove_panel(self._controls)
105 self._controls = None
107 self._parent.update_view()
108 self._parent = None
110 def bind_state(self, state):
111 base.Element.bind_state(self, state)
113 self.talkie_connect(
114 state,
115 ['visible', 'geometry', 'display_parameter', 'time',
116 'opacity', 'color', 'line_width'],
117 self.update)
119 self.cpt_handler.bind_state(state.cpt, self.update)
121 def unbind_state(self):
122 self.cpt_handler.unbind_state()
123 base.Element.unbind_state(self)
125 def update_cpt(self, state):
127 if len(state.display_parameter) != 0:
128 values = state.geometry.get_property(state.display_parameter)
129 # TODO Check
130 # if values.ndim == 2:
131 # values = values.sum(1)
133 self.cpt_handler._values = values
134 self.cpt_handler.update_cpt()
136 def get_name(self):
137 return 'Geometry'
139 def open_file_load_dialog(self):
140 caption = 'Select one file containing a geometry to open'
141 fns, _ = qw.QFileDialog.getOpenFileNames(
142 self._parent, caption, options=common.qfiledialog_options)
144 if fns:
145 self.load_file(str(fns[0]))
146 else:
147 return
149 def load_file(self, path):
151 loaded_geometry = load(filename=path)
152 props = loaded_geometry.properties.get_col_names(sub_headers=False)
154 if props:
155 if self._state.display_parameter not in props:
156 self._state.display_parameter = props[0]
158 self._parent.remove_panel(self._controls)
159 self._controls = None
160 self._state.geometry = loaded_geometry
162 self._parent.add_panel(
163 self.get_title_label(),
164 self._get_controls(),
165 visible=True,
166 title_controls=[
167 self.get_title_control_remove(),
168 self.get_title_control_visible()])
170 self.update()
172 def get_values(self, geom):
173 values = geom.get_property(self._state.display_parameter)
175 if geom.event is not None:
176 ref_time = geom.event.time
177 else:
178 ref_time = 0.
180 if len(values.shape) == 2:
181 tmin = self._parent.state.tmin
182 tmax = self._parent.state.tmax
183 if tmin is not None:
184 ref_tmin = tmin - ref_time
185 ref_idx_min = geom.time2idx(ref_tmin)
186 else:
187 ref_idx_min = geom.time2idx(self._state.time)
189 if tmax is not None:
190 ref_tmax = tmax - ref_time
191 ref_idx_max = geom.time2idx(ref_tmax)
192 else:
193 ref_idx_max = geom.time2idx(self._state.time)
195 if ref_idx_min == ref_idx_max:
196 out = values[:, ref_idx_min]
197 elif ref_idx_min > ref_idx_max:
198 out = values[:, ref_idx_min]
199 elif ref_idx_max < ref_idx_min:
200 out = values[:, ref_idx_max]
201 else:
202 # TODO CHECK
203 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
204 out = values[:, ref_idx_max]
205 else:
206 out = values.ravel()
207 return out
209 def update_view(self, *args):
210 pstate = self._parent.state
211 geom = self._state.geometry
213 if geom.no_faces() > 0:
214 latlon = geom.get_vertices('latlon')
215 pstate.lat, pstate.lon = geographic_midpoint(
216 latlon[:, 0],
217 latlon[:, 1])
218 elif geom.outlines:
219 latlon = num.concatenate([
220 outline.get_col('latlon') for outline in geom.outlines
221 ])
222 pstate.lat, pstate.lon = geographic_midpoint(
223 latlon[:, 0],
224 latlon[:, 1])
225 elif geom.event:
226 pstate.lat = geom.event.lat
227 pstate.lon = geom.event.lon
228 else:
229 raise ValueError('Geometry Element has no location information.')
231 self.update()
233 def clear(self):
234 self._parent.remove_panel(self._controls)
235 self._controls = None
236 self._state.geometry = None
238 self._parent.add_panel(
239 self.get_title_label(),
240 self._get_controls(),
241 visible=True,
242 title_controls=[
243 self.get_title_control_remove(),
244 self.get_title_control_visible()])
246 self.update()
248 def update_outlines(self, geo):
249 state = self._state
250 if len(self._outlines_pipe) == 0:
251 for cs in ['latlondepth']:
252 outline_pipe = OutlinesPipe(
253 geo, color=state.color, cs=cs)
254 outline_pipe.set_line_width(state.line_width)
255 self._outlines_pipe.append(outline_pipe)
256 self._parent.add_actor(
257 self._outlines_pipe[-1].actor)
259 else:
260 for outline_pipe in self._outlines_pipe:
261 outline_pipe.set_color(state.color)
262 outline_pipe.set_line_width(state.line_width)
264 def update(self, *args):
266 state = self._state
268 if state.geometry and self._controls:
269 self._update_controls()
270 self.update_cpt(state)
272 if state.visible:
273 geo = state.geometry
274 lut = self.cpt_handler._lookuptable
275 no_faces = geo.no_faces()
276 if no_faces:
277 values = self.get_values(geo)
278 if not isinstance(self._pipe, TrimeshPipe):
279 vertices = arr_vertices(geo.get_vertices('xyz'))
280 faces = arr_faces(geo.get_faces())
281 self._pipe = TrimeshPipe(
282 vertices, faces,
283 values=values,
284 lut=lut,
285 backface_culling=False)
286 self._cbar_pipe = ColorbarPipe(
287 lut=lut, cbar_title=state.display_parameter)
288 self._parent.add_actor(self._pipe.actor)
289 self._parent.add_actor(self._cbar_pipe.actor)
290 else:
291 self._pipe.set_values(values)
292 self._pipe.set_lookuptable(lut)
293 self._pipe.set_opacity(self._state.opacity)
295 self._cbar_pipe.set_lookuptable(lut)
296 self._cbar_pipe.set_title(state.display_parameter)
298 if geo.outlines:
299 self.update_outlines(geo)
300 else:
301 self.remove_pipes()
303 else:
304 self.remove_pipes()
306 self._parent.update_view()
308 def _get_controls(self):
309 state = self._state
310 if not self._controls:
311 from ..state import state_bind_combobox, \
312 state_bind_slider, state_bind_combobox_color
314 frame = qw.QFrame()
315 layout = qw.QGridLayout()
316 layout.setAlignment(qc.Qt.AlignTop)
317 frame.setLayout(layout)
319 # load geometry
320 il = 0
321 if not state.geometry:
322 pb = qw.QPushButton('Load')
323 layout.addWidget(pb, il, 0)
324 pb.clicked.connect(self.open_file_load_dialog)
326 # property choice
327 else:
328 props = []
329 for prop in state.geometry.properties.get_col_names(
330 sub_headers=False):
331 props.append(prop)
333 layout.addWidget(qw.QLabel('Display Parameter'), il, 0)
334 cb = qw.QComboBox()
336 unique_props = list(set(props))
337 for i, s in enumerate(unique_props):
338 cb.insertItem(i, s)
340 layout.addWidget(cb, il, 1)
341 state_bind_combobox(self, state, 'display_parameter', cb)
343 if state.geometry.no_faces != 0:
344 # color maps
345 self.cpt_handler.cpt_controls(
346 self._parent, self._state.cpt, layout)
348 il += 1
349 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
351 self.cpt_handler._update_cpt_combobox()
352 self.cpt_handler._update_cptscale_lineedit()
354 # times slider
355 if state.geometry.times is not None:
356 il = layout.rowCount() + 1
357 slider = qw.QSlider(qc.Qt.Horizontal)
358 slider.setSizePolicy(
359 qw.QSizePolicy(
360 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
362 def iround(x):
363 return int(round(x))
365 slider.setMinimum(iround(state.geometry.times.min()))
366 slider.setMaximum(iround(state.geometry.times.max()))
367 slider.setSingleStep(iround(state.geometry.deltat))
368 slider.setPageStep(iround(state.geometry.deltat))
370 time_label = qw.QLabel('Time')
371 layout.addWidget(time_label, il, 0)
372 layout.addWidget(slider, il, 1)
374 state_bind_slider(
375 self, state, 'time', slider, dtype=int)
377 self._time_label = time_label
378 self._time_slider = slider
380 il = layout.rowCount() + 1
381 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
382 slider_opacity.setSizePolicy(
383 qw.QSizePolicy(
384 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
385 slider_opacity.setMinimum(0)
386 slider_opacity.setMaximum(1000)
388 opacity_label = qw.QLabel('Opacity')
389 layout.addWidget(opacity_label, il, 0)
390 layout.addWidget(slider_opacity, il, 1)
392 state_bind_slider(
393 self, state, 'opacity', slider_opacity, factor=0.001)
395 self._opacity_label = opacity_label
396 self._opacity_slider = slider_opacity
398 # color
399 il += 1
400 layout.addWidget(qw.QLabel('Color'), il, 0)
402 cb = common.strings_to_combobox(
403 ['black', 'white', 'blue', 'red'])
405 layout.addWidget(cb, il, 1)
406 state_bind_combobox_color(self, state, 'color', cb)
408 # linewidth outline
409 il += 1
410 layout.addWidget(qw.QLabel('Line Width'), il, 0)
412 slider = qw.QSlider(qc.Qt.Horizontal)
413 slider.setSizePolicy(
414 qw.QSizePolicy(
415 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
416 slider.setMinimum(0)
417 slider.setMaximum(100)
418 layout.addWidget(slider, il, 1)
419 state_bind_slider(
420 self, state, 'line_width', slider, factor=0.1)
422 # Clear scene
423 il += 1
424 pb = qw.QPushButton('Clear')
425 layout.addWidget(pb, il, 1)
426 pb.clicked.connect(self.clear)
428 # Change view to source
429 pb = qw.QPushButton('Move To')
430 layout.addWidget(pb, il, 2)
431 pb.clicked.connect(self.update_view)
433 self._controls = frame
435 self._update_controls()
437 return self._controls
439 def _update_controls(self):
440 state = self._state
441 if state.geometry:
442 if len(state.display_parameter) != 0:
443 values = state.geometry.get_property(state.display_parameter)
445 if values.ndim == 2:
446 self._time_label.setVisible(True)
447 self._time_slider.setVisible(True)
448 self._opacity_label.setVisible(True)
449 self._opacity_slider.setVisible(True)
450 else:
451 self._time_label.setVisible(False)
452 self._time_slider.setVisible(False)
453 self._opacity_label.setVisible(False)
454 self._opacity_slider.setVisible(False)
457__all__ = [
458 'GeometryElement',
459 'GeometryState'
460]