Error when number of vertical dimensions is not 1 or 4

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)

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)
 

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

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

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