SCH-1073: Add face-centered subset support

suxarray/grid/grid.py

@@ -106,7 +106,9 @@ def isel(self, **dim_kwargs):
     def sgrid_isel(self, **kwargs):
         from uxarray.constants import GRID_DIMS
 
-        kwargs_no_grid_dim = {k: v for k, v in kwargs.items() if k not in GRID_DIMS}
+        # Filter out grid dimensions and inverse_indices (not a valid xarray dimension)
+        excluded_keys = set(GRID_DIMS) | {"inverse_indices"}
+        kwargs_no_grid_dim = {k: v for k, v in kwargs.items() if k not in excluded_keys}
         if kwargs_no_grid_dim:
             sgrid_info = self.sgrid_info.copy().isel(kwargs_no_grid_dim)
         else:

suxarray/io/_schismgrid.py

@@ -2,6 +2,7 @@
 from typing import Optional, Union
 import pyproj
 import xarray as xr
+import numpy as np
 from uxarray.conventions.ugrid import (
     CARTESIAN_NODE_COORDINATES,
     CARTESIAN_FACE_COORDINATES,
@@ -73,7 +74,10 @@ def _read_schism_grid(
         }
         ds_zcoords = ds_zcoords.swap_dims(dim_dict_zcoords)
         if "nSCHISM_vgrid_layers" in ds_zcoords.dims:
-            ds_zcoords = ds_zcoords.swap_dims({"nSCHISM_vgrid_layers": "n_layer"})
+            ds_zcoords = ds_zcoords.swap_dims({"nSCHISM_vgrid_layers":
+                                               "n_layer"})
+
+        ds_out2d = _assign_z_coords(ds_out2d, ds_zcoords)
 
     ds_sgrid_info = xr.Dataset()
     for varname in SCHISM_GRID_TIME_VARIABLES:
@@ -83,6 +87,127 @@ def _read_schism_grid(
 
     return ds_out2d, ds_sgrid_info
 
+def _assign_z_coords(ds_out2d, ds_zcoords) -> xr.Dataset:
+    """Assign z-coordinates to the grid dataset.
+
+    Creates both static 1D zCoordinates (bottom layer) for subsetting operations
+    and full 3D zCoordinates as data variables for actual z-coordinate access.
+
+    Parameters
+    ----------
+    ds_out2d : xr.Dataset
+        Dataset containing grid connectivity information
+    ds_zcoords : xr.Dataset
+        Dataset containing zCoordinates variable with dims (time,n_node,n_layer)
+
+    Returns
+    -------
+    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.
+    """
+    # Calculate full 3D z-coordinates (transposed to have spatial dim first)
+    node_z_3d = ds_zcoords.zCoordinates.values.transpose(1, 0, 2)
+    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,)
+
+    # 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),
+    })
+
+    # 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),
+    })
+
+    return ds_out2d
+
+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.
+    """
+    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)
+    return edge_z
+
+
+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)
+
+    Returns
+    -------
+    xr.Dataset
+        ds_sgrid_info with face_z coordinate added, shape (time,n_edge,n_layer)
+    """
+
+    # TODO: Make this dask-compatible - same approach as _calculate_edge_z.
+    face_ids = ds_out2d.face_node_connectivity.values  # shape (n_face, 4)
+
+    # Identify triangular and quad elements for correct face calculation
+    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)
+    return face_z
 
 def _find_grid_topology_varname(ds: xr.Dataset) -> str:
     """Find the UGRID grid topology variable name in the dataset

tests/test_suxarray.py

@@ -81,6 +81,20 @@ def test_subset_bounding_polygon(sxds_test_dask):
     assert da_subset.uxgrid.n_face == 10
 
 
+def test_subset_elements_bounding_polygon(sxds_element_test):
+    """Test subsetting dataarray with data at elements with a polygon."""
+    grid = sxds_element_test.uxgrid
+
+    lon_bounds = (grid.node_lon.min().values, grid.node_lon.max().values)
+    lat_bounds = (grid.node_lat.min().values, 0)
+
+    vertVel = sxds_element_test["verticalVelAtElement"]
+
+    da_subset = vertVel.subset.bounding_box(lon_bounds, lat_bounds, [['face']])
+    assert da_subset.uxgrid.n_node == 1319
+    assert da_subset.uxgrid.n_face == 2221
+
+
 def test_subset_bounding_box(sxds_test_dask):
     """Test find_element_at"""
     da_subset = sxds_test_dask["salinity"].subset.bounding_box_xy(

tests/testfixtures.py

@@ -29,3 +29,22 @@ def sxds_test_dask():
     ).astype(np.float64)
     sxds = sx.read_schism_nc(sxgrid, ds_salinity)
     return sxds
+
+
+@pytest.fixture(scope="module")
+def sxds_element_test():
+    """Test output at element fixture"""
+    p_cur = Path(__file__).parent.absolute()
+    path_out2d = [str(p_cur / "testdata/out2d_alt_{}.nc".format(i)) for i in range(1, 3)]
+    path_zcoords = [
+        str(p_cur / "testdata/zCoordinates_alt_{}.nc".format(i)) for i in range(1, 3)
+    ]
+    chunks = {"time": 12}
+    sxgrid = sx.core.api.open_grid(path_out2d, path_zcoords, chunks=chunks)
+    path_var = [str(p_cur / "testdata/verticalVelAtElement_{}.nc".format(i))
+                for i in range(1, 3)]
+    ds_vertVel = xr.open_mfdataset(
+        path_var, mask_and_scale=False, data_vars="minimal"
+    ).astype(np.float64)
+    sxds = sx.read_schism_nc(sxgrid, ds_vertVel)
+    return sxds