From a02cfe0fd412f13a8bb833b07c6784adce127935 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 26 Oct 2026 10:47:06 +0100 Subject: [PATCH 1/6] Moving check on number of vertical dimensions from index_search to fieldset creation --- src/parcels/_core/index_search.py | 2 -- src/parcels/_core/model.py | 12 ++++++++++++ tests/test_index_search.py | 8 ++++---- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py index 095a43b12..bf44debb9 100644 --- a/src/parcels/_core/index_search.py +++ b/src/parcels/_core/index_search.py @@ -42,8 +42,6 @@ def _search_1d_array( Barycentric coordinate. """ # TODO v4: We probably rework this to deal with 0D arrays before this point (as we already know field dimensionality) - if arr.ndim > 1: - raise ValueError(f"Expected 1D array for {arr}, got {arr.ndim}D array") if len(arr) < 2: return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x) index = np.clip(np.searchsorted(arr, x, side="left") - 1, 0, len(arr) - 2) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index cb8594cff..2041dd143 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -257,6 +257,7 @@ def from_sgrid_conventions( # ds["lon"] = ds[node_dimensions[0]] # ds["lat"] = ds[node_dimensions[1]] + assert_1D_or4D_vertical_dimensions(ds) assert_metadata_ds_consistency(ds, ds.sgrid.metadata) vector_fields = resolve_vector_fields(ds, vector_fields) @@ -304,6 +305,15 @@ def assert_valid_vector_fields(ds: xr.Dataset, vector_fields: ptyping.VectorFiel return +def assert_1D_or4D_vertical_dimensions(ds: xr.Dataset) -> None: + # Assert that vertical dimension has either length 1 or is full 4D + if (vertical_dims := ds.cf.axes.get("Z")) is not None: + for vertical_dim in vertical_dims: + if 1 < ds[vertical_dim].ndim < 4: + raise ValueError(f"Vertical dimension {vertical_dim} must have length 1 or be part of a full 4D field.") + return + + def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: ptyping.VectorFields) -> None: for components in vector_fields.values(): for c in components: @@ -400,6 +410,8 @@ def from_ugrid_conventions( if mesh is None: mesh = _infer_mesh_from_ugrid_data(ds) + assert_1D_or4D_vertical_dimensions(ds) + grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh=mesh) ds = _discover_ux_U_and_V(ds) diff --git a/tests/test_index_search.py b/tests/test_index_search.py index a1a56c614..223d206c5 100644 --- a/tests/test_index_search.py +++ b/tests/test_index_search.py @@ -3,11 +3,11 @@ from parcels._core.fieldset import FieldSet from parcels._core.index_search import ( - _search_1d_array, _search_indices_curvilinear_2d, curvilinear_point_in_cell, uxgrid_point_in_cell, ) +from parcels._datasets.structured.generated import simple_UV_dataset from parcels._datasets.structured.generic import datasets from tests.utils import create_lonlat_grid, sample_points_inside_faces @@ -20,10 +20,10 @@ def field_cone(): def test_multi_array_errors_in_search_1d_array(): - arr = np.array([[1, 2, 3]]) - x = np.array([[1], [2]]) + ds = simple_UV_dataset() + ds["depth"] = (("depth_a", "depth_b"), np.array([[1, 2]])) with pytest.raises(ValueError): - _search_1d_array(arr, x) + FieldSet.from_sgrid_conventions(ds) def test_grid_indexing_fpoints(field_cone): From a423c45f8b15a96aa5bb69d37d816e7896360e7f Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 26 Oct 2026 12:40:27 +0100 Subject: [PATCH 2/6] Moving and renaming unit test --- tests/test_fieldset.py | 9 ++++++++- tests/test_index_search.py | 8 -------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 10f7c7045..cfdc74bf9 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -17,7 +17,7 @@ from parcels._core.mesh import SphericalMesh from parcels._core.model import _default_vector_field_components, _infer_mesh_from_sgrid_data from parcels._core.utils.time import TimeInterval -from parcels._datasets.structured.generated import simple_UV_sigma_dataset +from parcels._datasets.structured.generated import simple_UV_dataset, simple_UV_sigma_dataset from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.structured.generic import datasets_sgrid from parcels._datasets.structured.strategies import reduced_dims_dataset, sgrid_datasets_varying_dims @@ -462,6 +462,13 @@ def test_fieldset_add_context_values(): assert fset.context["c2"] == 2.0 +def test_fieldset_error_on_unsupported_depth_dimensionality(): + ds = simple_UV_dataset() + ds["depth"] = (("depth_a", "depth_b"), np.array([[1, 2]])) + with pytest.raises(ValueError): + FieldSet.from_sgrid_conventions(ds) + + def _field_name_of(ds): (name,) = [name for name in ds.data_vars if ds[name].attrs.get("cf_role") != "grid_topology"] return name diff --git a/tests/test_index_search.py b/tests/test_index_search.py index 223d206c5..1785fb402 100644 --- a/tests/test_index_search.py +++ b/tests/test_index_search.py @@ -7,7 +7,6 @@ curvilinear_point_in_cell, uxgrid_point_in_cell, ) -from parcels._datasets.structured.generated import simple_UV_dataset from parcels._datasets.structured.generic import datasets from tests.utils import create_lonlat_grid, sample_points_inside_faces @@ -19,13 +18,6 @@ def field_cone(): return fieldset.data_g -def test_multi_array_errors_in_search_1d_array(): - ds = simple_UV_dataset() - ds["depth"] = (("depth_a", "depth_b"), np.array([[1, 2]])) - with pytest.raises(ValueError): - FieldSet.from_sgrid_conventions(ds) - - def test_grid_indexing_fpoints(field_cone): grid = field_cone.grid From 045ae468ad184e38ea4b781134f267fce0b6c327 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 26 Oct 2026 20:16:59 +0100 Subject: [PATCH 3/6] Generalising vertical dimensionality check for uxgrid And fixing unit test to properly set Z axis --- src/parcels/_core/model.py | 16 ++++++++++------ src/parcels/_datasets/unstructured/generated.py | 2 ++ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 2041dd143..d9149f7e3 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Self import cf_xarray # noqa: F401 +import numpy as np import xarray as xr from dask import is_dask_collection @@ -257,7 +258,7 @@ def from_sgrid_conventions( # ds["lon"] = ds[node_dimensions[0]] # ds["lat"] = ds[node_dimensions[1]] - assert_1D_or4D_vertical_dimensions(ds) + assert_vertical_dimensionality(ds, allowed_dimensionality=[1, 4]) assert_metadata_ds_consistency(ds, ds.sgrid.metadata) vector_fields = resolve_vector_fields(ds, vector_fields) @@ -305,12 +306,15 @@ def assert_valid_vector_fields(ds: xr.Dataset, vector_fields: ptyping.VectorFiel return -def assert_1D_or4D_vertical_dimensions(ds: xr.Dataset) -> None: - # Assert that vertical dimension has either length 1 or is full 4D +def assert_vertical_dimensionality(ds: xr.DataArray, allowed_dimensionality: list[int]): + # Assert that vertical dimensions have allowed dimensionality. if (vertical_dims := ds.cf.axes.get("Z")) is not None: for vertical_dim in vertical_dims: - if 1 < ds[vertical_dim].ndim < 4: - raise ValueError(f"Vertical dimension {vertical_dim} must have length 1 or be part of a full 4D field.") + ndim = np.atleast_1d(ds[vertical_dim]).ndim # cast 0D to 1D + if ndim not in allowed_dimensionality: + raise ValueError( + f"Vertical dimension {vertical_dim} has dimensionality {ndim}, but allowed dimensionality is {allowed_dimensionality}." + ) return @@ -410,7 +414,7 @@ def from_ugrid_conventions( if mesh is None: mesh = _infer_mesh_from_ugrid_data(ds) - assert_1D_or4D_vertical_dimensions(ds) + assert_vertical_dimensionality(ds, allowed_dimensionality=[1, 3]) grid = UxGrid(ds.uxgrid, z=ds.coords["zf"], mesh=mesh) ds = _discover_ux_U_and_V(ds) diff --git a/src/parcels/_datasets/unstructured/generated.py b/src/parcels/_datasets/unstructured/generated.py index 2a514cf5c..38a8b6bee 100644 --- a/src/parcels/_datasets/unstructured/generated.py +++ b/src/parcels/_datasets/unstructured/generated.py @@ -403,4 +403,6 @@ def sigma_coordinate_lattice_dataset(nx, x_range, y_range, nz, bottom_depth, eta ds = ux.UxDataset(xr.Dataset(velocities, coords={"time": time, **vertical_coords}), uxgrid=uxgrid) ds.uxgrid.node_lon.attrs["units"] = "m" ds.uxgrid.node_lat.attrs["units"] = "m" + ds.variables["zf"].attrs["axis"] = "Z" + ds.variables["zc"].attrs["axis"] = "Z" return ds From 8fb60953942306272c48b5fa40162dde2ff04f36 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Mon, 26 Oct 2026 20:20:00 +0100 Subject: [PATCH 4/6] Fixing typing error --- src/parcels/_core/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index d9149f7e3..3c6aabf8c 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -306,7 +306,7 @@ def assert_valid_vector_fields(ds: xr.Dataset, vector_fields: ptyping.VectorFiel return -def assert_vertical_dimensionality(ds: xr.DataArray, allowed_dimensionality: list[int]): +def assert_vertical_dimensionality(ds: xr.Dataset, allowed_dimensionality: list[int]): # Assert that vertical dimensions have allowed dimensionality. if (vertical_dims := ds.cf.axes.get("Z")) is not None: for vertical_dim in vertical_dims: From 25312990cd66150d49fcb5e5096105215bf61f61 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 27 Oct 2026 07:27:48 +0100 Subject: [PATCH 5/6] Add uxgrid test for unsupported depth dimensionality --- tests/test_fieldset.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index cfdc74bf9..84cc462e4 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -21,6 +21,7 @@ from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.structured.generic import datasets_sgrid from parcels._datasets.structured.strategies import reduced_dims_dataset, sgrid_datasets_varying_dims +from parcels._datasets.unstructured.generated import simple_small_delaunay from parcels._datasets.unstructured.generic import datasets as datasets_unstructured from parcels._sgrid.accessor import SGridDatasetInconsistency @@ -462,13 +463,20 @@ def test_fieldset_add_context_values(): assert fset.context["c2"] == 2.0 -def test_fieldset_error_on_unsupported_depth_dimensionality(): +def test_xgrid_fieldset_error_on_unsupported_depth_dimensionality(): ds = simple_UV_dataset() ds["depth"] = (("depth_a", "depth_b"), np.array([[1, 2]])) with pytest.raises(ValueError): FieldSet.from_sgrid_conventions(ds) +def test_ugrid_fieldset_error_on_unsupported_depth_dimensionality(): + ds = simple_small_delaunay() + ds = ds.drop_vars("zf").assign_coords(zf=(("depth_a", "depth_b"), np.array([[1, 2]]))) + with pytest.raises(ValueError, match="z must be a 1D or 3D array"): + FieldSet.from_ugrid_conventions(ds) + + def _field_name_of(ds): (name,) = [name for name in ds.data_vars if ds[name].attrs.get("cf_role") != "grid_topology"] return name From d34c367dab97639c9fe555567bd154de4d2b7ee7 Mon Sep 17 00:00:00 2001 From: Erik van Sebille Date: Tue, 27 Oct 2026 07:30:45 +0100 Subject: [PATCH 6/6] Avoiding dask computation From the LLM: np.atleast_1d(DataArray) converts to a numpy array, which computes lazy data. .ndim is metadata on the DataArray and costs nothing --- src/parcels/_core/model.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 3c6aabf8c..3bb4272a4 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Any, Self import cf_xarray # noqa: F401 -import numpy as np import xarray as xr from dask import is_dask_collection @@ -310,7 +309,7 @@ def assert_vertical_dimensionality(ds: xr.Dataset, allowed_dimensionality: list[ # Assert that vertical dimensions have allowed dimensionality. if (vertical_dims := ds.cf.axes.get("Z")) is not None: for vertical_dim in vertical_dims: - ndim = np.atleast_1d(ds[vertical_dim]).ndim # cast 0D to 1D + ndim = max(ds[vertical_dim].ndim, 1) # cast 0D to 1D; metadata only, so lazy data is not computed if ndim not in allowed_dimensionality: raise ValueError( f"Vertical dimension {vertical_dim} has dimensionality {ndim}, but allowed dimensionality is {allowed_dimensionality}."