Add windowarray ugrid tests (#2760)

* broke test_windowed_arrays_wraps_dask_but_not_numpy into structured and unstructured tests

* Parameterized struc/unstruc tests and added unstructured advection test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Wyatt Sieminski <wyatt_sieminski@brown.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Joe Schoonover <11430768+fluidnumericsJoe@users.noreply.github.com>

Author: 4 people

tests/test_windowed_array.py

@@ -8,6 +8,7 @@
 from parcels import FieldSet, ParticleSet
 from parcels._core._windowed_array import WindowedArray, maybe_windowed
 from parcels._datasets.structured.generated import simple_UV_dataset
+from parcels._datasets.unstructured.generic import _ux_constant_flow_face_centered_2D
 from parcels.kernels import AdvectionRK2
 
 
@@ -71,10 +72,17 @@ def test_windowed_isel_backward_clock_loads_once_and_evicts():
     assert max_cache <= 2  # only the bracketing levels resident
 
 
-def test_to_windowed_arrays_wraps_dask_but_not_numpy():
-    ds = simple_UV_dataset(mesh="flat")
-    fset_np = FieldSet.from_sgrid_conventions(ds, mesh="flat")
-    fset_dk = FieldSet.from_sgrid_conventions(ds.chunk({"time": 1}), mesh="flat")
+@pytest.mark.parametrize(
+    "fset_convention, ds",
+    [
+        (FieldSet.from_sgrid_conventions, simple_UV_dataset(mesh="flat")),
+        (FieldSet.from_ugrid_conventions, _ux_constant_flow_face_centered_2D()),
+    ],
+    ids=["structured", "unstructured"],
+)
+def test_windowed_arrays_wraps_dask_but_not_numpy(fset_convention: callable, ds: xr.Dataset):
+    fset_np = fset_convention(ds, mesh="flat")
+    fset_dk = fset_convention(ds.chunk({"time": 1}), mesh="flat")
 
     # construction is never windowing -- it is opt-in via the fieldset method
     assert not isinstance(fset_np.U.data, WindowedArray)
@@ -92,9 +100,16 @@ def test_to_windowed_arrays_wraps_dask_but_not_numpy():
     assert fset_dk.U.data.shape == fset_np.U.data.shape
 
 
-def test_to_windowed_arrays_is_idempotent_and_forwards_max_levels():
-    ds = simple_UV_dataset(mesh="flat")
-    fs = FieldSet.from_sgrid_conventions(ds.chunk({"time": 1}), mesh="flat")
+@pytest.mark.parametrize(
+    "fset_convention, ds",
+    [
+        (FieldSet.from_sgrid_conventions, simple_UV_dataset(mesh="flat")),
+        (FieldSet.from_ugrid_conventions, _ux_constant_flow_face_centered_2D()),
+    ],
+    ids=["structured", "unstructured"],
+)
+def test_to_windowed_arrays_is_idempotent_and_forwards_max_levels(fset_convention: callable, ds: xr.Dataset):
+    fs = fset_convention(ds.chunk({"time": 1}), mesh="flat")
 
     fs.to_windowed_arrays(max_levels=3)
     first = fs.U.data
@@ -141,10 +156,10 @@ def test_maybe_windowed_passthrough_for_non_time_leading():
 
 @pytest.mark.parametrize("mesh", ["flat", "spherical"])
 @pytest.mark.parametrize("dt_minutes", [15, -15], ids=["forward", "backward"])
-def test_dask_advection_matches_numpy(mesh, dt_minutes):
+def test_dask_advection_matches_numpy_on_structured_grids(mesh, dt_minutes):
     """An identical advection must give identical trajectories whether the field
     is numpy-backed or dask-backed (windowed) -- for both forward (dt > 0) and
-    backward (dt < 0) integration.
+    backward (dt < 0) integration on structured grids.
     """
     ds = simple_UV_dataset(mesh=mesh)
     ds["U"].data[:] = 1.0  # steady zonal flow -> in-bounds, deterministic
@@ -162,3 +177,27 @@ def run(chunked):
     x_dk, y_dk = run(True)
     np.testing.assert_allclose(x_dk, x_np, atol=1e-9)
     np.testing.assert_allclose(y_dk, y_np, atol=1e-9)
+
+
+@pytest.mark.parametrize("mesh", ["flat", "spherical"])
+@pytest.mark.parametrize("dt_minutes", [15, -15], ids=["forward", "backward"])
+def test_dask_advection_matches_numpy_on_unstructured_grids(mesh, dt_minutes):
+    """An identical advection must give identical trajectories whether the field
+    is numpy-backed or dask-backed (windowed) -- for both forward (dt > 0) and
+    backward (dt < 0) integration on unstructured grids.
+    """
+    ds = _ux_constant_flow_face_centered_2D()
+
+    def run(chunked):
+        d = ds.chunk({"time": 1}) if chunked else ds
+        fs = FieldSet.from_ugrid_conventions(d, mesh=mesh)
+        if chunked:
+            fs.to_windowed_arrays()
+        pset = ParticleSet(fs, x=10 * np.ones(10), y=np.linspace(5, 15, 10))
+        pset.execute(AdvectionRK2, runtime=3600, dt=np.timedelta64(dt_minutes, "m"))
+        return np.array(pset.x), np.array(pset.y)
+
+    x_np, y_np = run(False)
+    x_dk, y_dk = run(True)
+    np.testing.assert_allclose(x_dk, x_np, atol=1e-9)
+    np.testing.assert_allclose(y_dk, y_np, atol=1e-9)