Repository: OceanParcels/Parcels
Pull request #2082 (2025-07-11 14:59:45)

XGrid localization
<!-- Feel free to remove list items that are not relevant for your changes. -->

Currently we have a search method which returns a particle position relative to the F points. However, since our data can be defined on a staggered grid, it's important to "localize" this particle position to the grid for the array of interest. This mainly applies for C grids and when working with MITgcm and NEMO where their F points and C points are defined differently relative to each other (see diagram in docs or in #2037) .

This PR introduces this grid localization. This is really just the first draft to get feedback (code is a bit more messy than I would like - and there are no tests). This good localization will help with writing interpolators.

- [ ] Chose the correct base branch (`v4-dev` for v4 changes)
- [x] Fixes None
- [x] Added tests (not yet)
- [x] Added documentation


VeckoTheGecko: Below is a small testing script along with output. `data_c` with the -1 index seems a bit strange, but I guess it makes sense (if vertical positions are defined on the cell centers and the particle is at the surface then -1 with bcoord 0.5 is expected). Thoughts?
```python
# %%
import numpy as np
import xarray as xr
from pprint import pprint

from parcels import xgcm
from parcels._datasets.structured.generic import X, Y, Z
from parcels.xgrid import XGrid

T = 2
Z = Y = X = 3
TIME = xr.date_range("2000", "2001", T)

ds_mitgcm = xr.Dataset(
    {
        "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
        "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
        "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
    },
    coords={
        "XG": (
            ["XG"],
            np.arange(0, X),
            {"axis": "X", "c_grid_axis_shift": -0.5},
        ),
        "XC": (["XC"], np.arange(0, X) + 0.5, {"axis": "X"}),
        "YG": (
            ["YG"],
            np.arange(0, Y),
            {"axis": "Y", "c_grid_axis_shift": -0.5},
        ),
        "YC": (
            ["YC"],
            np.arange(0, Y) + 0.5,
            {"axis": "Y"},
        ),
        "ZG": (
            ["ZG"],
            np.arange(Z),
            {"axis": "Z", "c_grid_axis_shift": -0.5},
        ),
        "ZC": (
            ["ZC"],
            np.arange(Z) + 0.5,
            {"axis": "Z"},
        ),
        "lon": (["XG"], np.arange(0, X)),
        "lat": (["YG"], np.arange(0, Y)),
        "depth": (["ZG"], np.arange(Z)),
        "time": (["time"], TIME, {"axis": "T"}),
    },
)

ds_nemo = xr.Dataset(
    {
        "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
        "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
        "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
        "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
    },
    coords={
        "XG": (
            ["XG"],
            np.arange(0, X),
            {"axis": "X", "c_grid_axis_shift": 0.5},
        ),
        "XC": (["XC"], np.arange(0, X) - 0.5, {"axis": "X"}),
        "YG": (
            ["YG"],
            np.arange(0, Y),
            {"axis": "Y", "c_grid_axis_shift": 0.5},
        ),
        "YC": (
            ["YC"],
            np.arange(0, Y) - 0.5,
            {"axis": "Y"},
        ),
        "ZG": (
            ["ZG"],
            np.arange(Z),
            {"axis": "Z", "c_grid_axis_shift": -0.5},
        ),
        "ZC": (
            ["ZC"],
            np.arange(Z) + 0.5,
            {"axis": "Z"},
        ),
        "lon": (["XG"], np.arange(0, X)),
        "lat": (["YG"], np.arange(0, Y)),
        "depth": (["ZG"], np.arange(Z)),
        "time": (["time"], TIME, {"axis": "T"}),
    },
)


grid_mitgcm = XGrid(xgcm.Grid(ds_mitgcm, periodic=False))
grid_nemo = XGrid(xgcm.Grid(ds_nemo, periodic=False))

print("XGCM repr of MITgcm grid:")
print(grid_mitgcm.xgcm_grid)
print("\nXGCM repr of NEMO grid:")
print(grid_nemo.xgcm_grid)


# %%
def show_point_on_grid(grid, z, y, x):
    """Pretty printing of some info"""

    position = grid.search(z, y, x)
    print(f"Position wrt. fpoints (lon/lat grid):")
    pprint(position)

    for da in grid.xgcm_grid._ds.data_vars.values():
        local_position = grid.localize(position, da.dims)
        print(f"On {da.name=} with {da.dims=}, local position:")
        pprint(local_position)


print("----Working with MITgcm grid----")
show_point_on_grid(grid_mitgcm, 0, 0.8, 0.8)

print("\n----Working with NEMO grid----")
show_point_on_grid(grid_nemo, 0, 0.8, 0.8)
```

output:
```
XGCM repr of MITgcm grid:
<parcels.Grid>
Y Axis (not periodic, boundary=None):
  * center   YC --> left
  * left     YG --> center
X Axis (not periodic, boundary=None):
  * center   XC --> left
  * left     XG --> center
Z Axis (not periodic, boundary=None):
  * center   ZC --> left
  * left     ZG --> center
T Axis (not periodic, boundary=None):
  * center   time

XGCM repr of NEMO grid:
<parcels.Grid>
Y Axis (not periodic, boundary=None):
  * center   YC --> right
  * right    YG --> center
X Axis (not periodic, boundary=None):
  * center   XC --> right
  * right    XG --> center
Z Axis (not periodic, boundary=None):
  * center   ZC --> left
  * left     ZG --> center
T Axis (not periodic, boundary=None):
  * center   time
----Working with MITgcm grid----
Position wrt. fpoints (lon/lat grid):
{'X': (np.int64(0), np.float64(0.8)),
 'Y': (np.int64(0), np.float64(0.8)),
 'Z': (np.int64(0), np.float64(0.0))}
On da.name='data_g' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='data_c' with da.dims=('time', 'ZC', 'YC', 'XC'), local position:
{'XC': (np.int64(0), np.float64(0.30000000000000004)),
 'YC': (np.int64(0), np.float64(0.30000000000000004)),
 'ZC': (np.int64(-1), np.float64(0.5))}
On da.name='U (A grid)' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='V (A grid)' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='U (C grid)' with da.dims=('time', 'ZG', 'YC', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YC': (np.int64(0), np.float64(0.30000000000000004)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='V (C grid)' with da.dims=('time', 'ZG', 'YG', 'XC'), local position:
{'XC': (np.int64(0), np.float64(0.30000000000000004)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}

----Working with NEMO grid----
Position wrt. fpoints (lon/lat grid):
{'X': (np.int64(0), np.float64(0.8)),
 'Y': (np.int64(0), np.float64(0.8)),
 'Z': (np.int64(0), np.float64(0.0))}
On da.name='data_g' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='data_c' with da.dims=('time', 'ZC', 'YC', 'XC'), local position:
{'XC': (np.int64(1), np.float64(0.30000000000000004)),
 'YC': (np.int64(1), np.float64(0.30000000000000004)),
 'ZC': (np.int64(-1), np.float64(0.5))}
On da.name='U (A grid)' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='V (A grid)' with da.dims=('time', 'ZG', 'YG', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='U (C grid)' with da.dims=('time', 'ZG', 'YC', 'XG'), local position:
{'XG': (np.int64(0), np.float64(0.8)),
 'YC': (np.int64(1), np.float64(0.30000000000000004)),
 'ZG': (np.int64(0), np.float64(0.0))}
On da.name='V (C grid)' with da.dims=('time', 'ZG', 'YG', 'XC'), local position:
{'XC': (np.int64(1), np.float64(0.30000000000000004)),
 'YG': (np.int64(0), np.float64(0.8)),
 'ZG': (np.int64(0), np.float64(0.0))}

```
VeckoTheGecko: if we're comfortable with this approach, I'll do a self review to clear up some stuff (e.g., variable naming since it might be a bit confusing)
VeckoTheGecko: I realised that this localization step isn't actually needed.

<img width="786" height="458" alt="image" src="https://github.com/user-attachments/assets/df99bd3f-ee4f-41f3-b7fa-e6077a1e6660" />

If a point is defined at (in X) `index, bcoord = 0,0.8` wrt. the f-points, that will be at `0, 0.3` wrt. the C points. However, since the dual grid already has dimension coordinates that are offset by 0.5 (i.e., 0.5, 1.5 ...) we can just use xarray's interpolation functionality with `0.8` on this grid in order to get this point no matter the grid.

Still we need to build tooling around getting the dimension names that correspond to the axes (i.e, for `da.dims == ["XC", "YC"]` the grid needs to tell us `"X" -> "XC"` and `"Y" -> "YC"`

VeckoTheGecko: Also wrapped into this `numpydoc!=1.9.0` so that our docs build again on `v4-dev` (didn't bother with v3 since I assume they'll issue a fix before we need to update v3 docs again)

Patch:
diff --git a/environment.yml b/environment.yml
index 4b8694f0..2a470405 100644
--- a/environment.yml
+++ b/environment.yml
@@ -41,7 +41,7 @@ dependencies: #! Keep in sync with [tool.pixi.dependencies] in pyproject.toml
 
   # Docs
   - ipython
-  - numpydoc
+  - numpydoc!=1.9.0
   - nbsphinx
   - sphinx
   - pandoc
diff --git a/parcels/_datasets/structured/generic.py b/parcels/_datasets/structured/generic.py
index 2432f079..1b1cd7d8 100644
--- a/parcels/_datasets/structured/generic.py
+++ b/parcels/_datasets/structured/generic.py
@@ -136,7 +136,7 @@ def _unrolled_cone_curvilinear_grid():
 
 datasets = {
     "2d_left_rotated": _rotated_curvilinear_grid(),
-    "ds_2d_left": xr.Dataset(
+    "ds_2d_left": xr.Dataset(  # MITgcm indexing style
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
@@ -178,5 +178,47 @@ datasets = {
             "time": (["time"], TIME, {"axis": "T"}),
         },
     ),
+    "ds_2d_right": xr.Dataset(  # NEMO indexing style
+        {
+            "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
+            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+        },
+        coords={
+            "XG": (
+                ["XG"],
+                2 * np.pi / X * np.arange(0, X),
+                {"axis": "X", "c_grid_axis_shift": 0.5},
+            ),
+            "XC": (["XC"], 2 * np.pi / X * (np.arange(0, X) - 0.5), {"axis": "X"}),
+            "YG": (
+                ["YG"],
+                2 * np.pi / (Y) * np.arange(0, Y),
+                {"axis": "Y", "c_grid_axis_shift": 0.5},
+            ),
+            "YC": (
+                ["YC"],
+                2 * np.pi / (Y) * (np.arange(0, Y) - 0.5),
+                {"axis": "Y"},
+            ),
+            "ZG": (
+                ["ZG"],
+                np.arange(Z),
+                {"axis": "Z", "c_grid_axis_shift": 0.5},
+            ),
+            "ZC": (
+                ["ZC"],
+                np.arange(Z) - 0.5,
+                {"axis": "Z"},
+            ),
+            "lon": (["XG"], 2 * np.pi / X * np.arange(0, X)),
+            "lat": (["YG"], 2 * np.pi / (Y) * np.arange(0, Y)),
+            "depth": (["ZG"], np.arange(Z)),
+            "time": (["time"], TIME, {"axis": "T"}),
+        },
+    ),
     "2d_left_unrolled_cone": _unrolled_cone_curvilinear_grid(),
 }
diff --git a/parcels/xgrid.py b/parcels/xgrid.py
index c3c11048..2cfda507 100644
--- a/parcels/xgrid.py
+++ b/parcels/xgrid.py
@@ -129,6 +129,53 @@ class XGrid(BaseGrid):
 
         return get_cell_count_along_dim(self.xgcm_grid.axes[axis])
 
+    def localize(self, position: dict[_XGRID_AXES, tuple[int, float]], dims: list[str]) -> dict[str, tuple[int, float]]:
+        """
+        Uses the grid context (i.e., the staggering of the grid) to convert a position relative
+        to the F-points in the grid to a position relative to the staggered grid the array
+        of interest is defined on.
+
+        Uses dimensions of the DataArray to determine the staggered grid.
+
+        WARNING: This API is unstable and subject to change in future versions.
+
+        Parameters
+        ----------
+        position : dict
+            A mapping of the axis to a tuple of (index, barycentric coordinate) for the
+            F-points in the grid.
+        dims : list[str]
+            A list of dimension names that the DataArray is defined on. This is used to determine
+            the staggering of the grid and which axis each dimension corresponds to.
+
+        Returns
+        -------
+        dict[str, tuple[int, float]]
+            A mapping of the dimension names to a tuple of (index, barycentric coordinate) for
+            the staggered grid the DataArray is defined on.
+
+        Example
+        -------
+        >>> position = {'X': (5, 0.51), 'Y': (
+            10, 0.25), 'Z': (3, 0.75)}
+        >>> dims = ['time', 'depth', 'YC', 'XC']
+        >>> grid.localize(position, dims)
+        {'depth': (3, 0.75), 'YC': (9, 0.75), 'XC': (5, 0.01)}
+        """
+        axis_to_var = {get_axis_from_dim_name(self.xgcm_grid.axes, dim): dim for dim in dims}
+        var_positions = {
+            axis: get_xgcm_position_from_dim_name(self.xgcm_grid.axes, dim) for axis, dim in axis_to_var.items()
+        }
+        return {
+            axis_to_var[axis]: _convert_center_pos_to_fpoint(
+                index=index,
+                bcoord=bcoord,
+                xgcm_position=var_positions[axis],
+                f_points_xgcm_position=self._fpoint_info[axis],
+            )
+            for axis, (index, bcoord) in position.items()
+        }
+
     @property
     def _z4d(self) -> Literal[0, 1]:
         """
@@ -185,6 +232,20 @@ class XGrid(BaseGrid):
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
+    @cached_property
+    def _fpoint_info(self):
+        """Returns a mapping of the spatial axes in the Grid to their XGCM positions."""
+        xgcm_axes = self.xgcm_grid.axes
+        f_point_positions = ["left", "right", "inner", "outer"]
+        axis_position_mapping = {}
+        for axis in self.axes:
+            coords = xgcm_axes[axis].coords
+            edge_positions = [pos for pos in coords.keys() if pos in f_point_positions]
+            assert len(edge_positions) == 1, f"Axis {axis} has multiple edge positions: {edge_positions}"
+            axis_position_mapping[axis] = edge_positions[0]
+
+        return axis_position_mapping
+
 
 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."""
@@ -337,3 +398,28 @@ def _search_1d_array(
     i = np.argmin(arr <= x) - 1
     bcoord = (x - arr[i]) / (arr[i + 1] - arr[i])
     return i, bcoord
+
+
+def _convert_center_pos_to_fpoint(
+    *, index: int, bcoord: float, xgcm_position: _XGCM_AXIS_POSITION, f_points_xgcm_position: _XGCM_AXIS_POSITION
+) -> tuple[int, float]:
+    """Converts a physical position relative to the cell edges defined in the grid to be relative to the center point.
+
+    This is used to "localize" a position to be relative to the staggered grid at which the field is defined, so that
+    it can be easily interpolated.
+
+    This also handles different model input cell edges and centers are staggered in different directions (e.g., with NEMO and MITgcm).
+    """
+    if xgcm_position != "center":  # Data is already defined on the F points
+        return index, bcoord
+
+    bcoord = bcoord - 0.5
+    if bcoord < 0:
+        bcoord += 1.0
+        index -= 1
+
+    # Correct relative to the f-point position
+    if f_points_xgcm_position in ["inner", "right"]:
+        index += 1
+
+    return index, bcoord
diff --git a/pyproject.toml b/pyproject.toml
index 8fde067a..318f5851 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -96,7 +96,7 @@ pre_commit = "*"
 
 # Docs
 ipython = "*"
-numpydoc = "*"
+numpydoc = "!=1.9.0"
 nbsphinx = "*"
 sphinx = "*"
 pandoc = "*"


Test patch:
diff --git a/tests/v4/test_datasets.py b/tests/v4/test_datasets.py
new file mode 100644
index 00000000..e5201e58
--- /dev/null
+++ b/tests/v4/test_datasets.py
@@ -0,0 +1,21 @@
+from parcels._datasets.structured.generic import datasets
+from parcels.xgcm import Grid
+
+
+def test_left_indexed_dataset():
+    """Checks that 'ds_2d_left' is right indexed on all variables."""
+    ds = datasets["ds_2d_left"]
+    grid = Grid(ds)
+
+    for _axis_name, axis in grid.axes.items():
+        for pos, _dim_name in axis.coords.items():
+            assert pos in ["left", "center"]
+
+
+def test_right_indexed_dataset():
+    """Checks that 'ds_2d_right' is right indexed on all variables."""
+    ds = datasets["ds_2d_right"]
+    grid = Grid(ds)
+    for _axis_name, axis in grid.axes.items():
+        for pos, _dim_name in axis.coords.items():
+            assert pos in ["center", "right"]
diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py
index 2af1e1f0..210ca6b7 100644
--- a/tests/v4/test_xgrid.py
+++ b/tests/v4/test_xgrid.py
@@ -145,3 +145,57 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi):
     xi, xsi = _search_1d_array(array, x)
     assert xi == expected_xi
     assert np.isclose(xsi, expected_xsi)
+
+
+@pytest.mark.parametrize(
+    "grid, da_name, expected",
+    [
+        pytest.param(
+            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            "U (C grid)",
+            {
+                "XG": (np.int64(0), np.float64(0.0)),
+                "YC": (np.int64(-1), np.float64(0.5)),
+                "ZG": (np.int64(0), np.float64(0.0)),
+            },
+            id="MITgcm indexing style U (C grid)",
+        ),
+        pytest.param(
+            XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)),
+            "V (C grid)",
+            {
+                "XC": (np.int64(-1), np.float64(0.5)),
+                "YG": (np.int64(0), np.float64(0.0)),
+                "ZG": (np.int64(0), np.float64(0.0)),
+            },
+            id="MITgcm indexing style V (C grid)",
+        ),
+        pytest.param(
+            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            "U (C grid)",
+            {
+                "XG": (np.int64(0), np.float64(0.0)),
+                "YC": (np.int64(0), np.float64(0.5)),
+                "ZG": (np.int64(0), np.float64(0.0)),
+            },
+            id="NEMO indexing style U (C grid)",
+        ),
+        pytest.param(
+            XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)),
+            "V (C grid)",
+            {
+                "XC": (np.int64(0), np.float64(0.5)),
+                "YG": (np.int64(0), np.float64(0.0)),
+                "ZG": (np.int64(0), np.float64(0.0)),
+            },
+            id="NEMO indexing style V (C grid)",
+        ),
+    ],
+)
+def test_xgrid_localize_zero_position(grid, da_name, expected):
+    """Test localize function using left and right datasets."""
+    position = grid.search(0, 0, 0)
+    da = grid.xgcm_grid._ds[da_name]
+
+    local_position = grid.localize(position, da.dims)
+    assert local_position == expected, f"Expected {expected}, got {local_position}"