Preserve coords and attrs in topological aggregations - #1666
Conversation
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.
cmdupuis3
left a comment
There was a problem hiding this comment.
Tests pass, and the code isn't complicated and it looks fine to me
erogluorhan
left a comment
There was a problem hiding this comment.
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?
erogluorhan
left a comment
There was a problem hiding this comment.
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.
|
Makes sense — done here. It was actually eight sites, not three: the three in |
|
@cmdupuis3 this grew since your approval — please check once. |
|
@rajeeja Looks fine to me, all tests pass too. |
erogluorhan
left a comment
There was a problem hiding this comment.
It looks great to me! Thanks for taking care of the broader refactor!
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=...)raisedKeyErroron the result, andunits/long_namewere lost for plotting and CF output. The fix carries over any coordinate that
does not span the reduced
n_nodedimension, plus the attrs. Coordinates alongn_nodeare still dropped, since after the reduction they would be the wronglength.
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_coordsfromremap/structured.pyand moves it touxarray/utils/coords.pywith optionaloutput_dimsandexcludearguments soevery site can share it. It has no uxarray imports, so
core,remapandcross_sectionscan 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 previouslydiscarded. Nothing that was present before is removed.