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 

55 self.cpt_handler = base.CPTHandler() 

56 

57 def remove(self): 

58 if self._parent and self._state: 

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

60 

61 def init_pipeslots(self): 

62 if not self._pipe: 

63 self._pipe.append([]) 

64 

65 def remove_pipes(self): 

66 if self._pipe is not None: 

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

68 

69 if self._cbar_pipe is not None: 

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

71 

72 if len(self._outlines_pipe) > 0: 

73 for pipe in self._outlines_pipe: 

74 self._parent.remove_actor(pipe.actor) 

75 

76 self._pipe = None 

77 self._cbar_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'], 

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

135 

136 def get_name(self): 

137 return 'Geometry' 

138 

139 def open_file_load_dialog(self): 

140 caption = 'Select one file containing a geometry to open' 

141 fns, _ = qw.QFileDialog.getOpenFileNames( 

142 self._parent, caption, options=common.qfiledialog_options) 

143 

144 if fns: 

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

146 else: 

147 return 

148 

149 def load_file(self, path): 

150 

151 loaded_geometry = load(filename=path) 

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

153 

154 if props: 

155 if self._state.display_parameter not in props: 

156 self._state.display_parameter = props[0] 

157 

158 self._parent.remove_panel(self._controls) 

159 self._controls = None 

160 self._state.geometry = loaded_geometry 

161 

162 self._parent.add_panel( 

163 self.get_title_label(), 

164 self._get_controls(), 

165 visible=True, 

166 title_controls=[ 

167 self.get_title_control_remove(), 

168 self.get_title_control_visible()]) 

169 

170 self.update() 

171 

172 def get_values(self, geom): 

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

174 

175 if geom.event is not None: 

176 ref_time = geom.event.time 

177 else: 

178 ref_time = 0. 

179 

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

181 tmin = self._parent.state.tmin 

182 tmax = self._parent.state.tmax 

183 if tmin is not None: 

184 ref_tmin = tmin - ref_time 

185 ref_idx_min = geom.time2idx(ref_tmin) 

186 else: 

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

188 

189 if tmax is not None: 

190 ref_tmax = tmax - ref_time 

191 ref_idx_max = geom.time2idx(ref_tmax) 

192 else: 

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

194 

195 if ref_idx_min == ref_idx_max: 

196 out = values[:, ref_idx_min] 

197 elif ref_idx_min > ref_idx_max: 

198 out = values[:, ref_idx_min] 

199 elif ref_idx_max < ref_idx_min: 

200 out = values[:, ref_idx_max] 

201 else: 

202 # TODO CHECK 

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

204 out = values[:, ref_idx_max] 

205 else: 

206 out = values.ravel() 

207 return out 

208 

209 def update_view(self, *args): 

210 pstate = self._parent.state 

211 geom = self._state.geometry 

212 

213 if geom.no_faces() > 0: 

214 latlon = geom.get_vertices('latlon') 

215 pstate.lat, pstate.lon = geographic_midpoint( 

216 latlon[:, 0], 

217 latlon[:, 1]) 

218 elif geom.outlines: 

219 latlon = num.concatenate([ 

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

221 ]) 

222 pstate.lat, pstate.lon = geographic_midpoint( 

223 latlon[:, 0], 

224 latlon[:, 1]) 

225 elif geom.event: 

226 pstate.lat = geom.event.lat 

227 pstate.lon = geom.event.lon 

228 else: 

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

230 

231 self.update() 

232 

233 def clear(self): 

234 self._parent.remove_panel(self._controls) 

235 self._controls = None 

236 self._state.geometry = None 

237 

238 self._parent.add_panel( 

239 self.get_title_label(), 

240 self._get_controls(), 

241 visible=True, 

242 title_controls=[ 

243 self.get_title_control_remove(), 

244 self.get_title_control_visible()]) 

245 

246 self.update() 

247 

248 def update_outlines(self, geo): 

249 state = self._state 

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

251 for cs in ['latlondepth']: 

252 outline_pipe = OutlinesPipe( 

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

254 outline_pipe.set_line_width(state.line_width) 

255 self._outlines_pipe.append(outline_pipe) 

256 self._parent.add_actor( 

257 self._outlines_pipe[-1].actor) 

258 

259 else: 

260 for outline_pipe in self._outlines_pipe: 

261 outline_pipe.set_color(state.color) 

262 outline_pipe.set_line_width(state.line_width) 

263 

264 def update(self, *args): 

265 

266 state = self._state 

267 

268 if state.geometry and self._controls: 

269 self._update_controls() 

270 self.update_cpt(state) 

271 

272 if state.visible: 

273 geo = state.geometry 

274 lut = self.cpt_handler._lookuptable 

275 no_faces = geo.no_faces() 

276 if no_faces: 

277 values = self.get_values(geo) 

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

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

280 faces = arr_faces(geo.get_faces()) 

281 self._pipe = TrimeshPipe( 

282 vertices, faces, 

283 values=values, 

284 lut=lut, 

285 backface_culling=False) 

286 self._cbar_pipe = ColorbarPipe( 

287 lut=lut, cbar_title=state.display_parameter) 

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

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

290 else: 

291 self._pipe.set_values(values) 

292 self._pipe.set_lookuptable(lut) 

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

294 

295 self._cbar_pipe.set_lookuptable(lut) 

296 self._cbar_pipe.set_title(state.display_parameter) 

297 

298 if geo.outlines: 

299 self.update_outlines(geo) 

300 else: 

301 self.remove_pipes() 

302 

303 else: 

304 self.remove_pipes() 

305 

306 self._parent.update_view() 

307 

308 def _get_controls(self): 

309 state = self._state 

310 if not self._controls: 

311 from ..state import state_bind_combobox, \ 

312 state_bind_slider, state_bind_combobox_color 

313 

314 frame = qw.QFrame() 

315 layout = qw.QGridLayout() 

316 layout.setAlignment(qc.Qt.AlignTop) 

317 frame.setLayout(layout) 

318 

319 # load geometry 

320 il = 0 

321 if not state.geometry: 

322 pb = qw.QPushButton('Load') 

323 layout.addWidget(pb, il, 0) 

324 pb.clicked.connect(self.open_file_load_dialog) 

325 

326 # property choice 

327 else: 

328 props = [] 

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

330 sub_headers=False): 

331 props.append(prop) 

332 

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

334 cb = qw.QComboBox() 

335 

336 unique_props = list(set(props)) 

337 for i, s in enumerate(unique_props): 

338 cb.insertItem(i, s) 

339 

340 layout.addWidget(cb, il, 1) 

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

342 

343 if state.geometry.no_faces != 0: 

344 # color maps 

345 self.cpt_handler.cpt_controls( 

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

347 

348 il += 1 

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

350 

351 self.cpt_handler._update_cpt_combobox() 

352 self.cpt_handler._update_cptscale_lineedit() 

353 

354 # times slider 

355 if state.geometry.times is not None: 

356 il = layout.rowCount() + 1 

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

358 slider.setSizePolicy( 

359 qw.QSizePolicy( 

360 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

361 

362 def iround(x): 

363 return int(round(x)) 

364 

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

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

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

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

369 

370 time_label = qw.QLabel('Time') 

371 layout.addWidget(time_label, il, 0) 

372 layout.addWidget(slider, il, 1) 

373 

374 state_bind_slider( 

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

376 

377 self._time_label = time_label 

378 self._time_slider = slider 

379 

380 il = layout.rowCount() + 1 

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

382 slider_opacity.setSizePolicy( 

383 qw.QSizePolicy( 

384 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

385 slider_opacity.setMinimum(0) 

386 slider_opacity.setMaximum(1000) 

387 

388 opacity_label = qw.QLabel('Opacity') 

389 layout.addWidget(opacity_label, il, 0) 

390 layout.addWidget(slider_opacity, il, 1) 

391 

392 state_bind_slider( 

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

394 

395 self._opacity_label = opacity_label 

396 self._opacity_slider = slider_opacity 

397 

398 # color 

399 il += 1 

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

401 

402 cb = common.strings_to_combobox( 

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

404 

405 layout.addWidget(cb, il, 1) 

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

407 

408 # linewidth outline 

409 il += 1 

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

411 

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

413 slider.setSizePolicy( 

414 qw.QSizePolicy( 

415 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

416 slider.setMinimum(0) 

417 slider.setMaximum(100) 

418 layout.addWidget(slider, il, 1) 

419 state_bind_slider( 

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

421 

422 # Clear scene 

423 il += 1 

424 pb = qw.QPushButton('Clear') 

425 layout.addWidget(pb, il, 1) 

426 pb.clicked.connect(self.clear) 

427 

428 # Change view to source 

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

430 layout.addWidget(pb, il, 2) 

431 pb.clicked.connect(self.update_view) 

432 

433 self._controls = frame 

434 

435 self._update_controls() 

436 

437 return self._controls 

438 

439 def _update_controls(self): 

440 state = self._state 

441 if state.geometry: 

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

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

444 

445 if values.ndim == 2: 

446 self._time_label.setVisible(True) 

447 self._time_slider.setVisible(True) 

448 self._opacity_label.setVisible(True) 

449 self._opacity_slider.setVisible(True) 

450 else: 

451 self._time_label.setVisible(False) 

452 self._time_slider.setVisible(False) 

453 self._opacity_label.setVisible(False) 

454 self._opacity_slider.setVisible(False) 

455 

456 

457__all__ = [ 

458 'GeometryElement', 

459 'GeometryState' 

460]