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._parent = None
48 self._state = None
49 self._controls = None
51 self._pipe = None
52 self._cbar_pipe = None
53 self._outlines_pipe = []
55 self.cpt_handler = base.CPTHandler()
57 def remove(self):
58 if self._parent and self._state:
59 self._parent.state.elements.remove(self._state)
61 def init_pipeslots(self):
62 if not self._pipe:
63 self._pipe.append([])
65 def remove_pipes(self):
66 if self._pipe is not None:
67 self._parent.remove_actor(self._pipe.actor)
69 if self._cbar_pipe is not None:
70 self._parent.remove_actor(self._cbar_pipe.actor)
72 if len(self._outlines_pipe) > 0:
73 for pipe in self._outlines_pipe:
74 self._parent.remove_actor(pipe.actor)
76 self._pipe = None
77 self._cbar_pipe = None
78 self._outlines_pipe = []
80 def set_parent(self, parent):
81 self._parent = parent
82 self._parent.add_panel(
83 self.get_title_label(),
84 self._get_controls(),
85 visible=True,
86 title_controls=[
87 self.get_title_control_remove(),
88 self.get_title_control_visible()])
90 self.talkie_connect(
91 self._parent.state,
92 ['tmin', 'tmax', 'lat', 'lon'],
93 self.update)
95 self.update()
97 def unset_parent(self):
98 self.unbind_state()
99 if self._parent:
100 if self._pipe:
101 self.remove_pipes()
103 if self._controls:
104 self._parent.remove_panel(self._controls)
105 self._controls = None
107 self._parent.update_view()
108 self._parent = None
110 def bind_state(self, state):
111 base.Element.bind_state(self, state)
113 self.talkie_connect(
114 state,
115 ['visible', 'geometry', 'display_parameter', 'time', 'opacity'],
116 self.update)
118 self.cpt_handler.bind_state(state.cpt, self.update)
120 def unbind_state(self):
121 self.cpt_handler.unbind_state()
122 base.Element.unbind_state(self)
124 def get_cpt_name(self, cpt, display_parameter):
125 return '{}_{}'.format(cpt, display_parameter)
127 def update_cpt(self, state):
129 values = state.geometry.get_property(state.display_parameter)
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]
158 else:
159 raise ValueError(
160 'Imported geometry contains no property to be displayed!')
162 self._parent.remove_panel(self._controls)
163 self._controls = None
164 self._state.geometry = loaded_geometry
166 self._parent.add_panel(
167 self.get_title_label(),
168 self._get_controls(),
169 visible=True,
170 title_controls=[
171 self.get_title_control_remove(),
172 self.get_title_control_visible()])
174 self.update()
176 def get_values(self, geom):
177 values = geom.get_property(self._state.display_parameter)
179 if geom.event is not None:
180 ref_time = geom.event.time
181 else:
182 ref_time = 0.
184 if len(values.shape) == 2:
185 tmin = self._parent.state.tmin
186 tmax = self._parent.state.tmax
187 if tmin is not None:
188 ref_tmin = tmin - ref_time
189 ref_idx_min = geom.time2idx(ref_tmin)
190 else:
191 ref_idx_min = geom.time2idx(self._state.time)
193 if tmax is not None:
194 ref_tmax = tmax - ref_time
195 ref_idx_max = geom.time2idx(ref_tmax)
196 else:
197 ref_idx_max = geom.time2idx(self._state.time)
199 if ref_idx_min == ref_idx_max:
200 out = values[:, ref_idx_min]
201 elif ref_idx_min > ref_idx_max:
202 out = values[:, ref_idx_min]
203 elif ref_idx_max < ref_idx_min:
204 out = values[:, ref_idx_max]
205 else:
206 # TODO CHECK
207 # out = values[:, ref_idx_min:ref_idx_max].sum(1)
208 out = values[:, ref_idx_max]
209 else:
210 out = values.ravel()
211 return out
213 def update_view(self, *args):
214 pstate = self._parent.state
215 geom = self._state.geometry
217 if geom.event:
218 pstate.lat = geom.event.lat
219 pstate.lon = geom.event.lon
221 self.update()
223 def update(self, *args):
225 state = self._state
227 if state.geometry and self._controls:
228 self._update_controls()
229 # base.update_cpt(self)
230 self.update_cpt(state)
232 if state.visible:
233 # cpt_name = self.get_cpt_name(
234 # state.cpt, state.display_parameter)
235 geo = state.geometry
236 values = self.get_values(geo)
237 lut = self.cpt_handler._lookuptable
238 if not isinstance(self._pipe, TrimeshPipe):
239 vertices = arr_vertices(geo.get_vertices('xyz'))
240 faces = arr_faces(geo.get_faces())
241 self._pipe = TrimeshPipe(
242 vertices, faces,
243 values=values,
244 lut=lut,
245 backface_culling=False)
246 self._cbar_pipe = ColorbarPipe(
247 lut=lut, cbar_title=state.display_parameter)
248 self._parent.add_actor(self._pipe.actor)
249 self._parent.add_actor(self._cbar_pipe.actor)
251 if geo.outlines:
252 self._outlines_pipe.append(OutlinesPipe(
253 geo, color=(1., 1., 1.), cs='latlondepth'))
254 self._parent.add_actor(
255 self._outlines_pipe[-1].actor)
256 self._outlines_pipe.append(OutlinesPipe(
257 geo, color=(0.6, 0.6, 0.6), cs='latlon'))
258 self._parent.add_actor(
259 self._outlines_pipe[-1].actor)
261 else:
262 self._pipe.set_values(values)
263 self._pipe.set_lookuptable(lut)
264 self._pipe.set_opacity(self._state.opacity)
266 self._cbar_pipe.set_lookuptable(lut)
267 self._cbar_pipe.set_title(state.display_parameter)
268 else:
269 if self._pipe:
270 self.remove_pipes()
272 self._parent.update_view()
274 def _get_controls(self):
275 state = self._state
276 if not self._controls:
277 from ..state import state_bind_combobox, \
278 state_bind_slider
280 frame = qw.QFrame()
281 layout = qw.QGridLayout()
282 layout.setAlignment(qc.Qt.AlignTop)
283 frame.setLayout(layout)
285 # load geometry
286 pb = qw.QPushButton('Load')
287 layout.addWidget(pb, 0, 0)
289 pb.clicked.connect(self.open_file_load_dialog)
291 # property choice
292 il = 1
293 if state.geometry:
295 pb = qw.QPushButton('Move to')
296 layout.addWidget(pb, 0, 1)
297 pb.clicked.connect(self.update_view)
299 props = []
300 for prop in state.geometry.properties.get_col_names(
301 sub_headers=False):
302 props.append(prop)
304 layout.addWidget(qw.QLabel('Display parameter'), il, 0)
305 cb = qw.QComboBox()
307 unique_props = list(set(props))
308 for i, s in enumerate(unique_props):
309 cb.insertItem(i, s)
311 layout.addWidget(cb, il, 1)
312 state_bind_combobox(self, state, 'display_parameter', cb)
314 # color maps
315 self.cpt_handler.cpt_controls(
316 self._parent, self._state.cpt, layout)
318 # times slider
319 il = layout.rowCount() + 1
320 slider = qw.QSlider(qc.Qt.Horizontal)
321 slider.setSizePolicy(
322 qw.QSizePolicy(
323 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
325 def iround(x):
326 return int(round(x))
328 slider.setMinimum(iround(state.geometry.times.min()))
329 slider.setMaximum(iround(state.geometry.times.max()))
330 slider.setSingleStep(iround(state.geometry.deltat))
331 slider.setPageStep(iround(state.geometry.deltat))
333 time_label = qw.QLabel('Time')
334 layout.addWidget(time_label, il, 0)
335 layout.addWidget(slider, il, 1)
337 state_bind_slider(
338 self, state, 'time', slider, dtype=int)
340 self._time_label = time_label
341 self._time_slider = slider
343 il += 1
344 slider_opacity = qw.QSlider(qc.Qt.Horizontal)
345 slider_opacity.setSizePolicy(
346 qw.QSizePolicy(
347 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))
348 slider_opacity.setMinimum(0)
349 slider_opacity.setMaximum(1000)
351 opacity_label = qw.QLabel('Opacity')
352 layout.addWidget(opacity_label, il, 0)
353 layout.addWidget(slider_opacity, il, 1)
355 state_bind_slider(
356 self, state, 'opacity', slider_opacity, factor=0.001)
358 self._opacity_label = opacity_label
359 self._opacity_slider = slider_opacity
361 il += 1
362 layout.addWidget(qw.QFrame(), il, 0, 1, 3)
364 self.cpt_handler._update_cpt_combobox()
365 self.cpt_handler._update_cptscale_lineedit()
367 self._controls = frame
369 self._update_controls()
371 return self._controls
373 def _update_controls(self):
374 state = self._state
375 if state.geometry:
376 values = state.geometry.get_property(state.display_parameter)
378 if values.ndim == 2:
379 self._time_label.setVisible(True)
380 self._time_slider.setVisible(True)
381 self._opacity_label.setVisible(True)
382 self._opacity_slider.setVisible(True)
383 else:
384 self._time_label.setVisible(False)
385 self._time_slider.setVisible(False)
386 self._opacity_label.setVisible(False)
387 self._opacity_slider.setVisible(False)
390__all__ = [
391 'GeometryElement',
392 'GeometryState'
393]