Support for time varying z coordinates

[wyatt-fluidnumerics]

Description

The goal for this PR is to provide support for sigma coordinate models.

In nearly all sigma coordinate models, it is possible to output the z value of each node or cell face, this is fairly essential for analysis that uses a z coordinate unless the modeler chooses to plot entirely in sigma space, or reconstruct z themself. In this case, z has dimensions of (time, zf, n_node) for unstructured horizontal grids, and (time, z, y, x) for structured grids. We can utilize this field to advect particles in z coordinates for which the native model is in sigma coordinates (I believe this approach would work for ALE as well as long as a user has z coordinate output file, though I have not looked into this at all).

The primary change here is to optionally propagate ti and tau from the time search into the search(..., ti=None, tau=None) function itself, which allows for time and vertical column indexing of the z grid during the vertical grid search (i.e. at every time step each vertical column has unique z levels). On each of these columns the vertical search runs exactly the same as before. This allows _get_positions to return barycentric coordinates back to interpolators in the exact same format as parcels does currently. Nothing changes for classic z grids that are 1D.

Current Assumptions and Interpolations:

• TimeVarying UxGrid Z coordinate data must be stored on nodes.
• TimeVarying XGrid Z coordinate data must be stored on cell centres at layer interfaces
• For XGrids, horizontal Z coordinate interpolation is piecewise constant (i.e. data is taken directly from the face value for a particle)
• For UxGrids. horizontal Z coordinate interpolation is barycentric
• For both types of grid temporal Z coordinate interpolation is linear
• There is no default particle depths when using a time varying Z coordinate, z must be provided during particleset construction

The changes to each file are the following:

• _core/uxgrid.py: allow for 3D z dimensions in addition to 1d z dimensions. Change search to handle both cases, with time and node indexing for the 3D case.
• _core/xgrid.py: allow a 4D depth with dims (time, depth, YC, XC). For a 4D depth, the new _search_time_varying_depth takes the column at the centre of each particle's cell, interpolates it linearly in time with ti/tau, and then searches it's columns
• _core/index_search.py: a vectorized version of _search_1d_array for searching each individual column that gets time/node indexed into.
• interpolators/_uxinterpolators.py: for interpolators that are linear in z, use the z calculated barycentric coordinate directly to avoid indexes into the z grid (which would now need if branching for the 1d vs time varying case).
• _core/field.py: _get_positions passes the time index into grid.search
• _core/basegrid.py: search takes in ti=None
• _core/model.py: UnstructuredModelData fails whenever a 3D z coordinate's time coordinate differs from the data time coordinate
• _core/particleset.py: Whenever a particleset is created on a time varying zgrid that doesn't provide an initial particle depth, this errors. This needs updating to compute the closest z to zero while remaining on the vertical grid, similar to case for 1d z coordinates.
• _datasets/unstructured/generated.py: new dataset generator functions for testing and tutorials

There is also a new tutorial tutorial_sigma_coordinates.ipynb which demonstrates the utility for both structured and unstructured grids. Hopefully we will be able to replace the synthetic dataset with a model dataset (I am hoping to put together a schism run for this).

This PR is just an initial stab at the implementation for the sake of initiating discussion, @fluidnumericsJoe and I have a whiteboard with ideas to polish things and improve the scalability, we'll have an issue open soon but the biggest thing is the following:

• The main problem of this design is that the z coordinate can carry a large memory footprint. I believe that holding z as a field has been previously investigated (Re-evaluating support for timevarying depth dimension #1914) however the implementation seemed quite complicated and it introduces a bit of a circular dependency since fields require Grid objects, and z would also somehow need to be a part of the grid. We think it could be a good idea to give vertical coordinates their own class. This would likely be very field like, and could even implement a system similar to to_windowed_array/'to_cached_chunked_array` so that the memory foot print is reduced.
• A second benefit to a unique vertical coordinate class is the alleviation of hard coding z coordinate interpolators. Similar to the way a field object is given interpolators, a vertical grid class could be given interpolators which could be user customizable, for more freedom over defining interpolation methods. Currently, piecewise constant and barycentric are hardcoded during search for both XGrids and UxGrids respectively.

Checklist

• Closes Add native support for terrain following (sigma) coordinates #2907
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

AI Disclosure

• This PR contains AI-generated content.
  • I have tested any AI-generated content in my PR.
  • I take responsibility for any AI-generated content in my PR.
  • Describe how you used it (e.g., by pasting your prompt): Parts of this implementation were done by guiding opus 5.5. All code written by the llm was thoroughly reviewed and I made modifications as needed for scope, accuracy, and readability.

[erikvansebille]

Very cool, @wyatt-fluidnumerics! I love the final animation in the tutorial! Amazing that we can get time-varying depth grid support in v4!

Let me know when you want a full review

[fluidnumericsJoe]

@wyatt-fluidnumerics - this looks awesome. One thing I noticed is that the z interpolation in time does not do any blending between z[ti] and z[ti+1] - For a constant U field, like you have in the tutorial, this is not an issue. As soon as we have a field that depends on \sigma (or has any spatial variation), the errors are quite pronounced. In this PR, it might be good to add in the temporal blending.

Alternatively, I could be convinced to treat this as POC and work on pairing interpolators with coordinates, much like we do with Field objects and handling coordinate interpolation in interpolators.py .