Coverage for C:\src\imod-python\imod\visualize\pyvista.py: 12%
307 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
1"""
2This module creates unstructured grids from DataArrays.
3Directly creating pyvista.UnstructuredGrids is advantageous for several reasons.
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.
102. Cells are rectangular, and they jump in the z dimension from one cell to the
11other:
12 __
13 __ |2 |
14 |1 ||__|
15 |__|
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.
22The grid is constructed with:
23offset, cells, cell_type, points, values
25As presented here:
26https://github.com/pyvista/pyvista/blob/0.23.0/examples/00-load/create-unstructured-surface.py
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.
34The methods below construct pyvista.UnstructuredGrids for voxel models (z1d),
35"layer models" (z3d), and two dimensional data (e.g. a DEM).
36"""
38from pathlib import Path
39from typing import Optional, Tuple
41import numba
42import numpy as np
43import pandas as pd
44import scipy.ndimage
45import tqdm
46import xarray as xr
48import imod
49from imod.select import points_values
50from imod.util.imports import MissingOptionalModule
52try:
53 import pyvista as pv
54 import vtk
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")
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
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.
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)
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
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)
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
146 return indices, offset, cells, cell_type, points, values
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.
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)
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
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)
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
227 return indices, offset, cells, cell_type, points, values
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.
241 These arrays are used to create a pyvista.UnstructuredGrid.
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
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
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)
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
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
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
349 return offset, cells, cell_type, points, values
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
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:
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"})``.
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
397 Returns
398 -------
399 pyvista.UnstructuredGrid
401 Examples
402 --------
404 >>> grid = imod.visualize.grid_3d(da)
406 To plot the grid, call the ``.plot()`` method.
408 >>> grid.plot()
410 Use ``.assign_coords`` to assign tops and bottoms to layer models:
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()
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
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])
447 if exterior_only:
448 da = da.where(exterior(da, exterior_depth))
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)
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)
470 if exterior_only:
471 da = da.where(exterior(da, exterior_depth))
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 )
490 grid = pv.UnstructuredGrid(cells, cell_type, points)
491 grid.points[:, -1] *= vertical_exaggeration
492 grid.cell_data["values"] = values
494 if return_index:
495 return grid, indices
496 else:
497 return grid
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
508def line_3d(polygon, z=0.0):
509 """
510 Returns the exterior line of a shapely polygon.
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.
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)
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.
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.
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.
561 Examples
562 --------
564 Initialize the animation:
566 >>> animation = imod.visualize.GridAnimation3D(concentration, mesh_kwargs=dict(cmap="jet"))
568 Check what it looks like (if a window pops up: press "q" instead of the X to
569 return):
571 >>> animation.peek()
573 Change the camera position, add bounding box, and check the result:
575 >>> animation.plotter.camera_position = (2, 1, 0.5)
576 >>> animation.plotter.add_bounding_box()
577 >>> animation.peek()
579 When it looks good, write to a file:
581 >>> animation.write("example.mp4")
583 If you've made some changes that don't look good, call ``.reset()`` to start
584 over:
586 >>> animation.reset()
588 Note that ``.reset()`` is automatically called when the animation has
589 finished writing.
591 You can use "stitle" in mesh_kwargs in conjunction with the "timestamp"
592 keyword to print a formatted timestamp in the animation:
594 >>> animation = imod.visualize.GridAnimation3D(concentration, mesh_kwargs=dict(stitle="Concentration on {timestamp:%Y-%m-%d}"))
595 """
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 )
607 self.mesh_actor = self.plotter.add_mesh(self.mesh, **mesh_kwargs)
609 def _update(self, da):
610 self.plotter.remove_actor(self.mesh_actor)
611 self._initialize(da)
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))
625 def peek(self):
626 """
627 Display the current state of the animation plotter.
628 """
629 self.plotter.show(auto_close=False)
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))
638 def write(self, filename, framerate=24):
639 """
640 Write the animation to a video or gif.
642 Resets the plotter when finished animating.
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()
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()
664 # Close and reinitialize
665 self.plotter.close()
666 self.reset()
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``.
675 Refer to examples of ``imod.visualize.GridAnimation3D``.
676 """
678 def _initialize(self, da):
679 self.mesh, self.indices = grid_3d(
680 da, vertical_exaggeration=self.vertical_exaggeration, return_index=True
681 )
683 def _update(self, da):
684 self.mesh.cell_data["values"] = da.values.ravel()[self.indices]
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
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 )
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.")
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
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 )