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

1"""

2 This module creates unstructured grids from DataArrays.

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

4

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

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

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

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

9

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

11 other:

12 __

13 __ |2 |

14 |1 ||__|

15 |__|

16

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

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

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

20 wasteful.

21

22 The grid is constructed with:

23 offset, cells, cell_type, points, values

24

25 As presented here:

26 https://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

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

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

36"""

37

38 from pathlib import Path

39 from typing import Optional, Tuple

40

41 import numba

42 import numpy as np

43 import pandas as pd

44 import scipy.ndimage

45 import tqdm

46 import xarray as xr

47

48 import imod

49 from imod.select import points_values

50 from imod.util.imports import MissingOptionalModule

51

52 try:

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

58 except ImportError:

59 pv = MissingOptionalModule("pyvista")

60 vtk = MissingOptionalModule("vtk")

61

62

63 def 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

70 def _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

150 def _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

231 def _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

352 def 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

367 def 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

508 def 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

532 class 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

669 class 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

687 def 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 )