Preserve coords and attrs in topological aggregations

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

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"}

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"})
 

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
 

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)
 

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],

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)

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],