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, OutlinesPipe, Color 

13from pyrocko.orthodrome import geographic_midpoint 

14 

15from pyrocko.model import Geometry 

16 

17from . import base 

18from .. import common 

19from ..state import state_bind_combobox, state_bind_slider, \ 

20 state_bind_combobox_color 

21 

22 

23logger = logging.getLogger('geometry') 

24 

25guts_prefix = 'sparrow' 

26 

27km = 1e3 

28 

29 

30class GeometryState(base.ElementState): 

31 opacity = Float.T(default=1.0) 

32 visible = Bool.T(default=True) 

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

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

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

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

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

38 line_width = Float.T(default=1.0) 

39 

40 def create(self): 

41 element = GeometryElement() 

42 return element 

43 

44 

45class GeometryElement(base.Element): 

46 

47 def __init__(self): 

48 base.Element.__init__(self) 

49 self._parent = None 

50 self._state = None 

51 self._controls = None 

52 

53 self._pipe = None 

54 self._cbar_pipe = None 

55 self._outlines_pipe = [] 

56 self._time_label = None 

57 self._time_slider = None 

58 

59 self.cpt_handler = base.CPTHandler() 

60 

61 def remove(self): 

62 if self._parent and self._state: 

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

64 

65 def init_pipeslots(self): 

66 if not self._pipe: 

67 self._pipe.append([]) 

68 

69 def remove_pipes(self): 

70 if self._pipe is not None: 

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

72 

73 if len(self._outlines_pipe) > 0: 

74 for pipe in self._outlines_pipe: 

75 self._parent.remove_actor(pipe.actor) 

76 

77 self._pipe = None 

78 self._outlines_pipe = [] 

79 

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()]) 

89 

90 self.talkie_connect( 

91 self._parent.state, 

92 ['tmin', 'tmax', 'lat', 'lon', 'tmax_effective'], 

93 self.update) 

94 

95 self.update() 

96 

97 def unset_parent(self): 

98 self.unbind_state() 

99 if self._parent: 

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

101 self.remove_pipes() 

102 

103 if self._controls: 

104 self._parent.remove_panel(self._controls) 

105 self._controls = None 

106 

107 self._parent.update_view() 

108 self._parent = None 

109 

110 def bind_state(self, state): 

111 base.Element.bind_state(self, state) 

112 

113 self.talkie_connect( 

114 state, 

115 ['visible', 'geometry', 'display_parameter', 'time', 

116 'opacity', 'color', 'line_width'], 

117 self.update) 

118 

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

120 

121 def unbind_state(self): 

122 self.cpt_handler.unbind_state() 

123 base.Element.unbind_state(self) 

124 

125 def update_cpt(self, state): 

126 

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

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

129 # TODO Check 

130 # if values.ndim == 2: 

131 # values = values.sum(1) 

132 

133 self.cpt_handler._values = values 

134 self.cpt_handler.update_cpt(mask_zeros=True) 

135 self.cpt_handler.update_cbar(state.display_parameter) 

136 

137 def get_name(self): 

138 return 'Geometry' 

139 

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) 

144 

145 if fns: 

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

147 else: 

148 return 

149 

150 def load_file(self, path): 

151 

152 loaded_geometry = load(filename=path) 

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

154 

155 if props: 

156 if self._state.display_parameter not in props: 

157 self._state.display_parameter = props[0] 

158 

159 self._parent.remove_panel(self._controls) 

160 self._controls = None 

161 self._state.geometry = loaded_geometry 

162 

163 self._parent.add_panel( 

164 self.get_title_label(), 

165 self._get_controls(), 

166 visible=True, 

167 title_controls=[ 

168 self.get_title_control_remove(), 

169 self.get_title_control_visible()]) 

170 

171 self.update() 

172 

173 def get_values(self, geom): 

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

175 

176 if geom.event is not None: 

177 ref_time = geom.event.time 

178 else: 

179 ref_time = 0. 

180 

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

182 tmin = self._parent.state.tmin 

183 tmax = self._parent.state.tmax_effective 

184 

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 # for cumulative display 

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._parent.add_actor(self._pipe.actor) 

289 else: 

290 self._pipe.set_values(values) 

291 self._pipe.set_lookuptable(lut) 

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

293 

294 if geo.outlines: 

295 self.update_outlines(geo) 

296 else: 

297 self.remove_pipes() 

298 

299 else: 

300 self.remove_pipes() 

301 

302 self._parent.update_view() 

303 

304 def _get_controls(self): 

305 state = self._state 

306 if not self._controls: 

307 

308 frame = qw.QFrame() 

309 layout = qw.QGridLayout() 

310 layout.setAlignment(qc.Qt.AlignTop) 

311 frame.setLayout(layout) 

312 

313 # load geometry 

314 il = 0 

315 if not state.geometry: 

316 pb = qw.QPushButton('Load') 

317 layout.addWidget(pb, il, 0) 

318 pb.clicked.connect(self.open_file_load_dialog) 

319 

320 # property choice 

321 else: 

322 props = [] 

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

324 sub_headers=False): 

325 props.append(prop) 

326 

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

328 cb = qw.QComboBox() 

329 

330 unique_props = list(set(props)) 

331 for i, s in enumerate(unique_props): 

332 cb.insertItem(i, s) 

333 

334 layout.addWidget(cb, il, 1) 

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

336 

337 if state.geometry.no_faces != 0: 

338 # color maps 

339 self.cpt_handler.cpt_controls( 

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

341 

342 il += 1 

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

344 

345 self.cpt_handler._update_cpt_combobox() 

346 self.cpt_handler._update_cptscale_lineedit() 

347 

348 # times slider 

349 if state.geometry.times is not None: 

350 il = layout.rowCount() + 1 

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

352 slider.setSizePolicy( 

353 qw.QSizePolicy( 

354 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

355 

356 def iround(x): 

357 return int(round(x)) 

358 

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

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

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

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

363 

364 time_label = qw.QLabel('Time') 

365 layout.addWidget(time_label, il, 0) 

366 layout.addWidget(slider, il, 1) 

367 

368 state_bind_slider( 

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

370 

371 self._time_label = time_label 

372 self._time_slider = slider 

373 

374 il = layout.rowCount() + 1 

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

376 slider_opacity.setSizePolicy( 

377 qw.QSizePolicy( 

378 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

379 slider_opacity.setMinimum(0) 

380 slider_opacity.setMaximum(1000) 

381 

382 opacity_label = qw.QLabel('Opacity') 

383 layout.addWidget(opacity_label, il, 0) 

384 layout.addWidget(slider_opacity, il, 1) 

385 

386 state_bind_slider( 

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

388 

389 self._opacity_label = opacity_label 

390 self._opacity_slider = slider_opacity 

391 

392 # color 

393 il += 1 

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

395 

396 cb = common.strings_to_combobox( 

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

398 

399 layout.addWidget(cb, il, 1) 

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

401 

402 # linewidth outline 

403 il += 1 

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

405 

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

407 slider.setSizePolicy( 

408 qw.QSizePolicy( 

409 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

410 slider.setMinimum(0) 

411 slider.setMaximum(100) 

412 layout.addWidget(slider, il, 1) 

413 state_bind_slider( 

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

415 

416 # Clear scene 

417 il += 1 

418 pb = qw.QPushButton('Clear') 

419 layout.addWidget(pb, il, 1) 

420 pb.clicked.connect(self.clear) 

421 

422 # Change view to source 

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

424 layout.addWidget(pb, il, 2) 

425 pb.clicked.connect(self.update_view) 

426 

427 self._controls = frame 

428 

429 self._update_controls() 

430 

431 return self._controls 

432 

433 def _update_controls(self): 

434 state = self._state 

435 if state.geometry: 

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

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

438 

439 if values.ndim == 2: 

440 if self._time_label: 

441 self._time_label.setVisible(True) 

442 if self._time_slider: 

443 self._time_slider.setVisible(True) 

444 self._opacity_label.setVisible(True) 

445 self._opacity_slider.setVisible(True) 

446 else: 

447 if self._time_label: 

448 self._time_label.setVisible(False) 

449 if self._time_slider: 

450 self._time_slider.setVisible(False) 

451 self._opacity_label.setVisible(False) 

452 self._opacity_slider.setVisible(False) 

453 

454 

455__all__ = [ 

456 'GeometryElement', 

457 'GeometryState' 

458]