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

uxarray/core/aggregation.py

@@ -18,6 +18,22 @@
 }
 
 
+def _non_source_coords(uxda, source_dim):
+    """Coordinates that survive a topological aggregation.
+
+    The source dimension is reduced away, so any coordinate spanning it (the
+    grid dimension itself, or auxiliary coordinates like ``node_lon``) cannot be
+    carried over. Everything else -- most importantly the leading dimensions
+    such as ``time`` or ``lev`` -- is untouched by the aggregation and must be
+    preserved so that label-based indexing keeps working on the result.
+    """
+    return {
+        name: coord
+        for name, coord in uxda.coords.items()
+        if source_dim not in coord.dims
+    }
+
+
 def _uxda_grid_aggregate(uxda, destination, aggregation, **kwargs):
     """Applies a desired aggregation on the data stored in the provided
     UxDataArray."""
@@ -96,6 +112,8 @@ def _node_to_face_aggregation(uxda, aggregation, aggregation_func_kwargs):
         uxgrid=uxda.uxgrid,
         data=aggregated_var,
         dims=uxda.dims,
+        coords=_non_source_coords(uxda, "n_node"),
+        attrs=uxda.attrs,
         name=uxda.name,
     ).rename({"n_node": "n_face"})
 
@@ -164,6 +182,8 @@ def _node_to_edge_aggregation(uxda, aggregation, aggregation_func_kwargs):
         uxgrid=uxda.uxgrid,
         data=aggregation_var,
         dims=uxda.dims,
+        coords=_non_source_coords(uxda, "n_node"),
+        attrs=uxda.attrs,
         name=uxda.name,
     ).rename({"n_node": "n_edge"})