Automatic units handling with UDUNITS inference and strict mode (closes #28) - #32
Merged
Conversation
Every derived variable now carries an inferred `units` attribute, composed
through the typed tree with real UDUNITS arithmetic via cf-units: products
multiply units, sums require compatible summands, difference preserves,
reciprocal inverts, lateral_divergence carries the flux unit. Constants may
declare units via a {value, units} mapping; a bare number is dimensionless.
`collect_budgets(strict=True)` enforces that each budget declaring a top-level
`units` reconciles its lhs/rhs roots to it (else UnitError); non-strict warns.
BudgetQuery gains .units() and .budget_units(). All shipped recipes are
annotated and declare budget units; verified strict passes on the MOM6 example.
psu is treated as dimensionless (PSS-78); degC is dimensionalized in products
so heat-content terms reconcile to W.
Closes #28 (mechanism; ECCO pipeline units follow).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Attach UDUNITS units to the derived budget inputs the ECCO notebooks build (dt, volcello, WVELMASS_interior, boundary_forcing_*) so unit inference has a complete picture, and switch the three budget notebooks (MOM6 + both ECCO) to collect_budgets(..., strict=True). Verified: strict passes on both the MOM6 and ECCO V4r4 (LLC90) datasets; derived variables now display their inferred units. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
View / edit / reply to this conversation on ReviewNB hdrake commented on 2026-07-28T18:53:32Z The recipe's |
|
View / edit / reply to this conversation on ReviewNB hdrake commented on 2026-07-28T18:53:32Z Are there any |
|
View / edit / reply to this conversation on ReviewNB hdrake commented on 2026-07-28T18:53:34Z Line #34. # to-do: add comparable chunking inside the difference and convergence routines Either do this or remove the comment. |
hdrake
commented
Jul 28, 2026
…the repr
Five changes, one per review comment on the units-handling PR.
- Shared constants at the recipe root. A new top-level `constants` key declares
a scalar once (`seawater_density: {value: 1035., units: "kg m-3"}`) and terms
refer to it by name (`density: "seawater_density"`), instead of repeating the
three-line mapping at every use — 36 of them in ECCOV4r4_native.yaml alone.
References are substituted during parsing, so nothing downstream knows the
table existed; a string operand resolves against it first and otherwise stays
a diagnostic name.
- `strict` is now `assert_unit_consistency`, which says what it asserts.
- The recipe display annotates every node with its units in square brackets
(`density: 1035. [kg m-3]`, `lateral [×] [W]`). Units enter at the leaves, so
`show_recipe` takes an optional dataset/grid to read raw diagnostics' units
from and the whole tree resolves; a `BudgetQuery` shown in a notebook prefers
the units its run actually stamped. Backed by `units.infer_units`, which runs
the evaluator's arithmetic statically over the typed tree.
- Dropped the "legacy comment" in the MOM6 notebook: MOM6.yaml has no `var:
null` entries left, so the text now explains what it is really showing.
- Did the ECCO notebook's chunking to-do rather than deleting it.
`lateral_divergence` now gets the same `allow_rechunk` treatment `difference`
has always had, via three helpers the two now share. Verified identical to the
byte on the real ECCO LLC90 budgets, and roughly half the wall clock.
`lateral_divergence` also gains its first CI test, on a two-tile face-connected
synthetic grid; previously it ran only against the data-gated ECCO notebooks.
Verified: full suite green (136 -> 157) with the MOM6 example dataset present,
so the characterization golden is included; `assert_unit_consistency=True`
passes on both the real MOM6 and ECCO V4r4 datasets; all three affected
notebooks re-executed; docs build clean under `-W`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Merged
6 tasks
Review feedback on #33: `_divergence` was vaguer than the operation it wraps. Its body calls the bare `collect.lateral_divergence`, which the docstring now says explicitly so the two are not mistaken for each other. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Address review of #32: shared recipe constants, assert_unit_consistency, units in the repr
…ference
An independent reviewer went through the units feature hunting for bugs and
tried to write a recipe for their own model from the docs alone. Four findings
were worth fixing, plus the doc gap that stopped them cold.
Correctness
- The stamped `units` on the ECCO salt budget were wrong by 1000x. ECCO's salt
fluxes carry the salinity scale in their own units string (`1e-3 m3 s-1`), and
the recipe's compensating `unit_conversion: 0.001` was a bare number — so the
scale survived into the composed units while the values were already kg/s.
Anyone reading the attribute got a thousandth of the real flux. The conversion
constants now declare the reciprocal scale (`{value: 0.001, units: "1e3"}`,
and `kg g-1` for the g->kg one), so value x units == 1 and the scale cancels.
Values are untouched; only the label changed.
- `assert_unit_consistency` now compares magnitude, not just dimension, so it
catches that class of error instead of passing it. `is_convertible` cannot see
a numeric scale, which is exactly what a mis-declared conversion factor leaves
behind. Both shipped models reconcile exactly under the stricter check. This
reverses the "dimension, not exact units" decision documented in #32; the
motivating case there (salinity's `1e-3`) is better handled by making the
recipe cancel the scale than by declining to look at it.
- `lateral_divergence` without an `xgcm.Grid` raised `AttributeError: 'NoneType'
object has no attribute 'axes'` — a regression from the rechunking in #33,
which reads the grid before the guarded helper is called. Both grid-needing
operations now share one `_require_grid`, and the divergence has the
grid-guard test `difference` already had.
- `show_recipe` blanked the units of every term above a missing diagnostic,
including the budget roots — on the raw ECCO dataset, all three rhs roots.
`infer_units` takes an optional `available` predicate and, given one, mirrors
what the evaluator builds (dropping operands it would drop) rather than what
the recipe describes. A term that genuinely will not be built is still blank.
Smaller fixes
- A `constants` entry missing `value` raised a bare `KeyError`; a budget named
`constants` raised the same. Both now say what is wrong, with the path.
- The `reciprocal` operand was the one leaf the HTML tree left unannotated.
- Missing diagnostics are reported before unit issues: an absent input causes
unknown units, so `on_missing="raise"` now names the diagnostic, not the
symptom.
- `units:` on a term gets its own warning instead of the generic "missing an
enclosing product" hint, which pointed at the wrong fix.
- Rechunking is a no-op on unchunked data rather than silently returning dask.
Docs
- `recipes.md` gained the `xgcm.Grid` section it never had. The reviewer got
through everything else from the docs and then had to read `tests/conftest.py`
to learn that xgcm >= 0.10 wants `padding=`, which every older tutorial spells
`boundary=`/`periodic=`.
- The "conversion factors are dimensionless" guidance led them straight into the
bug above; it now says the opposite, with the worked pattern.
- Documented what the assertion does not check, that constants substitute only
in sum/product operands, and that `aggregate()` returns metadata beside sides.
Verified: 178 tests pass (was 157) with both example datasets present;
`assert_unit_consistency=True` passes on real MOM6 and ECCO under the stricter
magnitude check, with ECCO closure residuals unchanged; three notebooks
re-executed; docs build clean under `-W`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Editable installs bind a conda env to a single checkout, so sharing docs_env_xbudget across worktrees silently tests the wrong tree. Require a dedicated docs_env_xbudget_<branch-or-worktree-name> env instead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Act on independent review: unit scale correctness, grid guard, lenient display inference
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Automatic units handling (closes #28)
Every derived variable xbudget materializes now carries an inferred
unitsattribute, and an opt-in
assert_unit_consistency=Truemode verifies that eachbudget reconciles to its declared units. Units compose through the typed tree
with real UDUNITS arithmetic via
cf-units.What's new
its inputs: a
productmultiplies operand units, asumtakes the commonunit of its (dimensionally compatible) summands, a
differencepreservesunits, a
reciprocalinverts them, and alateral_divergencecarries the fluxunit. Inputs' units come from each diagnostic's own
unitsattribute; unknown(missing/unparseable) units propagate rather than being guessed. Results are
stamped as a
unitsattribute on every emitted variable.constantstable.constantsis now the one reserved top-level recipekey: shared scalars are declared once with
{value:, units:}(e.g.seawater_density: {value: 1035., units: "kg m-3"}) and terms refer to them byname (
density: "seawater_density"). References are substituted duringparsing, so nothing downstream sees the table; a string operand resolves
against it first and otherwise stays a diagnostic name. A bare number inline
still works and stays dimensionless — only density and heat capacity needed
annotating, since signs, salinity, and conversion factors are genuinely
dimensionless.
assert_unit_consistency=True.collect_budgets(..., assert_unit_consistency=True)requires every budget that declares a top-level
units(all shipped recipesnow do: mass/salt →
kg s-1, heat →W) to have itslhs/rhsroots inferunits convertible to the declared ones, raising the new
xbudget.UnitErrorotherwise. Left at its
Falsedefault, a mismatch only warns.units.infer_units(node, lookup)answers theunits question straight off the typed tree, without evaluating — which is what
lets
show_recipe(recipe, budget, data)and a notebook-displayedBudgetQueryannotate every node with its
[units].BudgetQuery.units(term)returns the stamped units of amaterialized term;
.budget_units(budget)returns the declared target.psu), undefined in UDUNITS, is treatedas dimensionless (PSS-78); offset temperature (
degC) is dimensionalized tokelvin inside products, so heat-content terms reconcile to
W.cf-unitsis now a dependency (added topyproject.tomland both env files).Also on this branch
lateral_divergencerechunking. The chunking wrapper is now_lateral_divergence, anddifference/lateral_divergenceshare_dataset_chunks/_single_chunk_along/_restore_chunksinevaluate.py.On the ECCO LLC90 budgets this is numerically identical and roughly halves
wall clock.
test_collect.pycovers it on a two-tile face-connected syntheticgrid — the first CI coverage of face-connected differencing.
CLAUDE.mdnow requires developmentand testing in a git worktree to use a dedicated
docs_env_xbudget_<branch-or-worktree-name>env rather than the shareddocs_env_xbudget: an editable install binds an env to one checkout, so ashared env silently tests the wrong tree.
Verified
test_units.py; new unitstests covering the arithmetic, the
constantstable, parsing, staticinference, stamping,
assert_unit_consistencypass/fail, and the queryaccessors — all CI-safe on the synthetic grid).
assert_unit_consistency=Truepasses on the real MOM6 example dataset —all six sides reconcile (mass/salt →
kg s-1, heat →W).assert_unit_consistency=Truepasses on the real ECCO V4r4 (LLC90) dataset— mass/heat/salt all reconcile. The ECCO example notebooks attach units to the
handful of derived inputs they build (
dt,volcello,WVELMASS_interior,boundary_forcing_*); the helper functions ineccov4r4_budget_diagnostics.pyalready stamped theirs.budget notebooks now call
collect_budgets(..., assert_unit_consistency=True)and the derived variables display their inferred
units.Notes for review
degC→K andpsu→dimensionless conventions are physical choices encodedin
xbudget/units.py— worth a look to confirm they match intent.Scale factors: why the check is dimensional, not exact
Unit consistency is checked with
is_convertible(same physical dimension),not exact string equality. This is deliberate, and the salt budget is the
motivating case.
UDUNITS treats a leading numeric factor as a first-class part of a unit — a unit
spec is
[number] [unit-expression], so0.001,1e-3, and0.001 kg s-1areall valid units, where the number is a scale relative to the base unit. CF
encodes practical salinity exactly this way: ECCO's
SALT/SALT_boundscarryunits
1e-3(a salinity of 35 means 0.035 kg/kg), and UDUNITS knows theconversion —
Unit("1e-3").convert(35, "1") == 0.035.Because that
1e-3scale rides along through the salt-flux multiplication, theECCO salt roots infer units of
0.001 kg s-1, not a barekg s-1. That isthe same dimension (mass/time) as
kg s-1, differing only by a numeric scale(
Unit("0.001 kg s-1").is_convertible("kg s-1")isTrue), so the check acceptsit and the budget is confirmed to be a mass-per-time rate. The array values
are already correct SI
kg/s(the recipe's explicitunit_conversion: 0.001does the arithmetic); the
0.001only appears in the composed unit string.Requiring an exact
kg s-1string instead would false-fail on every legitimatescale factor in the recipes (salinity's
1e-3, theunit_conversion: 0.001constants), so the check is built to catch genuine dimensional errors (a term
that is
kg m-2 s-1orWwhen it should bekg s-1) while ignoring harmlessscale. MOM6 salinity carries
psu(aliased to dimensionless"1", no scale),so MOM6 salt infers a clean
kg s-1; both models pass.🤖 Generated with Claude Code