Automatic units handling with UDUNITS inference and strict mode (closes #28)

CHANGELOG.md

@@ -2,6 +2,85 @@
 
 ## Unreleased
 
+### Units handling (issue #28)
+
+- Every derived variable is now stamped with an inferred `units` attribute.
+  Units compose through the tree with real UDUNITS arithmetic (via `cf-units`):
+  a `product` multiplies its operands' units, a `sum` takes the common unit of
+  its (dimensionally compatible) summands, a `difference` preserves units, a
+  `reciprocal` inverts them, and a `lateral_divergence` carries the flux unit.
+  Inputs' units come from each diagnostic's own `units` attribute; hard-coded
+  constants declare theirs in the recipe.
+- Recipe constants may now carry units via a `{value:, units:}` mapping, e.g.
+  `density: {value: 1035., units: "kg m-3"}`. A bare number (`sign: -1.`) is
+  dimensionless. All shipped recipes are annotated (only `density` and
+  `specific_heat_capacity` needed it; salinity/sign/conversion factors are
+  dimensionless) and declare a top-level budget `units` (`mass`/`salt` →
+  `kg s-1`, `heat` → `W`).
+- Constants shared across a recipe can be declared once under a new top-level
+  `constants` key and referred to by name (`density: "seawater_density"`),
+  instead of repeating a `{value:, units:}` mapping at every use. `constants` is
+  the only top-level key that is not a budget; a string operand resolves against
+  it first and otherwise stays a diagnostic name. The shipped recipes use it for
+  their reference density and heat capacity (36 repeated blocks in
+  `ECCOV4r4_native.yaml` alone).
+- `collect_budgets(..., assert_unit_consistency=True)` enforces units: each
+  budget that declares a top-level `units` must have its side roots infer *the
+  same* units — same dimension **and** same magnitude — raising `UnitError`
+  (exported) otherwise. Left at its default of `False`, a reconciliation failure
+  only warns; inference and stamping happen either way. A side root that never
+  materialized is skipped, since absent inputs are `on_missing`'s business.
+- Magnitude is checked because a leading numeric factor is part of a UDUNITS
+  spec, and so part of what a stamped attribute claims: a root inferring
+  `0.001 kg s-1` while holding values already in kg/s is mislabelled by a
+  thousand. That is what a conversion factor written as a bare (dimensionless)
+  number does when it cancels a scale carried in an input's own `units` — the
+  case CF's `1e-3` practical salinity creates. Such a factor must declare the
+  reciprocal scale as its units (`{value: 0.001, units: "1e3"}`), which
+  `ECCOV4r4_native.yaml` now does; its salt budget previously stamped
+  `0.001 kg s-1` on correct kg/s values.
+- `BudgetQuery` gained `.units(term)` (the stamped units of a materialized
+  term) and `.budget_units(budget)` (the declared target).
+- The recipe display (`show_recipe`, and a `BudgetQuery` shown in a notebook)
+  annotates every node with its units in square brackets — `density: 1035.
+  [kg m-3]`, `lateral [×] [W]`. `show_recipe` takes an optional third argument,
+  the grid/dataset to read raw diagnostics' units from, so a recipe can be read
+  with its full unit arithmetic before anything is collected.
+- `units.infer_units(node, lookup, available=)` takes an optional
+  "which diagnostics does the dataset hold?" predicate, and with it answers what
+  a run will *build* rather than what the recipe *describes* — dropping the
+  operands the evaluator drops instead of letting one absent diagnostic leave
+  every term above it unannotated. The display passes it whenever it has data.
+- Practical salinity (`psu`), which UDUNITS does not define, is treated as
+  dimensionless (PSS-78); offset temperature (`degC`) is handled by cf-units'
+  in-product dimensionalization, so heat-content terms reconcile to `W`.
+- `cf-units` is now a dependency.
+
+### Engine
+
+- A recipe using `lateral_divergence` without an `xgcm.Grid` now raises the same
+  clear `ValueError` a grid-less `difference` does, instead of failing later on
+  the grid-dependent rechunking bookkeeping.
+- Missing diagnostics are reported before unit issues: an absent input *causes*
+  unknown units, so under `on_missing="raise"` the error now names the
+  diagnostic rather than the symptom.
+- Rechunking is a no-op on unchunked (in-memory) data, which is already one chunk
+  along everything — `difference` and `lateral_divergence` no longer hand back a
+  lazy result where the caller passed an eager one.
+- A budget-level key written on a term (`units:` on a term rather than beside
+  `lhs`/`rhs`) gets its own warning instead of the generic "you may be missing an
+  enclosing `product`" hint, which pointed at the wrong fix.
+- A `constants` entry with no `value` raises `BudgetParseError` naming the path,
+  not a bare `KeyError`; and a budget named `constants` says the key is reserved.
+- `lateral_divergence` now gets the same `allow_rechunk` treatment `difference`
+  has always had: its flux operands are put into a single chunk along the
+  differenced dimensions while xgcm works, and the caller's chunking is restored
+  afterwards. Previously only `difference` rechunked, so a horizontally chunked
+  flux made xgcm shuffle across chunk (and face) boundaries. Verified identical
+  to the byte on the ECCO V4r4 LLC90 budgets, roughly halving their wall clock.
+  `lateral_divergence` now also has a CI test, on a two-tile face-connected
+  synthetic grid.
+
 ### Packaging
 
 - `pyyaml` is now declared as a dependency. It is imported at package import

CLAUDE.md

@@ -28,13 +28,21 @@ conda activate docs_env_xbudget
 pip install -e .
 ```
 
+**Worktrees get their own env.** Any development or testing work in a git worktree must use a dedicated conda env named `docs_env_xbudget_<branch-or-worktree-name>` — never the shared `docs_env_xbudget`. The editable install (`pip install -e .`) points an env at a single checkout, so sharing one env across worktrees silently tests the wrong tree. For the `units-handling` branch:
+
+```bash
+conda env create -f docs/environment.yml -n docs_env_xbudget_units-handling
+conda activate docs_env_xbudget_units-handling
+pip install -e .   # run from inside the worktree
+```
+
 ## Core architecture
 
 The central abstraction is the **`recipe`** — a nested provenance tree (loaded from a YAML *recipe* file) describing how to build each budget term from raw diagnostics. It is the public input format. Internally it is parsed into a typed expression tree and evaluated.
 
 ### The recipe tree (input format — unchanged)
 
-Top-level keys are budgets (`mass`, `heat`, `salt`). Each budget has `lhs` and/or `rhs` sub-trees plus metadata keys (`lambda`, `thickness`, `surface_lambda`) that the engine does not interpret. Within a side, terms nest recursively. A node names a diagnostic with a `var` key (`var: "thetao"`); a derived term omits `var` entirely (the `var: null` placeholders are optional as of 0.7.0) and carries one or more **operation** keys instead:
+Top-level keys are budgets (`mass`, `heat`, `salt`) — except `constants`, the one reserved key, which declares scalars shared across the recipe (`seawater_density: {value: 1035., units: "kg m-3"}`) that terms then refer to by name (`density: "seawater_density"`). References are substituted during parsing, so nothing downstream sees the table; a string operand resolves against it first and otherwise stays a diagnostic name. Each budget has `lhs` and/or `rhs` sub-trees plus metadata keys (`units`, `lambda`, `thickness`, `surface_lambda`) that the engine does not interpret as terms. Within a side, terms nest recursively. A node names a diagnostic with a `var` key (`var: "thetao"`); a derived term omits `var` entirely (the `var: null` placeholders are optional as of 0.7.0) and carries one or more **operation** keys instead:
 
 - `sum` — add the child terms together
 - `product` — multiply child terms (scalar numbers allowed as factors, e.g. `density: 1035.`, `sign: -1.`)
@@ -63,20 +71,23 @@ recipe ──parse_budgets──▶ typed tree (nodes.py) ──evaluate_budgets
 
 - **Naming.** One variable per node/operation with operator infixes dropped; no duplicate "copy" variables (108 → 57 on the MOM6 example; 73 on the ECCOv4r4 native-grid recipe, whose volume budget is Eulerian-only — the GM bolus is exposed as `bolus` metadata, not a budget term). `CHANGELOG.md` has the old→new table for anyone migrating from 0.6.x.
 - **Missing diagnostics are skipped, not fatal** — a `sum` drops only the missing operand and builds from the rest (flagged `xbudget_incomplete`); a `product` with a missing *required* factor is dropped entirely (an unknown factor is not a zero one). `collect_budgets(on_missing="warn"|"raise"|"ignore")` sets the policy; `optional: true` on a term exempts its subtree. Query it back with `BudgetQuery.missing()`.
-- **`difference` rechunking:** `allow_rechunk=True` (default) temporarily rechunks the difference dimension into a single chunk (required by `grid.diff`) then restores chunking.
+- **Rechunking:** `allow_rechunk=True` (default) temporarily rechunks the differenced dimension(s) into a single chunk (required by `grid.diff`) then restores the caller's chunking. This applies to **both** `difference` and `lateral_divergence` — they share `_dataset_chunks`/`_single_chunk_along`/`_restore_chunks` in `evaluate.py`. On the ECCO LLC90 budgets it is numerically identical and roughly halves wall clock.
+- **Units.** Every emitted variable is stamped with a `units` attr inferred through the tree (`units.py`, UDUNITS via `cf-units`); `collect_budgets(..., assert_unit_consistency=True)` additionally requires each budget's side roots to reconcile to its declared top-level `units` — same dimension **and** same magnitude — raising `UnitError` (it only warns by default). A side root that never materialized is exempt (that is `on_missing`'s job). Magnitude matters because a UDUNITS spec may carry a leading numeric scale: a conversion factor cancelling a scale in an input's own units (CF writes practical salinity as `1e-3`) must declare the reciprocal scale as *its* units (`{value: 0.001, units: "1e3"}`), or the term stamps itself `0.001 kg s-1` over correct kg/s values.
+- **Static unit inference.** `units.infer_units(node, lookup, available=)` answers the units question off the typed tree, without evaluating — which is what lets `show_recipe(recipe, budget, data)` and a notebook-displayed `BudgetQuery` annotate every node with `[units]`. Passing `available` (a `name -> bool` over the dataset) makes it mirror the *run* rather than the recipe: operands the evaluator would drop are dropped here too, instead of one absent diagnostic leaving every ancestor unannotated. The `_DROPPED` sentinel is what distinguishes "not built" from "built but unknown".
 - **Lenient parser.** `parse.py` **warns and skips** unavailable-diagnostic placeholders (e.g. a `null`-source `difference`) and stray non-operation keys instead of failing, so real recipes with such terms still load. A stray key usually means an enclosing `product:` was omitted (its factors are then silently dropped rather than multiplied), so those warnings are worth reading rather than filtering.
 - **xgcm version:** requires **xgcm >= 0.10.0** — `lateral_divergence` needs native face-connected differencing, first shipped there. 0.10.0 also removed `xgcm.Grid`'s `periodic` argument and renamed `boundary` to `padding` (both now raise `ValueError`), so every grid construction in `examples/` and `tests/` uses `padding=`; `periodic=False` translates to `padding="fill"`.
 
 ### Tests
 
 - `test_parse.py` — parser units + validation; asserts all shipped recipes parse; covers the tolerated-malformation path.
 - `test_characterization.py` (+ `characterization_MOM6.json`) — golden snapshot of the engine's absolute MOM6 output (data-gated; exercises reciprocal/difference-of-sub-term only via the real recipe locally).
-- `test_collect.py` — `collect_budgets` behavior (no recipe mutation, lhs/rhs, the grid-guard on `difference`).
+- `test_collect.py` — `collect_budgets` behavior (no recipe mutation, lhs/rhs, the grid-guard on `difference`, and the `lateral_divergence` rechunking on a two-tile face-connected synthetic grid — the only CI coverage of face-connected differencing).
 - `test_query.py` — the query layer, all on the synthetic grid so it actually runs in CI.
-- `test_missing_handling.py` / `test_optional_var.py` / `test_broadcast_warning.py` / `test_display.py` — missing-diagnostic handling, `var: null` optionality, the issue-#11 broadcast guard, and the tree display; all synthetic-grid, CI-safe.
+- `test_units.py` — the unit algebra, the `constants` table, static inference, stamping, and `assert_unit_consistency`; all pure or synthetic-grid.
+- `test_missing_handling.py` / `test_optional_var.py` / `test_broadcast_warning.py` / `test_display.py` — missing-diagnostic handling, `var: null` optionality, the issue-#11 broadcast guard, and the tree display (including its `[units]` annotation); all synthetic-grid, CI-safe.
 - `conftest.py` — `synthetic_grid` / `synthetic_preset` fixtures, plus `SYNTHETIC_PRESET_SKIPS` (a preset whose structurally-primary op is skipped at run time — the regression test for name resolution being a runtime fact).
 
-**CI has neither example dataset**, so every data-gated test skips there. Only the synthetic-grid tests actually protect the engine in CI; `lateral_divergence` and `reciprocal` on a real face-connected grid are exercised *only* by the data-gated characterization / ECCO notebooks locally.
+**CI has neither example dataset**, so every data-gated test skips there. Only the synthetic-grid tests actually protect the engine in CI; `reciprocal` on a real face-connected grid is exercised *only* by the data-gated characterization / ECCO notebooks locally (`lateral_divergence` now has the miniature face-connected test in `test_collect.py`, but the real LLC90 topology still only runs locally).
 
 ## Data & examples
 

ci/environment.yml

@@ -4,6 +4,7 @@ channels:
   - nodefaults
 dependencies:
   - python>=3.11
+  - cf-units
   - cftime
   - dask
   - netcdf4

docs/environment.yml

@@ -3,6 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - python=3.12
+  - cf-units
   - cftime
   - dask
   - ipython

docs/quickstart.md

@@ -96,6 +96,14 @@ Each term is named after its path through the recipe, joined with underscores 
 `heat` / `rhs` / `advection` becomes `heat_rhs_advection`. Nothing to memorize,
 and nothing to parse back out.
 
+```{note}
+A budget's metadata sits alongside its sides in that result, so a recipe
+declaring `units: "kg s-1"` or `lambda: "thetao"` gives
+`{'mass': {'units': 'kg s-1', 'lhs': {...}, 'rhs': {...}}}`. Iterate with
+`q.aggregate()["mass"]["rhs"]` rather than over every value, or you will hit a
+metadata string where you expected a dict of terms.
+```
+
 ## Does it close?
 
 This is the question xbudget exists to answer:
@@ -205,7 +213,8 @@ Two things this page skipped, both of which real budgets need:
 
 - **Grids.** Terms built by differencing a staggered flux (or by a lateral
   divergence) need an `xgcm.Grid` instead of a plain `Dataset`, so xbudget knows
-  the discretization. Pass the grid to `collect_budgets` in exactly the same way.
+  the discretization. Pass the grid to `collect_budgets` in exactly the same way;
+  [Writing a recipe](the-grid) shows how to build one.
 - **Scale.** The examples below run against real model output, with dask.
 
 For a full worked budget, see the