Merge pull request #2363 from Parcels-code/support_1d_fields

Support for 1D fields without lon and/or lat

Author: erikvansebille

src/parcels/_core/field.py

@@ -206,6 +206,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)
 
@@ -289,6 +292,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

@@ -84,10 +84,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

@@ -6,6 +6,7 @@
 import xarray as xr
 from numpy.testing import assert_allclose
 
+from parcels import Field
 from parcels._core.index_search import (
     LEFT_OUT_OF_BOUNDS,
     RIGHT_OUT_OF_BOUNDS,
@@ -16,6 +17,7 @@
     _transpose_xfield_data_to_tzyx,
 )
 from parcels._datasets.structured.generic import X, Y, Z, datasets
+from parcels.interpolators import XLinear
 from tests import utils
 
 GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"])
@@ -134,6 +136,30 @@ def test_invalid_depth():
         XGrid.from_dataset(ds)
 
 
+def test_vertical1D_field():
+    nz = 11
+    ds = xr.Dataset(
+        {"z1d": (["depth"], np.linspace(0, 10, nz))},
+        coords={"depth": (["depth"], np.linspace(0, 1, nz), {"axis": "Z"})},
+    )
+    grid = XGrid.from_dataset(ds)
+    field = Field("z1d", ds["z1d"], grid, XLinear)
+
+    assert field.eval(np.timedelta64(0, "s"), 0.45, 0, 0) == 4.5
+
+
+def test_time1D_field():
+    timerange = xr.date_range("2000-01-01", "2000-01-20")
+    ds = xr.Dataset(
+        {"t1d": (["time"], np.arange(0, len(timerange)))},
+        coords={"time": (["time"], timerange, {"axis": "T"})},
+    )
+    grid = XGrid.from_dataset(ds)
+    field = Field("t1d", ds["t1d"], grid, XLinear)
+
+    assert field.eval(np.datetime64("2000-01-10T12:00:00"), -20, 5, 6) == 9.5
+
+
 @pytest.mark.parametrize(
     "ds",
     [