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

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import math 

7import numpy as num 

8 

9from pyrocko.guts import Bool, Float, StringChoice 

10from pyrocko import cake, plot 

11from pyrocko.plot import automap 

12from pyrocko.dataset import topo 

13from pyrocko.gui.qt_compat import qw, qc 

14 

15from pyrocko.gui.vtk_util import TrimeshPipe, faces_to_cells, \ 

16 cpt_to_vtk_lookuptable 

17 

18from .base import Element, ElementState 

19from .. import state as vstate 

20from pyrocko import geometry 

21 

22from .. import common 

23 

24guts_prefix = 'sparrow' 

25 

26 

27def ticks(vmin, vmax, vstep): 

28 vmin = num.floor(vmin / vstep) * vstep 

29 vmax = num.ceil(vmax / vstep) * vstep 

30 n = int(round((vmax - vmin) / vstep)) 

31 return vmin + num.arange(n+1) * vstep 

32 

33 

34def nice_value_circle(step): 

35 step = plot.nice_value(step) 

36 if step > 30.: 

37 return 30. 

38 

39 return step 

40 

41 

42class TopoMeshPipe(TrimeshPipe): 

43 

44 def __init__( 

45 self, tile, cells_cache=None, 

46 mask_ocean=False, mask_land=False, **kwargs): 

47 

48 vertices, faces = geometry.topo_to_mesh( 

49 tile.y(), tile.x(), tile.data, 

50 cake.earthradius) 

51 

52 self._exaggeration = 1.0 

53 self._tile = tile 

54 self._raw_vertices = vertices 

55 

56 centers = geometry.face_centers(vertices, faces) 

57 

58 altitudes = (geometry.vnorm(centers) - 1.0) * cake.earthradius 

59 if mask_ocean: 

60 mask = num.all(tile.data.flatten()[faces] == 0, axis=1) 

61 altitudes[mask] = None 

62 

63 if mask_land: 

64 mask = num.all(tile.data.flatten()[faces] >= 0, axis=1) 

65 altitudes[mask] = None 

66 

67 if cells_cache is not None: 

68 if id(faces) not in cells_cache: 

69 cells_cache[id(faces)] = faces_to_cells(faces) 

70 

71 cells = cells_cache[id(faces)] 

72 else: 

73 cells = faces_to_cells(faces) 

74 

75 TrimeshPipe.__init__( 

76 self, self._raw_vertices, cells=cells, values=altitudes, **kwargs) 

77 

78 def set_exaggeration(self, exaggeration): 

79 if self._exaggeration != exaggeration: 

80 factors = \ 

81 (cake.earthradius + self._tile.data.flatten()*exaggeration) \ 

82 / (cake.earthradius + self._tile.data.flatten()) 

83 self.set_vertices(self._raw_vertices * factors[:, num.newaxis]) 

84 self._exaggeration = exaggeration 

85 

86 

87class TopoCPTChoice(StringChoice): 

88 choices = ['light', 'uniform'] 

89 

90 

91class TopoState(ElementState): 

92 visible = Bool.T(default=True) 

93 exaggeration = Float.T(default=1.0) 

94 opacity = Float.T(default=1.0) 

95 smooth = Bool.T(default=False) 

96 cpt = TopoCPTChoice.T(default='light') 

97 shading = vstate.ShadingChoice.T(default='phong') 

98 resolution_max_factor = Float.T(default=1.0) 

99 resolution_min_factor = Float.T(default=1.0) 

100 coverage_factor = Float.T(default=1.0) 

101 

102 def create(self): 

103 element = TopoElement() 

104 return element 

105 

106 

107class TopoElement(Element): 

108 

109 def __init__(self): 

110 Element.__init__(self) 

111 self._parent = None 

112 self.mesh = None 

113 self._controls = None 

114 

115 # region = (-180., 180, -90, 90) 

116 # self._tile = topo.get('ETOPO1_D8', region) 

117 

118 self._visible = False 

119 self._active_meshes = {} 

120 self._meshes = {} 

121 self._cells = {} 

122 self._cpt_name = None 

123 self._lookuptables = {} 

124 

125 def get_name(self): 

126 return 'Topography' 

127 

128 def bind_state(self, state): 

129 Element.bind_state(self, state) 

130 self.talkie_connect( 

131 state, 

132 ['visible', 'exaggeration', 'opacity', 'smooth', 'shading', 'cpt', 

133 'resolution_min_factor', 'resolution_max_factor', 

134 'coverage_factor'], 

135 self.update) 

136 

137 self._state = state 

138 

139 def set_parent(self, parent): 

140 self._parent = parent 

141 

142 for var in ['distance', 'lat', 'lon']: 

143 self.talkie_connect( 

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

145 

146 self._parent.add_panel( 

147 self.get_title_label(), 

148 self._get_controls(), 

149 visible=True, 

150 title_controls=[ 

151 self.get_title_control_remove(), 

152 self.get_title_control_visible()]) 

153 

154 self.update() 

155 

156 def unset_parent(self): 

157 self.unbind_state() 

158 if self._parent: 

159 for k in self._active_meshes: 

160 self._parent.remove_actor(self._active_meshes[k].actor) 

161 

162 self._active_meshes.clear() 

163 self._meshes.clear() 

164 self._cells.clear() 

165 

166 if self._controls: 

167 self._parent.remove_panel(self._controls) 

168 self._controls = None 

169 

170 self._parent.update_view() 

171 self._parent = None 

172 

173 def select_dems(self, delta, region): 

174 if not self._parent: 

175 return [], [] 

176 

177 _, size_y = self._parent.renwin.GetSize() 

178 

179 dmin = 2.0 * delta * 1.0 / float(size_y) \ 

180 / self._state.resolution_max_factor # [deg] 

181 dmax = 2.0 * delta * 20.0 \ 

182 * self._state.resolution_min_factor / float(size_y) 

183 

184 result = [ 

185 topo.select_dem_names( 

186 k, dmin, dmax, topo.positive_region(region), mode='highest') 

187 

188 for k in ['ocean', 'land']] 

189 

190 if not any(result) and delta > 20.: 

191 return ['ETOPO1_D8'], ['ETOPO1_D8'] 

192 else: 

193 return result 

194 

195 def update_cpt(self, cpt_name): 

196 if cpt_name not in self._lookuptables: 

197 if cpt_name == 'light': 

198 topo_cpt_wet = 'light_sea' 

199 topo_cpt_dry = 'light_land' 

200 

201 elif cpt_name == 'uniform': 

202 topo_cpt_wet = 'light_sea_uniform' 

203 topo_cpt_dry = 'light_land_uniform' 

204 

205 cpt_wet = automap.read_cpt(topo.cpt(topo_cpt_wet)) 

206 cpt_dry = automap.read_cpt(topo.cpt(topo_cpt_dry)) 

207 cpt_combi = automap.cpt_merge_wet_dry(cpt_wet, cpt_dry) 

208 

209 lut_combi = cpt_to_vtk_lookuptable(cpt_combi) 

210 lut_combi.SetNanColor(0.0, 0.0, 0.0, 0.0) 

211 

212 self._lookuptables[cpt_name] = lut_combi 

213 

214 def update(self, *args): 

215 

216 pstate = self._parent.state 

217 delta = pstate.distance * geometry.r2d * 0.5 

218 visible = self._state.visible 

219 

220 self.update_cpt(self._state.cpt) 

221 

222 step = nice_value_circle( 

223 max(1./8., min(2**round(math.log(delta) / math.log(2.)), 10.))) 

224 

225 lat_min, lat_max, lon_min, lon_max, lon_closed = common.cover_region( 

226 pstate.lat, pstate.lon, delta*self._state.coverage_factor, step) 

227 

228 if lon_closed: 

229 lon_max = 180. 

230 

231 region = lon_min, lon_max, lat_min, lat_max 

232 

233 dems_ocean, dems_land = self.select_dems(delta, region) 

234 

235 lat_majors = ticks(lat_min, lat_max-step, step) 

236 lon_majors = ticks(lon_min, lon_max-step, step) 

237 

238 wanted = set() 

239 if visible: 

240 

241 show_progress = [False] 

242 abort = [False] 

243 ntiles = lat_majors.size * lon_majors.size 

244 itile = 0 

245 

246 def download_progress_update(): 

247 abort[0] = self._parent.progressbars.set_status( 

248 'Downloading topography data', 

249 itile/ntiles * 100., can_abort=False) 

250 

251 show_progress[0] = True 

252 self._parent.update() 

253 

254 self._parent.download_progress_update.connect( 

255 download_progress_update) 

256 

257 for ilat, lat in enumerate(lat_majors): 

258 for ilon, lon in enumerate(lon_majors): 

259 itile = ilon + ilat * lon_majors.size 

260 lon = ((lon + 180.) % 360.) - 180. 

261 

262 region = topo.positive_region( 

263 (lon, lon+step, lat, lat+step)) 

264 

265 for demname in dems_land[:1] + dems_ocean[:1]: 

266 mask_ocean = demname.startswith('SRTM') \ 

267 or demname.startswith('Iceland') 

268 

269 mask_land = demname.startswith('ETOPO') \ 

270 and (dems_land and dems_land[0] != demname) 

271 

272 k = (step, demname, region, mask_ocean, mask_land) 

273 if k not in self._meshes: 

274 tile = topo.get(demname, region) 

275 if not tile: 

276 continue 

277 

278 self._meshes[k] = TopoMeshPipe( 

279 tile, 

280 cells_cache=self._cells, 

281 mask_ocean=mask_ocean, 

282 mask_land=mask_land, 

283 smooth=self._state.smooth, 

284 lut=self._lookuptables[self._state.cpt]) 

285 

286 wanted.add(k) 

287 

288 # prevent adding both, SRTM and ETOPO because 

289 # vtk produces artifacts when showing the two masked 

290 # meshes (like this, mask_land is never used): 

291 break 

292 

293 if show_progress[0]: 

294 download_progress_update() 

295 

296 if abort[0]: 

297 break 

298 

299 if abort[0]: 

300 break 

301 

302 itile = ntiles 

303 if show_progress[0]: 

304 download_progress_update() 

305 

306 self._parent.download_progress_update.disconnect( 

307 download_progress_update) 

308 

309 if abort[0]: 

310 self._state.visible = False 

311 

312 unwanted = set() 

313 for k in self._active_meshes: 

314 if k not in wanted or not visible: 

315 unwanted.add(k) 

316 

317 for k in unwanted: 

318 self._parent.remove_actor(self._active_meshes[k].actor) 

319 del self._active_meshes[k] 

320 

321 for k in wanted: 

322 if k not in self._active_meshes: 

323 m = self._meshes[k] 

324 self._active_meshes[k] = m 

325 self._parent.add_actor(m.actor) 

326 

327 self._active_meshes[k].set_exaggeration(self._state.exaggeration) 

328 self._active_meshes[k].set_opacity(self._state.opacity) 

329 self._active_meshes[k].set_smooth(self._state.smooth) 

330 self._active_meshes[k].set_shading(self._state.shading) 

331 self._active_meshes[k].set_lookuptable( 

332 self._lookuptables[self._state.cpt]) 

333 

334 self._parent.update_view() 

335 

336 def _get_controls(self): 

337 state = self._state 

338 if not self._controls: 

339 from ..state import state_bind_slider, state_bind_checkbox, \ 

340 state_bind_combobox 

341 

342 frame = qw.QFrame() 

343 layout = qw.QGridLayout() 

344 frame.setLayout(layout) 

345 

346 iy = 0 

347 

348 # exaggeration 

349 

350 layout.addWidget(qw.QLabel('Exaggeration'), iy, 0) 

351 

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

353 slider.setSizePolicy( 

354 qw.QSizePolicy( 

355 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

356 slider.setMinimum(0) 

357 slider.setMaximum(2000) 

358 layout.addWidget(slider, iy, 1) 

359 

360 state_bind_slider(self, state, 'exaggeration', slider, factor=0.01) 

361 

362 iy += 1 

363 

364 # opacity 

365 

366 layout.addWidget(qw.QLabel('Opacity'), iy, 0) 

367 

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

369 slider.setSizePolicy( 

370 qw.QSizePolicy( 

371 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

372 slider.setMinimum(0) 

373 slider.setMaximum(1000) 

374 layout.addWidget(slider, iy, 1) 

375 

376 state_bind_slider(self, state, 'opacity', slider, factor=0.001) 

377 

378 iy += 1 

379 

380 # high resolution 

381 

382 layout.addWidget(qw.QLabel('High-Res Factor'), iy, 0) 

383 

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

385 slider.setSizePolicy( 

386 qw.QSizePolicy( 

387 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

388 slider.setMinimum(500) 

389 slider.setMaximum(4000) 

390 layout.addWidget(slider, iy, 1) 

391 

392 state_bind_slider( 

393 self, state, 'resolution_max_factor', slider, factor=0.001) 

394 

395 iy += 1 

396 

397 # low resolution 

398 

399 layout.addWidget(qw.QLabel('Low-Res Factor'), iy, 0) 

400 

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

402 slider.setSizePolicy( 

403 qw.QSizePolicy( 

404 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

405 slider.setMinimum(500) 

406 slider.setMaximum(4000) 

407 layout.addWidget(slider, iy, 1) 

408 

409 state_bind_slider( 

410 self, state, 'resolution_min_factor', slider, factor=0.001) 

411 

412 iy += 1 

413 

414 # low resolution 

415 

416 layout.addWidget(qw.QLabel('Coverage Factor'), iy, 0) 

417 

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

419 slider.setSizePolicy( 

420 qw.QSizePolicy( 

421 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

422 slider.setMinimum(500) 

423 slider.setMaximum(4000) 

424 layout.addWidget(slider, iy, 1) 

425 

426 state_bind_slider( 

427 self, state, 'coverage_factor', slider, factor=0.001) 

428 

429 iy += 1 

430 

431 cb = common.string_choices_to_combobox(TopoCPTChoice) 

432 layout.addWidget(qw.QLabel('CPT'), iy, 0) 

433 layout.addWidget(cb, iy, 1) 

434 state_bind_combobox(self, state, 'cpt', cb) 

435 

436 iy += 1 

437 

438 cb = qw.QCheckBox('Smooth') 

439 layout.addWidget(cb, iy, 0) 

440 state_bind_checkbox(self, state, 'smooth', cb) 

441 

442 cb = common.string_choices_to_combobox(vstate.ShadingChoice) 

443 layout.addWidget(cb, iy, 1) 

444 state_bind_combobox(self, state, 'shading', cb) 

445 

446 iy += 1 

447 

448 layout.addWidget(qw.QFrame(), iy, 0, 1, 2) 

449 

450 self._controls = frame 

451 

452 return self._controls 

453 

454 

455__all__ = [ 

456 'TopoElement', 

457 'TopoState', 

458]