From 56d7190a81e49a0910af87c3d3bfa843d85eb46a Mon Sep 17 00:00:00 2001 From: Rajeev Jain Date: Mon, 10 Aug 2026 12:51:04 -0500 Subject: [PATCH] Preserve coords and attrs in topological aggregations Node-to-face and node-to-edge aggregations rebuilt the output UxDataArray from data/dims/name only, so every coordinate and all variable metadata were dropped. A result keeps its 'time' dimension but loses the 'time' coordinate, which breaks label-based indexing downstream: .sel(time=...), groupby('time.season') and .resample(time=...) all raise KeyError, and units/long_name are lost for plotting and CF output. Carry over any coordinate that does not span the reduced node dimension, along with the variable attrs. Coordinates along n_node are still dropped, since they no longer match the length of the output dimension. --- test/core/test_topological_agg.py | 63 +++++++++++++++++++++++++++++++ uxarray/core/aggregation.py | 20 ++++++++++ 2 files changed, 83 insertions(+) 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/uxarray/core/aggregation.py b/uxarray/core/aggregation.py index b78bd6dbe..277f97d00 100644 --- a/uxarray/core/aggregation.py +++ b/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"})