Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/sparrow/state.py: 77%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import logging 

7 

8import numpy as num 

9 

10from pyrocko import util 

11from pyrocko.guts import StringChoice, Float, List, Bool, Timestamp, Tuple, \ 

12 Duration, Object, get_elements, set_elements, path_to_str, clone 

13 

14from pyrocko.color import Color, interpolate as interpolate_color 

15 

16from pyrocko.gui import talkie 

17from pyrocko.gui import util as gui_util 

18from . import common, light 

19 

20guts_prefix = 'sparrow' 

21 

22logger = logging.getLogger('pyrocko.gui.sparrow.state') 

23 

24 

25class FocalPointChoice(StringChoice): 

26 choices = ['center', 'target'] 

27 

28 

29class ShadingChoice(StringChoice): 

30 choices = ['flat', 'gouraud', 'phong', 'pbr'] 

31 

32 

33class LightingChoice(StringChoice): 

34 choices = light.get_lighting_theme_names() 

35 

36 

37class ViewerGuiState(talkie.TalkieRoot): 

38 panels_visible = Bool.T(default=True) 

39 size = Tuple.T(2, Float.T(), default=(100., 100.)) 

40 fixed_size = Tuple.T(2, Float.T(), optional=True) 

41 focal_point = FocalPointChoice.T(default='center') 

42 detached = Bool.T(default=False) 

43 tcursor = Timestamp.T(optional=True) 

44 

45 def next_focal_point(self): 

46 choices = FocalPointChoice.choices 

47 ii = choices.index(self.focal_point) 

48 self.focal_point = choices[(ii+1) % len(choices)] 

49 

50 

51class Background(Object): 

52 color = Color.T(default=Color.D('black')) 

53 

54 def vtk_apply(self, ren): 

55 ren.GradientBackgroundOff() 

56 ren.SetBackground(*self.color.rgb) 

57 

58 def __str__(self): 

59 return str(self.color) 

60 

61 @property 

62 def color_top(self): 

63 return self.color 

64 

65 @property 

66 def color_bottom(self): 

67 return self.color 

68 

69 # def __eq__(self, other): 

70 # print('in==', self.color.rgb, other.color.rgb) 

71 # return type(self) is type(other) and self.color == other.color 

72 

73 

74class BackgroundGradient(Background): 

75 color_top = Color.T(default=Color.D('skyblue1')) 

76 color_bottom = Color.T(default=Color.D('white')) 

77 

78 def vtk_apply(self, ren): 

79 ren.GradientBackgroundOn() 

80 ren.SetBackground(*self.color_bottom.rgb) 

81 ren.SetBackground2(*self.color_top.rgb) 

82 

83 def __str__(self): 

84 return '%s - %s' % (self.color_top, self.color_bottom) 

85 

86 # def __eq__(self, other): 

87 # return type(self) is type(other) and \ 

88 # self.color_top == other.color_top and \ 

89 # self.color_bottom == other.color_bottom 

90 

91 

92def interpolate_background(a, b, blend): 

93 if type(a) is Background and type(b) is Background: 

94 return Background(color=interpolate_color(a.color, b.color, blend)) 

95 else: 

96 return BackgroundGradient( 

97 color_top=interpolate_color( 

98 a.color_top, b.color_top, blend), 

99 color_bottom=interpolate_color( 

100 a.color_bottom, b.color_bottom, blend)) 

101 

102 

103@talkie.has_computed 

104class ViewerState(talkie.TalkieRoot): 

105 lat = Float.T(default=0.0) 

106 lon = Float.T(default=0.0) 

107 depth = Float.T(default=0.0) 

108 strike = Float.T(default=90.0) 

109 dip = Float.T(default=0.0) 

110 distance = Float.T(default=3.0) 

111 elements = List.T(talkie.Talkie.T()) 

112 tmin = Timestamp.T(optional=True) 

113 tmax = Timestamp.T(optional=True) 

114 tduration = Duration.T(optional=True) 

115 tposition = Float.T(default=0.0) 

116 lighting = LightingChoice.T(default=LightingChoice.choices[0]) 

117 background = Background.T(default=Background.D(color=Color('black'))) 

118 

119 @talkie.computed(['tmin', 'tmax', 'tduration', 'tposition']) 

120 def tmin_effective(self): 

121 return common.tmin_effective( 

122 self.tmin, self.tmax, self.tduration, self.tposition) 

123 

124 @talkie.computed(['tmin', 'tmax', 'tduration', 'tposition']) 

125 def tmax_effective(self): 

126 return common.tmax_effective( 

127 self.tmin, self.tmax, self.tduration, self.tposition) 

128 

129 def sort_elements(self): 

130 self.elements.sort(key=lambda el: el.element_id) 

131 

132 

133def state_bind( 

134 owner, state, paths, update_state, 

135 widget, signals, update_widget, attribute=None): 

136 

137 def make_wrappers(widget): 

138 def wrap_update_widget(*args): 

139 if attribute: 

140 update_widget(state, attribute, widget) 

141 else: 

142 update_widget(state, widget) 

143 common.de_errorize(widget) 

144 

145 def wrap_update_state(*args): 

146 try: 

147 if attribute: 

148 update_state(widget, state, attribute) 

149 else: 

150 update_state(widget, state) 

151 common.de_errorize(widget) 

152 except Exception as e: 

153 logger.warn('Caught exception: %s' % e) 

154 common.errorize(widget) 

155 

156 return wrap_update_widget, wrap_update_state 

157 

158 wrap_update_widget, wrap_update_state = make_wrappers(widget) 

159 

160 for sig in signals: 

161 sig.connect(wrap_update_state) 

162 

163 for path in paths: 

164 owner.talkie_connect(state, path, wrap_update_widget) 

165 

166 wrap_update_widget() 

167 

168 

169def state_bind_slider( 

170 owner, state, path, widget, factor=1., 

171 dtype=float, 

172 min_is_none=False, 

173 max_is_none=False): 

174 

175 viewer = common.get_viewer() 

176 widget.sliderPressed.connect(viewer.disable_capture) 

177 widget.sliderReleased.connect(viewer.enable_capture) 

178 

179 def make_funcs(): 

180 def update_state(widget, state): 

181 val = widget.value() 

182 if (min_is_none and val == widget.minimum()) \ 

183 or (max_is_none and val == widget.maximum()): 

184 state.set(path, None) 

185 else: 

186 viewer.status('%g' % (val * factor)) 

187 state.set(path, dtype(val * factor)) 

188 

189 def update_widget(state, widget): 

190 val = state.get(path) 

191 widget.blockSignals(True) 

192 if min_is_none and val is None: 

193 widget.setValue(widget.minimum()) 

194 elif max_is_none and val is None: 

195 widget.setValue(widget.maximum()) 

196 else: 

197 widget.setValue(int(state.get(path) * 1. / factor)) 

198 widget.blockSignals(False) 

199 

200 return update_state, update_widget 

201 

202 update_state, update_widget = make_funcs() 

203 

204 state_bind( 

205 owner, state, [path], update_state, widget, [widget.valueChanged], 

206 update_widget) 

207 

208 

209def state_bind_slider_float( 

210 owner, state, path, widget, 

211 min_is_none=False, 

212 max_is_none=False): 

213 

214 assert isinstance(widget, gui_util.QSliderFloat) 

215 

216 viewer = common.get_viewer() 

217 widget.sliderPressed.connect(viewer.disable_capture) 

218 widget.sliderReleased.connect(viewer.enable_capture) 

219 

220 def make_funcs(): 

221 def update_state(widget, state): 

222 val = widget.valueFloat() 

223 if (min_is_none and val == widget.minimumFloat()) \ 

224 or (max_is_none and val == widget.maximumFloat()): 

225 state.set(path, None) 

226 else: 

227 viewer.status('%g' % (val)) 

228 state.set(path, val) 

229 

230 def update_widget(state, widget): 

231 val = state.get(path) 

232 widget.blockSignals(True) 

233 if min_is_none and val is None: 

234 widget.setValueFloat(widget.minimumFloat()) 

235 elif max_is_none and val is None: 

236 widget.setValueFloat(widget.maximumFloat()) 

237 else: 

238 widget.setValueFloat(state.get(path)) 

239 widget.blockSignals(False) 

240 

241 return update_state, update_widget 

242 

243 update_state, update_widget = make_funcs() 

244 

245 state_bind( 

246 owner, state, [path], update_state, widget, [widget.valueChanged], 

247 update_widget) 

248 

249 

250def state_bind_spinbox(owner, state, path, widget, factor=1., dtype=float): 

251 return state_bind_slider(owner, state, path, widget, factor, dtype) 

252 

253 

254def state_bind_combobox(owner, state, path, widget): 

255 

256 def make_funcs(): 

257 def update_state(widget, state): 

258 state.set(path, str(widget.currentText())) 

259 

260 def update_widget(state, widget): 

261 widget.blockSignals(True) 

262 val = state.get(path) 

263 for i in range(widget.count()): 

264 if str(widget.itemText(i)) == val: 

265 widget.setCurrentIndex(i) 

266 widget.blockSignals(False) 

267 

268 return update_state, update_widget 

269 

270 update_state, update_widget = make_funcs() 

271 

272 state_bind( 

273 owner, state, [path], update_state, widget, [widget.activated], 

274 update_widget) 

275 

276 

277def state_bind_combobox_background(owner, state, path, widget): 

278 

279 def make_funcs(): 

280 def update_state(widget, state): 

281 values = str(widget.currentText()).split(' - ') 

282 if len(values) == 1: 

283 state.set( 

284 path, 

285 Background(color=Color(values[0]))) 

286 

287 elif len(values) == 2: 

288 state.set( 

289 path, 

290 BackgroundGradient( 

291 color_top=Color(values[0]), 

292 color_bottom=Color(values[1]))) 

293 

294 def update_widget(state, widget): 

295 widget.blockSignals(True) 

296 val = str(state.get(path)) 

297 for i in range(widget.count()): 

298 if str(widget.itemText(i)) == val: 

299 widget.setCurrentIndex(i) 

300 widget.blockSignals(False) 

301 

302 return update_state, update_widget 

303 

304 update_state, update_widget = make_funcs() 

305 

306 state_bind( 

307 owner, state, [path], update_state, widget, [widget.activated], 

308 update_widget) 

309 

310 

311def state_bind_combobox_color(owner, state, path, widget): 

312 

313 def make_funcs(): 

314 def update_state(widget, state): 

315 value = str(widget.currentText()) 

316 state.set(path, Color(value)) 

317 

318 def update_widget(state, widget): 

319 widget.blockSignals(True) 

320 val = str(state.get(path)) 

321 for i in range(widget.count()): 

322 if str(widget.itemText(i)) == val: 

323 widget.setCurrentIndex(i) 

324 widget.blockSignals(False) 

325 

326 return update_state, update_widget 

327 

328 update_state, update_widget = make_funcs() 

329 

330 state_bind( 

331 owner, state, [path], update_state, widget, [widget.activated], 

332 update_widget) 

333 

334 

335def state_bind_checkbox(owner, state, path, widget): 

336 

337 def make_funcs(): 

338 def update_state(widget, state): 

339 state.set(path, bool(widget.isChecked())) 

340 

341 def update_widget(state, widget): 

342 widget.blockSignals(True) 

343 widget.setChecked(state.get(path)) 

344 widget.blockSignals(False) 

345 

346 return update_state, update_widget 

347 

348 update_state, update_widget = make_funcs() 

349 

350 state_bind( 

351 owner, state, [path], update_state, widget, [widget.toggled], 

352 update_widget) 

353 

354 

355def state_bind_lineedit( 

356 owner, state, path, widget, from_string=str, to_string=str): 

357 

358 def make_funcs(): 

359 

360 def update_state(widget, state): 

361 state.set(path, from_string(widget.text())) 

362 

363 def update_widget(state, widget): 

364 widget.blockSignals(True) 

365 widget.setText(to_string(state.get(path))) 

366 widget.blockSignals(False) 

367 

368 return update_state, update_widget 

369 

370 update_state, update_widget = make_funcs() 

371 

372 state_bind( 

373 owner, 

374 state, [path], update_state, 

375 widget, [widget.editingFinished, widget.returnPressed], update_widget) 

376 

377 

378def interpolateables(state_a, state_b): 

379 

380 animate = [] 

381 for tag, path, values in state_b.diff(state_a): 

382 if tag == 'set': 

383 ypath = path_to_str(path) 

384 v_new = get_elements(state_b, ypath)[0] 

385 v_old = values 

386 for type in [float, Color, Background]: 

387 if isinstance(v_old, type) and isinstance(v_new, type): 

388 animate.append((ypath, v_old, v_new)) 

389 

390 return animate 

391 

392 

393def interpolate(times, states, times_inter): 

394 

395 assert len(times) == len(states) 

396 

397 states_inter = [] 

398 for i in range(len(times) - 1): 

399 

400 state_a = states[i] 

401 state_b = states[i+1] 

402 time_a = times[i] 

403 time_b = times[i+1] 

404 

405 animate = interpolateables(state_a, state_b) 

406 

407 if i == 0: 

408 times_inter_this = times_inter[num.logical_and( 

409 time_a <= times_inter, times_inter <= time_b)] 

410 else: 

411 times_inter_this = times_inter[num.logical_and( 

412 time_a < times_inter, times_inter <= time_b)] 

413 

414 for time_inter in times_inter_this: 

415 state = clone(state_b) 

416 if time_b == time_a: 

417 blend = 0. 

418 else: 

419 blend = (time_inter - time_a) / (time_b - time_a) 

420 

421 for ypath, v_old, v_new in animate: 

422 if isinstance(v_old, float) and isinstance(v_new, float): 

423 if ypath == 'strike': 

424 if v_new - v_old > 180.: 

425 v_new -= 360. 

426 elif v_new - v_old < -180.: 

427 v_new += 360. 

428 

429 if ypath != 'distance': 

430 v_inter = v_old + blend * (v_new - v_old) 

431 else: 

432 v_old = num.log(v_old) 

433 v_new = num.log(v_new) 

434 v_inter = v_old + blend * (v_new - v_old) 

435 v_inter = num.exp(v_inter) 

436 

437 set_elements(state, ypath, v_inter) 

438 else: 

439 set_elements(state, ypath, v_new) 

440 

441 states_inter.append(state) 

442 

443 return states_inter 

444 

445 

446class Interpolator(object): 

447 

448 def __init__(self, times, states, fps=25.): 

449 

450 assert len(times) == len(states) 

451 

452 self.dt = 1.0 / fps 

453 self.tmin = times[0] 

454 self.tmax = times[-1] 

455 times_inter = util.arange2( 

456 self.tmin, self.tmax, self.dt, error='floor') 

457 times_inter[-1] = times[-1] 

458 

459 if times_inter.size == 1: 

460 self._states_inter = [clone(states[-1])] 

461 return 

462 

463 states_inter = [] 

464 for i in range(len(times) - 1): 

465 

466 state_a = states[i] 

467 state_b = states[i+1] 

468 time_a = times[i] 

469 time_b = times[i+1] 

470 

471 animate = interpolateables(state_a, state_b) 

472 

473 if i == 0: 

474 times_inter_this = times_inter[num.logical_and( 

475 time_a <= times_inter, times_inter <= time_b)] 

476 else: 

477 times_inter_this = times_inter[num.logical_and( 

478 time_a < times_inter, times_inter <= time_b)] 

479 

480 for time_inter in times_inter_this: 

481 state = clone(state_b) 

482 

483 if time_b == time_a: 

484 blend = 0. 

485 else: 

486 blend = (time_inter - time_a) / (time_b - time_a) 

487 

488 for ypath, v_old, v_new in animate: 

489 if isinstance(v_old, float) and isinstance(v_new, float): 

490 if ypath in ('lon', 'strike'): 

491 if v_new - v_old > 180.: 

492 v_new -= 360. 

493 elif v_new - v_old < -180.: 

494 v_new += 360. 

495 

496 if ypath != 'distance': 

497 v_inter = v_old + blend * (v_new - v_old) 

498 else: 

499 v_old = num.log(v_old) 

500 v_new = num.log(v_new) 

501 v_inter = v_old + blend * (v_new - v_old) 

502 v_inter = num.exp(v_inter) 

503 

504 set_elements(state, ypath, v_inter) 

505 

506 elif isinstance(v_old, Color) and isinstance(v_new, Color): 

507 v_inter = interpolate_color(v_old, v_new, blend) 

508 set_elements(state, ypath, v_inter) 

509 

510 elif isinstance(v_old, Background) \ 

511 and isinstance(v_new, Background): 

512 v_inter = interpolate_background(v_old, v_new, blend) 

513 set_elements(state, ypath, v_inter) 

514 

515 else: 

516 set_elements(state, ypath, v_new) 

517 

518 states_inter.append(state) 

519 

520 self._states_inter = states_inter 

521 

522 def __call__(self, t): 

523 itime = int(round((t - self.tmin) / self.dt)) 

524 itime = min(max(0, itime), len(self._states_inter)-1) 

525 return self._states_inter[itime]