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 self._listeners = []
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_name(),
85 self._get_controls(),
86 visible=True,
87 title_controls=[
88 self.get_title_control_remove(),
89 self.get_title_control_visible()])
91 for var in ['tmin', 'tmax', 'lat', 'lon']:
92 self.register_state_listener3(self.update, self._parent.state, var)
94 self.update()
96 def unset_parent(self):
97 self.unbind_state()
98 if self._parent:
99 if self._pipe:
100 self.remove_pipes()
102 if self._controls:
103 self._parent.remove_panel(self._controls)
104 self._controls = None
106 self._parent.update_view()
107 self._parent = None
109 def bind_state(self, state):
110 base.Element.bind_state(self, state)
111 for var in [
112 'visible', 'geometry', 'display_parameter', 'time', 'opacity']:
114 self.register_state_listener3(self.update, state, var)
116 self.cpt_handler.bind_state(state.cpt, self.update)
118 def unbind_state(self):
119 for listener in self._listeners:
120 try:
121 listener.release()
122 except Exception:
123 pass
125 self.cpt_handler.unbind_state()
126 self._state = None
128 def get_cpt_name(self, cpt, display_parameter):
129 return '{}_{}'.format(cpt, display_parameter)
131 def update_cpt(self, state):
133 values = state.geometry.get_property(state.display_parameter)
134 # TODO Check
135 # if values.ndim == 2:
136 # values = values.sum(1)
138 self.cpt_handler._values = values
139 self.cpt_handler.update_cpt()
141 def get_name(self):
142 return 'Geometry'
144 def open_file_load_dialog(self):
145 caption = 'Select one file containing a geometry to open'
146 fns, _ = qw.QFileDialog.getOpenFileNames(
147 self._parent, caption, options=common.qfiledialog_options)
149 if fns:
150 self.load_file(str(fns[0]))
151 else:
152 return
154 def load_file(self, path):
156 loaded_geometry = load(filename=path)
157 props = loaded_geometry.properties.get_col_names(sub_headers=False)
159 if props:
160 if self._state.display_parameter not in props:
161 self._state.display_parameter = props[0]
162 else:
163 raise ValueError(
164 'Imported geometry contains no property to be displayed!')
166 self._parent.remove_panel(self._controls)
167 self._controls = None
168 self._state.geometry = loaded_geometry
170 self._parent.add_panel(
171 self.get_name(),
172 self._get_controls(),
173 visible=True,
174 title_controls=[
175 self.get_title_control_remove(),
176 self.get_title_control_visible()])
178 self.update()
180 def get_values(self, geom):
181 values = geom.get_property(self._state.display_parameter)
183 if geom.event is not None:
184 ref_time = geom.event.time
185 else:
186 ref_time = 0.
188 if len(values.shape) == 2:
189 tmin = self._parent.state.tmin
190 tmax = self._parent.state.tmax
191 if tmin is not None:
192 ref_tmin = tmin - ref_time
193 ref_idx_min = geom.time2idx(ref_tmin)
194 else:
195 ref_idx_min = geom.time2idx(self._state.time)
197 if tmax is not None:
198 ref_tmax = tmax - ref_time
199 ref_idx_max = geom.time2idx(ref_tmax)
200 else:
201 ref_idx_max = geom.time2idx(self._state.time)
203 if ref_idx_min == ref_idx_max:
204 out = values[:, ref_idx_min]
205 elif ref_idx_min > ref_idx_max:
206 out = values[:, ref_idx_min]
207 elif ref_idx_max < ref_idx_min:
208 out = values[:, ref_idx_max]
209 else:
210 # TODO CHECK
211 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
212 out = values[:, ref_idx_max]
213 else:
214 out = values.ravel()
215 return out
217 def update_view(self, *args):
218 pstate = self._parent.state
219 geom = self._state.geometry
221 if geom.event:
222 pstate.lat = geom.event.lat
223 pstate.lon = geom.event.lon
225 self.update()
227 def update(self, *args):
229 state = self._state
231 if state.geometry and self._controls:
232 self._update_controls()
233 # base.update_cpt(self)
234 self.update_cpt(state)
236 if state.visible:
237 # cpt_name = self.get_cpt_name(
238 # state.cpt, state.display_parameter)
239 geo = state.geometry
240 values = self.get_values(geo)
241 lut = self.cpt_handler._lookuptable
242 if not isinstance(self._pipe, TrimeshPipe):
243 vertices = arr_vertices(geo.get_vertices('xyz'))
244 faces = arr_faces(geo.get_faces())
245 self._pipe = TrimeshPipe(
246 vertices, faces,
247 values=values,
248 lut=lut,
249 backface_culling=False)
250 self._cbar_pipe = ColorbarPipe(
251 lut=lut, cbar_title=state.display_parameter)
252 self._parent.add_actor(self._pipe.actor)
253 self._parent.add_actor(self._cbar_pipe.actor)
255 if geo.outlines:
256 self._outlines_pipe.append(OutlinesPipe(
257 geo, color=(1., 1., 1.), cs='latlondepth'))
258 self._parent.add_actor(
259 self._outlines_pipe[-1].actor)
260 self._outlines_pipe.append(OutlinesPipe(
261 geo, color=(0.6, 0.6, 0.6), cs='latlon'))
262 self._parent.add_actor(
263 self._outlines_pipe[-1].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)
272 else:
273 if self._pipe:
274 self.remove_pipes()
276 self._parent.update_view()
278 def _get_controls(self):
279 state = self._state
280 if not self._controls:
281 from ..state import state_bind_combobox, \
282 state_bind_slider
284 frame = qw.QFrame()
285 layout = qw.QGridLayout()
286 layout.setAlignment(qc.Qt.AlignTop)
287 frame.setLayout(layout)
289 # load geometry
290 pb = qw.QPushButton('Load')
291 layout.addWidget(pb, 0, 0)
293 pb.clicked.connect(self.open_file_load_dialog)
295 # property choice
296 il = 1
297 if state.geometry:
299 pb = qw.QPushButton('Move to')
300 layout.addWidget(pb, 0, 1)
301 pb.clicked.connect(self.update_view)
303 props = []
304 for prop in state.geometry.properties.get_col_names(
305 sub_headers=False):
306 props.append(prop)
308 layout.addWidget(qw.QLabel('Display parameter'), il, 0)
309 cb = qw.QComboBox()
311 unique_props = list(set(props))
312 for i, s in enumerate(unique_props):
313 cb.insertItem(i, s)
315 layout.addWidget(cb, il, 1)
316 state_bind_combobox(self, state, 'display_parameter', cb)
318 # color maps
319 self.cpt_handler.cpt_controls(
320 self._parent, self._state.cpt, layout)
322 # times slider
323 il = layout.rowCount() + 1
324 slider = qw.QSlider(qc.Qt.Horizontal)
325 slider.setSizePolicy(
326 qw.QSizePolicy(
327 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
329 def iround(x):
330 return int(round(x))
332 slider.setMinimum(iround(state.geometry.times.min()))
333 slider.setMaximum(iround(state.geometry.times.max()))
334 slider.setSingleStep(iround(state.geometry.deltat))
335 slider.setPageStep(iround(state.geometry.deltat))
337 time_label = qw.QLabel('Time')
338 layout.addWidget(time_label, il, 0)
339 layout.addWidget(slider, il, 1)
341 state_bind_slider(
342 self, state, 'time', slider, dtype=int)
344 self._time_label = time_label
345 self._time_slider = slider
347 il += 1
348 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
349 slider_opacity.setSizePolicy(
350 qw.QSizePolicy(
351 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
352 slider_opacity.setMinimum(0)
353 slider_opacity.setMaximum(1000)
355 opacity_label = qw.QLabel('Opacity')
356 layout.addWidget(opacity_label, il, 0)
357 layout.addWidget(slider_opacity, il, 1)
359 state_bind_slider(
360 self, state, 'opacity', slider_opacity, factor=0.001)
362 self._opacity_label = opacity_label
363 self._opacity_slider = slider_opacity
365 il += 1
366 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
368 self.cpt_handler._update_cpt_combobox()
369 self.cpt_handler._update_cptscale_lineedit()
371 self._controls = frame
373 self._update_controls()
375 return self._controls
377 def _update_controls(self):
378 state = self._state
379 if state.geometry:
380 values = state.geometry.get_property(state.display_parameter)
382 if values.ndim == 2:
383 self._time_label.setVisible(True)
384 self._time_slider.setVisible(True)
385 self._opacity_label.setVisible(True)
386 self._opacity_slider.setVisible(True)
387 else:
388 self._time_label.setVisible(False)
389 self._time_slider.setVisible(False)
390 self._opacity_label.setVisible(False)
391 self._opacity_slider.setVisible(False)
394__all__ = [
395 'GeometryElement',
396 'GeometryState'
397]