Merge pull request #2284 from OceanParcels/bugfix/2260-vectorize-pwlinear-node

Bugfix/2260 vectorize pwlinear node

Author: fluidnumerics-joe

parcels/interpolators.py

@@ -653,12 +653,12 @@ def UXPiecewiseLinearNode(
     """
     k, fi = position["Z"][0], position["FACE"][0]
     bcoords = position["FACE"][1]
-    node_ids = field.grid.uxgrid.face_node_connectivity[fi, :]
+    node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values
     # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels.
     # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1.
     # First, do barycentric interpolation in the lateral direction for each interface level
-    fzk = np.sum(field.data.values[ti, k, node_ids] * bcoords, axis=-1)
-    fzkp1 = np.sum(field.data.values[ti, k + 1, node_ids] * bcoords, axis=-1)
+    fzk = np.sum(field.data.values[ti[:, None], k[:, None], node_ids] * bcoords, axis=-1)
+    fzkp1 = np.sum(field.data.values[ti[:, None], k[:, None] + 1, node_ids] * bcoords, axis=-1)
 
     # Then, do piecewise linear interpolation in the vertical direction
     zk = field.grid.z.values[k]

tests/test_particleset_execute.py

@@ -22,8 +22,8 @@
 from parcels._datasets.structured.generated import simple_UV_dataset
 from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
-from parcels.interpolators import UXPiecewiseConstantFace
-from parcels.kernels import AdvectionEE, AdvectionRK4
+from parcels.interpolators import UXPiecewiseConstantFace, UXPiecewiseLinearNode
+from parcels.kernels import AdvectionEE, AdvectionRK4, AdvectionRK4_3D
 from tests import utils
 from tests.common_kernels import DoNothing
 
@@ -479,6 +479,50 @@ def test_uxstommelgyre_pset_execute():
     assert utils.round_and_hash_float_array([p.lat for p in pset]) == 1142124776
 
 
+def test_uxstommelgyre_multiparticle_pset_execute():
+    ds = datasets_unstructured["stommel_gyre_delaunay"]
+    grid = UxGrid(grid=ds.uxgrid, z=ds.coords["nz"], mesh="spherical")
+    U = Field(
+        name="U",
+        data=ds.U,
+        grid=grid,
+        interp_method=UXPiecewiseConstantFace,
+    )
+    V = Field(
+        name="V",
+        data=ds.V,
+        grid=grid,
+        interp_method=UXPiecewiseConstantFace,
+    )
+    W = Field(
+        name="W",
+        data=ds.W,
+        grid=grid,
+        interp_method=UXPiecewiseLinearNode,
+    )
+    P = Field(
+        name="P",
+        data=ds.p,
+        grid=grid,
+        interp_method=UXPiecewiseConstantFace,
+    )
+    UVW = VectorField(name="UVW", U=U, V=V, W=W)
+    fieldset = FieldSet([UVW, UVW.U, UVW.V, UVW.W, P])
+    pset = ParticleSet(
+        fieldset,
+        lon=[30.0, 32.0],
+        lat=[5.0, 5.0],
+        depth=[50.0, 50.0],
+        time=[np.timedelta64(0, "s")],
+        pclass=Particle,
+    )
+    pset.execute(
+        runtime=np.timedelta64(10, "m"),
+        dt=np.timedelta64(60, "s"),
+        pyfunc=AdvectionRK4_3D,
+    )
+
+
 @pytest.mark.xfail(reason="Output file not implemented yet")
 def test_uxstommelgyre_pset_execute_output():
     ds = datasets_unstructured["stommel_gyre_delaunay"]