Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/sparrow/elements/geometry.py: 16%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6import logging 

7import numpy as num 

8 

9from pyrocko.guts import Bool, String, load, Float 

10from pyrocko.geometry import arr_vertices, arr_faces 

11from pyrocko.gui.qt_compat import qw, qc 

12from pyrocko.gui.vtk_util import TrimeshPipe, ColorbarPipe, OutlinesPipe, Color 

13from pyrocko.orthodrome import geographic_midpoint 

14 

15from pyrocko.model import Geometry 

16 

17from . import base 

18from .. import common 

19 

20 

21logger = logging.getLogger('geometry') 

22 

23guts_prefix = 'sparrow' 

24 

25km = 1e3 

26 

27 

28class GeometryState(base.ElementState): 

29 opacity = Float.T(default=1.0) 

30 visible = Bool.T(default=True) 

31 geometry = Geometry.T(default=None, optional=True) 

32 display_parameter = String.T(default="") 

33 time = Float.T(default=0., optional=True) 

34 cpt = base.CPTState.T(default=base.CPTState.D()) 

35 color = Color.T(default=Color.D('white')) 

36 line_width = Float.T(default=1.0) 

37 

38 def create(self): 

39 element = GeometryElement() 

40 return element 

41 

42 

43class GeometryElement(base.Element): 

44 

45 def __init__(self): 

46 base.Element.__init__(self) 

47 self._parent = None 

48 self._state = None 

49 self._controls = None 

50 

51 self._pipe = None 

52 self._cbar_pipe = None 

53 self._outlines_pipe = [] 

54 self._time_label = None 

55 self._time_slider = None 

56 

57 self.cpt_handler = base.CPTHandler() 

58 

59 def remove(self): 

60 if self._parent and self._state: 

61 self._parent.state.elements.remove(self._state) 

62 

63 def init_pipeslots(self): 

64 if not self._pipe: 

65 self._pipe.append([]) 

66 

67 def remove_pipes(self): 

68 if self._pipe is not None: 

69 self._parent.remove_actor(self._pipe.actor) 

70 

71 if self._cbar_pipe is not None: 

72 self._parent.remove_actor(self._cbar_pipe.actor) 

73 

74 if len(self._outlines_pipe) > 0: 

75 for pipe in self._outlines_pipe: 

76 self._parent.remove_actor(pipe.actor) 

77 

78 self._pipe = None 

79 self._cbar_pipe = None 

80 self._outlines_pipe = [] 

81 

82 def set_parent(self, parent): 

83 self._parent = parent 

84 self._parent.add_panel( 

85 self.get_title_label(), 

86 self._get_controls(), 

87 visible=True, 

88 title_controls=[ 

89 self.get_title_control_remove(), 

90 self.get_title_control_visible()]) 

91 

92 self.talkie_connect( 

93 self._parent.state, 

94 ['tmin', 'tmax', 'lat', 'lon'], 

95 self.update) 

96 

97 self.update() 

98 

99 def unset_parent(self): 

100 self.unbind_state() 

101 if self._parent: 

102 if self._pipe or self._cbar_pipe or self._outlines_pipe: 

103 self.remove_pipes() 

104 

105 if self._controls: 

106 self._parent.remove_panel(self._controls) 

107 self._controls = None 

108 

109 self._parent.update_view() 

110 self._parent = None 

111 

112 def bind_state(self, state): 

113 base.Element.bind_state(self, state) 

114 

115 self.talkie_connect( 

116 state, 

117 ['visible', 'geometry', 'display_parameter', 'time', 

118 'opacity', 'color', 'line_width'], 

119 self.update) 

120 

121 self.cpt_handler.bind_state(state.cpt, self.update) 

122 

123 def unbind_state(self): 

124 self.cpt_handler.unbind_state() 

125 base.Element.unbind_state(self) 

126 

127 def update_cpt(self, state): 

128 

129 if len(state.display_parameter) != 0: 

130 values = state.geometry.get_property(state.display_parameter) 

131 # TODO Check 

132 # if values.ndim == 2: 

133 # values = values.sum(1) 

134 

135 self.cpt_handler._values = values 

136 self.cpt_handler.update_cpt() 

137 

138 def get_name(self): 

139 return 'Geometry' 

140 

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) 

145 

146 if fns: 

147 self.load_file(str(fns[0])) 

148 else: 

149 return 

150 

151 def load_file(self, path): 

152 

153 loaded_geometry = load(filename=path) 

154 props = loaded_geometry.properties.get_col_names(sub_headers=False) 

155 

156 if props: 

157 if self._state.display_parameter not in props: 

158 self._state.display_parameter = props[0] 

159 

160 self._parent.remove_panel(self._controls) 

161 self._controls = None 

162 self._state.geometry = loaded_geometry 

163 

164 self._parent.add_panel( 

165 self.get_title_label(), 

166 self._get_controls(), 

167 visible=True, 

168 title_controls=[ 

169 self.get_title_control_remove(), 

170 self.get_title_control_visible()]) 

171 

172 self.update() 

173 

174 def get_values(self, geom): 

175 values = geom.get_property(self._state.display_parameter) 

176 

177 if geom.event is not None: 

178 ref_time = geom.event.time 

179 else: 

180 ref_time = 0. 

181 

182 if len(values.shape) == 2: 

183 tmin = self._parent.state.tmin 

184 tmax = self._parent.state.tmax 

185 if tmin is not None: 

186 ref_tmin = tmin - ref_time 

187 ref_idx_min = geom.time2idx(ref_tmin) 

188 else: 

189 ref_idx_min = geom.time2idx(self._state.time) 

190 

191 if tmax is not None: 

192 ref_tmax = tmax - ref_time 

193 ref_idx_max = geom.time2idx(ref_tmax) 

194 else: 

195 ref_idx_max = geom.time2idx(self._state.time) 

196 

197 if ref_idx_min == ref_idx_max: 

198 out = values[:, ref_idx_min] 

199 elif ref_idx_min > ref_idx_max: 

200 out = values[:, ref_idx_min] 

201 elif ref_idx_max < ref_idx_min: 

202 out = values[:, ref_idx_max] 

203 else: 

204 # TODO CHECK 

205 # out = values[:, ref_idx_min:ref_idx_max].sum(1) 

206 out = values[:, ref_idx_max] 

207 else: 

208 out = values.ravel() 

209 return out 

210 

211 def update_view(self, *args): 

212 pstate = self._parent.state 

213 geom = self._state.geometry 

214 

215 if geom.no_faces() > 0: 

216 latlon = geom.get_vertices('latlon') 

217 pstate.lat, pstate.lon = geographic_midpoint( 

218 latlon[:, 0], 

219 latlon[:, 1]) 

220 elif geom.outlines: 

221 latlon = num.concatenate([ 

222 outline.get_col('latlon') for outline in geom.outlines 

223 ]) 

224 pstate.lat, pstate.lon = geographic_midpoint( 

225 latlon[:, 0], 

226 latlon[:, 1]) 

227 elif geom.event: 

228 pstate.lat = geom.event.lat 

229 pstate.lon = geom.event.lon 

230 else: 

231 raise ValueError('Geometry Element has no location information.') 

232 

233 self.update() 

234 

235 def clear(self): 

236 self._parent.remove_panel(self._controls) 

237 self._controls = None 

238 self._state.geometry = None 

239 

240 self._parent.add_panel( 

241 self.get_title_label(), 

242 self._get_controls(), 

243 visible=True, 

244 title_controls=[ 

245 self.get_title_control_remove(), 

246 self.get_title_control_visible()]) 

247 

248 self.update() 

249 

250 def update_outlines(self, geo): 

251 state = self._state 

252 if len(self._outlines_pipe) == 0: 

253 for cs in ['latlondepth']: 

254 outline_pipe = OutlinesPipe( 

255 geo, color=state.color, cs=cs) 

256 outline_pipe.set_line_width(state.line_width) 

257 self._outlines_pipe.append(outline_pipe) 

258 self._parent.add_actor( 

259 self._outlines_pipe[-1].actor) 

260 

261 else: 

262 for outline_pipe in self._outlines_pipe: 

263 outline_pipe.set_color(state.color) 

264 outline_pipe.set_line_width(state.line_width) 

265 

266 def update(self, *args): 

267 

268 state = self._state 

269 

270 if state.geometry and self._controls: 

271 self._update_controls() 

272 self.update_cpt(state) 

273 

274 if state.visible: 

275 geo = state.geometry 

276 lut = self.cpt_handler._lookuptable 

277 no_faces = geo.no_faces() 

278 if no_faces: 

279 values = self.get_values(geo) 

280 if not isinstance(self._pipe, TrimeshPipe): 

281 vertices = arr_vertices(geo.get_vertices('xyz')) 

282 faces = arr_faces(geo.get_faces()) 

283 self._pipe = TrimeshPipe( 

284 vertices, faces, 

285 values=values, 

286 lut=lut, 

287 backface_culling=False) 

288 self._cbar_pipe = ColorbarPipe( 

289 lut=lut, cbar_title=state.display_parameter) 

290 self._parent.add_actor(self._pipe.actor) 

291 self._parent.add_actor(self._cbar_pipe.actor) 

292 else: 

293 self._pipe.set_values(values) 

294 self._pipe.set_lookuptable(lut) 

295 self._pipe.set_opacity(self._state.opacity) 

296 

297 self._cbar_pipe.set_lookuptable(lut) 

298 self._cbar_pipe.set_title(state.display_parameter) 

299 

300 if geo.outlines: 

301 self.update_outlines(geo) 

302 else: 

303 self.remove_pipes() 

304 

305 else: 

306 self.remove_pipes() 

307 

308 self._parent.update_view() 

309 

310 def _get_controls(self): 

311 state = self._state 

312 if not self._controls: 

313 from ..state import state_bind_combobox, \ 

314 state_bind_slider, state_bind_combobox_color 

315 

316 frame = qw.QFrame() 

317 layout = qw.QGridLayout() 

318 layout.setAlignment(qc.Qt.AlignTop) 

319 frame.setLayout(layout) 

320 

321 # load geometry 

322 il = 0 

323 if not state.geometry: 

324 pb = qw.QPushButton('Load') 

325 layout.addWidget(pb, il, 0) 

326 pb.clicked.connect(self.open_file_load_dialog) 

327 

328 # property choice 

329 else: 

330 props = [] 

331 for prop in state.geometry.properties.get_col_names( 

332 sub_headers=False): 

333 props.append(prop) 

334 

335 layout.addWidget(qw.QLabel('Display Parameter'), il, 0) 

336 cb = qw.QComboBox() 

337 

338 unique_props = list(set(props)) 

339 for i, s in enumerate(unique_props): 

340 cb.insertItem(i, s) 

341 

342 layout.addWidget(cb, il, 1) 

343 state_bind_combobox(self, state, 'display_parameter', cb) 

344 

345 if state.geometry.no_faces != 0: 

346 # color maps 

347 self.cpt_handler.cpt_controls( 

348 self._parent, self._state.cpt, layout) 

349 

350 il += 1 

351 layout.addWidget(qw.QFrame(), il, 0, 1, 3) 

352 

353 self.cpt_handler._update_cpt_combobox() 

354 self.cpt_handler._update_cptscale_lineedit() 

355 

356 # times slider 

357 if state.geometry.times is not None: 

358 il = layout.rowCount() + 1 

359 slider = qw.QSlider(qc.Qt.Horizontal) 

360 slider.setSizePolicy( 

361 qw.QSizePolicy( 

362 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

363 

364 def iround(x): 

365 return int(round(x)) 

366 

367 slider.setMinimum(iround(state.geometry.times.min())) 

368 slider.setMaximum(iround(state.geometry.times.max())) 

369 slider.setSingleStep(iround(state.geometry.deltat)) 

370 slider.setPageStep(iround(state.geometry.deltat)) 

371 

372 time_label = qw.QLabel('Time') 

373 layout.addWidget(time_label, il, 0) 

374 layout.addWidget(slider, il, 1) 

375 

376 state_bind_slider( 

377 self, state, 'time', slider, dtype=int) 

378 

379 self._time_label = time_label 

380 self._time_slider = slider 

381 

382 il = layout.rowCount() + 1 

383 slider_opacity = qw.QSlider(qc.Qt.Horizontal) 

384 slider_opacity.setSizePolicy( 

385 qw.QSizePolicy( 

386 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

387 slider_opacity.setMinimum(0) 

388 slider_opacity.setMaximum(1000) 

389 

390 opacity_label = qw.QLabel('Opacity') 

391 layout.addWidget(opacity_label, il, 0) 

392 layout.addWidget(slider_opacity, il, 1) 

393 

394 state_bind_slider( 

395 self, state, 'opacity', slider_opacity, factor=0.001) 

396 

397 self._opacity_label = opacity_label 

398 self._opacity_slider = slider_opacity 

399 

400 # color 

401 il += 1 

402 layout.addWidget(qw.QLabel('Color'), il, 0) 

403 

404 cb = common.strings_to_combobox( 

405 ['black', 'white', 'blue', 'red']) 

406 

407 layout.addWidget(cb, il, 1) 

408 state_bind_combobox_color(self, state, 'color', cb) 

409 

410 # linewidth outline 

411 il += 1 

412 layout.addWidget(qw.QLabel('Line Width'), il, 0) 

413 

414 slider = qw.QSlider(qc.Qt.Horizontal) 

415 slider.setSizePolicy( 

416 qw.QSizePolicy( 

417 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

418 slider.setMinimum(0) 

419 slider.setMaximum(100) 

420 layout.addWidget(slider, il, 1) 

421 state_bind_slider( 

422 self, state, 'line_width', slider, factor=0.1) 

423 

424 # Clear scene 

425 il += 1 

426 pb = qw.QPushButton('Clear') 

427 layout.addWidget(pb, il, 1) 

428 pb.clicked.connect(self.clear) 

429 

430 # Change view to source 

431 pb = qw.QPushButton('Move To') 

432 layout.addWidget(pb, il, 2) 

433 pb.clicked.connect(self.update_view) 

434 

435 self._controls = frame 

436 

437 self._update_controls() 

438 

439 return self._controls 

440 

441 def _update_controls(self): 

442 state = self._state 

443 if state.geometry: 

444 if len(state.display_parameter) != 0: 

445 values = state.geometry.get_property(state.display_parameter) 

446 

447 if values.ndim == 2: 

448 if self._time_label: 

449 self._time_label.setVisible(True) 

450 if self._time_slider: 

451 self._time_slider.setVisible(True) 

452 self._opacity_label.setVisible(True) 

453 self._opacity_slider.setVisible(True) 

454 else: 

455 if self._time_label: 

456 self._time_label.setVisible(False) 

457 if self._time_slider: 

458 self._time_slider.setVisible(False) 

459 self._opacity_label.setVisible(False) 

460 self._opacity_slider.setVisible(False) 

461 

462 

463__all__ = [ 

464 'GeometryElement', 

465 'GeometryState' 

466]