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, StringChoice, 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 CPTChoices(StringChoice):
28 choices = ['slip_colors', 'seismic', 'jet', 'hot_r', 'gist_earth_r']
31class GeometryState(base.ElementState):
32 opacity = Float.T(default=1.0)
33 visible = Bool.T(default=True)
34 geometry = Geometry.T(default=None, optional=True)
35 display_parameter = String.T(default='slip')
36 time = Float.T(default=0., optional=True)
37 cpt = base.CPTState.T(default=base.CPTState.D())
39 def create(self):
40 element = GeometryElement()
41 return element
44class GeometryElement(base.Element):
46 def __init__(self):
47 base.Element.__init__(self)
48 self._parent = None
49 self._state = None
50 self._controls = None
52 self._pipe = None
53 self._cbar_pipe = None
54 self._outlines_pipe = []
56 self.cpt_handler = base.CPTHandler()
58 def remove(self):
59 if self._parent and self._state:
60 self._parent.state.elements.remove(self._state)
62 def init_pipeslots(self):
63 if not self._pipe:
64 self._pipe.append([])
66 def remove_pipes(self):
67 if self._pipe is not None:
68 self._parent.remove_actor(self._pipe.actor)
70 if self._cbar_pipe is not None:
71 self._parent.remove_actor(self._cbar_pipe.actor)
73 if len(self._outlines_pipe) > 0:
74 for pipe in self._outlines_pipe:
75 self._parent.remove_actor(pipe.actor)
77 self._pipe = None
78 self._cbar_pipe = None
79 self._outlines_pipe = []
81 def set_parent(self, parent):
82 self._parent = parent
83 self._parent.add_panel(
84 self.get_title_label(),
85 self._get_controls(),
86 visible=True,
87 title_controls=[
88 self.get_title_control_remove(),
89 self.get_title_control_visible()])
91 self.talkie_connect(
92 self._parent.state,
93 ['tmin', 'tmax', 'lat', 'lon'],
94 self.update)
96 self.update()
98 def unset_parent(self):
99 self.unbind_state()
100 if self._parent:
101 if self._pipe:
102 self.remove_pipes()
104 if self._controls:
105 self._parent.remove_panel(self._controls)
106 self._controls = None
108 self._parent.update_view()
109 self._parent = None
111 def bind_state(self, state):
112 base.Element.bind_state(self, state)
114 self.talkie_connect(
115 state,
116 ['visible', 'geometry', 'display_parameter', 'time', 'opacity'],
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 values = state.geometry.get_property(state.display_parameter)
131 # TODO Check
132 # if values.ndim == 2:
133 # values = values.sum(1)
135 self.cpt_handler._values = values
136 self.cpt_handler.update_cpt()
138 def get_name(self):
139 return 'Geometry'
141 def open_file_load_dialog(self):
142 caption = 'Select one file containing a geometry to open'
143 fns, _ = qw.QFileDialog.getOpenFileNames(
144 self._parent, caption, options=common.qfiledialog_options)
146 if fns:
147 self.load_file(str(fns[0]))
148 else:
149 return
151 def load_file(self, path):
153 loaded_geometry = load(filename=path)
154 props = loaded_geometry.properties.get_col_names(sub_headers=False)
156 if props:
157 if self._state.display_parameter not in props:
158 self._state.display_parameter = props[0]
159 else:
160 raise ValueError(
161 'Imported geometry contains no property to be displayed!')
163 self._parent.remove_panel(self._controls)
164 self._controls = None
165 self._state.geometry = loaded_geometry
167 self._parent.add_panel(
168 self.get_title_label(),
169 self._get_controls(),
170 visible=True,
171 title_controls=[
172 self.get_title_control_remove(),
173 self.get_title_control_visible()])
175 self.update()
177 def get_values(self, geom):
178 values = geom.get_property(self._state.display_parameter)
180 if geom.event is not None:
181 ref_time = geom.event.time
182 else:
183 ref_time = 0.
185 if len(values.shape) == 2:
186 tmin = self._parent.state.tmin
187 tmax = self._parent.state.tmax
188 if tmin is not None:
189 ref_tmin = tmin - ref_time
190 ref_idx_min = geom.time2idx(ref_tmin)
191 else:
192 ref_idx_min = geom.time2idx(self._state.time)
194 if tmax is not None:
195 ref_tmax = tmax - ref_time
196 ref_idx_max = geom.time2idx(ref_tmax)
197 else:
198 ref_idx_max = geom.time2idx(self._state.time)
200 if ref_idx_min == ref_idx_max:
201 out = values[:, ref_idx_min]
202 elif ref_idx_min > ref_idx_max:
203 out = values[:, ref_idx_min]
204 elif ref_idx_max < ref_idx_min:
205 out = values[:, ref_idx_max]
206 else:
207 # TODO CHECK
208 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
209 out = values[:, ref_idx_max]
210 else:
211 out = values.ravel()
212 return out
214 def update_view(self, *args):
215 pstate = self._parent.state
216 geom = self._state.geometry
218 if geom.event:
219 pstate.lat = geom.event.lat
220 pstate.lon = geom.event.lon
222 self.update()
224 def update(self, *args):
226 state = self._state
228 if state.geometry and self._controls:
229 self._update_controls()
230 # base.update_cpt(self)
231 self.update_cpt(state)
233 if state.visible:
234 # cpt_name = self.get_cpt_name(
235 # state.cpt, state.display_parameter)
236 geo = state.geometry
237 values = self.get_values(geo)
238 lut = self.cpt_handler._lookuptable
239 if not isinstance(self._pipe, TrimeshPipe):
240 vertices = arr_vertices(geo.get_vertices('xyz'))
241 faces = arr_faces(geo.get_faces())
242 self._pipe = TrimeshPipe(
243 vertices, faces,
244 values=values,
245 lut=lut,
246 backface_culling=False)
247 self._cbar_pipe = ColorbarPipe(
248 lut=lut, cbar_title=state.display_parameter)
249 self._parent.add_actor(self._pipe.actor)
250 self._parent.add_actor(self._cbar_pipe.actor)
252 if geo.outlines:
253 self._outlines_pipe.append(OutlinesPipe(
254 geo, color=(1., 1., 1.), cs='latlondepth'))
255 self._parent.add_actor(
256 self._outlines_pipe[-1].actor)
257 self._outlines_pipe.append(OutlinesPipe(
258 geo, color=(0.6, 0.6, 0.6), cs='latlon'))
259 self._parent.add_actor(
260 self._outlines_pipe[-1].actor)
262 else:
263 self._pipe.set_values(values)
264 self._pipe.set_lookuptable(lut)
265 self._pipe.set_opacity(self._state.opacity)
267 self._cbar_pipe.set_lookuptable(lut)
268 self._cbar_pipe.set_title(state.display_parameter)
269 else:
270 if self._pipe:
271 self.remove_pipes()
273 self._parent.update_view()
275 def _get_controls(self):
276 state = self._state
277 if not self._controls:
278 from ..state import state_bind_combobox, \
279 state_bind_slider
281 frame = qw.QFrame()
282 layout = qw.QGridLayout()
283 layout.setAlignment(qc.Qt.AlignTop)
284 frame.setLayout(layout)
286 # load geometry
287 pb = qw.QPushButton('Load')
288 layout.addWidget(pb, 0, 0)
290 pb.clicked.connect(self.open_file_load_dialog)
292 # property choice
293 il = 1
294 if state.geometry:
296 pb = qw.QPushButton('Move to')
297 layout.addWidget(pb, 0, 1)
298 pb.clicked.connect(self.update_view)
300 props = []
301 for prop in state.geometry.properties.get_col_names(
302 sub_headers=False):
303 props.append(prop)
305 layout.addWidget(qw.QLabel('Display parameter'), il, 0)
306 cb = qw.QComboBox()
308 unique_props = list(set(props))
309 for i, s in enumerate(unique_props):
310 cb.insertItem(i, s)
312 layout.addWidget(cb, il, 1)
313 state_bind_combobox(self, state, 'display_parameter', cb)
315 # color maps
316 self.cpt_handler.cpt_controls(
317 self._parent, self._state.cpt, layout)
319 # times slider
320 il = layout.rowCount() + 1
321 slider = qw.QSlider(qc.Qt.Horizontal)
322 slider.setSizePolicy(
323 qw.QSizePolicy(
324 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
326 def iround(x):
327 return int(round(x))
329 slider.setMinimum(iround(state.geometry.times.min()))
330 slider.setMaximum(iround(state.geometry.times.max()))
331 slider.setSingleStep(iround(state.geometry.deltat))
332 slider.setPageStep(iround(state.geometry.deltat))
334 time_label = qw.QLabel('Time')
335 layout.addWidget(time_label, il, 0)
336 layout.addWidget(slider, il, 1)
338 state_bind_slider(
339 self, state, 'time', slider, dtype=int)
341 self._time_label = time_label
342 self._time_slider = slider
344 il += 1
345 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
346 slider_opacity.setSizePolicy(
347 qw.QSizePolicy(
348 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
349 slider_opacity.setMinimum(0)
350 slider_opacity.setMaximum(1000)
352 opacity_label = qw.QLabel('Opacity')
353 layout.addWidget(opacity_label, il, 0)
354 layout.addWidget(slider_opacity, il, 1)
356 state_bind_slider(
357 self, state, 'opacity', slider_opacity, factor=0.001)
359 self._opacity_label = opacity_label
360 self._opacity_slider = slider_opacity
362 il += 1
363 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
365 self.cpt_handler._update_cpt_combobox()
366 self.cpt_handler._update_cptscale_lineedit()
368 self._controls = frame
370 self._update_controls()
372 return self._controls
374 def _update_controls(self):
375 state = self._state
376 if state.geometry:
377 values = state.geometry.get_property(state.display_parameter)
379 if values.ndim == 2:
380 self._time_label.setVisible(True)
381 self._time_slider.setVisible(True)
382 self._opacity_label.setVisible(True)
383 self._opacity_slider.setVisible(True)
384 else:
385 self._time_label.setVisible(False)
386 self._time_slider.setVisible(False)
387 self._opacity_label.setVisible(False)
388 self._opacity_slider.setVisible(False)
391__all__ = [
392 'GeometryElement',
393 'GeometryState'
394]