Skip to content

Preserve coords and attrs in topological aggregations - #1666

Merged
erogluorhan merged 3 commits into
mainfrom
rajeeja/fix_agg_coords_attrs
Aug 14, 2026
Merged

Preserve coords and attrs in topological aggregations#1666
erogluorhan merged 3 commits into
mainfrom
rajeeja/fix_agg_coords_attrs

Conversation

@rajeeja

@rajeeja rajeeja commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Overview

Fixes #1665.

Node-to-face and node-to-edge aggregations dropped every coordinate and all
variable metadata, so .sel(time=...), groupby('time.season') and
.resample(time=...) raised KeyError on the result, and units/long_name
were lost for plotting and CF output. The fix carries over any coordinate that
does not span the reduced n_node dimension, plus the attrs. Coordinates along
n_node are still dropped, since after the reduction they would be the wrong
length.

Per @erogluorhan's review, the same coord-filtering logic was duplicated across
several other sites: both zonal means, the azimuthal mean, both rectilinear
reshape paths, the constant-latitude cross section, and the two remap backends.
Rather than add a new helper, this generalizes the existing
_preserve_valid_coords from remap/structured.py and moves it to
uxarray/utils/coords.py with optional output_dims and exclude arguments so
every site can share it. It has no uxarray imports, so core, remap and
cross_sections can all use it without import cycles.

One behavior change worth flagging: the YAC weights path keyed its coord dict on
dimension names, so it silently dropped non-dimension coordinates such as an
auxiliary time reference. That is the same bug as #1665 in a second location and
is now fixed, so remap.yac(...) output may carry coordinates it previously
discarded. Nothing that was present before is removed.

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.
@rajeeja
rajeeja requested a review from cmdupuis3 August 10, 2026 21:11
@rajeeja rajeeja added the bug Something isn't working label Aug 10, 2026

@cmdupuis3 cmdupuis3 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tests pass, and the code isn't complicated and it looks fine to me

@erogluorhan erogluorhan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks good to me; thanks for the fix; however, this exact pattern was already duplicated at least three times before (what I can remember from my recent work): in conservative and non-conservative zonal mean and azimuthal mean, see PR #1460 and the later PR #1471 . There might be more.

Would you be interested in coming up with a single shared helper and have zonal_mean--conservative & non-conservative, azimuthal_mean, and topological aggregations all call it?

@rajeeja
rajeeja requested a review from erogluorhan August 12, 2026 19:44

@erogluorhan erogluorhan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd already reviewed this, leaving this comment.

Please let me know if it makes sense.

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.
@rajeeja

rajeeja commented Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

Makes sense — done here. It was actually eight sites, not three: the three in dataarray.py, the aggregations, both rectilinear reshape paths, the constant-latitude cross section, and the two remap backends. I reused the existing _preserve_valid_coords name from remap/structured.py instead of inventing a new one, moved it to uxarray/utils/coords.py, and gave it optional output_dims/exclude args so every site fits; worth flagging that the YAC weights path was keyed on dimension names and silently dropped non-dimension coords, which is the same bug as #1665 and is now fixed too.

@rajeeja

rajeeja commented Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

@cmdupuis3 this grew since your approval — please check once.

@rajeeja
rajeeja requested a review from erogluorhan August 13, 2026 21:49
@cmdupuis3

Copy link
Copy Markdown
Collaborator

@rajeeja Looks fine to me, all tests pass too.

@erogluorhan erogluorhan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks great to me! Thanks for taking care of the broader refactor!

@erogluorhan
erogluorhan merged commit 8159b25 into main Aug 14, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Topological aggregations drop coordinates and attributes

3 participants