From 56d7190a81e49a0910af87c3d3bfa843d85eb46a Mon Sep 17 00:00:00 2001
From: Rajeev Jain <rajeeja@gmail.com>
Date: Mon, 10 Aug 2026 12:51:04 -0500
Subject: [PATCH 1/2] 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"})
 

From 65aa4cde1f7d2e12d5d872986451b66d8590ee51 Mon Sep 17 00:00:00 2001
From: Rajeev Jain <rajeeja@gmail.com>
Date: Thu, 13 Aug 2026 15:37:37 -0500
Subject: [PATCH 2/2] Consolidate coordinate preservation into a single shared
 helper

The dict comprehension that filters out coordinates spanning a dimension
consumed by an operation was duplicated across eight call sites: the three
zonal/azimuthal means, the topological aggregations added here, both
rectilinear reshape paths, the constant-latitude cross section, and the two
remap backends. Four were byte-identical; the rest differed only by also
excluding coordinates by name or restricting to the output dimensions.

Generalize the existing remap helper into uxarray/utils/coords.py with
optional output_dims and exclude arguments so every site can share it, and
drop the local copies. The module has no uxarray imports, so it is safe to
use from core, remap, and cross_sections alike.

This also fixes the YAC weights path, which keyed its coordinate dict on
dimension names and so silently dropped non-dimension coordinates such as an
auxiliary time reference; it now preserves them like every other path.
---
 test/utils/test_coords.py                    | 82 ++++++++++++++++++++
 uxarray/core/aggregation.py                  | 21 +----
 uxarray/core/dataarray.py                    | 17 +---
 uxarray/cross_sections/dataarray_accessor.py |  7 +-
 uxarray/remap/apply_weights.py               |  8 +-
 uxarray/remap/structured.py                  | 38 +++------
 uxarray/remap/yac.py                         |  4 +-
 uxarray/utils/coords.py                      | 54 +++++++++++++
 8 files changed, 160 insertions(+), 71 deletions(-)
 create mode 100644 test/utils/test_coords.py
 create mode 100644 uxarray/utils/coords.py

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 a090ed088..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,
@@ -18,22 +19,6 @@
 }
 
 
-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."""
@@ -112,7 +97,7 @@ 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"),
+        coords=_preserve_valid_coords(uxda, "n_node"),
         attrs=uxda.attrs,
         name=uxda.name,
     ).rename({"n_node": "n_face"})
@@ -182,7 +167,7 @@ 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"),
+        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))
+    }