Make edge/face z-coordinate calculations dask-compatible#8
Conversation
… use lazy xarray operations instead of eager numpy, making them compatible with dask-backed arrays.
There was a problem hiding this comment.
Pull request overview
This PR updates SCHISM grid z-coordinate computation in suxarray.io._schismgrid to use lazy xarray operations (better compatibility with dask-backed arrays) and adds regression tests validating face/edge z-averaging behavior on real SCHISM test fixtures.
Changes:
- Reworked
_calculate_edge_zand_calculate_face_zto useisel(...).mean(...)with masking for triangle fill nodes, avoiding eager NumPy indexing. - Updated
_assign_z_coordsto work withDataArrayoutputs from the new calculation functions. - Added
tests/test_schismgrid_io.pyto validate output shapes and numeric correctness for tri/quad faces and edges.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
suxarray/io/_schismgrid.py |
Replaces eager NumPy z-calculations with xarray-based indexing/masking and updates z-coordinate assignment. |
tests/test_schismgrid_io.py |
Adds tests to verify face/edge z-coordinate calculations against known SCHISM test data. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Assign 1D static coordinates (for uxarray subsetting) | ||
| ds_out2d = ds_out2d.assign_coords({ | ||
| "node_z": (("n_node",), node_z), | ||
| "edge_z": (("n_edge",), edge_z), | ||
| "face_z": (("n_face",), face_z), | ||
| "node_z": (("n_node",), node_z.values), | ||
| "edge_z": (("n_edge",), edge_z.values), | ||
| "face_z": (("n_face",), face_z.values), | ||
| }) |
There was a problem hiding this comment.
These values are 1D static coordinates (from the first timestep in bottom layer). They should be light enough to call eagerly. But as we're moving towards lazy loading, dropping .values brings consistency.
There was a problem hiding this comment.
Fixed in commit that removes the .values calls — assign_coords now receives the DataArray slices directly, keeping the z-coordinate calculations lazy.
| def _calculate_edge_z( | ||
| ds_zcoords: xr.Dataset, ds_out2d: xr.Dataset | ||
| ) -> xr.Dataset: | ||
| """Calculate z-coordinates at edge centers by averaging node z-coordinates. | ||
|
|
||
| For each edge, computes the mean z-coordinate of its two endpoint nodes | ||
| across all time steps and vertical layers. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| ds_zcoords : xr.Dataset | ||
| Dataset containing zCoordinates variable with dims (time,n_node,n_layer) | ||
| ds_out2d : xr.Dataset | ||
| Dataset containing edge_node_connectivity with shape (n_edge, 2) | ||
|
|
||
| Returns | ||
| ------- | ||
| xr.Dataset | ||
| ds_sgrid_info with edge_z coordinate added, shape (time,n_edge,n_layer) | ||
|
|
||
| TODO: Make this dask-compatible to avoid loading large arrays in memory. | ||
| Using xr.apply_ufunc with dask="parallelized" or da.map_blocks. | ||
| ds_sgrid_info with edge_z coordinate added, shape (n_edge,time,n_layer) | ||
| """ |
There was a problem hiding this comment.
Fixed — updated the return type annotation from xr.Dataset to xr.DataArray and updated the docstring to describe the returned DataArray with shape (n_edge, time, n_layer).
| def _calculate_face_z( | ||
| ds_zcoords: xr.Dataset, ds_out2d: xr.Dataset | ||
| ) -> xr.Dataset: | ||
| """Calculate z-coordinates at face centers by averaging node z-coordinates. | ||
|
|
||
| For each face, computes the mean z-coordinate of its vertices. Making sure | ||
| to select three or four vertices depending on triangular or quad elements. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| ds_zcoords : xr.Dataset | ||
| Dataset containing zCoordinates variable with dims (time,n_node,n_layer) | ||
| ds_out2d : xr.Dataset | ||
| Dataset containing edge_node_connectivity with shape (n_edge, 2) | ||
| Dataset containing face_node_connectivity with shape (n_face, 4) | ||
|
|
||
| Returns | ||
| ------- | ||
| xr.Dataset | ||
| ds_sgrid_info with face_z coordinate added, shape (time,n_edge,n_layer) | ||
| ds_sgrid_info with face_z coordinate added, shape (n_face,time,n_layer) | ||
| """ |
There was a problem hiding this comment.
Fixed in the latest commit — updated the return type annotation from xr.Dataset to xr.DataArray and updated the docstring to describe the returned DataArray with shape (n_face, time, n_layer).
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
…and _calculate_face_z functions
Rewrites
_calculate_edge_zand_calculate_face_zin_schismgrid.pyto use lazy xarray operations instead of eager numpy to make them compatible with dask-based arrays. Adds tests for these z calculations.Changes:
_calculate_face_z: replaced the eager numpy loop (separate tri/quad branches + np.zeros allocation) with xarrayisel(...).where(mask).mean(), which correctly handles the fill-value mask for triangular cells without materialising the full array._assign_z_coords: updated callers to useisel()and.valueswith the new DataArray return types.test_schismgrid_io.pycovering output shape, correct 3-node averaging for tri cells, fill-node exclusion guard, 4-node averaging for quad cells, and edge z-value correctness against the test mesh (out2d_1.nc / zCoordinates_1.nc).