parcels icon indicating copy to clipboard operation
parcels copied to clipboard

Create a Grid class to consolidate Parcels specific grid operations

Open fluidnumerics-joe opened this issue 9 months ago • 1 comments

At the moment, in v4-dev we have a rather clunky setup that I put in for handling Field.eval .

The call structure is currently as follows

Field.eval()
│
├──> Field._interpolate()
│     │
│     ├──> Field._search_indices()
│     │     │
│     │     ├──> _search_time_index(self, time)
│     │     │
│     │     ├──> (type check on self.data)
│     │     │     │
│     │     │     ├──> If unstructured:
│     │     │     │     └──> Field._search_indices_unstructured()
│     │     │     │           ├──> If ei is None:
│     │     │     │           │     └──> grid.get_spatial_hash().query([[x, y]])
│     │     │     │           │           └──> Return face id, barycentric coords
│     │     │     │           └──> Else:
│     │     │     │                 ├──> Field.unravel_index()
│     │     │     │                 ├──> Field._get_ux_barycentric_coordinates()
│     │     │     │                 ├──> If failed: iterate neighbors
│     │     │     │                 └──> Fallback to spatial hash
│     │     │     │
│     │     │     └──> If structured:
│     │     │           └──> Field._search_indices_structured()
│     │     │                 ├──> Check grid type
│     │     │                 ├──> Call _search_indices_rectilinear()
│     │     │                 └──> Return (zeta, eta, xsi, zi, yi, xi)
│     │     │
│     │     └──> Return bcoords, ei, tau, ti
│     │
│     └──> Field._interp_method(...)
│           └──> Compute interpolated value at (ti, ei, bcoords, tau, time, z, y, x)
│
├──> Check for NaN in interpolated value
└──> Apply unit conversion (if applyConversion is True)

I'd like to propose that the Grid class have a method called search with the following API

def search(self, field, z, y, x, ei=None, search2D=False):
    """
    Perform a spatial (and optionally vertical) search to locate the grid element
    that contains a given point (x, y, z).

    This method delegates to grid-type-specific logic (e.g., structured or unstructured)
    to determine the appropriate indices and interpolation coordinates for evaluating a field.

    Parameters
    ----------
    field : Field
        The field requesting the search. Used to extract grid-specific parameters,
        unravel index metadata, or perform coordinate system-specific operations.
    z : float
        Vertical coordinate of the query point. If `search2D=True`, this may be ignored.
    y : float
        Latitude or vertical index, depending on grid type and projection.
    x : float
        Longitude or horizontal index, depending on grid type and projection.
    ei : int, optional
        A previously computed encoded index (e.g., raveled face or cell index). If provided,
        the search will first attempt to validate and reuse it before falling back to
        a global or local search strategy.
    search2D : bool, default=False
        If True, perform only a 2D search (x, y), ignoring the vertical component z.

    Returns
    -------
    bcoords : np.ndarray or tuple
        Interpolation weights or barycentric coordinates within the containing cell/face.
        The interpretation of `bcoords` depends on the grid type.
    ei : int
        Encoded index of the identified grid cell or face. This value can be cached for
        future lookups to accelerate repeated searches.

    Raises
    ------
    FieldOutOfBoundError
        Raised when the queried point lies outside the bounds of the grid.
    NotImplementedError
        Raised if the search method is not implemented for the current grid type.
    """

Additionally, I'd propose adding a new UxGrid (parcels.UxGrid) that can be initialized from a ux.Grid ; this would be a type extension of ux.Grid so that we inherit all the useful bits of uxarray. The purpose of the UXGrid is to provide a search API that leverage the ux.Grid.spatial_hash.query() or a nearest neighbor search (that needs to be implemented).

With this, the Field.__init__ method API would change to

def __init__(
        self,
        name: str,
        data: xr.DataArray | ux.UxDataArray,
        grid: UxGrid | Grid,                                 # Formerly this was  ux.Grid | Grid
        mesh_type: Mesh = "flat",
        interp_method: Callable | None = None,
    ):

With this proposed change, the Field.eval call structure can be simplified

Field.eval()
│
├──> Field._search_time_index(self, time)
│     └──> Return tau, ti
|
├──> Field.grid.search()
│     └──> Return bcoords, ei
│
└──> Field._interp_method(...)
│     └──> Compute interpolated value at (ti, ei, bcoords, tau, time, z, y, x)
│
├──> Check for NaN in interpolated value
└──> Apply unit conversion (if applyConversion is True)

This would also allow us to remove the _index_search.py module and push us to isolate index searching for each type of grid. In this case, I'm also proposing to push the _search_time_index to be an internal API function for the Field class

@erikvansebille @VeckoTheGecko - Would be good to hear your thoughts on this. I'm happy to take a crack at restructuring code a bit to give us something "tangible" to look at.

fluidnumerics-joe avatar May 29 '25 18:05 fluidnumerics-joe

I think this sounds good. It would be really helpful to have all the grid specific tasks isolated in the grid classes. Some comments:

I'd like to propose that the Grid class have a method called search with the following API

def search(self, field, z, y, x, ei=None, search2D=False):

This looks good for the most part.

  1. I don't think field is necessary as a param since the grid already has access to the grid information
    • After the meeting, decided that we'd need to move ravel and unravel index onto the grid itself as well
  2. search2D=True/False , what impact does that have on the return type?
  3. Would this mean that we should have information on the grid class for encoding/decoding to and from ei?

VeckoTheGecko avatar Jun 04 '25 11:06 VeckoTheGecko

This has been resolved via the PRs xref'ed above, and also clarified in #2048

VeckoTheGecko avatar Jul 04 '25 11:07 VeckoTheGecko