Fixing support for 1D depth grids (without lon or lat)

Author: erikvansebille

src/parcels/_core/field.py

@@ -217,6 +217,9 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
             _ei = None
         else:
             _ei = particles.ei[:, self.igrid]
+        z = np.atleast_1d(z)
+        y = np.atleast_1d(y)
+        x = np.atleast_1d(x)
 
         particle_positions, grid_positions = _get_positions(self, time, z, y, x, particles, _ei)
 
@@ -300,6 +303,9 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
             _ei = None
         else:
             _ei = particles.ei[:, self.igrid]
+        z = np.atleast_1d(z)
+        y = np.atleast_1d(y)
+        x = np.atleast_1d(x)
 
         particle_positions, grid_positions = _get_positions(self.U, time, z, y, x, particles, _ei)
 

src/parcels/_core/xgrid.py

@@ -289,22 +289,13 @@ def search(self, z, y, x, ei=None):
         else:
             zi, zeta = np.zeros(z.shape, dtype=int), np.zeros(z.shape, dtype=float)
 
-        if ds.lon.ndim == 1:
-            yi, eta = _search_1d_array(ds.lat.values, y)
-            xi, xsi = _search_1d_array(ds.lon.values, x)
-            return {
-                "Z": {"index": zi, "bcoord": zeta},
-                "Y": {"index": yi, "bcoord": eta},
-                "X": {"index": xi, "bcoord": xsi},
-            }
+        if "X" in self.axes and "Y" in self.axes and ds.lon.ndim == 2:
+            yi, xi = None, None
+            if ei is not None:
+                axis_indices = self.unravel_index(ei)
+                xi = axis_indices.get("X")
+                yi = axis_indices.get("Y")
 
-        yi, xi = None, None
-        if ei is not None:
-            axis_indices = self.unravel_index(ei)
-            xi = axis_indices.get("X")
-            yi = axis_indices.get("Y")
-
-        if ds.lon.ndim == 2:
             yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
             return {
@@ -313,7 +304,24 @@ def search(self, z, y, x, ei=None):
                 "X": {"index": xi, "bcoord": xsi},
             }
 
-        raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
+        if "X" in self.axes and ds.lon.ndim > 2:
+            raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
+
+        if "Y" in self.axes:
+            yi, eta = _search_1d_array(ds.lat.values, y)
+        else:
+            yi, eta = np.zeros(y.shape, dtype=int), np.zeros(y.shape, dtype=float)
+
+        if "X" in self.axes:
+            xi, xsi = _search_1d_array(ds.lon.values, x)
+        else:
+            xi, xsi = np.zeros(x.shape, dtype=int), np.zeros(x.shape, dtype=float)
+
+        return {
+            "Z": {"index": zi, "bcoord": zeta},
+            "Y": {"index": yi, "bcoord": eta},
+            "X": {"index": xi, "bcoord": xsi},
+        }
 
     @cached_property
     def _fpoint_info(self):

src/parcels/interpolators.py

@@ -83,10 +83,11 @@ def _get_corner_data_Agrid(
     xi = np.tile(np.array([xi, xi_1]).flatten(), lenT * lenZ * 2)
 
     # Create DataArrays for indexing
-    selection_dict = {
-        axis_dim["X"]: xr.DataArray(xi, dims=("points")),
-        axis_dim["Y"]: xr.DataArray(yi, dims=("points")),
-    }
+    selection_dict = {}
+    if "X" in axis_dim:
+        selection_dict[axis_dim["X"]] = xr.DataArray(xi, dims=("points"))
+    if "Y" in axis_dim:
+        selection_dict[axis_dim["Y"]] = xr.DataArray(yi, dims=("points"))
     if "Z" in axis_dim:
         selection_dict[axis_dim["Z"]] = xr.DataArray(zi, dims=("points"))
     if "time" in data.dims:

tests/test_xgrid.py

@@ -139,7 +139,10 @@ def test_invalid_depth():
 
 def test_vertical1D_field():
     nz = 11
-    ds = xr.Dataset({"z1d": (["depth"], np.linspace(0, 10, nz))}, coords={"depth": np.linspace(0, 1, nz)})
+    ds = xr.Dataset(
+        {"z1d": (["depth"], np.linspace(0, 10, nz))},
+        coords={"depth": (["depth"], np.linspace(0, 1, nz), {"axis": "Z"})},
+    )
     grid = XGrid(xgcm.Grid(ds, **_DEFAULT_XGCM_KWARGS))
     field = Field("z1d", ds["z1d"], grid)