Structured grids barycentric coordinates

[VeckoTheGecko]

• Chose the correct base branch (v4-dev for v4 changes)
• Fixes Structured grid index searching and computing barycentric coordinates #2035
• Added tests
• Added documentation

Changes:

• Require dataset to have lon lat grid of F points (and validates that this is the right format)
• Index search and barycentric coordinate calculation on these F-points (both for 1D and 2D lon and lat)
  • Assumes that depth is 1D
  • Uses algorithm from v3 (in future we're sure this can be optimised to be more robust and faster - but this is good for now)
• ravel and unravel method

Here is a diagram to help visualise what's happening here:

Diagram on NEMO/MITgcm indexing in general

Note that result from the search() method is wrt. the red points (i.e., yi, xi = 0, 0, eta, xsi = 0.5, 0.5 means the particle is in the cell centre with lower left being red point at (0, 0).

One thing to note: Particles on the edge of the model grid are not representable in the F-points (in the diagram, this is evident in the NEMO model for the lower left edge of the grid). In the periodic case, the grid has a halo anyway so this is a non-issue. The grid is not patched to be a "full grid" on initialisation - we instead work close to the original dataset.

If we have a particle in tracer cell (yi, xi) = 0,0, then the relevant velocities for interpolation are:

• NEMO:
  • U[y, x] -> U[1,0] and U[1,1]
  • V[y, x] -> V[0,1] and V[1,1]
• MITgcm
  • U[y, x] -> U[0,0] and U[0,1]
  • V[y, x] -> V[0,0] and V[1,0]

The lining up of these indices with on-disk representation will be handled in the interpolator itself (in a future PR).

[VeckoTheGecko]

@erikvansebille I'm not sure what the code at da87f9b does (in particular the if grid.mesh == 'spherical': block), something to do with the international dateline? Could you provide some insight here?

Also insight on _reconnect_bnd_indices would be helpful (first mention: 812be0c)

Happy to discuss irl, and will update here accordingly.

[VeckoTheGecko]

@fluidnumerics-joe getting some failing tests now with FieldSet.add_constant_field: ValueError: Longitude DataArray 'lon' with dims ('XG',) is defined on the center of the grid, but must be defined on the F points.

This error makes sense, the current implementation is:

         ...
         da = xr.DataArray(
            data=np.full((1, 1, 1, 1), value),
            dims=["time", "ZG", "YG", "XG"],
            coords={
                "ZG": (["ZG"], np.arange(1), {"axis": "Z"}),
                "YG": (["YG"], np.arange(1), {"axis": "Y"}),
                "XG": (["XG"], np.arange(1), {"axis": "X"}),
                "lon": (["XG"], np.arange(1), {"axis": "X"}),
                "lat": (["YG"], np.arange(1), {"axis": "Y"}),
                "depth": (["ZG"], np.arange(1), {"axis": "Z"}),
                "time": (["time"], np.arange(1), {"axis": "T"}),
            },
        )
        breakpoint()
        grid = XGrid(xgcm.Grid(da))
        ...

Basically making a 1 point structured grid and then going forward with that.

I think this implementation should change (it doesn't make sense for a constant field to have its own grid, especially now since grids now have the concepts of searching, raveling, and "out of bounds"). Not sure exactly how yet, having a think...

[fluidnumerics-joe]

@VeckoTheGecko - at a minimum, a single point grid would be a single tracer point, which would give 4 vorticity points that form the boundary of the tracer cell. If you have only one vorticity point, there are no tracer cells.

[VeckoTheGecko]

at a minimum, a single point grid would be a single tracer point, which would give 4 vorticity points that form the boundary of the tracer cell. If you have only one vorticity point, there are no tracer cells.

Yes, that would get past the initialiser - but would fail on eval due to the index searching etc.. I'm also thinking about futureproofing (does this mean for a constant Field that an extra ei for each particle causing a n_particles * 8 byte memory overhead?). Going a different route avoids this memory overhead. I might just xfail for this PR and fix in a different one

[fluidnumerics-joe]

at a minimum, a single point grid would be a single tracer point, which would give 4 vorticity points that form the boundary of the tracer cell. If you have only one vorticity point, there are no tracer cells.

Yes, that would get past the initialiser - but would fail on eval due to the index searching etc.. I'm also thinking about futureproofing (does this mean for a constant Field that an extra ei for each particle causing a n_particles * 8 byte memory overhead?). Going a different route avoids this memory overhead. I might just xfail for this PR and fix in a different one

I don't understand why it fails on index searching. An Arakawa-C grid is not completely defined in the example you've provided with just an XG,YG point. It's also not completely defined with just an XC,YC point; both sets of points are needed to define a grid. A single finite volume cell is defined completely by specifying the four corner points of the volume (the four XG,YG F-points) and the cell center (which can be inferred as the average of the four corner points - this is the XC,YC point).

[erikvansebille]

I don't understand why it fails on index searching. An Arakawa-C grid is not completely defined in the example you've provided with just an XG,YG point. It's also not completely defined with just an XC,YC point; both sets of points are needed to define a grid. A single finite volume cell is defined completely by specifying the four corner points of the volume (the four XG,YG F-points) and the cell center (which can be inferred as the average of the four corner points - this is the XC,YC point).

Not sure if this is super-pertinent to this discussion, but note that in Parcels v3, we use 'one-dimensional fields' to represent e.g. horizontally uniform diffusivity or variables that would only depend on depth. See for an example e.g.
https://github.com/OceanParcels/Parcels/blob/ae2d508fa8608ab81f3b5e9e1c34acc220b2c1eb/docs/examples/example_brownian.py#L26-L30
These fields have only one longitude and latitude point, and by definition cannot throw an OutOfBounds error.

It's quite useful (albeit a bit hacky?) to keep support for these one-dimensional fields?

[VeckoTheGecko]

@fluidnumerics-joe full traceback

Details

============================= test session starts ==============================
platform darwin -- Python 3.12.10, pytest-8.3.5, pluggy-1.5.0 -- /Users/Hodgs004/miniforge3/envs/parcels-dev/bin/python3.12
cachedir: .pytest_cache
metadata: {'Python': '3.12.10', 'Platform': 'macOS-15.3.1-arm64-arm-64bit', 'Packages': {'pytest': '8.3.5', 'pluggy': '1.5.0'}, 'Plugins': {'anyio': '4.9.0', 'html': '4.1.1', 'metadata': '3.1.1', 'hypothesis': '6.131.9', 'nbval': '0.11.0', 'reportlog': '0.1.2'}}
hypothesis profile 'default' -> database=DirectoryBasedExampleDatabase(PosixPath('/Users/Hodgs004/coding/repos/parcels/.hypothesis/examples'))
rootdir: /Users/Hodgs004/coding/repos/parcels
configfile: pyproject.toml
plugins: anyio-4.9.0, html-4.1.1, metadata-3.1.1, hypothesis-6.131.9, nbval-0.11.0, reportlog-0.1.2
collecting ... collected 1 item

tests/v4/test_fieldset.py::test_fieldset_add_constant_field FAILED       [100%]

=================================== FAILURES ===================================
_______________________ test_fieldset_add_constant_field _______________________

fieldset = <parcels.fieldset.FieldSet object at 0x1771c0980>

    def test_fieldset_add_constant_field(fieldset):
>       fieldset.add_constant_field("test_constant_field", 1.0)

tests/v4/test_fieldset.py:42: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
parcels/fieldset.py:177: in add_constant_field
    grid = XGrid(xgcm.Grid(da))
parcels/xgrid.py:55: in __init__
    assert_valid_lat_lon(ds["lat"], ds["lon"], grid.axes)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

da_lat = <xarray.DataArray 'lat' (YG: 1)> Size: 8B
array([0])
Coordinates:
  * YG       (YG) int64 8B 0
    lat      (YG) int64 8B 0
Attributes:
    axis:     Y
da_lon = <xarray.DataArray 'lon' (XG: 1)> Size: 8B
array([0])
Coordinates:
  * XG       (XG) int64 8B 0
    lon      (XG) int64 8B 0
Attributes:
    axis:     X
axes = OrderedDict({'Y': <parcels.Axis 'Y' (periodic, boundary=None)>
Axis Coordinates:
  * center   YG, 'T': <parcels.Axis '...xis Coordinates:
  * center   XG, 'Z': <parcels.Axis 'Z' (periodic, boundary=None)>
Axis Coordinates:
  * center   ZG})

    def assert_valid_lat_lon(da_lat, da_lon, axes: _XGCM_AXES):
        """
        Asserts that the provided longitude and latitude DataArrays are defined appropriately
        on the F points to match the internal representation in Parcels.
    
        - Longitude and latitude must be 1D or 2D (both must have the same dimensionality)
        - Both are defined on the left points (i.e., not the centers)
        - If 1D:
          - Longitude is associated with the X axis
          - Latitude is associated with the Y axis
        - If 2D:
          - Lon and lat are defined on the same dimensions
          - Lon and lat are transposed such they're Y, X
        """
        assert_all_dimensions_correspond_with_axis(da_lon, axes)
        assert_all_dimensions_correspond_with_axis(da_lat, axes)
    
        dim_to_position = {dim: get_position_from_dim_name(axes, dim) for dim in da_lon.dims}
        dim_to_position.update({dim: get_position_from_dim_name(axes, dim) for dim in da_lat.dims})
    
        for dim in da_lon.dims:
            if get_position_from_dim_name(axes, dim) == "center":
>               raise ValueError(
                    f"Longitude DataArray {da_lon.name!r} with dims {da_lon.dims} is defined on the center of the grid, but must be defined on the F points."
                )
E               ValueError: Longitude DataArray 'lon' with dims ('XG',) is defined on the center of the grid, but must be defined on the F points.

parcels/xgrid.py:328: ValueError
=============================== warnings summary ===============================
../../../miniforge3/envs/parcels-dev/lib/python3.12/site-packages/geopandas/_compat.py:7
  /Users/Hodgs004/miniforge3/envs/parcels-dev/lib/python3.12/site-packages/geopandas/_compat.py:7: DeprecationWarning: The 'shapely.geos' module is deprecated, and will be removed in a future version. All attributes of 'shapely.geos' are available directly from the top-level 'shapely' namespace (since shapely 2.0.0).
    import shapely.geos

tests/v4/test_fieldset.py::test_fieldset_add_constant_field
  /Users/Hodgs004/coding/repos/parcels/tests/v4/test_fieldset.py:21: DeprecationWarning: The `periodic` argument will be deprecated. To preserve previous behavior supply `boundary = 'periodic'.
    grid = XGrid(xgcm.Grid(ds))

tests/v4/test_fieldset.py::test_fieldset_add_constant_field
  /Users/Hodgs004/coding/repos/parcels/parcels/fieldset.py:177: DeprecationWarning: The `periodic` argument will be deprecated. To preserve previous behavior supply `boundary = 'periodic'.
    grid = XGrid(xgcm.Grid(da))

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================== short test summary info ============================
FAILED tests/v4/test_fieldset.py::test_fieldset_add_constant_field - ValueError: Longitude DataArray 'lon' with dims ('XG',) is defined on the center of the grid, but must be defined on the F points.
======================== 1 failed, 3 warnings in 4.26s =========================

[VeckoTheGecko]

Let's postpone this a couple days.... Working on a proposal that affects this

[VeckoTheGecko]

I've just rebased this and updated the API to line up with #2048 - ready for review.

[VeckoTheGecko]

@fluidnumerics-joe getting some failing tests now with FieldSet.add_constant_field: ValueError: Longitude DataArray 'lon' with dims ('XG',) is defined on the center of the grid, but must be defined on the F points.

Resolved with #2058 - all v4 tests are passing in that PR

[erikvansebille]

Is this needed if we don't do periodic grids anymore?

[VeckoTheGecko]

I was wondering if this wasn't for periodic grids, but instead was to do with tripolar grids - would that make sense? I thought the edges/wrapping in that grid would result in the index search needing to be "reconnected" to the grid like this.

[erikvansebille]

Yes, you may indeed be right. Let's leave in for now and we can see whether this is still necessary when we test curvilinear NEMO grid in the Arctic

[VeckoTheGecko]

Just implemented the review changes. Still a couple pending items above that we didn't discuss during our meeting @erikvansebille - after that I think we're good to merge.

[erikvansebille]

Good to merge, as far as I'm concerned