fix(vtk): transpose the vector after the point scalar expansion

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

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)