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 get_cpt_name(self, cpt, display_parameter):
126 return '{}_{}'.format(cpt, display_parameter)
128 def update_cpt(self, state):
130 if len(state.display_parameter) != 0:
131 values = state.geometry.get_property(state.display_parameter)
132 # TODO Check
133 # if values.ndim == 2:
134 # values = values.sum(1)
136 self.cpt_handler._values = values
137 self.cpt_handler.update_cpt()
139 def get_name(self):
140 return 'Geometry'
142 def open_file_load_dialog(self):
143 caption = 'Select one file containing a geometry to open'
144 fns, _ = qw.QFileDialog.getOpenFileNames(
145 self._parent, caption, options=common.qfiledialog_options)
147 if fns:
148 self.load_file(str(fns[0]))
149 else:
150 return
152 def load_file(self, path):
154 loaded_geometry = load(filename=path)
155 props = loaded_geometry.properties.get_col_names(sub_headers=False)
157 if props:
158 if self._state.display_parameter not in props:
159 self._state.display_parameter = props[0]
161 self._parent.remove_panel(self._controls)
162 self._controls = None
163 self._state.geometry = loaded_geometry
165 self._parent.add_panel(
166 self.get_title_label(),
167 self._get_controls(),
168 visible=True,
169 title_controls=[
170 self.get_title_control_remove(),
171 self.get_title_control_visible()])
173 self.update()
175 def get_values(self, geom):
176 values = geom.get_property(self._state.display_parameter)
178 if geom.event is not None:
179 ref_time = geom.event.time
180 else:
181 ref_time = 0.
183 if len(values.shape) == 2:
184 tmin = self._parent.state.tmin
185 tmax = self._parent.state.tmax
186 if tmin is not None:
187 ref_tmin = tmin - ref_time
188 ref_idx_min = geom.time2idx(ref_tmin)
189 else:
190 ref_idx_min = geom.time2idx(self._state.time)
192 if tmax is not None:
193 ref_tmax = tmax - ref_time
194 ref_idx_max = geom.time2idx(ref_tmax)
195 else:
196 ref_idx_max = geom.time2idx(self._state.time)
198 if ref_idx_min == ref_idx_max:
199 out = values[:, ref_idx_min]
200 elif ref_idx_min > ref_idx_max:
201 out = values[:, ref_idx_min]
202 elif ref_idx_max < ref_idx_min:
203 out = values[:, ref_idx_max]
204 else:
205 # TODO CHECK
206 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
207 out = values[:, ref_idx_max]
208 else:
209 out = values.ravel()
210 return out
212 def update_view(self, *args):
213 pstate = self._parent.state
214 geom = self._state.geometry
216 if geom.no_faces() > 0:
217 latlon = geom.get_vertices('latlon')
218 pstate.lat, pstate.lon = geographic_midpoint(
219 latlon[:, 0],
220 latlon[:, 1])
221 elif geom.outlines:
222 latlon = num.concatenate([
223 outline.get_col('latlon') for outline in geom.outlines
224 ])
225 pstate.lat, pstate.lon = geographic_midpoint(
226 latlon[:, 0],
227 latlon[:, 1])
228 elif geom.event:
229 pstate.lat = geom.event.lat
230 pstate.lon = geom.event.lon
231 else:
232 raise ValueError('Geometry Element has no location information.')
234 self.update()
236 def clear(self):
237 self._parent.remove_panel(self._controls)
238 self._controls = None
239 self._state.geometry = None
241 self._parent.add_panel(
242 self.get_title_label(),
243 self._get_controls(),
244 visible=True,
245 title_controls=[
246 self.get_title_control_remove(),
247 self.get_title_control_visible()])
249 self.update()
251 def update_outlines(self, geo):
252 state = self._state
253 if len(self._outlines_pipe) == 0:
254 for cs in ['latlondepth']:
255 outline_pipe = OutlinesPipe(
256 geo, color=state.color, cs=cs)
257 outline_pipe.set_line_width(state.line_width)
258 self._outlines_pipe.append(outline_pipe)
259 self._parent.add_actor(
260 self._outlines_pipe[-1].actor)
262 else:
263 for outline_pipe in self._outlines_pipe:
264 outline_pipe.set_color(state.color)
265 outline_pipe.set_line_width(state.line_width)
267 def update(self, *args):
269 state = self._state
271 if state.geometry and self._controls:
272 self._update_controls()
273 # base.update_cpt(self)
274 self.update_cpt(state)
276 if state.visible:
277 # cpt_name = self.get_cpt_name(
278 # state.cpt, state.display_parameter)
279 geo = state.geometry
280 lut = self.cpt_handler._lookuptable
281 no_faces = geo.no_faces()
282 if no_faces:
283 values = self.get_values(geo)
284 if not isinstance(self._pipe, TrimeshPipe):
285 vertices = arr_vertices(geo.get_vertices('xyz'))
286 faces = arr_faces(geo.get_faces())
287 self._pipe = TrimeshPipe(
288 vertices, faces,
289 values=values,
290 lut=lut,
291 backface_culling=False)
292 self._cbar_pipe = ColorbarPipe(
293 lut=lut, cbar_title=state.display_parameter)
294 self._parent.add_actor(self._pipe.actor)
295 self._parent.add_actor(self._cbar_pipe.actor)
296 else:
297 self._pipe.set_values(values)
298 self._pipe.set_lookuptable(lut)
299 self._pipe.set_opacity(self._state.opacity)
301 self._cbar_pipe.set_lookuptable(lut)
302 self._cbar_pipe.set_title(state.display_parameter)
304 if geo.outlines:
305 self.update_outlines(geo)
306 else:
307 self.remove_pipes()
309 else:
310 self.remove_pipes()
312 self._parent.update_view()
314 def _get_controls(self):
315 state = self._state
316 if not self._controls:
317 from ..state import state_bind_combobox, \
318 state_bind_slider, state_bind_combobox_color
320 frame = qw.QFrame()
321 layout = qw.QGridLayout()
322 layout.setAlignment(qc.Qt.AlignTop)
323 frame.setLayout(layout)
325 # load geometry
326 il = 0
327 if not state.geometry:
328 pb = qw.QPushButton('Load')
329 layout.addWidget(pb, il, 0)
330 pb.clicked.connect(self.open_file_load_dialog)
332 # property choice
333 else:
334 props = []
335 for prop in state.geometry.properties.get_col_names(
336 sub_headers=False):
337 props.append(prop)
339 layout.addWidget(qw.QLabel('Display Parameter'), il, 0)
340 cb = qw.QComboBox()
342 unique_props = list(set(props))
343 for i, s in enumerate(unique_props):
344 cb.insertItem(i, s)
346 layout.addWidget(cb, il, 1)
347 state_bind_combobox(self, state, 'display_parameter', cb)
349 if state.geometry.no_faces != 0:
350 # color maps
351 self.cpt_handler.cpt_controls(
352 self._parent, self._state.cpt, layout)
354 il += 1
355 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
357 self.cpt_handler._update_cpt_combobox()
358 self.cpt_handler._update_cptscale_lineedit()
360 # times slider
361 if state.geometry.times is not None:
362 il = layout.rowCount() + 1
363 slider = qw.QSlider(qc.Qt.Horizontal)
364 slider.setSizePolicy(
365 qw.QSizePolicy(
366 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
368 def iround(x):
369 return int(round(x))
371 slider.setMinimum(iround(state.geometry.times.min()))
372 slider.setMaximum(iround(state.geometry.times.max()))
373 slider.setSingleStep(iround(state.geometry.deltat))
374 slider.setPageStep(iround(state.geometry.deltat))
376 time_label = qw.QLabel('Time')
377 layout.addWidget(time_label, il, 0)
378 layout.addWidget(slider, il, 1)
380 state_bind_slider(
381 self, state, 'time', slider, dtype=int)
383 self._time_label = time_label
384 self._time_slider = slider
386 il = layout.rowCount() + 1
387 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
388 slider_opacity.setSizePolicy(
389 qw.QSizePolicy(
390 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
391 slider_opacity.setMinimum(0)
392 slider_opacity.setMaximum(1000)
394 opacity_label = qw.QLabel('Opacity')
395 layout.addWidget(opacity_label, il, 0)
396 layout.addWidget(slider_opacity, il, 1)
398 state_bind_slider(
399 self, state, 'opacity', slider_opacity, factor=0.001)
401 self._opacity_label = opacity_label
402 self._opacity_slider = slider_opacity
404 # color
405 il += 1
406 layout.addWidget(qw.QLabel('Color'), il, 0)
408 cb = common.strings_to_combobox(
409 ['black', 'white', 'blue', 'red'])
411 layout.addWidget(cb, il, 1)
412 state_bind_combobox_color(self, state, 'color', cb)
414 # linewidth outline
415 il += 1
416 layout.addWidget(qw.QLabel('Line Width'), il, 0)
418 slider = qw.QSlider(qc.Qt.Horizontal)
419 slider.setSizePolicy(
420 qw.QSizePolicy(
421 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
422 slider.setMinimum(0)
423 slider.setMaximum(100)
424 layout.addWidget(slider, il, 1)
425 state_bind_slider(
426 self, state, 'line_width', slider, factor=0.1)
428 # Clear scene
429 il += 1
430 pb = qw.QPushButton('Clear')
431 layout.addWidget(pb, il, 1)
432 pb.clicked.connect(self.clear)
434 # Change view to source
435 pb = qw.QPushButton('Move To')
436 layout.addWidget(pb, il, 2)
437 pb.clicked.connect(self.update_view)
439 self._controls = frame
441 self._update_controls()
443 return self._controls
445 def _update_controls(self):
446 state = self._state
447 if state.geometry:
448 if len(state.display_parameter) != 0:
449 values = state.geometry.get_property(state.display_parameter)
451 if values.ndim == 2:
452 self._time_label.setVisible(True)
453 self._time_slider.setVisible(True)
454 self._opacity_label.setVisible(True)
455 self._opacity_slider.setVisible(True)
456 else:
457 self._time_label.setVisible(False)
458 self._time_slider.setVisible(False)
459 self._opacity_label.setVisible(False)
460 self._opacity_slider.setVisible(False)
463__all__ = [
464 'GeometryElement',
465 'GeometryState'
466]