Add XGrid.get_axis_dim_mapping

Author: VeckoTheGecko

parcels/application_kernels/interpolation.py

@@ -33,15 +33,16 @@ def XTriCurviLinear(
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
     data = field.data
+    axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
 
     return (
         (
-            (1 - xsi) * (1 - eta) * data.isel(YG=yi, XG=xi)
-            + xsi * (1 - eta) * data.isel(YG=yi, XG=xi + 1)
-            + xsi * eta * data.isel(YG=yi + 1, XG=xi + 1)
-            + (1 - xsi) * eta * data.isel(YG=yi + 1, XG=xi)
+            (1 - xsi) * (1 - eta) * data.isel({axis_dim["Y"]: yi, axis_dim["X"]: xi})
+            + xsi * (1 - eta) * data.isel({axis_dim["Y"]: yi, axis_dim["X"]: xi + 1})
+            + xsi * eta * data.isel({axis_dim["Y"]: yi + 1, axis_dim["X"]: xi + 1})
+            + (1 - xsi) * eta * data.isel({axis_dim["Y"]: yi + 1, axis_dim["X"]: xi})
         )
-        .interp(time=t, ZG=zi + zeta)
+        .interp(time=t, **{axis_dim["Z"]: zi + zeta})
         .values
     )
 
@@ -57,10 +58,13 @@ def XTriRectiLinear(
     x: np.float32 | np.float64,
 ):
     """Trilinear interpolation on a rectilinear grid."""
+    axis_dim = field.grid.get_axis_dim_mapping(field.data.dims)
+
     xi, xsi = position["X"]
     yi, eta = position["Y"]
     zi, zeta = position["Z"]
-    return field.data.interp(time=t, ZG=zi + zeta, YG=yi + eta, XG=xi + xsi).values
+    kwargs = {axis_dim["X"]: xi + xsi, axis_dim["Y"]: yi + eta, axis_dim["Z"]: zi + zeta}
+    return field.data.interp(time=t, **kwargs).values
 
 
 def UXPiecewiseConstantFace(

parcels/xgrid.py

@@ -185,6 +185,38 @@ def search(self, z, y, x, ei=None):
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
+    def get_axis_dim_mapping(self, dims: list[str]) -> dict[_XGRID_AXES, str]:
+        """
+        Maps xarray dimension names to their corresponding axis (X, Y, Z).
+
+        WARNING: This API is unstable and subject to change in future versions.
+
+        Parameters
+        ----------
+        dims : list[str]
+            List of xarray dimension names
+
+        Returns
+        -------
+        dict[_XGRID_AXES, str]
+            Dictionary mapping axes (X, Y, Z) to their corresponding dimension names
+
+        Examples
+        --------
+        >>> grid.get_axis_dim_mapping(['time', 'lat', 'lon'])
+        {'Y': 'lat', 'X': 'lon'}
+
+        Notes
+        -----
+        Only returns mappings for spatial axes (X, Y, Z) that are present in the grid.
+        """
+        result = {}
+        for dim in dims:
+            axis = get_axis_from_dim_name(self.xgcm_grid.axes, dim)
+            if axis in self.axes:  # Only include spatial axes (X, Y, Z)
+                result[cast(_XGRID_AXES, axis)] = dim
+        return result
+
 
 def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None:
     """For a given dimension name in a grid, returns the direction axis it is on."""