1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import logging
8from pyrocko.guts import Bool, String, load, Float
9from pyrocko.geometry import arr_vertices, arr_faces
10from pyrocko.gui.qt_compat import qw, qc
11from pyrocko.gui.vtk_util import TrimeshPipe, ColorbarPipe, OutlinesPipe, Color
13from pyrocko.model import Geometry
15from . import base
16from .. import common
19logger = logging.getLogger('geometry')
21guts_prefix = 'sparrow'
23km = 1e3
26class GeometryState(base.ElementState):
27 opacity = Float.T(default=1.0)
28 visible = Bool.T(default=True)
29 geometry = Geometry.T(default=None, optional=True)
30 display_parameter = String.T(default="")
31 time = Float.T(default=0., optional=True)
32 cpt = base.CPTState.T(default=base.CPTState.D())
33 color = Color.T(default=Color.D('white'))
34 line_width = Float.T(default=1.0)
36 def create(self):
37 element = GeometryElement()
38 return element
41class GeometryElement(base.Element):
43 def __init__(self):
44 base.Element.__init__(self)
45 self._parent = None
46 self._state = None
47 self._controls = None
49 self._pipe = None
50 self._cbar_pipe = None
51 self._outlines_pipe = []
53 self.cpt_handler = base.CPTHandler()
55 def remove(self):
56 if self._parent and self._state:
57 self._parent.state.elements.remove(self._state)
59 def init_pipeslots(self):
60 if not self._pipe:
61 self._pipe.append([])
63 def remove_pipes(self):
64 if self._pipe is not None:
65 self._parent.remove_actor(self._pipe.actor)
67 if self._cbar_pipe is not None:
68 self._parent.remove_actor(self._cbar_pipe.actor)
70 if len(self._outlines_pipe) > 0:
71 for pipe in self._outlines_pipe:
72 self._parent.remove_actor(pipe.actor)
74 self._pipe = None
75 self._cbar_pipe = None
76 self._outlines_pipe = []
78 def set_parent(self, parent):
79 self._parent = parent
80 self._parent.add_panel(
81 self.get_title_label(),
82 self._get_controls(),
83 visible=True,
84 title_controls=[
85 self.get_title_control_remove(),
86 self.get_title_control_visible()])
88 self.talkie_connect(
89 self._parent.state,
90 ['tmin', 'tmax', 'lat', 'lon'],
91 self.update)
93 self.update()
95 def unset_parent(self):
96 self.unbind_state()
97 if self._parent:
98 if self._pipe:
99 self.remove_pipes()
101 if self._controls:
102 self._parent.remove_panel(self._controls)
103 self._controls = None
105 self._parent.update_view()
106 self._parent = None
108 def bind_state(self, state):
109 base.Element.bind_state(self, state)
111 self.talkie_connect(
112 state,
113 ['visible', 'geometry', 'display_parameter', 'time',
114 'opacity', 'color', 'line_width'],
115 self.update)
117 self.cpt_handler.bind_state(state.cpt, self.update)
119 def unbind_state(self):
120 self.cpt_handler.unbind_state()
121 base.Element.unbind_state(self)
123 def get_cpt_name(self, cpt, display_parameter):
124 return '{}_{}'.format(cpt, display_parameter)
126 def update_cpt(self, state):
128 if len(state.display_parameter) != 0:
129 values = state.geometry.get_property()
130 # TODO Check
131 # if values.ndim == 2:
132 # values = values.sum(1)
134 self.cpt_handler._values = values
135 self.cpt_handler.update_cpt()
137 def get_name(self):
138 return 'Geometry'
140 def open_file_load_dialog(self):
141 caption = 'Select one file containing a geometry to open'
142 fns, _ = qw.QFileDialog.getOpenFileNames(
143 self._parent, caption, options=common.qfiledialog_options)
145 if fns:
146 self.load_file(str(fns[0]))
147 else:
148 return
150 def load_file(self, path):
152 loaded_geometry = load(filename=path)
153 props = loaded_geometry.properties.get_col_names(sub_headers=False)
155 if props:
156 if self._state.display_parameter not in props:
157 self._state.display_parameter = props[0]
159 self._parent.remove_panel(self._controls)
160 self._controls = None
161 self._state.geometry = loaded_geometry
163 self._parent.add_panel(
164 self.get_title_label(),
165 self._get_controls(),
166 visible=True,
167 title_controls=[
168 self.get_title_control_remove(),
169 self.get_title_control_visible()])
171 self.update()
173 def get_values(self, geom):
174 values = geom.get_property(self._state.display_parameter)
176 if geom.event is not None:
177 ref_time = geom.event.time
178 else:
179 ref_time = 0.
181 if len(values.shape) == 2:
182 tmin = self._parent.state.tmin
183 tmax = self._parent.state.tmax
184 if tmin is not None:
185 ref_tmin = tmin - ref_time
186 ref_idx_min = geom.time2idx(ref_tmin)
187 else:
188 ref_idx_min = geom.time2idx(self._state.time)
190 if tmax is not None:
191 ref_tmax = tmax - ref_time
192 ref_idx_max = geom.time2idx(ref_tmax)
193 else:
194 ref_idx_max = geom.time2idx(self._state.time)
196 if ref_idx_min == ref_idx_max:
197 out = values[:, ref_idx_min]
198 elif ref_idx_min > ref_idx_max:
199 out = values[:, ref_idx_min]
200 elif ref_idx_max < ref_idx_min:
201 out = values[:, ref_idx_max]
202 else:
203 # TODO CHECK
204 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
205 out = values[:, ref_idx_max]
206 else:
207 out = values.ravel()
208 return out
210 def update_view(self, *args):
211 pstate = self._parent.state
212 geom = self._state.geometry
214 if geom.event:
215 pstate.lat = geom.event.lat
216 pstate.lon = geom.event.lon
218 self.update()
220 def update_outlines(self, geo):
221 state = self._state
222 if len(self._outlines_pipe) == 0:
223 for cs in ['latlondepth']:
224 outline_pipe = OutlinesPipe(
225 geo, color=state.color, cs=cs)
226 outline_pipe.set_line_width(state.line_width)
227 self._outlines_pipe.append(outline_pipe)
228 self._parent.add_actor(
229 self._outlines_pipe[-1].actor)
231 else:
232 for outline_pipe in self._outlines_pipe:
233 outline_pipe.set_color(state.color)
234 outline_pipe.set_line_width(state.line_width)
236 def update(self, *args):
238 state = self._state
240 if state.geometry and self._controls:
241 self._update_controls()
242 # base.update_cpt(self)
243 self.update_cpt(state)
245 if state.visible:
246 # cpt_name = self.get_cpt_name(
247 # state.cpt, state.display_parameter)
248 geo = state.geometry
249 lut = self.cpt_handler._lookuptable
250 no_faces = geo.no_faces()
251 if no_faces:
252 values = self.get_values(geo)
253 if not isinstance(self._pipe, TrimeshPipe):
254 vertices = arr_vertices(geo.get_vertices('xyz'))
255 faces = arr_faces(geo.get_faces())
256 self._pipe = TrimeshPipe(
257 vertices, faces,
258 values=values,
259 lut=lut,
260 backface_culling=False)
261 self._cbar_pipe = ColorbarPipe(
262 lut=lut, cbar_title=state.display_parameter)
263 self._parent.add_actor(self._pipe.actor)
264 self._parent.add_actor(self._cbar_pipe.actor)
265 else:
266 self._pipe.set_values(values)
267 self._pipe.set_lookuptable(lut)
268 self._pipe.set_opacity(self._state.opacity)
270 self._cbar_pipe.set_lookuptable(lut)
271 self._cbar_pipe.set_title(state.display_parameter)
273 if geo.outlines:
274 self.update_outlines(geo)
275 else:
276 self.remove_pipes()
278 self._parent.update_view()
280 def _get_controls(self):
281 state = self._state
282 if not self._controls:
283 from ..state import state_bind_combobox, \
284 state_bind_slider, state_bind_combobox_color
286 frame = qw.QFrame()
287 layout = qw.QGridLayout()
288 layout.setAlignment(qc.Qt.AlignTop)
289 frame.setLayout(layout)
291 # load geometry
292 pb = qw.QPushButton('Load')
293 layout.addWidget(pb, 0, 0)
295 pb.clicked.connect(self.open_file_load_dialog)
297 # property choice
298 il = 1
299 if state.geometry:
301 pb = qw.QPushButton('Move to')
302 layout.addWidget(pb, 0, 1)
303 pb.clicked.connect(self.update_view)
305 props = []
306 for prop in state.geometry.properties.get_col_names(
307 sub_headers=False):
308 props.append(prop)
310 layout.addWidget(qw.QLabel('Display parameter'), il, 0)
311 cb = qw.QComboBox()
313 unique_props = list(set(props))
314 for i, s in enumerate(unique_props):
315 cb.insertItem(i, s)
317 layout.addWidget(cb, il, 1)
318 state_bind_combobox(self, state, 'display_parameter', cb)
320 if state.geometry.no_faces != 0:
321 # color maps
322 self.cpt_handler.cpt_controls(
323 self._parent, self._state.cpt, layout)
325 il += 1
326 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
328 self.cpt_handler._update_cpt_combobox()
329 self.cpt_handler._update_cptscale_lineedit()
331 # times slider
332 if state.geometry.times:
333 il = layout.rowCount() + 1
334 slider = qw.QSlider(qc.Qt.Horizontal)
335 slider.setSizePolicy(
336 qw.QSizePolicy(
337 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
339 def iround(x):
340 return int(round(x))
342 slider.setMinimum(iround(state.geometry.times.min()))
343 slider.setMaximum(iround(state.geometry.times.max()))
344 slider.setSingleStep(iround(state.geometry.deltat))
345 slider.setPageStep(iround(state.geometry.deltat))
347 time_label = qw.QLabel('Time')
348 layout.addWidget(time_label, il, 0)
349 layout.addWidget(slider, il, 1)
351 state_bind_slider(
352 self, state, 'time', slider, dtype=int)
354 self._time_label = time_label
355 self._time_slider = slider
357 il = layout.rowCount() + 1
358 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
359 slider_opacity.setSizePolicy(
360 qw.QSizePolicy(
361 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
362 slider_opacity.setMinimum(0)
363 slider_opacity.setMaximum(1000)
365 opacity_label = qw.QLabel('Opacity')
366 layout.addWidget(opacity_label, il, 0)
367 layout.addWidget(slider_opacity, il, 1)
369 state_bind_slider(
370 self, state, 'opacity', slider_opacity, factor=0.001)
372 self._opacity_label = opacity_label
373 self._opacity_slider = slider_opacity
375 # color
376 il += 1
377 layout.addWidget(qw.QLabel('Color'), il, 0)
379 cb = common.strings_to_combobox(
380 ['black', 'white', 'blue', 'red'])
382 layout.addWidget(cb, il, 1)
383 state_bind_combobox_color(self, state, 'color', cb)
385 # linewidth outline
386 il += 1
387 layout.addWidget(qw.QLabel('Line width'), il, 0)
389 slider = qw.QSlider(qc.Qt.Horizontal)
390 slider.setSizePolicy(
391 qw.QSizePolicy(
392 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
393 slider.setMinimum(0)
394 slider.setMaximum(100)
395 layout.addWidget(slider, il, 1)
396 state_bind_slider(
397 self, state, 'line_width', slider, factor=0.1)
399 self._controls = frame
401 self._update_controls()
403 return self._controls
405 def _update_controls(self):
406 state = self._state
407 if state.geometry:
408 if len(state.display_parameter) != 0:
409 values = state.geometry.get_property()
411 if values.ndim == 2:
412 self._time_label.setVisible(True)
413 self._time_slider.setVisible(True)
414 self._opacity_label.setVisible(True)
415 self._opacity_slider.setVisible(True)
416 else:
417 self._time_label.setVisible(False)
418 self._time_slider.setVisible(False)
419 self._opacity_label.setVisible(False)
420 self._opacity_slider.setVisible(False)
423__all__ = [
424 'GeometryElement',
425 'GeometryState'
426]