diff --git a/autotest/test_export.py b/autotest/test_export.py index 1bd9a9514..74f8e1884 100644 --- a/autotest/test_export.py +++ b/autotest/test_export.py @@ -1417,7 +1417,6 @@ def test_vtk_cbc(function_tmpdir, example_data_path): @requires_pkg("vtk") -@pytest.mark.slow def test_vtk_vector(function_tmpdir, example_data_path): # test mf 2005 freyberg mpth = example_data_path / "freyberg_multilayer_transient" @@ -1480,6 +1479,40 @@ def test_vtk_vector(function_tmpdir, example_data_path): assert info["pointdata_names"] == [] +@requires_pkg("vtk") +@pytest.mark.parametrize("point_scalars", [False, True]) +@pytest.mark.parametrize("size", ["nnodes", "ncpl"]) +def test_vtk_add_vector_components(size, point_scalars): + """Cell i must get (x[i], y[i], z[i]) for either input size""" + from vtk.util import numpy_support + + nlay, nrow, ncol = 2, 3, 4 + grid = StructuredGrid( + delr=np.full(ncol, 10.0), + delc=np.full(nrow, 10.0), + top=np.full((nrow, ncol), 10.0), + botm=np.array([np.full((nrow, ncol), 0.0), np.full((nrow, ncol), -10.0)]), + nlay=nlay, + ) + n = nlay * nrow * ncol if size == "nnodes" else nrow * ncol + + # a constant field is used because inverse distance weighting of a + # constant returns the constant, so the point and the cell case have the + # same expected value + vector = np.array([np.full(n, 1.0), np.full(n, 2.0), np.full(n, 3.0)]) + vtk = Vtk(modelgrid=grid, point_scalars=point_scalars) + vtk.add_vector(vector, "v") + + data = vtk.vtk_grid.GetPointData() if point_scalars else vtk.vtk_grid.GetCellData() + arr = numpy_support.vtk_to_numpy(data.GetVectors()) + assert arr.shape[1] == 3 + + # cells that the ncpl sized vector does not reach are filled with nan + finite = arr[np.isfinite(arr).all(axis=1)] + assert len(finite) > 0 + assert np.allclose(finite, [1.0, 2.0, 3.0]) + + @requires_pkg("vtk") def test_vtk_unstructured(function_tmpdir, unstructured_grid): from vtkmodules.util.numpy_support import vtk_to_numpy diff --git a/flopy/export/vtk.py b/flopy/export/vtk.py index 9fdb4741f..387e5d7e8 100644 --- a/flopy/export/vtk.py +++ b/flopy/export/vtk.py @@ -897,7 +897,7 @@ def add_vector(self, vector, name, masked_values=None): else: raise AssertionError("Size of vector must be 3 * nnodes or 3 * ncpl") else: - vector = np.reshape(vector, (3, self.nnodes)).T + vector = np.reshape(vector, (3, self.nnodes)) if self.point_scalars: tmp = [] @@ -907,8 +907,11 @@ def add_vector(self, vector, name, masked_values=None): vector = self._mask_values(vector, masked_values) + # a row per component is built above, but numpy_to_vtk flattens in + # row-major order and vtk reads each row as one tuple, so the array + # is transposed to give a row per cell or point vtk_arr = numpy_support.numpy_to_vtk( - num_array=vector, array_type=self.__vtk.VTK_FLOAT + num_array=vector.T, array_type=self.__vtk.VTK_FLOAT ) vtk_arr.SetName(name) vtk_arr.SetNumberOfComponents(3)