diff --git a/tests/test_windowed_array.py b/tests/test_windowed_array.py index 47a85c8eb..2e7fa963f 100644 --- a/tests/test_windowed_array.py +++ b/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)