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..3bb4272a4 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_vertical_dimensionality(ds, allowed_dimensionality=[1, 4]) assert_metadata_ds_consistency(ds, ds.sgrid.metadata) vector_fields = resolve_vector_fields(ds, vector_fields) @@ -304,6 +305,18 @@ def assert_valid_vector_fields(ds: xr.Dataset, vector_fields: ptyping.VectorFiel return +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: + 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}." + ) + 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 +413,8 @@ def from_ugrid_conventions( if mesh is None: mesh = _infer_mesh_from_ugrid_data(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 diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 10f7c7045..84cc462e4 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -17,10 +17,11 @@ 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 +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,6 +463,20 @@ def test_fieldset_add_context_values(): assert fset.context["c2"] == 2.0 +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 diff --git a/tests/test_index_search.py b/tests/test_index_search.py index a1a56c614..1785fb402 100644 --- a/tests/test_index_search.py +++ b/tests/test_index_search.py @@ -3,7 +3,6 @@ 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, @@ -19,13 +18,6 @@ def field_cone(): return fieldset.data_g -def test_multi_array_errors_in_search_1d_array(): - arr = np.array([[1, 2, 3]]) - x = np.array([[1], [2]]) - with pytest.raises(ValueError): - _search_1d_array(arr, x) - - def test_grid_indexing_fpoints(field_cone): grid = field_cone.grid