1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function, division
8import logging
10from pyrocko.guts import Bool, String, load, StringChoice, Float
11from pyrocko.geometry import arr_vertices, arr_faces
12from pyrocko.gui.qt_compat import qw, qc
13from pyrocko.gui.vtk_util import TrimeshPipe, ColorbarPipe, OutlinesPipe
15from pyrocko.model import Geometry
17from . import base
18from .. import common
21logger = logging.getLogger('geometry')
23guts_prefix = 'sparrow'
25km = 1e3
28class CPTChoices(StringChoice):
30 choices = ['slip_colors', 'seismic', 'jet', 'hot_r', 'gist_earth_r']
33class GeometryState(base.ElementState):
34 opacity = Float.T(default=1.0)
35 visible = Bool.T(default=True)
36 geometry = Geometry.T(default=None, optional=True)
37 display_parameter = String.T(default='slip')
38 time = Float.T(default=0., optional=True)
39 cpt = base.CPTState.T(default=base.CPTState.D())
41 def create(self):
42 element = GeometryElement()
43 return element
46class GeometryElement(base.Element):
48 def __init__(self):
49 self._listeners = []
50 self._parent = None
51 self._state = None
52 self._controls = None
54 self._pipe = None
55 self._cbar_pipe = None
56 self._outlines_pipe = []
58 self.cpt_handler = base.CPTHandler()
60 def remove(self):
61 if self._parent and self._state:
62 self._parent.state.elements.remove(self._state)
64 def init_pipeslots(self):
65 if not self._pipe:
66 self._pipe.append([])
68 def remove_pipes(self):
69 if self._pipe is not None:
70 self._parent.remove_actor(self._pipe.actor)
72 if self._cbar_pipe is not None:
73 self._parent.remove_actor(self._cbar_pipe.actor)
75 if len(self._outlines_pipe) > 0:
76 for pipe in self._outlines_pipe:
77 self._parent.remove_actor(pipe.actor)
79 self._pipe = None
80 self._cbar_pipe = None
81 self._outlines_pipe = []
83 def set_parent(self, parent):
84 self._parent = parent
85 self._parent.add_panel(
86 self.get_name(),
87 self._get_controls(),
88 visible=True,
89 title_controls=[
90 self.get_title_control_remove(),
91 self.get_title_control_visible()])
93 for var in ['tmin', 'tmax', 'lat', 'lon']:
94 self.register_state_listener3(self.update, self._parent.state, var)
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)
113 for var in [
114 'visible', 'geometry', 'display_parameter', 'time', 'opacity']:
116 self.register_state_listener3(self.update, state, var)
118 self.cpt_handler.bind_state(state.cpt, self.update)
120 def unbind_state(self):
121 for listener in self._listeners:
122 try:
123 listener.release()
124 except Exception:
125 pass
127 self.cpt_handler.unbind_state()
128 self._state = None
130 def get_cpt_name(self, cpt, display_parameter):
131 return '{}_{}'.format(cpt, display_parameter)
133 def update_cpt(self, state):
135 values = state.geometry.get_property(state.display_parameter)
136 # TODO Check
137 # if values.ndim == 2:
138 # values = values.sum(1)
140 self.cpt_handler._values = values
141 self.cpt_handler.update_cpt()
143 def get_name(self):
144 return 'Geometry'
146 def open_file_load_dialog(self):
147 caption = 'Select one file containing a geometry to open'
148 fns, _ = qw.QFileDialog.getOpenFileNames(
149 self._parent, caption, options=common.qfiledialog_options)
151 if fns:
152 self.load_file(str(fns[0]))
153 else:
154 return
156 def load_file(self, path):
158 loaded_geometry = load(filename=path)
159 props = loaded_geometry.properties.get_col_names(sub_headers=False)
161 if props:
162 if self._state.display_parameter not in props:
163 self._state.display_parameter = props[0]
164 else:
165 raise ValueError(
166 'Imported geometry contains no property to be displayed!')
168 self._parent.remove_panel(self._controls)
169 self._controls = None
170 self._state.geometry = loaded_geometry
172 self._parent.add_panel(
173 self.get_name(),
174 self._get_controls(),
175 visible=True,
176 title_controls=[
177 self.get_title_control_remove(),
178 self.get_title_control_visible()])
180 self.update()
182 def get_values(self, geom):
183 values = geom.get_property(self._state.display_parameter)
185 if geom.event is not None:
186 ref_time = geom.event.time
187 else:
188 ref_time = 0.
190 if len(values.shape) == 2:
191 tmin = self._parent.state.tmin
192 tmax = self._parent.state.tmax
193 if tmin is not None:
194 ref_tmin = tmin - ref_time
195 ref_idx_min = geom.time2idx(ref_tmin)
196 else:
197 ref_idx_min = geom.time2idx(self._state.time)
199 if tmax is not None:
200 ref_tmax = tmax - ref_time
201 ref_idx_max = geom.time2idx(ref_tmax)
202 else:
203 ref_idx_max = geom.time2idx(self._state.time)
205 if ref_idx_min == ref_idx_max:
206 out = values[:, ref_idx_min]
207 elif ref_idx_min > ref_idx_max:
208 out = values[:, ref_idx_min]
209 elif ref_idx_max < ref_idx_min:
210 out = values[:, ref_idx_max]
211 else:
212 # TODO CHECK
213 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
214 out = values[:, ref_idx_max]
215 else:
216 out = values.ravel()
217 return out
219 def update_view(self, *args):
220 pstate = self._parent.state
221 geom = self._state.geometry
223 if geom.event:
224 pstate.lat = geom.event.lat
225 pstate.lon = geom.event.lon
227 self.update()
229 def update(self, *args):
231 state = self._state
233 if state.geometry and self._controls:
234 self._update_controls()
235 # base.update_cpt(self)
236 self.update_cpt(state)
238 if state.visible:
239 # cpt_name = self.get_cpt_name(
240 # state.cpt, state.display_parameter)
241 geo = state.geometry
242 values = self.get_values(geo)
243 lut = self.cpt_handler._lookuptable
244 if not isinstance(self._pipe, TrimeshPipe):
245 vertices = arr_vertices(geo.get_vertices('xyz'))
246 faces = arr_faces(geo.get_faces())
247 self._pipe = TrimeshPipe(
248 vertices, faces,
249 values=values,
250 lut=lut,
251 backface_culling=False)
252 self._cbar_pipe = ColorbarPipe(
253 lut=lut, cbar_title=state.display_parameter)
254 self._parent.add_actor(self._pipe.actor)
255 self._parent.add_actor(self._cbar_pipe.actor)
257 if geo.outlines:
258 self._outlines_pipe.append(OutlinesPipe(
259 geo, color=(1., 1., 1.), cs='latlondepth'))
260 self._parent.add_actor(
261 self._outlines_pipe[-1].actor)
262 self._outlines_pipe.append(OutlinesPipe(
263 geo, color=(0.6, 0.6, 0.6), cs='latlon'))
264 self._parent.add_actor(
265 self._outlines_pipe[-1].actor)
267 else:
268 self._pipe.set_values(values)
269 self._pipe.set_lookuptable(lut)
270 self._pipe.set_opacity(self._state.opacity)
272 self._cbar_pipe.set_lookuptable(lut)
273 self._cbar_pipe.set_title(state.display_parameter)
274 else:
275 if self._pipe:
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
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 # color maps
321 self.cpt_handler.cpt_controls(
322 self._parent, self._state.cpt, layout)
324 # times slider
325 il = layout.rowCount() + 1
326 slider = qw.QSlider(qc.Qt.Horizontal)
327 slider.setSizePolicy(
328 qw.QSizePolicy(
329 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
331 def iround(x):
332 return int(round(x))
334 slider.setMinimum(iround(state.geometry.times.min()))
335 slider.setMaximum(iround(state.geometry.times.max()))
336 slider.setSingleStep(iround(state.geometry.deltat))
337 slider.setPageStep(iround(state.geometry.deltat))
339 time_label = qw.QLabel('Time')
340 layout.addWidget(time_label, il, 0)
341 layout.addWidget(slider, il, 1)
343 state_bind_slider(
344 self, state, 'time', slider, dtype=int)
346 self._time_label = time_label
347 self._time_slider = slider
349 il += 1
350 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
351 slider_opacity.setSizePolicy(
352 qw.QSizePolicy(
353 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
354 slider_opacity.setMinimum(0)
355 slider_opacity.setMaximum(1000)
357 opacity_label = qw.QLabel('Opacity')
358 layout.addWidget(opacity_label, il, 0)
359 layout.addWidget(slider_opacity, il, 1)
361 state_bind_slider(
362 self, state, 'opacity', slider_opacity, factor=0.001)
364 self._opacity_label = opacity_label
365 self._opacity_slider = slider_opacity
367 il += 1
368 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
370 self.cpt_handler._update_cpt_combobox()
371 self.cpt_handler._update_cptscale_lineedit()
373 self._controls = frame
375 self._update_controls()
377 return self._controls
379 def _update_controls(self):
380 state = self._state
381 if state.geometry:
382 values = state.geometry.get_property(state.display_parameter)
384 if values.ndim == 2:
385 self._time_label.setVisible(True)
386 self._time_slider.setVisible(True)
387 self._opacity_label.setVisible(True)
388 self._opacity_slider.setVisible(True)
389 else:
390 self._time_label.setVisible(False)
391 self._time_slider.setVisible(False)
392 self._opacity_label.setVisible(False)
393 self._opacity_slider.setVisible(False)
396__all__ = [
397 'GeometryElement',
398 'GeometryState'
399]