Coverage for C:\src\imod-python\imod\visualize\pyvista.py: 79%

307 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:41 +0200

1""" 

2This module creates unstructured grids from DataArrays. 

3Directly creating pyvista.UnstructuredGrids is advantageous for several reasons. 

4 

51. Structured (modflow) grids are rectilinear, so the pyvista.RectilinearGrid might 

6seem obvious. However, nodata values are also present, and have to be removed using 

7a .treshold(). This means constructing every cell, followed by throwing most away. 

8This also returns an unstructured grid, and is slower. 

9 

102. Cells are rectangular, and they jump in the z dimension from one cell to the 

11other: 

12 __ 

13 __ |2 | 

14 |1 ||__| 

15 |__| 

16 

17In case of a rectilinear grid, the z value on the edge between cell 1 and 2 is 

18linearly interpolated, rather than making a jump. To create jumps, double the number 

19of cells is required, with dummy cells in between with width 0. This is clearly 

20wasteful. 

21 

22The grid is constructed with: 

23offset, cells, cell_type, points, values 

24 

25As presented here: 

26https://github.com/pyvista/pyvista/blob/0.23.0/examples/00-load/create-unstructured-surface.py 

27 

28* offset: start of each cell in the cells array. 

29* cells: number of points in the cell, then point number; for every cell. 

30* cell_type: integers, informing vtk of the geometry type. 

31* points are the coordinates (x, y, z) for every cell corner. 

32* values: are the values from the data array. Can generally by used for coloring. 

33 

34The methods below construct pyvista.UnstructuredGrids for voxel models (z1d), 

35"layer models" (z3d), and two dimensional data (e.g. a DEM). 

36""" 

37 

38from pathlib import Path 

39from typing import Optional, Tuple 

40 

41import numba 

42import numpy as np 

43import pandas as pd 

44import scipy.ndimage 

45import tqdm 

46import xarray as xr 

47 

48import imod 

49from imod.select import points_values 

50from imod.util.imports import MissingOptionalModule 

51 

52try: 

53 import pyvista as pv 

54 import vtk 

55 

56 if vtk.vtkVersion().GetVTKMajorVersion() < 9: 

57 raise ImportError("VTK version of 9.0 or higher required") 

58except ImportError: 

59 pv = MissingOptionalModule("pyvista") 

60 vtk = MissingOptionalModule("vtk") 

61 

62 

63def exterior(da, n): 

64 has_data = da.notnull() 

65 eroded = da.copy(data=scipy.ndimage.binary_erosion(has_data.values, iterations=n)) 

66 return has_data & ~eroded 

67 

68 

69@numba.njit 

70def _create_hexahedra_z1d(data, x, y, z): 

71 """ 

72 This function creates the necessary arrays to create hexahedra, based on a 

73 one-dimensional z array: i.e. one that does not depend on x and y. 

74 These arrays are used to create a pyvista.UnstructuredGrid. 

75 

76 Parameters 

77 ---------- 

78 data : np.array of size (nz, ny, nx) 

79 x : np.array of size nx + 1 

80 y: np.array of size ny + 1 

81 z : np.array of size (nz + 1) 

82 

83 Returns 

84 ------- 

85 offset : np.array of int 

86 cells : np.array of int 

87 cell_type : np.array of vkt enum 

88 points : np.array of float 

89 values : np.array of float 

90 """ 

91 nz, ny, nx = data.shape 

92 # First pass: count valid values 

93 n = 0 

94 for i in range(nz): 

95 for j in range(ny): 

96 for k in range(nx): 

97 if ~np.isnan(data[i, j, k]): 

98 n += 1 

99 

100 # Allocate 

101 # VTK_HEXAHEDRON is just an enum 

102 offset = np.arange(0, 9 * (n + 1), 9) 

103 cells = np.empty(n * 9, dtype=np.int32) 

104 indices = np.empty(n, dtype=np.int32) 

105 cell_type = np.full(n, vtk.VTK_HEXAHEDRON) 

106 # A hexahedron has 8 corners 

107 points = np.empty((n * 8, 3)) 

108 values = np.empty(n, data.dtype) 

109 

110 ii = 0 

111 jj = 0 

112 kk = 0 

113 linear_index = 0 

114 for i in range(nz): 

115 for j in range(ny): 

116 for k in range(nx): 

117 if ~np.isnan(data[i, j, k]): 

118 # Set coordinates of points 

119 points[ii] = (x[k], y[j], z[i]) 

120 points[ii + 1] = (x[k + 1], y[j], z[i]) 

121 points[ii + 2] = (x[k + 1], y[j + 1], z[i]) 

122 points[ii + 3] = (x[k], y[j + 1], z[i]) 

123 points[ii + 4] = (x[k], y[j], z[i + 1]) 

124 points[ii + 5] = (x[k + 1], y[j], z[i + 1]) 

125 points[ii + 6] = (x[k + 1], y[j + 1], z[i + 1]) 

126 points[ii + 7] = (x[k], y[j + 1], z[i + 1]) 

127 # Set number of cells, and point number 

128 cells[jj] = 8 

129 cells[jj + 1] = ii 

130 cells[jj + 2] = ii + 1 

131 cells[jj + 3] = ii + 2 

132 cells[jj + 4] = ii + 3 

133 cells[jj + 5] = ii + 4 

134 cells[jj + 6] = ii + 5 

135 cells[jj + 7] = ii + 6 

136 cells[jj + 8] = ii + 7 

137 ii += 8 

138 jj += 9 

139 # Set values 

140 values[kk] = data[i, j, k] 

141 # Set indices 

142 indices[kk] = linear_index 

143 kk += 1 

144 linear_index += 1 

145 

146 return indices, offset, cells, cell_type, points, values 

147 

148 

149@numba.njit 

150def _create_hexahedra_z3d(data, x, y, z3d): 

151 """ 

152 This function creates the necessary arrays to create hexahedra, based on a 

153 three-dimensional z array: i.e. one that depends on x and y. 

154 These arrays are used to create a pyvista.UnstructuredGrid. 

155 

156 Parameters 

157 ---------- 

158 data : np.array of size (nz, ny, nx) 

159 x : np.array of size nx + 1 

160 y: np.array of size ny + 1 

161 z : np.array of size (nz + 1, ny, nx) 

162 

163 Returns 

164 ------- 

165 offset : np.array of int 

166 cells : np.array of int 

167 cell_type : np.array of vkt enum 

168 points : np.array of float 

169 values : np.array of float 

170 """ 

171 nz, ny, nx = data.shape 

172 # First pass: count valid values 

173 n = 0 

174 for i in range(nz): 

175 for j in range(ny): 

176 for k in range(nx): 

177 if ~np.isnan(data[i, j, k]): 

178 n += 1 

179 

180 # Allocate 

181 # VTK_HEXAHEDRON is just an enum 

182 offset = np.arange(0, 9 * (n + 1), 9) 

183 cells = np.empty(n * 9, dtype=np.int32) 

184 indices = np.empty(n, dtype=np.int32) 

185 cell_type = np.full(n, vtk.VTK_HEXAHEDRON) 

186 # A hexahedron has 8 corners 

187 points = np.empty((n * 8, 3)) 

188 values = np.empty(n, data.dtype) 

189 

190 ii = 0 

191 jj = 0 

192 kk = 0 

193 linear_index = 0 

194 for i in range(nz): 

195 for j in range(ny): 

196 for k in range(nx): 

197 v = data[i, j, k] 

198 if ~np.isnan(v): 

199 # Set coordinates of points 

200 points[ii] = (x[k], y[j], z3d[i, j, k]) 

201 points[ii + 1] = (x[k + 1], y[j], z3d[i, j, k]) 

202 points[ii + 2] = (x[k + 1], y[j + 1], z3d[i, j, k]) 

203 points[ii + 3] = (x[k], y[j + 1], z3d[i, j, k]) 

204 points[ii + 4] = (x[k], y[j], z3d[i + 1, j, k]) 

205 points[ii + 5] = (x[k + 1], y[j], z3d[i + 1, j, k]) 

206 points[ii + 6] = (x[k + 1], y[j + 1], z3d[i + 1, j, k]) 

207 points[ii + 7] = (x[k], y[j + 1], z3d[i + 1, j, k]) 

208 # Set number of cells, and point number 

209 cells[jj] = 8 

210 cells[jj + 1] = ii 

211 cells[jj + 2] = ii + 1 

212 cells[jj + 3] = ii + 2 

213 cells[jj + 4] = ii + 3 

214 cells[jj + 5] = ii + 4 

215 cells[jj + 6] = ii + 5 

216 cells[jj + 7] = ii + 6 

217 cells[jj + 8] = ii + 7 

218 ii += 8 

219 jj += 9 

220 # Set values 

221 values[kk] = v 

222 # Set index 

223 indices[kk] = linear_index 

224 kk += 1 

225 linear_index += 1 

226 

227 return indices, offset, cells, cell_type, points, values 

228 

229 

230@numba.njit 

231def _create_plane_surface(data, x, y): 

232 """ 

233 This function creates the necessary arrays to create quads, based on a 

234 two-dimensional data array. The data array is used for the z dimension. 

235 All the horizontal surfaces are connected by vertical faces, exactly 

236 representing the cell geometry, with a constant value per cell. 

237 The alternative is linear interpolation, which does not represent 

238 geometry exactly, or holes in the surface, with isolated quads floating 

239 in space. 

240 

241 These arrays are used to create a pyvista.UnstructuredGrid. 

242 

243 Parameters 

244 ---------- 

245 data : np.array of size (ny, nx) 

246 x : np.array of size nx + 1 

247 y: np.array of size ny + 1 

248 

249 Returns 

250 ------- 

251 offset : np.array of int 

252 cells : np.array of int 

253 cell_type : np.array of vkt enum 

254 points : np.array of float 

255 values : np.array of float 

256 """ 

257 ny, nx = data.shape 

258 # First pass: count valid values 

259 n = 0 

260 for i in range(ny): 

261 for j in range(nx): 

262 if ~np.isnan(data[i, j]): 

263 n += 1 

264 if j < nx - 1: 

265 if ~np.isnan(data[i, j + 1]): 

266 n += 1 

267 if i < ny - 1: 

268 if ~np.isnan(data[i + 1, j]): 

269 n += 1 

270 

271 # Allocate 

272 # VTK_QUAD is just an enum 

273 offset = np.arange(0, 5 * (n + 1), 5) 

274 cells = np.empty(n * 5, dtype=np.int32) 

275 cell_type = np.full(n, vtk.VTK_QUAD) 

276 # A hexahedron has r corners 

277 points = np.empty((n * 4, 3)) 

278 values = np.empty(n, dtype=data.dtype) 

279 

280 ii = 0 

281 jj = 0 

282 kk = 0 

283 for i in range(ny): 

284 for j in range(nx): 

285 v = data[i, j] 

286 # Create horizontal quad 

287 if ~np.isnan(v): 

288 # Set coordinates of points 

289 points[ii] = (x[j], y[i], v) 

290 points[ii + 1] = (x[j + 1], y[i], v) 

291 points[ii + 2] = (x[j + 1], y[i + 1], v) 

292 points[ii + 3] = (x[j], y[i + 1], v) 

293 # Set number of cells, and point number 

294 cells[jj] = 4 

295 cells[jj + 1] = ii 

296 cells[jj + 2] = ii + 1 

297 cells[jj + 3] = ii + 2 

298 cells[jj + 4] = ii + 3 

299 ii += 4 

300 jj += 5 

301 # Set values 

302 values[kk] = v 

303 # Set index 

304 kk += 1 

305 

306 if j < nx - 1: 

307 v01 = data[i, j + 1] 

308 # Create vertical quads 

309 if ~np.isnan(v01): 

310 # Set coordinates of points 

311 points[ii] = (x[j + 1], y[i], v) 

312 points[ii + 1] = (x[j + 1], y[i + 1], v) 

313 points[ii + 2] = (x[j + 1], y[i + 1], v01) 

314 points[ii + 3] = (x[j + 1], y[i], v01) 

315 # Set number of cells, and point number 

316 cells[jj] = 4 

317 cells[jj + 1] = ii 

318 cells[jj + 2] = ii + 1 

319 cells[jj + 3] = ii + 2 

320 cells[jj + 4] = ii + 3 

321 ii += 4 

322 jj += 5 

323 # Set values 

324 values[kk] = v if v > v01 else v01 

325 # Set index 

326 kk += 1 

327 

328 if i < ny - 1: 

329 v10 = data[i + 1, j] 

330 if ~np.isnan(v10): 

331 # Set coordinates of points 

332 points[ii] = (x[j], y[i + 1], v) 

333 points[ii + 1] = (x[j + 1], y[i + 1], v) 

334 points[ii + 2] = (x[j + 1], y[i + 1], v10) 

335 points[ii + 3] = (x[j], y[i + 1], v10) 

336 # Set number of cells, and point number 

337 cells[jj] = 4 

338 cells[jj + 1] = ii 

339 cells[jj + 2] = ii + 1 

340 cells[jj + 3] = ii + 2 

341 cells[jj + 4] = ii + 3 

342 ii += 4 

343 jj += 5 

344 # Set values 

345 values[kk] = v if v > v10 else v10 

346 # Set index 

347 kk += 1 

348 

349 return offset, cells, cell_type, points, values 

350 

351 

352def vertices_coords(dx, xmin, xmax, nx): 

353 """ 

354 Return the coordinates of the vertices. 

355 (xarray stores midpoints) 

356 """ 

357 if isinstance(dx, float): 

358 dx = np.full(nx, dx) 

359 if dx[0] > 0: 

360 x = np.full(nx + 1, xmin) 

361 else: 

362 x = np.full(nx + 1, xmax) 

363 x[1:] += dx.cumsum() 

364 return x 

365 

366 

367def grid_3d( 

368 da, 

369 vertical_exaggeration=30.0, 

370 exterior_only=True, 

371 exterior_depth=1, 

372 return_index=False, 

373): 

374 """ 

375 Constructs a 3D PyVista representation of a DataArray. 

376 DataArrays should be two-dimensional or three-dimensional: 

377 

378 * 2D: dimensions should be ``{"y", "x"}``. E.g. a DEM. 

379 * 3D: dimensions should be ``{"z", "y", "x"}``, for a voxel model. 

380 * 3D: dimensions should be ``{"layer", "y", "x"}``, with coordinates 

381 ``"top"({"layer", "y", "x"})`` and ``"bottom"({"layer", "y", "x"})``. 

382 

383 Parameters 

384 ---------- 

385 da : xr.DataArray 

386 vertical_exaggeration : float, default 30.0 

387 exterior_only : bool, default True 

388 Whether or not to only draw the exterior. Greatly speeds up rendering, 

389 but it means that pyvista slices and filters produce "hollowed out" 

390 results. 

391 exterior_depth : int, default 1 

392 How many cells to consider as exterior. In case of large jumps, holes 

393 can occur. By settings this argument to a higher value, more of the 

394 inner cells will be rendered, reducing the chances of gaps occurring. 

395 return_index : bool, default False 

396 

397 Returns 

398 ------- 

399 pyvista.UnstructuredGrid 

400 

401 Examples 

402 -------- 

403 

404 >>> grid = imod.visualize.grid_3d(da) 

405 

406 To plot the grid, call the ``.plot()`` method. 

407 

408 >>> grid.plot() 

409 

410 Use ``.assign_coords`` to assign tops and bottoms to layer models: 

411 

412 >>> top = imod.idf.open("top*.idf") 

413 >>> bottom = imod.idf.open("bot*.idf") 

414 >>> kd = imod.idf.open("kd*.idf") 

415 >>> kd = kd.assign_coords(top=(("layer", "y", "x"), top)) 

416 >>> kd = kd.assign_coords(bottom=(("layer", "y", "x"), bottom)) 

417 >>> grid = imod.visualize.grid_3d(kd) 

418 >>> grid.plot() 

419 

420 Refer to the PyVista documentation on how to customize plots: 

421 https://docs.pyvista.org/index.html 

422 """ 

423 # x and y dimension 

424 dx, xmin, xmax, dy, ymin, ymax = imod.util.spatial.spatial_reference(da) 

425 nx = da.coords["x"].size 

426 ny = da.coords["y"].size 

427 x = vertices_coords(dx, xmin, xmax, nx) 

428 y = vertices_coords(dy, ymin, ymax, ny) 

429 # Coordinates should always have dtype == np.float64 by now 

430 

431 # z dimension 

432 if "top" in da.coords and "bottom" in da.coords: 

433 if not len(da.shape) == 3: 

434 raise ValueError( 

435 'Coordinates "top" and "bottom" are present, but data is not 3D.' 

436 ) 

437 if not set(da.dims) == {"layer", "y", "x"}: 

438 raise ValueError( 

439 'Coordinates "top" and "bottom" are present, only dimensions allowed are: ' 

440 '{"layer", "y", "x"}' 

441 ) 

442 da = da.transpose("layer", "y", "x", transpose_coords=True) 

443 ztop = da.coords["top"] 

444 zbot = da.coords["bottom"] 

445 z3d = np.vstack([np.expand_dims(ztop.isel(layer=1).values, 0), zbot.values]) 

446 

447 if exterior_only: 

448 da = da.where(exterior(da, exterior_depth)) 

449 

450 indices, offset, cells, cell_type, points, values = _create_hexahedra_z3d( 

451 da.values, x, y, z3d 

452 ) 

453 elif "z" in da.coords: 

454 if not len(da.shape) == 3: 

455 raise ValueError('Coordinate "z" is present, but data is not 3D.') 

456 if not (set(da.dims) == {"layer", "y", "x"} or set(da.dims) == {"z", "y", "x"}): 

457 raise ValueError( 

458 'Coordinate "z" is present, only dimensions allowed are: ' 

459 '{"layer", "y", "x"} or {"z", "y", "x"}' 

460 ) 

461 if "z" not in da.dims: 

462 da = da.transpose("layer", "y", "x", transpose_coords=True) 

463 else: 

464 da = da.transpose("z", "y", "x", transpose_coords=True) 

465 

466 dz, zmin, zmax = imod.util.spatial.coord_reference(da["z"]) 

467 nz = da.coords["z"].size 

468 z = vertices_coords(dz, zmin, zmax, nz) 

469 

470 if exterior_only: 

471 da = da.where(exterior(da, exterior_depth)) 

472 

473 indices, offset, cells, cell_type, points, values = _create_hexahedra_z1d( 

474 da.values, x, y, z 

475 ) 

476 elif set(da.dims) == {"y", "x"}: 

477 if return_index: 

478 raise ValueError("Cannot return indices for a 2D dataarray.") 

479 da = da.transpose("y", "x", transpose_coords=True) 

480 offset, cells, cell_type, points, values = _create_plane_surface( 

481 da.values.astype(np.float64), x, y 

482 ) 

483 else: 

484 raise ValueError( 

485 'Incorrect coordinates and/or dimensions: Neither "z" nor "top" and ' 

486 '"bottom" is present in the DataArray, but dimension are not {"y", "x"} ' 

487 " either." 

488 ) 

489 

490 grid = pv.UnstructuredGrid(cells, cell_type, points) 

491 grid.points[:, -1] *= vertical_exaggeration 

492 grid.cell_data["values"] = values 

493 

494 if return_index: 

495 return grid, indices 

496 else: 

497 return grid 

498 

499 

500# For arrow plots: compute magnitude of vector 

501# Downsample, using max rule. 

502# Upsample again 

503# Select with .where? 

504# Create an array of (n, 3) shape 

505# Vectors are located somewhere 

506 

507 

508def line_3d(polygon, z=0.0): 

509 """ 

510 Returns the exterior line of a shapely polygon. 

511 

512 Parameters 

513 ---------- 

514 polygon : shapely.geometry.Polygon 

515 z : float or xr.DataArray 

516 z-coordinate to assign to line. If DataArray, assigns z-coordinate 

517 based on xy locations in DataArray. 

518 

519 Returns 

520 ------- 

521 pyvista.PolyData 

522 """ 

523 x, y = map(np.array, polygon.exterior.coords.xy) 

524 if isinstance(z, xr.DataArray): 

525 z = points_values(z, x=x, y=y).values 

526 else: 

527 z = np.full_like(x, z) 

528 coords = np.vstack([x, y, z]).transpose() 

529 return pv.lines_from_points(coords) 

530 

531 

532class GridAnimation3D: 

533 """ 

534 Class to easily setup 3D animations for transient data. Use the 

535 ``imod.visualize.StaticGridAnimation3D`` when the location of the displayed 

536 cells is constant over time: it will render much faster. 

537 

538 You can iteratively add or change settings to the plotter, until you're 

539 satisfied. Call the ``.peek()`` method to take a look. When satisfied, call 

540 ``.output()`` to write to a file. 

541 

542 

543 Parameters 

544 ---------- 

545 da : xr.DataArray 

546 The dataarray with transient data. Must contain a "time" dimension. 

547 vertical_exaggeration : float, defaults to 30.0 mesh_kwargs : dict 

548 keyword arguments that are forwarded to the pyvista mesh representing 

549 "da". If "stitle" is given as one of the arguments, the special keyword 

550 "timestamp" can be used to render the plotted time as part of the title. 

551 See example. For a full list of kwargs supported, see the 

552 `plotter.add_mesh 

553 <https://docs.pyvista.org/api/plotting/_autosummary/pyvista.Plotter.add_mesh.html#pyvista.Plotter.add_mesh>`_ 

554 method documentation. 

555 plotter_kwargs : dict 

556 keyword arguments that are forwarded to the pyvista plotter. For a full 

557 list of of kwargs supported, see the `Plotter constructor 

558 <https://docs.pyvista.org/api/plotting/_autosummary/pyvista.Plotter.html>`_ 

559 documention. 

560 

561 Examples 

562 -------- 

563 

564 Initialize the animation: 

565 

566 >>> animation = imod.visualize.GridAnimation3D(concentration, mesh_kwargs=dict(cmap="jet")) 

567 

568 Check what it looks like (if a window pops up: press "q" instead of the X to 

569 return): 

570 

571 >>> animation.peek() 

572 

573 Change the camera position, add bounding box, and check the result: 

574 

575 >>> animation.plotter.camera_position = (2, 1, 0.5) 

576 >>> animation.plotter.add_bounding_box() 

577 >>> animation.peek() 

578 

579 When it looks good, write to a file: 

580 

581 >>> animation.write("example.mp4") 

582 

583 If you've made some changes that don't look good, call ``.reset()`` to start 

584 over: 

585 

586 >>> animation.reset() 

587 

588 Note that ``.reset()`` is automatically called when the animation has 

589 finished writing. 

590 

591 You can use "stitle" in mesh_kwargs in conjunction with the "timestamp" 

592 keyword to print a formatted timestamp in the animation: 

593 

594 >>> animation = imod.visualize.GridAnimation3D(concentration, mesh_kwargs=dict(stitle="Concentration on {timestamp:%Y-%m-%d}")) 

595 """ 

596 

597 def _initialize(self, da): 

598 self.mesh = grid_3d( 

599 da, vertical_exaggeration=self.vertical_exaggeration, return_index=False 

600 ) 

601 mesh_kwargs = self.mesh_kwargs.copy() 

602 if "stitle" in mesh_kwargs and "{timestamp" in mesh_kwargs["stitle"]: 

603 mesh_kwargs["stitle"] = mesh_kwargs["stitle"].format( 

604 timestamp=pd.Timestamp(da.time.values) 

605 ) 

606 

607 self.mesh_actor = self.plotter.add_mesh(self.mesh, **mesh_kwargs) 

608 

609 def _update(self, da): 

610 self.plotter.remove_actor(self.mesh_actor) 

611 self._initialize(da) 

612 

613 def __init__( 

614 self, da, vertical_exaggeration=30.0, mesh_kwargs={}, plotter_kwargs={} 

615 ): 

616 # Store data 

617 self.da = da 

618 self.vertical_exaggeration = vertical_exaggeration 

619 self.mesh_kwargs = mesh_kwargs.copy() 

620 self.plotter_kwargs = plotter_kwargs 

621 self.plotter = pv.Plotter(**plotter_kwargs) 

622 # Initialize pyvista objects 

623 self._initialize(da.isel(time=0)) 

624 

625 def peek(self): 

626 """ 

627 Display the current state of the animation plotter. 

628 """ 

629 self.plotter.show(auto_close=False) 

630 

631 def reset(self): 

632 """ 

633 Reset the plotter to its base state. 

634 """ 

635 self.plotter = pv.Plotter(**self.plotter_kwargs) 

636 self._update(self.da.isel(time=0)) 

637 

638 def write(self, filename, framerate=24): 

639 """ 

640 Write the animation to a video or gif. 

641 

642 Resets the plotter when finished animating. 

643 

644 Parameters 

645 ---------- 

646 filename : str, pathlib.Path 

647 Filename to write the video to. Should be an .mp4 or .gif. 

648 framerate : int, optional 

649 Frames per second. Not honoured for gif. 

650 """ 

651 if Path(filename).suffix.lower() == ".gif": 

652 self.plotter.open_gif( 

653 Path(filename).with_suffix(".gif").as_posix() 

654 ) # only lowercase gif and no Path allowed 

655 else: 

656 self.plotter.open_movie(filename, framerate=framerate) 

657 self.plotter.show(auto_close=False, interactive=False) 

658 self.plotter.write_frame() 

659 

660 for itime in tqdm.tqdm(range(1, self.da.coords["time"].size)): 

661 self._update(self.da.isel(time=itime)) 

662 self.plotter.write_frame() 

663 

664 # Close and reinitialize 

665 self.plotter.close() 

666 self.reset() 

667 

668 

669class StaticGridAnimation3D(GridAnimation3D): 

670 """ 

671 Class to easily setup 3D animations for transient data; 

672 Should only be used when the location of the displayed cells is constant 

673 over time. It will render much faster than ``imod.visualize.GridAnimation3D``. 

674 

675 Refer to examples of ``imod.visualize.GridAnimation3D``. 

676 """ 

677 

678 def _initialize(self, da): 

679 self.mesh, self.indices = grid_3d( 

680 da, vertical_exaggeration=self.vertical_exaggeration, return_index=True 

681 ) 

682 

683 def _update(self, da): 

684 self.mesh.cell_data["values"] = da.values.ravel()[self.indices] 

685 

686 

687def velocity_field( 

688 vx: xr.DataArray, 

689 vy: xr.DataArray, 

690 vz: xr.DataArray, 

691 z: Optional[xr.DataArray] = None, 

692 vertical_exaggeration: Optional[float] = 30.0, 

693 scale_by_magnitude: Optional[bool] = True, 

694 factor: Optional[float] = 1.0, 

695 tolerance: Optional[float] = 0.0, 

696 absolute: Optional[bool] = False, 

697 clamping: Optional[bool] = False, 

698 rng: Optional[Tuple[float, float]] = None, 

699): 

700 if z is None: 

701 z = vx["z"].values 

702 

703 if not (vx.shape == vy.shape == vz.shape): 

704 raise ValueError("Shapes of velocity components vx, vy, vz do not match") 

705 if not (vx.dims == ("layer", "y", "x") or vx.dims == ("z", "y", "x")): 

706 raise ValueError( 

707 'Velocity components must have dimensions ("layer", "y", "x") ' 

708 'or ("z", "y", "x") exactly.\n' 

709 f"Received {vx.dims} instead." 

710 ) 

711 

712 # Start by generating the location of the velocity arrows 

713 # Ensure the points follow the memory layout of the v DataArrays (z, y, x) 

714 # otherwise, the result of np.ravel() doesn't match up 

715 if z.dim == 1: 

716 zz, yy, xx = np.meshgrid(z, vx["y"].values, vx["x"].values, indexing="ij") 

717 elif z.dim == 3: 

718 if not z.shape == vx.shape: 

719 raise ValueError("Shape of `z` does not match velocity components.") 

720 _, yy, xx = np.meshgrid( 

721 np.arange(z.shape[0]), vx["x"].values, vy["y"].avlues, indexing="ij" 

722 ) 

723 zz = z 

724 else: 

725 raise ValueError("z should be one or three dimensional.") 

726 

727 zz *= vertical_exaggeration 

728 cellcenters = pv.PolyData(np.stack([np.ravel(x) for x in (xx, yy, zz)], axis=1)) 

729 # Add the velocity components in x, y, z order 

730 cellcenters["velocity"] = np.stack([x.data.ravel() for x in (vx, vy, vz)], axis=1) 

731 scale = "velocity" if scale_by_magnitude else None 

732 

733 return cellcenters.glyph( 

734 scale=scale, 

735 orient="velocity", 

736 factor=factor, 

737 tolerance=tolerance, 

738 absolute=absolute, 

739 clamping=clamping, 

740 rng=rng, 

741 )