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

311 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2024-03-07 11:54 +0000

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import string 

7 

8import numpy as num 

9 

10from pyrocko.guts import Bool, Float, Object, String 

11 

12from pyrocko import cake, geometry, gf 

13from pyrocko.gui.qt_compat import qc, qw 

14from pyrocko.gui.talkie import TalkieRoot 

15 

16from pyrocko.gui.vtk_util import \ 

17 ArrowPipe, ColorbarPipe, PolygonPipe, ScatterPipe, OutlinesPipe 

18 

19from .. import state as vstate 

20from .. import common 

21from . import base 

22 

23guts_prefix = 'sparrow' 

24 

25 

26d2r = num.pi / 180. 

27 

28 

29map_anchor = { 

30 'center': (0.0, 0.0), 

31 'center_left': (-1.0, 0.0), 

32 'center_right': (1.0, 0.0), 

33 'top': (0.0, -1.0), 

34 'top_left': (-1.0, -1.0), 

35 'top_right': (1.0, -1.0), 

36 'bottom': (0.0, 1.0), 

37 'bottom_left': (-1.0, 1.0), 

38 'bottom_right': (1.0, 1.0)} 

39 

40 

41class ProxySource(TalkieRoot): 

42 pass 

43 

44 

45for source_cls in [gf.RectangularSource]: 

46 

47 cls_name = 'Proxy' + source_cls.__name__ 

48 

49 class proxy_source_cls(ProxySource): 

50 class_name = cls_name 

51 

52 def __init__(self, **kwargs): 

53 ProxySource.__init__(self) 

54 for key, value in self._ranges.items(): 

55 setattr(self, key, value['ini']) 

56 

57 if kwargs is not None: 

58 for it in kwargs.items(): 

59 setattr(self, it[0], it[1]) 

60 

61 proxy_source_cls.__name__ = cls_name 

62 vars()[cls_name] = proxy_source_cls 

63 

64 for prop in source_cls.T.properties: 

65 proxy_source_cls.T.add_property(prop.name, prop) 

66 

67ProxyRectangularSource = vars()['ProxyRectangularSource'] # silence flake8 

68 

69ProxyRectangularSource._name = 'RectangularSource' 

70 

71ProxyRectangularSource._ranges = { 

72 'lat': {'min': -90., 'max': 90., 'step': 1, 'ini': 0.}, 

73 'lon': {'min': -180., 'max': 180., 'step': 1, 'ini': 0.}, 

74 'depth': {'min': 0., 'max': 600000., 'step': 1000, 'ini': 10000.}, 

75 'width': {'min': 0.1, 'max': 500000., 'step': 1000, 'ini': 10000.}, 

76 'length': {'min': 0.1, 'max': 1000000., 'step': 1000, 'ini': 50000.}, 

77 'strike': {'min': -180., 'max': 180., 'step': 1, 'ini': 0.}, 

78 'dip': {'min': 0., 'max': 90., 'step': 1, 'ini': 45.}, 

79 'rake': {'min': -180., 'max': 180., 'step': 1, 'ini': 0.}, 

80 'nucleation_x': 

81 {'min': -100., 'max': 100., 'step': 1, 'ini': 0., 'fac': .01}, 

82 'nucleation_y': 

83 {'min': -100., 'max': 100., 'step': 1, 'ini': 0., 'fac': .01}, 

84 'slip': {'min': 0., 'max': 1000., 'step': 1, 'ini': 1., 'fac': .01}} 

85 

86 

87class ProxyConfig(Object): 

88 deltas = num.array([1000., 1000.]) 

89 deltat = Float.T(default=0.5) 

90 rho = Float.T(default=2800) 

91 vs = Float.T(default=3600) 

92 

93 def get_shear_moduli(self, *args, **kwargs): 

94 points = kwargs.get('points') 

95 return num.ones(len(points)) * num.power(self.vs, 2) * self.rho 

96 

97 

98class ProxyStore(Object): 

99 def __init__(self, **kwargs): 

100 config = ProxyConfig() 

101 if kwargs: 

102 config.deltas = kwargs.get('deltas', config.deltas) 

103 config.deltat = kwargs.get('deltat', config.deltat) 

104 config.rho = kwargs.get('rho', config.rho) 

105 config.vs = kwargs.get('vs', config.vs) 

106 

107 self.config = config 

108 self.mode = String.T(default='r') 

109 self._f_data = None 

110 self._f_index = None 

111 

112 

113parameter_label = { 

114 'time (s)': 'times', 

115 'slip (m)': 'slip', 

116 'moment (Nm)': 'moment' 

117} 

118 

119parameter_geometry = { 

120 'time (s)': 't_arrival', 

121 'slip (m)': 'slip', 

122 'moment (Nm)': 'moment' 

123} 

124 

125unit_label = { 

126 'lat': '(deg)', 

127 'lon': '(deg)', 

128 'depth': '(m)', 

129 'strike': '(deg)', 

130 'dip': '(deg)', 

131 'rake': '(deg)', 

132 'length': '(m)', 

133 'width': '(m)', 

134 'slip': '(m)' 

135} 

136 

137 

138class SourceState(base.ElementState): 

139 visible = Bool.T(default=True) 

140 source_selection = ProxySource.T(default=ProxyRectangularSource()) # noqa 

141 deltat = Float.T(default=0.5) 

142 display_parameter = String.T(default='time (s)') 

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

144 

145 @classmethod 

146 def get_name(self): 

147 return 'Rectangular Source' 

148 

149 def create(self): 

150 element = SourceElement() 

151 return element 

152 

153 

154class SourceElement(base.Element): 

155 

156 def __init__(self): 

157 base.Element.__init__(self) 

158 self._parent = None 

159 self._pipe = [] 

160 self._controls = None 

161 self._points = num.array([]) 

162 

163 self.cpt_handler = base.CPTHandler() 

164 

165 def _state_bind_source(self, *args, **kwargs): 

166 vstate.state_bind(self, self._state.source_selection, *args, **kwargs) 

167 

168 def _state_bind_store(self, *args, **kwargs): 

169 vstate.state_bind(self, self._state, *args, **kwargs) 

170 

171 def bind_state(self, state): 

172 base.Element.bind_state(self, state) 

173 self.talkie_connect( 

174 state, 

175 ['visible', 'source_selection', 'deltat', 'display_parameter'], 

176 self.update) 

177 

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

179 

180 def unbind_state(self): 

181 self.cpt_handler.unbind_state() 

182 base.Element.unbind_state(self) 

183 

184 def get_name(self): 

185 return 'Rectangular Source' 

186 

187 def set_parent(self, parent): 

188 self._parent = parent 

189 self._parent.add_panel( 

190 self.get_title_label(), 

191 self._get_controls(), 

192 visible=True, 

193 title_controls=[ 

194 self.get_title_control_remove(), 

195 self.get_title_control_visible()]) 

196 

197 self.update() 

198 

199 def unset_parent(self): 

200 self.unbind_state() 

201 if self._parent: 

202 if self._pipe: 

203 for pipe in self._pipe: 

204 if isinstance(pipe.actor, list): 

205 for act in pipe.actor: 

206 self._parent.remove_actor(act) 

207 else: 

208 self._parent.remove_actor(pipe.actor) 

209 self._pipe = [] 

210 

211 if self._controls: 

212 self._parent.remove_panel(self._controls) 

213 self._controls = None 

214 

215 self._parent.update_view() 

216 self._parent = None 

217 

218 def open_file_load_dialog(self): 

219 caption = 'Select one file to open' 

220 fns, _ = qw.QFileDialog.getOpenFileNames( 

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

222 

223 if fns: 

224 try: 

225 self.load_source_file(str(fns[0])) 

226 except FileNotFoundError as e: 

227 raise e 

228 

229 else: 

230 return 

231 

232 def load_source_file(self, path): 

233 loaded_source = gf.load(filename=path) 

234 source = ProxyRectangularSource( 

235 **{prop: getattr(loaded_source, prop) 

236 for prop in loaded_source.T.propnames 

237 if getattr(loaded_source, prop)}) 

238 

239 self._parent.remove_panel(self._controls) 

240 self._controls = None 

241 self._state.source_selection = source 

242 self._parent.add_panel( 

243 self.get_title_label(), 

244 self._get_controls(), 

245 visible=True, 

246 title_controls=[ 

247 self.get_title_control_remove(), 

248 self.get_title_control_visible()]) 

249 

250 self.update() 

251 

252 def open_file_save_dialog(self, fn=None): 

253 caption = 'Choose a file name to write source' 

254 if not fn: 

255 fn, _ = qw.QFileDialog.getSaveFileName( 

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

257 if fn: 

258 self.save_file(str(fn)) 

259 

260 def save_file(self, path): 

261 source = self._state.source_selection 

262 source2dump = gf.RectangularSource( 

263 **{prop: getattr(source, prop) for prop in source.T.propnames}) 

264 

265 if path.split('.')[-1].lower() in ['xml']: 

266 source2dump.dump_xml(filename=path) 

267 else: 

268 source2dump.dump(filename=path) 

269 

270 def update_loc(self, *args): 

271 pstate = self._parent.state 

272 state = self._state 

273 

274 source = state.source_selection 

275 source.lat = pstate.lat 

276 source.lon = pstate.lon 

277 

278 self._state.source_selection.source = source 

279 

280 self.update() 

281 

282 def update_source(self, store): 

283 state = self._state 

284 

285 source = state.source_selection 

286 source_list = gf.source_classes 

287 

288 for i, a in enumerate(source_list): 

289 if a.__name__ is source._name: 

290 fault = a( 

291 **{prop: source.__dict__[prop] 

292 for prop in source.T.propnames}) 

293 

294 source_geom = fault.geometry(store) 

295 

296 self._update_outlines(source_geom) 

297 self._update_scatter(source, fault) 

298 self._update_raster(source_geom, state.display_parameter) 

299 self._update_rake_arrow(fault) 

300 

301 def _update_outlines(self, source_geom): 

302 

303 if source_geom.outlines: 

304 self._pipe.append(OutlinesPipe( 

305 source_geom, color=(1., 1., 1.), cs='latlondepth')) 

306 self._parent.add_actor( 

307 self._pipe[-1].actor) 

308 

309 self._pipe.append(OutlinesPipe( 

310 source_geom, color=(0.6, 0.6, 0.6), cs='latlon')) 

311 self._parent.add_actor( 

312 self._pipe[-1].actor) 

313 

314 def _update_scatter(self, source, fault): 

315 for point, color in zip( 

316 ((source.nucleation_x, 

317 source.nucleation_y), 

318 map_anchor[source.anchor]), 

319 (num.array([[1., 0., 0.]]), 

320 num.array([[0., 0., 1.]]))): 

321 

322 points = geometry.latlondepth2xyz( 

323 fault.xy_to_coord( 

324 x=[point[0]], y=[point[1]], 

325 cs='latlondepth'), 

326 planetradius=cake.earthradius) 

327 

328 vertices = geometry.arr_vertices(points) 

329 p = ScatterPipe(vertices) 

330 p.set_symbol('sphere') 

331 p.set_colors(color) 

332 self._pipe.append(p) 

333 self._parent.add_actor(p.actor) 

334 

335 def _update_raster(self, source_geom, param): 

336 vertices = geometry.arr_vertices( 

337 source_geom.get_vertices(col='xyz')) 

338 

339 faces = source_geom.get_faces() 

340 

341 if param not in parameter_label: 

342 raise NameError('No parameter label given for %s', param) 

343 

344 if not source_geom.has_property(parameter_geometry[param]): 

345 raise AttributeError( 

346 'No property within source geometry called %s', 

347 parameter_geometry[param]) 

348 

349 self.cpt_handler._values = source_geom.get_property( 

350 parameter_geometry[param]) 

351 cbar_title = parameter_label[param] 

352 

353 self.cpt_handler.update_cpt() 

354 

355 poly_pipe = PolygonPipe( 

356 vertices, faces, 

357 values=self.cpt_handler._values, 

358 lut=self.cpt_handler._lookuptable) 

359 

360 if not source_geom.has_property(parameter_geometry[param]): 

361 raise AttributeError( 

362 'No property within source geometry called %s', 

363 parameter_geometry[param]) 

364 

365 tmin = self._parent.state.tmin_effective 

366 tmax = self._parent.state.tmax_effective 

367 

368 times = source_geom.get_property(parameter_geometry[param]) 

369 times += self._state.source_selection.time 

370 

371 if tmin is not None: 

372 m1 = times < tmin 

373 else: 

374 m1 = num.zeros(times.size, dtype=bool) 

375 

376 if tmax is not None: 

377 m3 = tmax < times 

378 else: 

379 m3 = num.zeros(times.size, dtype=bool) 

380 

381 m2 = num.logical_not(num.logical_or(m1, m3)) 

382 

383 if not any(m2): 

384 poly_pipe.set_alpha(0.) 

385 # print(m2.astype(num.float_)) 

386 # poly_pipe.set_alpha(m2.copy().astype(num.float_)) 

387 # print(self.cpt_handler._lookuptable.GetRange()) 

388 

389 self._pipe.append(poly_pipe) 

390 self._parent.add_actor(self._pipe[-1].actor) 

391 

392 if cbar_title is not None: 

393 cbar_pipe = ColorbarPipe( 

394 parent_pipe=poly_pipe, cbar_title=cbar_title, 

395 lut=self.cpt_handler._lookuptable) 

396 

397 self._pipe.append(cbar_pipe) 

398 self._parent.add_actor(self._pipe[-1].actor) 

399 

400 def _update_rake_arrow(self, fault): 

401 source = self._state.source_selection 

402 rake = source.rake * d2r 

403 

404 nucl_x = source.nucleation_x 

405 nucl_y = source.nucleation_y 

406 

407 wd_ln = source.width / source.length 

408 

409 endpoint = [None] * 2 

410 endpoint[0] = nucl_x + num.cos(rake) * wd_ln 

411 endpoint[1] = nucl_y + num.sin(-rake) 

412 

413 points = geometry.latlondepth2xyz( 

414 fault.xy_to_coord( 

415 x=[nucl_x, endpoint[0]], 

416 y=[nucl_y, endpoint[1]], 

417 cs='latlondepth'), 

418 planetradius=cake.earthradius) 

419 vertices = geometry.arr_vertices(points) 

420 

421 self._pipe.append(ArrowPipe(vertices[0], vertices[1])) 

422 self._parent.add_actor(self._pipe[-1].actor) 

423 

424 def update(self, *args): 

425 state = self._state 

426 

427 store = ProxyStore( 

428 deltat=state.deltat) 

429 store.config.deltas = num.array( 

430 [(store.config.deltat * store.config.vs) + 1] * 2) 

431 

432 if self._pipe: 

433 for pipe in self._pipe: 

434 self._parent.remove_actor(pipe.actor) 

435 

436 self._pipe = [] 

437 

438 if state.visible: 

439 self.update_source(store) 

440 

441 self._parent.update_view() 

442 

443 def _get_controls(self): 

444 if not self._controls: 

445 from ..state import \ 

446 state_bind_slider, state_bind_combobox 

447 from pyrocko import gf 

448 

449 source = self._state.source_selection 

450 

451 frame = qw.QFrame() 

452 layout = qw.QGridLayout() 

453 frame.setLayout(layout) 

454 

455 def state_to_lineedit(state, attribute, widget): 

456 sel = getattr(state, attribute) 

457 

458 widget.setText('%g' % sel) 

459 

460 def lineedit_to_state(widget, state, attribute): 

461 s = float(widget.text()) 

462 try: 

463 setattr(state, attribute, s) 

464 except Exception: 

465 raise ValueError( 

466 'Value of %s needs to be a float or integer' 

467 % string.capwords(attribute)) 

468 

469 il = 0 

470 

471 # Origin time controls 

472 layout.addWidget(qw.QLabel('Origin time'), il, 0) 

473 le_time = qw.QLineEdit() 

474 layout.addWidget(le_time, il, 1, 1, 2) 

475 

476 self._state_bind_source( 

477 ['time'], common.lineedit_to_time, le_time, 

478 [le_time.editingFinished, le_time.returnPressed], 

479 common.time_to_lineedit, 

480 attribute='time') 

481 

482 for var in ['tmin', 'tmax', 'tduration', 'tposition']: 

483 self.talkie_connect( 

484 self._parent.state, var, self.update) 

485 

486 # Source property controls 

487 for il, label in enumerate(source.T.propnames, start=il+1): 

488 if label in source._ranges.keys(): 

489 

490 unit = unit_label[label] if label in unit_label else '' 

491 

492 layout.addWidget(qw.QLabel( 

493 f'{string.capwords(label)} {unit}'), il, 0) 

494 

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

496 slider.setSizePolicy( 

497 qw.QSizePolicy( 

498 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

499 slider.setMinimum( 

500 int(round(source._ranges[label]['min']))) 

501 slider.setMaximum( 

502 int(round(source._ranges[label]['max']))) 

503 slider.setSingleStep( 

504 int(round(source._ranges[label]['step']))) 

505 slider.setPageStep( 

506 int(round(source._ranges[label]['step']))) 

507 

508 layout.addWidget(slider, il, 1) 

509 try: 

510 state_bind_slider( 

511 self, self._state.source_selection, label, slider, 

512 factor=source._ranges[label]['fac']) 

513 except Exception: 

514 state_bind_slider( 

515 self, self._state.source_selection, label, slider) 

516 

517 le = qw.QLineEdit() 

518 layout.addWidget(le, il, 2) 

519 

520 self._state_bind_source( 

521 [label], lineedit_to_state, le, 

522 [le.editingFinished, le.returnPressed], 

523 state_to_lineedit, attribute=label) 

524 

525 for label, name in zip( 

526 ['Sampling int. (s)'], ['deltat']): 

527 il += 1 

528 layout.addWidget(qw.QLabel(label), il, 0) 

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

530 slider.setSizePolicy( 

531 qw.QSizePolicy( 

532 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

533 slider.setMinimum(1) 

534 slider.setMaximum(1000) 

535 slider.setSingleStep(1) 

536 slider.setPageStep(1) 

537 layout.addWidget(slider, il, 1) 

538 state_bind_slider( 

539 self, self._state, name, slider, factor=0.01) 

540 

541 le = qw.QLineEdit() 

542 layout.addWidget(le, il, 2) 

543 

544 self._state_bind_store( 

545 [name], lineedit_to_state, le, 

546 [le.editingFinished, le.returnPressed], 

547 state_to_lineedit, attribute=name) 

548 

549 il += 1 

550 layout.addWidget(qw.QLabel('Anchor'), il, 0) 

551 

552 cb = qw.QComboBox() 

553 for i, s in enumerate(gf.RectangularSource.anchor.choices): 

554 cb.insertItem(i, s) 

555 layout.addWidget(cb, il, 1, 1, 2) 

556 state_bind_combobox( 

557 self, self._state.source_selection, 'anchor', cb) 

558 

559 il += 1 

560 layout.addWidget(qw.QLabel('Display Param.'), il, 0) 

561 

562 cb = qw.QComboBox() 

563 for i, s in enumerate(parameter_label.keys()): 

564 cb.insertItem(i, s) 

565 layout.addWidget(cb, il, 1) 

566 state_bind_combobox( 

567 self, self._state, 'display_parameter', cb) 

568 

569 self.cpt_handler.cpt_controls( 

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

571 

572 il = layout.rowCount() + 1 

573 pb = qw.QPushButton('Move Here') 

574 layout.addWidget(pb, il, 0) 

575 pb.clicked.connect(self.update_loc) 

576 

577 pb = qw.QPushButton('Load') 

578 layout.addWidget(pb, il, 1) 

579 pb.clicked.connect(self.open_file_load_dialog) 

580 

581 pb = qw.QPushButton('Save') 

582 layout.addWidget(pb, il, 2) 

583 pb.clicked.connect(self.open_file_save_dialog) 

584 

585 il += 1 

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

587 

588 self._controls = frame 

589 

590 self.cpt_handler._update_cpt_combobox() 

591 self.cpt_handler._update_cptscale_lineedit() 

592 

593 return self._controls 

594 

595 

596__all__ = [ 

597 'SourceElement', 

598 'SourceState', 

599]