Initial 3d z coordinate implementation for unstructured grids

[wyatt-fluidnumerics]

Description

TODO list:

• Add structured grid support, verify SGRID compliance
• Add linear interpolation of 3d z fields (pass tau through _get_positions into the grid search).
• Rewrite the new tutorial markdown cells now that the dataset is made with a built in generator.

The goal for this PR is to provide support for sigma coordinate models that utilize unstructured grids.

In nearly all sigma coordinate models, it is possible to output the z value of each node, 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. 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 from the time search into the search(..., ti=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.

Currently, this is only implemented for unstructured grids, though the only fundamental change for structured grids is that z would be 4d (t, x, y, z) instead of 3d (t, z, n_node). Additionally, nothing changes for standard z coordinate models that use a 1d vertical grid.

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/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 3d case).
• _core/field.py: _get_positions passes the time index into grid.search
• _core/basegrid.py: search takes in ti=None, XGrid search ignores this
• _core/model.py: UnstructuredModelData fails whenever a 3D z dimensions time coordinate differs from the data time coordinate
• _core/particleset.py: Whenever a particleset is created on a 3d 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_unstructured_sigma_coordinates.ipynb which demonstrates the utility. 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 draft 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.

Checklist

• Will eventually close 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 .