diff --git a/suxarray/grid/grid.py b/suxarray/grid/grid.py index 8f471d8..9e16075 100644 --- a/suxarray/grid/grid.py +++ b/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: diff --git a/suxarray/io/_schismgrid.py b/suxarray/io/_schismgrid.py index c49caa8..db609d1 100644 --- a/suxarray/io/_schismgrid.py +++ b/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 diff --git a/tests/test_suxarray.py b/tests/test_suxarray.py index bc5df3f..61dc8ee 100644 --- a/tests/test_suxarray.py +++ b/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( diff --git a/tests/testdata/out2d_alt_1.nc b/tests/testdata/out2d_alt_1.nc new file mode 100644 index 0000000..690df41 Binary files /dev/null and b/tests/testdata/out2d_alt_1.nc differ diff --git a/tests/testdata/out2d_alt_2.nc b/tests/testdata/out2d_alt_2.nc new file mode 100644 index 0000000..43b422b Binary files /dev/null and b/tests/testdata/out2d_alt_2.nc differ diff --git a/tests/testdata/verticalVelAtElement_1.nc b/tests/testdata/verticalVelAtElement_1.nc new file mode 100644 index 0000000..d7d81b7 Binary files /dev/null and b/tests/testdata/verticalVelAtElement_1.nc differ diff --git a/tests/testdata/verticalVelAtElement_2.nc b/tests/testdata/verticalVelAtElement_2.nc new file mode 100644 index 0000000..91253c4 Binary files /dev/null and b/tests/testdata/verticalVelAtElement_2.nc differ diff --git a/tests/testdata/zCoordinates_alt_1.nc b/tests/testdata/zCoordinates_alt_1.nc new file mode 100644 index 0000000..c97e57e Binary files /dev/null and b/tests/testdata/zCoordinates_alt_1.nc differ diff --git a/tests/testdata/zCoordinates_alt_2.nc b/tests/testdata/zCoordinates_alt_2.nc new file mode 100644 index 0000000..0e8c708 Binary files /dev/null and b/tests/testdata/zCoordinates_alt_2.nc differ diff --git a/tests/testfixtures.py b/tests/testfixtures.py index 0ff1970..1423df0 100644 --- a/tests/testfixtures.py +++ b/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 \ No newline at end of file