Make edge/face z-coordinate calculations dask-compatible

suxarray/io/_schismgrid.py

@@ -105,38 +105,40 @@ def _assign_z_coords(ds_out2d, ds_zcoords) -> xr.Dataset:
     xr.Dataset
         ds_out2d with z-coordinates added:
         - node_z, edge_z, face_z: 1D static coords (bottom layer) for subsetting
-        - node_z_3d, edge_z_3d, face_z_3d: 3D data variables for z-coordinates.
+        - node_z_3d: shape (time, n_node, n_layer)
+        - edge_z_3d: shape (time, n_edge, n_layer)
+        - face_z_3d: shape (time, n_face, n_layer)
     """
-    # Calculate full 3D z-coordinates (transposed to have spatial dim first)
-    node_z_3d = ds_zcoords.zCoordinates.values.transpose(1, 0, 2)
+    # Calculate full 3D z-coordinates
+    node_z_3d = ds_zcoords.zCoordinates.transpose("time", "n_node", "n_layer")
     edge_z_3d = _calculate_edge_z(ds_zcoords, ds_out2d)
     face_z_3d = _calculate_face_z(ds_zcoords, ds_out2d)
 
     # Extract bottom layer (index -1) as static 1D coordinates for subsetting
     # Bottom layer represents depth which doesn't change with time
-    node_z = node_z_3d[:, 0, -1]  # shape (n_node,)
-    edge_z = edge_z_3d[:, 0, -1]  # shape (n_edge,)
-    face_z = face_z_3d[:, 0, -1]  # shape (n_face,)
+    node_z = node_z_3d.isel(time=0, n_layer=-1)
+    edge_z = edge_z_3d.isel(time=0, n_layer=-1)
+    face_z = face_z_3d.isel(time=0, n_layer=-1)
 
     # 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": node_z,
+        "edge_z": edge_z,
+        "face_z": face_z,
     })
 
     # Assign 3D z as data variables (for actual z-coordinate access)
     ds_out2d = ds_out2d.assign({
-        "node_z_3d": (("n_node", "time", "n_layer"), node_z_3d),
-        "edge_z_3d": (("n_edge", "time", "n_layer"), edge_z_3d),
-        "face_z_3d": (("n_face", "time", "n_layer"), face_z_3d),
+        "node_z_3d": node_z_3d,
+        "edge_z_3d": edge_z_3d,
+        "face_z_3d": face_z_3d,
     })
 
     return ds_out2d
 
 def _calculate_edge_z(
     ds_zcoords: xr.Dataset, ds_out2d: xr.Dataset
-) -> xr.Dataset:
+) -> xr.DataArray:
     """Calculate z-coordinates at edge centers by averaging node z-coordinates.
 
     For each edge, computes the mean z-coordinate of its two endpoint nodes
@@ -151,21 +153,24 @@ def _calculate_edge_z(
 
     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.
+    xr.DataArray
+        DataArray of edge z-coordinates with shape (time, n_edge, n_layer)
     """
+
+    # node_connectivity is static and should be light enough to load with numpy
     edge_node_ids = ds_out2d.edge_node_connectivity.values  # shape (n_edge, 2)
-    edge_z = ds_zcoords.zCoordinates.values[:, edge_node_ids, :].mean(axis=2)
-    edge_z = edge_z.transpose(1, 0, 2)
+    # Roll edge_node_ids into a DataArray
+    node_idx = xr.DataArray(edge_node_ids, dims=["n_edge", "two"])
+    # Lazy xarray operations for indexing and mean
+    edge_z = ds_zcoords.zCoordinates.isel(n_node=node_idx).mean(dim="two")
+    edge_z = edge_z.transpose("time", "n_edge", "n_layer")
+
     return edge_z
 
 
 def _calculate_face_z(
     ds_zcoords: xr.Dataset, ds_out2d: xr.Dataset
-) -> xr.Dataset:
+) -> xr.DataArray:
     """Calculate z-coordinates at face centers by averaging node z-coordinates.
 
     For each face, computes the mean z-coordinate of its vertices. Making sure
@@ -176,37 +181,33 @@ def _calculate_face_z(
     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)
+    xr.DataArray
+        face_z DataArray with shape (time, n_face, n_layer)
     """
 
-    # TODO: Make this dask-compatible - same approach as _calculate_edge_z.
+    # node_connectivity is static and should be light enough to load with numpy
     face_ids = ds_out2d.face_node_connectivity.values  # shape (n_face, 4)
 
-    # Identify triangular and quad elements for correct face calculation
+    # Grab fill value for face_node_connectivity to identify triangular elements
     fill_value = ds_out2d.face_node_connectivity._FillValue
-    tri_el_ids = np.where(face_ids[:, -1] == fill_value)[0]
-    quad_el_ids = np.where(face_ids[:, -1] != fill_value)[0]
-
-    # TODO: use Dataset.sizes instead of DataArray.dims for futureproofing
-    face_z = np.zeros((ds_zcoords.dims["time"],
-                       len(face_ids),
-                       ds_zcoords.dims["n_layer"]))
-
-    # Create face_z, distinguishing triangular and quad elements.
-    zCoords = ds_zcoords.zCoordinates.values
-    face_z[:, tri_el_ids, :] = zCoords[:,
-                                       face_ids[tri_el_ids, :-1],
-                                       :].mean(axis=2)
-    # TODO: Implement more rigorous face z-coord calculation method.
-    face_z[:, quad_el_ids, :] = zCoords[:,
-                                        face_ids[quad_el_ids, :],
-                                        :].mean(axis=2)
-    face_z = face_z.transpose(1, 0, 2)
+    valid_mask = face_ids != fill_value
+    # Replace fill_value with 0 for indexing
+    face_ids_masked = np.where(valid_mask, face_ids, 0)
+
+    # Roll face_ids into a DataArray
+    face_idx = xr.DataArray(face_ids_masked, dims=["n_face", "four"])
+    mask_da = xr.DataArray(valid_mask, dims=["n_face", "four"])
+
+    # Lazy xarray operations for indexing and mean
+    face_z = ds_zcoords.zCoordinates.isel(n_node=face_idx)
+    # Use the mask to ignore fill values when calculating the mean
+    face_z = face_z.where(mask_da).mean(dim="four")
+    face_z = face_z.transpose("time", "n_face", "n_layer")
+
     return face_z
 
 def _find_grid_topology_varname(ds: xr.Dataset) -> str:

tests/test_schismgrid_io.py

@@ -0,0 +1,157 @@
+"""Tests for suxarray.io._schismgrid face/edge z-coordinate calculations.
+
+Uses real SCHISM test data (out2d_1.nc / zCoordinates_1.nc).
+
+Focuses on:
+- _calculate_face_z: correct averaging for triangular (3-node) and
+  quadrilateral (4-node) cells, with fill-value masking for tri cells.
+- _calculate_edge_z: correct averaging over the two endpoint nodes.
+"""
+from pathlib import Path
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from suxarray.io._schismgrid import (
+    _calculate_edge_z,
+    _calculate_face_z,
+    _rename_coords,
+)
+from uxarray.io._ugrid import _read_ugrid
+
+TESTDATA = Path(__file__).parent / "testdata"
+
+# Expected face/edge counts for out2d_1.nc
+N_NODES = 823
+N_FACES = 791
+N_EDGES = 1613
+N_TRI = 87
+N_QUAD = 704
+N_LAYERS = 23
+N_TIME = 48
+
+# First tri face in the test mesh: index 64, nodes 63/80/79 (0-based)
+TRI_IDX = 64
+TRI_NODES = [63, 80, 79]
+
+# First quad face in the test mesh: index 0, nodes 0/1/6/3 (0-based)
+QUAD_IDX = 0
+QUAD_NODES = [0, 1, 6, 3]
+
+# First edge: index 0, nodes 0/1 (0-based)
+EDGE_IDX = 0
+EDGE_NODES = [0, 1]
+
+
+@pytest.fixture(scope="module")
+def calc_inputs():
+    """Prepare processed ds_out2d and ds_zcoords ready for the calculation
+    functions, mirroring the pre-processing done inside _read_schism_grid."""
+    ds_out2d_raw = xr.open_dataset(str(TESTDATA / "out2d_1.nc"))
+    ds_zcoords_raw = xr.open_dataset(str(TESTDATA / "zCoordinates_1.nc"))
+
+    # Process out2d the same way _read_schism_grid does
+    ds_out2d = ds_out2d_raw.drop_dims("time")
+    ds_out2d, _ = _read_ugrid(ds_out2d)
+    ds_out2d = _rename_coords(ds_out2d)
+
+    # Rename zcoords dims the same way _read_schism_grid does
+    ds_zcoords = ds_zcoords_raw.swap_dims(
+        {"nSCHISM_hgrid_node": "n_node", "nSCHISM_vgrid_layers": "n_layer"}
+    )
+
+    return ds_out2d, ds_zcoords, ds_zcoords_raw
+
+
+# Connectivity – sanity checks on the test mesh
+
+def test_mesh_tri_and_quad_counts(calc_inputs):
+    """Test mesh has the expected 87 tri and 704 quad cells."""
+    ds_out2d, _, _ = calc_inputs
+    fnc = ds_out2d.face_node_connectivity.values
+    fill = ds_out2d.face_node_connectivity._FillValue
+    assert int(np.sum(fnc[:, -1] == fill)) == N_TRI
+    assert int(np.sum(fnc[:, -1] != fill)) == N_QUAD
+    assert ds_out2d.face_node_connectivity.sizes["n_face"] == N_FACES
+
+
+# _calculate_face_z
+
+def test_calculate_face_z_output_shape(calc_inputs):
+    """_calculate_face_z returns (n_face, time, n_layer)."""
+    ds_out2d, ds_zcoords, _ = calc_inputs
+    face_z = _calculate_face_z(ds_zcoords, ds_out2d)
+    assert face_z.dims == ("n_face", "time", "n_layer")
+    assert face_z.sizes == {"n_face": N_FACES,
+                            "time": N_TIME,
+                            "n_layer": N_LAYERS}
+
+
+def test_calculate_face_z_tri_cell(calc_inputs):
+    """face_z for a tri cell equals the mean of exactly its 3 node z-values."""
+    ds_out2d, ds_zcoords, ds_zcoords_raw = calc_inputs
+    face_z = _calculate_face_z(ds_zcoords, ds_out2d)
+    z_raw = ds_zcoords_raw.zCoordinates.values  # (time, n_node, n_layer)
+    expected = float(z_raw[0, TRI_NODES, -1].mean())
+    assert float(face_z.values[TRI_IDX, 0, -1]) == pytest.approx(expected,
+                                                                 rel=1e-5)
+
+
+def test_calculate_face_z_tri_ignores_fill_node(calc_inputs):
+    """Guard that the fill node is NOT included in the tri-cell average.
+
+    If node index -1 (wrapping to the last node) were naively included,
+    the result would match a 4-node mean instead of the correct 3-node mean.
+    """
+    ds_out2d, ds_zcoords, ds_zcoords_raw = calc_inputs
+    face_z = _calculate_face_z(ds_zcoords, ds_out2d)
+    z_raw = ds_zcoords_raw.zCoordinates.values
+    # Wrong result if fill wraps to node[-1] = node[N_NODES - 1]
+    wrong_mean = float(z_raw[0, TRI_NODES + [N_NODES - 1], -1].mean())
+    assert float(face_z.values[TRI_IDX, 0, -1]) != pytest.approx(wrong_mean,
+                                                                rel=1e-5)
+
+
+def test_calculate_face_z_quad_cell(calc_inputs):
+    """face_z for a quad cell equals the mean of all 4 node z-values."""
+    ds_out2d, ds_zcoords, ds_zcoords_raw = calc_inputs
+    face_z = _calculate_face_z(ds_zcoords, ds_out2d)
+    z_raw = ds_zcoords_raw.zCoordinates.values
+    expected = float(z_raw[0, QUAD_NODES, -1].mean())
+    assert float(face_z.values[QUAD_IDX, 0, -1]) == pytest.approx(expected,
+                                                                  rel=1e-5)
+
+
+def test_calculate_face_z_tri_and_quad_differ(calc_inputs):
+    """
+    Tri and quad cells produce distinct z-values,
+    confirming correct node counts.
+    """
+    ds_out2d, ds_zcoords, _ = calc_inputs
+    face_z = _calculate_face_z(ds_zcoords, ds_out2d)
+    assert float(face_z.values[TRI_IDX, 0, -1]) != pytest.approx(
+        float(face_z.values[QUAD_IDX, 0, -1]), rel=1e-5
+    )
+
+
+# _calculate_edge_z
+
+def test_calculate_edge_z_output_shape(calc_inputs):
+    """_calculate_edge_z returns (n_edge, time, n_layer)."""
+    ds_out2d, ds_zcoords, _ = calc_inputs
+    edge_z = _calculate_edge_z(ds_zcoords, ds_out2d)
+    assert edge_z.dims == ("n_edge", "time", "n_layer")
+    assert edge_z.sizes == {"n_edge": N_EDGES,
+                            "time": N_TIME,
+                            "n_layer": N_LAYERS}
+
+
+def test_calculate_edge_z_values(calc_inputs):
+    """edge_z equals the two-node mean at every time step."""
+    ds_out2d, ds_zcoords, ds_zcoords_raw = calc_inputs
+    edge_z = _calculate_edge_z(ds_zcoords, ds_out2d)
+    z_raw = ds_zcoords_raw.zCoordinates.values  # (time, n_node, n_layer)
+    expected = z_raw[:, EDGE_NODES, -1].mean(axis=1)  # (time,)
+    np.testing.assert_allclose(edge_z.values[EDGE_IDX, :, -1],
+                               expected, rtol=1e-5)