Vertical Coordinate Interpolators

docs/user_guide/index.md

@@ -56,6 +56,11 @@ group_convert_concepts.md
 :caption: Work with FieldSets
 :name: work-with-fieldsets
 :titlesonly:
+examples/explanation_grids.ipynb
+examples/tutorial_time_varying_z_coordinates.ipynb
+examples/tutorial_velocityconversion.ipynb
+examples/tutorial_nestedgrids.ipynb
+examples/tutorial_manipulating_field_data.ipynb
 group_fieldsets.md
 ```
 

src/parcels/_core/particleset.py

@@ -86,7 +86,7 @@ def __init__(
             minz = None
             for field in self.fieldset.fields.values():
                 has_time_varying_z = (isinstance(field.grid, UxGrid) and field.grid.z.ndim == 3) or (
-                    isinstance(field.grid, XGrid) and "Z" in field.grid.axes and field.grid._ds["depth"].ndim == 4
+                    isinstance(field.grid, XGrid) and field.grid.depth.ndim == 4
                 )
                 if has_time_varying_z:
                     raise ValueError(

src/parcels/_core/uxgrid.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
-import warnings
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-import xarray as xr
 from dask import is_dask_collection
 
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, _search_1d_columns, uxgrid_point_in_cell
 from parcels._core.mesh import BaseMesh, TMesh
-from parcels._core.warnings import FieldSetWarning
+from parcels.interpolators._base import VerticalCoordinateInterpolator
+from parcels.interpolators._uxinterpolators import UxVerticalCoordinateConstantFace, UxVerticalCoordinateLinearNode
 
 if TYPE_CHECKING:
     import uxarray as ux
@@ -35,7 +34,7 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh: TMesh) -> None:
         z : ux.UxDataArray
             Vertical coordinates (depths) of the layer interface heights (not the mid-layer depths). Either a 1D array,
             constant in time and space (flat bottom topography, no moving vertical grid), or a 3D array with dims
-            ("time", "zf", "n_node"), varying in time and across the mesh nodes.
+            ("time", "zf", "n_node") or ("time", "zf", "n_face"), varying in time and across the mesh nodes or faces.
         mesh : FlatMesh or SphericalMesh
             The mesh of the grid, e.g. ``FlatMesh()`` or ``SphericalMesh(radius=...)``.
         """
@@ -50,24 +49,30 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh: TMesh) -> None:
             raise TypeError("z must be an instance of ux.UxDataArray")
         if z.ndim not in (1, 3):
             raise ValueError(f"z must be a 1D or 3D array of vertical coordinates, got {z.ndim}D")
-        if z.ndim == 3 and z.dims != ("time", "zf", "n_node"):
-            raise ValueError(f"A 3D z must have dims ('time', 'zf', 'n_node'), got {z.dims}")
-        if z.ndim == 3:
-            warnings.warn(
-                "Time-varying (3D) z coordinates are experimental and may cause significant memory overhead that "
-                f"leads to OOM errors. This z coordinate has sizes {dict(z.sizes)} ({z.nbytes / 1e9:.3g} GB). "
-                "Assumptions: z is defined at the layer interfaces ('zf') on the mesh nodes ('n_node') and is strictly "
-                "increasing along 'zf'; each particle's z column is interpolated barycentrically from its face's "
-                "nodes and linearly in time between z snapshots.",
-                FieldSetWarning,
-                stacklevel=4,
-            )
+
         self.z = z
         self.mesh = mesh
         self._spatialhash = None
+        self._z_interp_method = None
+        if z.ndim == 3:
+            self._init_time_varying_z()
         if is_dask_collection(self.uxgrid.face_node_connectivity.data):
             self.uxgrid.face_node_connectivity.load()
 
+    def _init_time_varying_z(self) -> None:
+        """Checks the dims of a time-varying (3D) z and assigns its default z_interp_method."""
+        z = self.z
+
+        if z.dims not in (("time", "zf", "n_node"), ("time", "zf", "n_face")):
+            raise ValueError(
+                f"A 3D z must have dims ('time', 'zf', 'n_node') or ('time', 'zf', 'n_face'), got {z.dims}"
+            )
+
+        if z.dims[-1] == "n_node":
+            self.z_interp_method = UxVerticalCoordinateLinearNode()
+        else:
+            self.z_interp_method = UxVerticalCoordinateConstantFace()
+
     @property
     def depth(self):
         """
@@ -76,16 +81,31 @@ def depth(self):
         Included for compatibility with v3 codebase. May be removed in future.
         TODO v4: Evaluate
         """
-        try:
-            _ = self.z.values
-        except KeyError:
-            return np.zeros(1)
+        # Load time varying z lazily
+        if self.z.ndim == 3:
+            return self.z
         return self.z.values
 
     @property
     def axes(self) -> list[_UXGRID_AXES]:
         return ["Z", "FACE"]
 
+    @property
+    def z_interp_method(self) -> VerticalCoordinateInterpolator | None:
+        return self._z_interp_method
+
+    @z_interp_method.setter
+    def z_interp_method(self, value: VerticalCoordinateInterpolator):
+        if not isinstance(value, VerticalCoordinateInterpolator):
+            raise ValueError(f"z_interp_method must be a `VerticalCoordinateInterpolator` object. Got {type(value)=!r}")
+        if self.depth.ndim != 3:
+            raise ValueError("z_interp_method can only be set on a grid with a time-varying (3D) z")
+        if isinstance(value, UxVerticalCoordinateLinearNode) and self.z.dims[-1] != "n_node":
+            raise ValueError(f"UxVerticalCoordinateLinearNode needs z on the mesh nodes, got dims {self.z.dims}")
+        if isinstance(value, UxVerticalCoordinateConstantFace) and self.z.dims[-1] != "n_face":
+            raise ValueError(f"UxVerticalCoordinateConstantFace needs z on the mesh faces, got dims {self.z.dims}")
+        self._z_interp_method = value
+
     def get_axis_dim(self, axis: _UXGRID_AXES) -> int:
         if axis not in self.axes:
             raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}")
@@ -163,32 +183,16 @@ def search(self, z, y, x, ei=None, ti=None, tau=None, tol=1e-6):
         return {"Z": {"index": zi, "bcoord": zeta}, "FACE": {"index": fi, "bcoord": coords}}
 
     def _search_time_varying_z(self, z, fi, coords, ti, tau):
-        """Vertical search of each particle in the z column interpolated barycentrically from the nodes of its face,
-        using temporal linear interpolation between the snapshots ti and ti + 1.
-        """
+        """Vertical search of each particle in the z column built by the grid's z_interp_method."""
         found = fi >= 0
-        cols_ti = self.z.isel(
-            time=xr.DataArray(np.broadcast_to(ti, fi.shape)[found], dims="points"),
-            n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
-            ignore_grid=True,
-        ).compute()
-
-        if self.z.shape[0] == 1:
-            node_columns = cols_ti
-        else:
-            cols_tnext = self.z.isel(
-                time=xr.DataArray(np.broadcast_to(ti + 1, fi.shape)[found], dims="points"),
-                n_node=xr.DataArray(self.uxgrid.face_node_connectivity[fi[found], :].values, dims=("points", "nodes")),
-                ignore_grid=True,
-            ).compute()
-
-            tau_points = xr.DataArray(np.broadcast_to(tau, fi.shape)[found], dims="points")
-            node_columns = (1 - tau_points) * cols_ti + tau_points * cols_tnext
+        grid_positions = {
+            "T": {"index": np.broadcast_to(ti, fi.shape)[found], "bcoord": np.broadcast_to(tau, fi.shape)[found]},
+            "FACE": {"index": fi[found], "bcoord": coords[found]},
+        }
 
-        bcoords = xr.DataArray(coords[found], dims=("points", "nodes"))
         # Particles outside the mesh are given a NaN z column for vertical searching
         columns = np.full((fi.size, self.z.sizes["zf"]), np.nan)
-        columns[found] = xr.dot(node_columns, bcoords, dim="nodes").values
+        columns[found] = self.z_interp_method.columns(grid_positions, self)
         zi, zeta = _search_1d_columns(columns, z)
 
         # Particles outside the mesh are given a 0 vertical index and a NaN vertical barycentric coordinate

src/parcels/_core/xgrid.py

@@ -1,4 +1,3 @@
-import warnings
 from collections.abc import Hashable, Sequence
 from functools import cached_property
 from typing import Literal, cast
@@ -13,9 +12,9 @@
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import _search_1d_array, _search_1d_columns, _search_indices_curvilinear_2d
 from parcels._core.mesh import BaseMesh
-from parcels._core.warnings import FieldSetWarning
 from parcels._sgrid.accessor import _get_dim_to_axis_mapping, get_dim_position
-from parcels.interpolators._xinterpolators import _get_offsets_dictionary
+from parcels.interpolators._base import VerticalCoordinateInterpolator
+from parcels.interpolators._xinterpolators import XVerticalCoordinateConstant, XVerticalCoordinateLinear
 
 _FIELD_DATA_ORDERING: Sequence[ptyping.XgcmAxisDirection] = "TZYX"
 _XGRID_AXES_ORDERING: Sequence[ptyping.XgridAxis] = "ZYX"
@@ -129,6 +128,7 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh):
         self._ds = model_data
         self.mesh = mesh
         self._spatialhash = None
+        self._z_interp_method = None
         ds = model_data
 
         axes = self.axes
@@ -149,34 +149,56 @@ def __init__(self, model_data: xr.Dataset, mesh: ptyping.TMesh):
             if depth.ndim not in (1, 4):
                 raise ValueError(f"depth must be a 1D or 4D array of vertical coordinates, got {depth.ndim}D")
             if depth.ndim == 4:
-                assert self.sgrid_metadata.vertical_dimensions is not None
-                vertical_dim = self.sgrid_metadata.vertical_dimensions[0].node
-                fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
-                expected_dims = ("time", vertical_dim, fnp_y.face, fnp_x.face)
-                if depth.dims != expected_dims:
-                    raise ValueError(f"A 4D depth must have dims {expected_dims}, got {depth.dims}")
-                warnings.warn(
-                    "Time-varying (4D) z coordinates are experimental and may cause significant memory overhead that "
-                    f"leads to OOM errors. This z coordinate has sizes {dict(depth.sizes)} ({depth.nbytes / 1e9:.3g} GB). "
-                    f"Assumptions: z has dims {expected_dims}, i.e. it is "
-                    f"defined at the layer interfaces on the cell centres, and is strictly increasing along "
-                    f"{vertical_dim!r}; each particle uses the z column at the centre of its grid cell, so z is constant "
-                    "within a cell and jumps between cells, it is linearly interpolated in time between z snapshots.",
-                    FieldSetWarning,
-                    stacklevel=5,
-                )
+                self._init_time_varying_depth()
             else:
                 assert_valid_depth(depth)
 
         self._ds = ds
 
+    def _init_time_varying_depth(self) -> None:
+        """Checks the dims of a time-varying (4D) depth and assigns its default z_interp_method."""
+        depth = self.depth
+
+        assert self.sgrid_metadata.vertical_dimensions is not None
+        vertical_dim = self.sgrid_metadata.vertical_dimensions[0].node
+
+        fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
+        face_registered_dims = ("time", vertical_dim, fnp_y.face, fnp_x.face)
+        node_registered_dims = ("time", vertical_dim, fnp_y.node, fnp_x.node)
+
+        if depth.dims not in (face_registered_dims, node_registered_dims):
+            raise ValueError(
+                f"A 4D depth must have dims {face_registered_dims} (cell centres) or {node_registered_dims} "
+                f"(cell corners), got {depth.dims}"
+            )
+
+        self.z_interp_method = XVerticalCoordinateLinear()
+
     # def __repr__(self):
     #     return xgrid_repr(self)
 
     @property
     def axes(self) -> list[ptyping.XgridAxis]:
         return _get_xgrid_axes(self.sgrid_metadata, set(self._ds.dims))
 
+    @property
+    def z_interp_method(self) -> VerticalCoordinateInterpolator | None:
+        return self._z_interp_method
+
+    @z_interp_method.setter
+    def z_interp_method(self, value: VerticalCoordinateInterpolator):
+        if not isinstance(value, VerticalCoordinateInterpolator):
+            raise ValueError(f"z_interp_method must be a `VerticalCoordinateInterpolator` object. Got {type(value)=!r}")
+        if self.depth.ndim != 4:
+            raise ValueError("z_interp_method can only be set on a grid with a time-varying (4D) depth")
+        depth_dims = self.depth.dims
+        if (
+            isinstance(value, XVerticalCoordinateConstant)
+            and get_dim_position(self.sgrid_metadata, depth_dims[-1]) != "face"
+        ):
+            raise ValueError(f"XVerticalCoordinateConstant needs depth on the cell centres, got dims {depth_dims}")
+        self._z_interp_method = value
+
     @property
     def lon(self):
         """
@@ -219,6 +241,9 @@ def depth(self):
         """
         if "Z" not in self.axes:
             return np.zeros(1)
+        # Load time varying z lazily
+        if self._ds["depth"].ndim == 4:
+            return self._ds["depth"]
         return self._ds["depth"].values
 
     @property
@@ -379,7 +404,7 @@ def search(self, z, y, x, ei=None, ti=None, tau=None):
             raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
         if "Z" in self.axes and ds.depth.ndim == 4:
-            zi, zeta = self._search_time_varying_depth(z, yi, xi, ti, tau)
+            zi, zeta = self._search_time_varying_depth(z, yi, eta, xi, xsi, ti, tau)
         elif "Z" in self.axes:
             zi, zeta = _search_1d_array(ds.depth.values, z)
         else:
@@ -391,35 +416,18 @@ def search(self, z, y, x, ei=None, ti=None, tau=None):
             "X": {"index": xi, "bcoord": xsi},
         }
 
-    def _search_time_varying_depth(self, z, yi, xi, ti, tau):
-        """Vertical search of each particle in the z column at the centre of its cell using temporal linear
-        interpolation between the snapshots ti and ti + 1.
-        """
+    def _search_time_varying_depth(self, z, yi, eta, xi, xsi, ti, tau):
+        """Vertical search of each particle in the z column built by the grid's z_interp_method."""
         found = (yi >= 0) & (xi >= 0)
-        offsets = _get_offsets_dictionary(self)
-        yi_centre = yi[found] + offsets["Y"]
-        xi_centre = xi[found] + offsets["X"]
-
-        depth = self._ds["depth"]
-        fnp_x, fnp_y = self.sgrid_metadata.face_dimensions
-        cell_centres = {
-            fnp_y.face: xr.DataArray(yi_centre, dims="points"),
-            fnp_x.face: xr.DataArray(xi_centre, dims="points"),
+        grid_positions = {
+            "T": {"index": np.broadcast_to(ti, yi.shape)[found], "bcoord": np.broadcast_to(tau, yi.shape)[found]},
+            "Y": {"index": yi[found], "bcoord": eta[found]},
+            "X": {"index": xi[found], "bcoord": xsi[found]},
         }
 
-        ti_found = np.broadcast_to(ti, yi.shape)[found]
-        columns_ti = depth.isel({"time": xr.DataArray(ti_found, dims="points"), **cell_centres})
-
-        if depth.sizes["time"] == 1:
-            cell_columns = columns_ti
-        else:
-            columns_tnext = depth.isel({"time": xr.DataArray(ti_found + 1, dims="points"), **cell_centres})
-            tau_found = xr.DataArray(np.broadcast_to(tau, yi.shape)[found], dims="points")
-            cell_columns = (1 - tau_found) * columns_ti + tau_found * columns_tnext
-
         # Particles outside the grid are given a NaN z column for vertical searching
-        columns = np.full((yi.size, cell_columns.shape[1]), np.nan)
-        columns[found] = cell_columns.values
+        columns = np.full((yi.size, self.depth.shape[1]), np.nan)
+        columns[found] = self.z_interp_method.columns(grid_positions, self)
         zi, zeta = _search_1d_columns(columns, z)
 
         # Particles outside the grid are given a 0 vertical index and a NaN vertical barycentric coordinate

src/parcels/interpolators/__init__.py

@@ -6,6 +6,8 @@
     UxConstantFaceLinearZF,
     UxLinearNodeConstantZC,
     UxLinearNodeLinearZF,
+    UxVerticalCoordinateConstantFace,
+    UxVerticalCoordinateLinearNode,
 )
 from ._xinterpolators import (
     CGrid_Tracer,
@@ -17,6 +19,9 @@
     XLinearInvdistLandTracer,
     XNearest,
     XPartialslip,
+    XVerticalCoordinateConstant,
+    XVerticalCoordinateLinear,
+    XVerticalCoordinateNearest,
 )
 
 __all__ = [  # noqa: RUF022
@@ -30,6 +35,9 @@
     "XLinear_Velocity",
     "XNearest",
     "XPartialslip",
+    "XVerticalCoordinateConstant",
+    "XVerticalCoordinateLinear",
+    "XVerticalCoordinateNearest",
     # uxinterpolators
     "Ux_Velocity",
     "UxConstantFaceConstantZC",
@@ -38,4 +46,6 @@
     "UxConstantEdgeLinearZF",
     "UxLinearNodeConstantZC",
     "UxLinearNodeLinearZF",
+    "UxVerticalCoordinateConstantFace",
+    "UxVerticalCoordinateLinearNode",
 ]