diff --git a/test/core/test_topological_agg.py b/test/core/test_topological_agg.py index 94ce1f58d..264d103e3 100644 --- a/test/core/test_topological_agg.py +++ b/test/core/test_topological_agg.py @@ -1,5 +1,8 @@ import uxarray as ux +import numpy as np +import numpy.testing as nt +import pandas as pd import pytest @@ -32,3 +35,63 @@ def test_node_to_edge_aggs(gridpath): grid_reduction = getattr(uxds['areaTriangle'], agg_func)(destination='edge') assert 'n_edge' in grid_reduction.dims + + +def _timeseries_uxda(gridpath): + """Node-centered data with a labelled time axis and CF-style attributes.""" + uxgrid = ux.open_grid(gridpath("mpas", "QU", "oQU480.231010.nc")) + rng = np.random.default_rng(0) + return ux.UxDataArray( + rng.random((6, uxgrid.n_node)), + dims=("time", "n_node"), + coords={"time": pd.date_range("2000-01-01", periods=6, freq="MS")}, + uxgrid=uxgrid, + name="var", + attrs={"units": "m", "long_name": "sea surface height"}, + ) + + +@pytest.mark.parametrize("destination", ["face", "edge"]) +def test_agg_preserves_leading_coords_and_attrs(gridpath, destination): + """Aggregating over the node dimension must not discard the leading + coordinates or the variable metadata. Regression test for topological + aggregations returning a coordinate-less result, which broke label-based + indexing (``.sel``/``.groupby``/``.resample``) on the output. + """ + uxda = _timeseries_uxda(gridpath) + + for agg_func in AGGS: + result = getattr(uxda, agg_func)(destination=destination) + + assert "time" in result.coords + nt.assert_array_equal(result.time.values, uxda.time.values) + assert result.attrs == uxda.attrs + + +@pytest.mark.parametrize("destination", ["face", "edge"]) +def test_agg_result_supports_label_based_indexing(gridpath, destination): + """The preserved time axis must actually be usable downstream.""" + result = _timeseries_uxda(gridpath).topological_mean(destination=destination) + + grid_dim = f"n_{destination}" + assert result.sel(time="2000-03-01").dims == (grid_dim,) + assert ( + result.groupby("time.season").mean().sizes[grid_dim] == result.sizes[grid_dim] + ) + assert result.resample(time="QS").mean().sizes["time"] == 2 + + +@pytest.mark.parametrize("destination", ["face", "edge"]) +def test_agg_drops_node_spanning_coords(gridpath, destination): + """Coordinates along the reduced dimension cannot be carried over, since + they no longer match the length of the output dimension. + """ + uxda = _timeseries_uxda(gridpath) + rng = np.random.default_rng(1) + uxda = uxda.assign_coords(node_lon=("n_node", rng.random(uxda.uxgrid.n_node))) + + result = uxda.topological_mean(destination=destination) + + assert "node_lon" not in result.coords + assert "n_node" not in result.dims + assert "time" in result.coords diff --git a/test/utils/test_coords.py b/test/utils/test_coords.py new file mode 100644 index 000000000..6130de55d --- /dev/null +++ b/test/utils/test_coords.py @@ -0,0 +1,82 @@ +import numpy as np +import pytest +import xarray as xr + +from uxarray.utils.coords import _preserve_valid_coords + + +@pytest.fixture +def da(): + """A DataArray carrying every kind of coordinate the helper must classify.""" + return xr.DataArray( + np.zeros((2, 3)), + dims=("time", "n_face"), + coords={ + "time": [1, 2], + "n_face": [0, 1, 2], + "lat": ("n_face", [10.0, 20.0, 30.0]), + "scalar": 5, + }, + ) + + +def test_drops_coords_spanning_dropped_dim(da): + coords = _preserve_valid_coords(da, "n_face") + + assert set(coords) == {"time", "scalar"} + + +def test_keeps_everything_when_no_filters_given(da): + coords = _preserve_valid_coords(da) + + assert set(coords) == set(da.coords) + + +def test_output_dims_drops_coords_on_absent_dims(da): + """A coordinate on a dimension missing from the result cannot be carried.""" + coords = _preserve_valid_coords(da, output_dims={"time"}) + + assert set(coords) == {"time", "scalar"} + + +def test_scalar_coords_always_survive(da): + """Dimensionless coords span nothing, so no filter can invalidate them.""" + coords = _preserve_valid_coords(da, "n_face", output_dims=set()) + + assert set(coords) == {"scalar"} + + +def test_exclude_drops_by_name_regardless_of_dims(da): + coords = _preserve_valid_coords(da, "n_face", exclude={"scalar"}) + + assert set(coords) == {"time"} + + +def test_dropped_dim_and_output_dims_compose(da): + """Both filters apply; a coord must satisfy each one to survive.""" + coords = _preserve_valid_coords(da, "n_face", output_dims={"n_face", "n_lat"}) + + assert set(coords) == {"scalar"} + + +def test_returns_the_original_coordinate_objects(da): + coords = _preserve_valid_coords(da, "n_face") + + assert coords["time"].equals(da.coords["time"]) + + +def test_result_is_accepted_by_the_dataarray_constructor(da): + """The mapping must be usable directly as a ``coords`` argument.""" + coords = _preserve_valid_coords(da, "n_face") + + result = xr.DataArray(np.zeros((2, 4)), dims=("time", "n_edge"), coords=coords) + + assert result.sel(time=1).sizes == {"n_edge": 4} + + +def test_works_on_datasets(da): + ds = da.to_dataset(name="v") + + coords = _preserve_valid_coords(ds, "n_face") + + assert set(coords) == {"time", "scalar"} diff --git a/uxarray/core/aggregation.py b/uxarray/core/aggregation.py index 4f75aaeb8..b7ffb5b40 100644 --- a/uxarray/core/aggregation.py +++ b/uxarray/core/aggregation.py @@ -3,6 +3,7 @@ import uxarray.core.dataarray from uxarray.errors import DataCenteringError from uxarray.grid.connectivity import get_face_node_partitions +from uxarray.utils.coords import _preserve_valid_coords NUMPY_AGGREGATIONS = { "mean": np.mean, @@ -96,6 +97,8 @@ def _node_to_face_aggregation(uxda, aggregation, aggregation_func_kwargs): uxgrid=uxda.uxgrid, data=aggregated_var, dims=uxda.dims, + coords=_preserve_valid_coords(uxda, "n_node"), + attrs=uxda.attrs, name=uxda.name, ).rename({"n_node": "n_face"}) @@ -164,6 +167,8 @@ def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs): uxgrid=uxda.uxgrid, data=aggregation_var, dims=uxda.dims, + coords=_preserve_valid_coords(uxda, "n_node"), + attrs=uxda.attrs, name=uxda.name, ).rename({"n_node": "n_edge"}) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index c3203989f..f265ad53b 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -39,6 +39,7 @@ from uxarray.plot.accessor import UxDataArrayPlotAccessor from uxarray.remap.accessor import RemapAccessor from uxarray.subset import DataArraySubsetAccessor +from uxarray.utils.coords import _preserve_valid_coords if TYPE_CHECKING: import cartopy.crs as ccrs @@ -742,11 +743,7 @@ def zonal_mean(self, lat=(-90, 90, 10), conservative: bool = False, **kwargs): dims[face_axis] = "latitudes" # Assign coords from `self` to the result except one that corresponds to `dims[face_axis]` - new_coords = { - k: v - for k, v in self.coords.items() - if self.dims[face_axis] not in v.dims - } + new_coords = _preserve_valid_coords(self, "n_face") # Add latitudes to the resulting coords new_coords["latitudes"] = latitudes @@ -796,11 +793,7 @@ def zonal_mean(self, lat=(-90, 90, 10), conservative: bool = False, **kwargs): dims[face_axis] = "latitudes" # Assign coords from `self` to the result except one that corresponds to `dims[face_axis]` - new_coords = { - k: v - for k, v in self.coords.items() - if self.dims[face_axis] not in v.dims - } + new_coords = _preserve_valid_coords(self, "n_face") # Add latitudes to the resulting coords new_coords["latitudes"] = centers @@ -995,9 +988,7 @@ def azimuthal_mean( ) # Assign coords from `self` to the result except one that corresponds to `dims[face_axis]` - new_coords = { - k: v for k, v in self.coords.items() if self.dims[face_axis] not in v.dims - } + new_coords = _preserve_valid_coords(self, "n_face") # Add radii_deg to the resulting coords new_coords["radius"] = radii_deg diff --git a/uxarray/cross_sections/dataarray_accessor.py b/uxarray/cross_sections/dataarray_accessor.py index be39a9af4..801f71a71 100644 --- a/uxarray/cross_sections/dataarray_accessor.py +++ b/uxarray/cross_sections/dataarray_accessor.py @@ -7,6 +7,7 @@ from uxarray.constants import INT_DTYPE from uxarray.errors import DataCenteringError +from uxarray.utils.coords import _preserve_valid_coords from .sample import ( _fill_numba, @@ -133,11 +134,7 @@ def __call__( data = np.moveaxis(filled, -1, dim_axis) # Build coords dict: keep everything except 'n_face' - coords = { - name: self.uxda.coords[name] - for name in self.uxda.coords - if name != "n_face" and "n_face" not in self.uxda.coords[name].dims - } + coords = _preserve_valid_coords(self.uxda, "n_face") # index along the arc coords[new_dim] = np.arange(steps) diff --git a/uxarray/remap/apply_weights.py b/uxarray/remap/apply_weights.py index 1c4e6efa7..8d68b7814 100644 --- a/uxarray/remap/apply_weights.py +++ b/uxarray/remap/apply_weights.py @@ -7,6 +7,7 @@ import uxarray.core.dataarray from uxarray.errors import DimensionError +from uxarray.utils.coords import _preserve_valid_coords from .utils import ( LABEL_TO_COORD, @@ -90,12 +91,7 @@ def _apply_weights( da_t = da.transpose(*other_dims, variable_source_dim) remapped_values = weights_obj._apply(np.asarray(da_t.values)) - other_dims_set = set(other_dims) - coords = { - coord_name: coord - for coord_name, coord in da.coords.items() - if set(coord.dims).issubset(other_dims_set) - } + coords = _preserve_valid_coords(da, variable_source_dim, other_dims) da_out = uxarray.core.dataarray.UxDataArray( remapped_values, dims=other_dims + [destination_dim], diff --git a/uxarray/remap/structured.py b/uxarray/remap/structured.py index 13064cd63..b9a616b65 100644 --- a/uxarray/remap/structured.py +++ b/uxarray/remap/structured.py @@ -6,6 +6,7 @@ import xarray as xr from uxarray.errors import DimensionError +from uxarray.utils.coords import _preserve_valid_coords @dataclass(frozen=True) @@ -127,21 +128,6 @@ def _normalize_rectilinear_target(lon, lat) -> RectilinearGridSpec: ) -def _preserve_valid_coords( - da: xr.DataArray, - dropped_dim: str, - output_dims: tuple[str, ...] | list[str], -) -> dict[str, xr.DataArray]: - """Keep only coords that remain valid after replacing ``dropped_dim``.""" - - output_dims = set(output_dims) - return { - name: coord - for name, coord in da.coords.items() - if dropped_dim not in coord.dims and set(coord.dims).issubset(output_dims) - } - - def _reshape_array_to_rectilinear( da: xr.DataArray, spec: RectilinearGridSpec ) -> xr.DataArray: @@ -157,12 +143,11 @@ def _reshape_array_to_rectilinear( shape = da.shape[:axis] + spec.shape + da.shape[axis + 1 :] dims = da.dims[:axis] + (spec.lat_dim, spec.lon_dim) + da.dims[axis + 1 :] - coords = { - name: coord - for name, coord in da.coords.items() - if "n_face" not in coord.dims - and name not in {spec.lat_name, spec.lon_name, spec.lat_dim, spec.lon_dim} - } + coords = _preserve_valid_coords( + da, + "n_face", + exclude={spec.lat_name, spec.lon_name, spec.lat_dim, spec.lon_dim}, + ) coords[spec.lat_name] = spec.lat coords[spec.lon_name] = spec.lon @@ -191,12 +176,11 @@ def _reshape_to_rectilinear(obj, spec: RectilinearGridSpec): name: _reshape_array_to_rectilinear(da, spec) for name, da in xr_obj.data_vars.items() } - coords = { - name: coord - for name, coord in xr_obj.coords.items() - if "n_face" not in coord.dims - and name not in {spec.lat_name, spec.lon_name, spec.lat_dim, spec.lon_dim} - } + coords = _preserve_valid_coords( + xr_obj, + "n_face", + exclude={spec.lat_name, spec.lon_name, spec.lat_dim, spec.lon_dim}, + ) coords[spec.lat_name] = spec.lat coords[spec.lon_name] = spec.lon return xr.Dataset(data_vars=data_vars, coords=coords, attrs=xr_obj.attrs) diff --git a/uxarray/remap/yac.py b/uxarray/remap/yac.py index 93d1d4fea..094a9da10 100644 --- a/uxarray/remap/yac.py +++ b/uxarray/remap/yac.py @@ -17,7 +17,6 @@ from uxarray.remap.structured import ( RectilinearGridSpec, _normalize_rectilinear_target, - _preserve_valid_coords, _reshape_to_rectilinear, ) from uxarray.remap.utils import ( @@ -27,6 +26,7 @@ _get_remap_dims, _to_dataset, ) +from uxarray.utils.coords import _preserve_valid_coords @dataclass @@ -448,7 +448,7 @@ def _yac_remap(source, destination_grid, remap_to: str, yac_method: str, yac_kwa out_shape = src_values.shape[:-1] + (remapper._tgt_size,) out_values = out_flat.reshape(out_shape) - coords = {dim: da.coords[dim] for dim in other_dims if dim in da.coords} + coords = _preserve_valid_coords(da, src_dim, other_dims) da_out = uxarray.core.dataarray.UxDataArray( out_values, dims=other_dims + [destination_dim], diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py new file mode 100644 index 000000000..d82214d05 --- /dev/null +++ b/uxarray/utils/coords.py @@ -0,0 +1,54 @@ +"""Utilities for carrying coordinates across operations that change dimensions.""" + +from __future__ import annotations + +from typing import Hashable, Iterable, Mapping + +import xarray as xr + + +def _preserve_valid_coords( + obj: xr.DataArray | xr.Dataset, + dropped_dim: str | None = None, + output_dims: Iterable[Hashable] | None = None, + exclude: Iterable[Hashable] | None = None, +) -> Mapping[Hashable, xr.DataArray]: + """Keep only the coordinates that remain valid on the result of an operation. + + Operations such as topological aggregations, zonal and azimuthal means, and + remapping consume one dimension and replace it with another. Any coordinate + spanning the consumed dimension no longer matches the output shape and has to + be dropped, but every other coordinate -- most importantly the leading ones + such as ``time`` or ``lev`` -- is untouched and must be carried over so that + label-based indexing keeps working on the result. + + Parameters + ---------- + obj : xr.DataArray or xr.Dataset + Object whose coordinates are being filtered. + dropped_dim : str, optional + Dimension consumed by the operation. Coordinates spanning it are dropped. + output_dims : iterable of hashable, optional + Dimensions present on the result. Coordinates spanning any dimension not + in this set are dropped. Useful when the operation also removes or + reshapes dimensions other than ``dropped_dim``. + exclude : iterable of hashable, optional + Coordinate names to drop regardless of their dimensions, for cases where + the caller supplies its own replacement under the same name. + + Returns + ------- + dict + Mapping of coordinate name to coordinate, suitable for passing straight + to the ``coords`` argument of a DataArray or Dataset constructor. + """ + output_dims = None if output_dims is None else set(output_dims) + exclude = frozenset() if exclude is None else frozenset(exclude) + + return { + name: coord + for name, coord in obj.coords.items() + if name not in exclude + and (dropped_dim is None or dropped_dim not in coord.dims) + and (output_dims is None or set(coord.dims).issubset(output_dims)) + }