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
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='slip')
31 time = Float.T(default=0., optional=True)
32 cpt = base.CPTState.T(default=base.CPTState.D())
34 def create(self):
35 element = GeometryElement()
36 return element
39class GeometryElement(base.Element):
41 def __init__(self):
42 base.Element.__init__(self)
43 self._parent = None
44 self._state = None
45 self._controls = None
47 self._pipe = None
48 self._cbar_pipe = None
49 self._outlines_pipe = []
51 self.cpt_handler = base.CPTHandler()
53 def remove(self):
54 if self._parent and self._state:
55 self._parent.state.elements.remove(self._state)
57 def init_pipeslots(self):
58 if not self._pipe:
59 self._pipe.append([])
61 def remove_pipes(self):
62 if self._pipe is not None:
63 self._parent.remove_actor(self._pipe.actor)
65 if self._cbar_pipe is not None:
66 self._parent.remove_actor(self._cbar_pipe.actor)
68 if len(self._outlines_pipe) > 0:
69 for pipe in self._outlines_pipe:
70 self._parent.remove_actor(pipe.actor)
72 self._pipe = None
73 self._cbar_pipe = None
74 self._outlines_pipe = []
76 def set_parent(self, parent):
77 self._parent = parent
78 self._parent.add_panel(
79 self.get_title_label(),
80 self._get_controls(),
81 visible=True,
82 title_controls=[
83 self.get_title_control_remove(),
84 self.get_title_control_visible()])
86 self.talkie_connect(
87 self._parent.state,
88 ['tmin', 'tmax', 'lat', 'lon'],
89 self.update)
91 self.update()
93 def unset_parent(self):
94 self.unbind_state()
95 if self._parent:
96 if self._pipe:
97 self.remove_pipes()
99 if self._controls:
100 self._parent.remove_panel(self._controls)
101 self._controls = None
103 self._parent.update_view()
104 self._parent = None
106 def bind_state(self, state):
107 base.Element.bind_state(self, state)
109 self.talkie_connect(
110 state,
111 ['visible', 'geometry', 'display_parameter', 'time', 'opacity'],
112 self.update)
114 self.cpt_handler.bind_state(state.cpt, self.update)
116 def unbind_state(self):
117 self.cpt_handler.unbind_state()
118 base.Element.unbind_state(self)
120 def get_cpt_name(self, cpt, display_parameter):
121 return '{}_{}'.format(cpt, display_parameter)
123 def update_cpt(self, state):
125 values = state.geometry.get_property(state.display_parameter)
126 # TODO Check
127 # if values.ndim == 2:
128 # values = values.sum(1)
130 self.cpt_handler._values = values
131 self.cpt_handler.update_cpt()
133 def get_name(self):
134 return 'Geometry'
136 def open_file_load_dialog(self):
137 caption = 'Select one file containing a geometry to open'
138 fns, _ = qw.QFileDialog.getOpenFileNames(
139 self._parent, caption, options=common.qfiledialog_options)
141 if fns:
142 self.load_file(str(fns[0]))
143 else:
144 return
146 def load_file(self, path):
148 loaded_geometry = load(filename=path)
149 props = loaded_geometry.properties.get_col_names(sub_headers=False)
151 if props:
152 if self._state.display_parameter not in props:
153 self._state.display_parameter = props[0]
154 else:
155 raise ValueError(
156 'Imported geometry contains no property to be displayed!')
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.event:
214 pstate.lat = geom.event.lat
215 pstate.lon = geom.event.lon
217 self.update()
219 def update(self, *args):
221 state = self._state
223 if state.geometry and self._controls:
224 self._update_controls()
225 # base.update_cpt(self)
226 self.update_cpt(state)
228 if state.visible:
229 # cpt_name = self.get_cpt_name(
230 # state.cpt, state.display_parameter)
231 geo = state.geometry
232 values = self.get_values(geo)
233 lut = self.cpt_handler._lookuptable
234 if not isinstance(self._pipe, TrimeshPipe):
235 vertices = arr_vertices(geo.get_vertices('xyz'))
236 faces = arr_faces(geo.get_faces())
237 self._pipe = TrimeshPipe(
238 vertices, faces,
239 values=values,
240 lut=lut,
241 backface_culling=False)
242 self._cbar_pipe = ColorbarPipe(
243 lut=lut, cbar_title=state.display_parameter)
244 self._parent.add_actor(self._pipe.actor)
245 self._parent.add_actor(self._cbar_pipe.actor)
247 if geo.outlines:
248 self._outlines_pipe.append(OutlinesPipe(
249 geo, color=(1., 1., 1.), cs='latlondepth'))
250 self._parent.add_actor(
251 self._outlines_pipe[-1].actor)
252 self._outlines_pipe.append(OutlinesPipe(
253 geo, color=(0.6, 0.6, 0.6), cs='latlon'))
254 self._parent.add_actor(
255 self._outlines_pipe[-1].actor)
257 else:
258 self._pipe.set_values(values)
259 self._pipe.set_lookuptable(lut)
260 self._pipe.set_opacity(self._state.opacity)
262 self._cbar_pipe.set_lookuptable(lut)
263 self._cbar_pipe.set_title(state.display_parameter)
264 else:
265 if self._pipe:
266 self.remove_pipes()
268 self._parent.update_view()
270 def _get_controls(self):
271 state = self._state
272 if not self._controls:
273 from ..state import state_bind_combobox, \
274 state_bind_slider
276 frame = qw.QFrame()
277 layout = qw.QGridLayout()
278 layout.setAlignment(qc.Qt.AlignTop)
279 frame.setLayout(layout)
281 # load geometry
282 pb = qw.QPushButton('Load')
283 layout.addWidget(pb, 0, 0)
285 pb.clicked.connect(self.open_file_load_dialog)
287 # property choice
288 il = 1
289 if state.geometry:
291 pb = qw.QPushButton('Move to')
292 layout.addWidget(pb, 0, 1)
293 pb.clicked.connect(self.update_view)
295 props = []
296 for prop in state.geometry.properties.get_col_names(
297 sub_headers=False):
298 props.append(prop)
300 layout.addWidget(qw.QLabel('Display parameter'), il, 0)
301 cb = qw.QComboBox()
303 unique_props = list(set(props))
304 for i, s in enumerate(unique_props):
305 cb.insertItem(i, s)
307 layout.addWidget(cb, il, 1)
308 state_bind_combobox(self, state, 'display_parameter', cb)
310 # color maps
311 self.cpt_handler.cpt_controls(
312 self._parent, self._state.cpt, layout)
314 # times slider
315 il = layout.rowCount() + 1
316 slider = qw.QSlider(qc.Qt.Horizontal)
317 slider.setSizePolicy(
318 qw.QSizePolicy(
319 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
321 def iround(x):
322 return int(round(x))
324 slider.setMinimum(iround(state.geometry.times.min()))
325 slider.setMaximum(iround(state.geometry.times.max()))
326 slider.setSingleStep(iround(state.geometry.deltat))
327 slider.setPageStep(iround(state.geometry.deltat))
329 time_label = qw.QLabel('Time')
330 layout.addWidget(time_label, il, 0)
331 layout.addWidget(slider, il, 1)
333 state_bind_slider(
334 self, state, 'time', slider, dtype=int)
336 self._time_label = time_label
337 self._time_slider = slider
339 il += 1
340 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
341 slider_opacity.setSizePolicy(
342 qw.QSizePolicy(
343 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
344 slider_opacity.setMinimum(0)
345 slider_opacity.setMaximum(1000)
347 opacity_label = qw.QLabel('Opacity')
348 layout.addWidget(opacity_label, il, 0)
349 layout.addWidget(slider_opacity, il, 1)
351 state_bind_slider(
352 self, state, 'opacity', slider_opacity, factor=0.001)
354 self._opacity_label = opacity_label
355 self._opacity_slider = slider_opacity
357 il += 1
358 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
360 self.cpt_handler._update_cpt_combobox()
361 self.cpt_handler._update_cptscale_lineedit()
363 self._controls = frame
365 self._update_controls()
367 return self._controls
369 def _update_controls(self):
370 state = self._state
371 if state.geometry:
372 values = state.geometry.get_property(state.display_parameter)
374 if values.ndim == 2:
375 self._time_label.setVisible(True)
376 self._time_slider.setVisible(True)
377 self._opacity_label.setVisible(True)
378 self._opacity_slider.setVisible(True)
379 else:
380 self._time_label.setVisible(False)
381 self._time_slider.setVisible(False)
382 self._opacity_label.setVisible(False)
383 self._opacity_slider.setVisible(False)
386__all__ = [
387 'GeometryElement',
388 'GeometryState'
389]