diff --git a/suxarray/io/_schismgrid.py b/suxarray/io/_schismgrid.py index 5d30dc2..d14aa81 100644 --- a/suxarray/io/_schismgrid.py +++ b/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: diff --git a/tests/test_schismgrid_io.py b/tests/test_schismgrid_io.py new file mode 100644 index 0000000..348427e --- /dev/null +++ b/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)