Address review of #32: shared recipe constants, assert_unit_consistency, units in the repr

CHANGELOG.md

@@ -17,17 +17,41 @@
   `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`).
-- `collect_budgets(..., strict=True)` enforces units: each budget that declares
-  a top-level `units` must have its side roots infer units convertible to the
-  declared ones, raising `UnitError` (exported) otherwise. Without `strict`, a
-  reconciliation failure only warns; inference and stamping happen either way.
+- 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 units
+  convertible to the declared ones, raising `UnitError` (exported) otherwise.
+  Left at its default of `False`, a reconciliation failure only warns; inference
+  and stamping happen either way.
 - `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.
 - 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
+
+- `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

@@ -34,7 +34,7 @@ The central abstraction is the **`recipe`** — a nested provenance tree (loaded
 
 ### 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 +63,22 @@ 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`, raising `UnitError` (it only warns by default). `units.infer_units(node, lookup)` answers the same question statically, off the typed tree, which is what lets `show_recipe(recipe, budget, data)` and a notebook-displayed `BudgetQuery` annotate every node with `[units]`.
 - **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
 

docs/recipes.md

@@ -90,6 +90,35 @@ Only genuinely dimensional constants need this (`density`, `specific_heat_capaci
 `sign`, salinity-to-mass conversion factors, and the like are dimensionless and
 stay bare. See the *Units* section below.
 
+**Shared constants live at the root.** A reference density recurs at nearly every
+term of a budget, and repeating the mapping at each use is both tedious and a
+place for the copies to drift apart. Declare it once under the top-level
+`constants` key and refer to it by name:
+
+```yaml
+constants:
+  seawater_density:
+    value: 1035.
+    units: "kg m-3"
+
+mass:
+  units: "kg s-1"
+  lhs:
+    sum:
+      Eulerian_tendency:
+        product:
+          thickness_tendency: "dhdt"
+          density: "seawater_density"   # -> 1035. kg m-3
+          area: "areacello"
+```
+
+Entries take the same two forms an inline constant does (a bare number, or a
+`{value:, units:}` mapping). `constants` is the one top-level key that is not a
+budget. A string operand is looked up there first and otherwise treated as a
+diagnostic name, so a constant sharing a name with a dataset variable shadows it
+— name them distinctly. All the shipped recipes use this for their reference
+density and heat capacity.
+
 [udunits]: https://www.unidata.ucar.edu/software/udunits/
 
 Budget-level keys that are not `lhs`/`rhs` are metadata, carried through
@@ -143,13 +172,13 @@ mass:
   ...
 ```
 
-With `strict=True`, this is enforced: each budget that declares `units` must have
-its `lhs`/`rhs` roots infer units convertible to the declared ones, or
-`collect_budgets` raises `xbudget.UnitError`. Without `strict` (the default) a
-mismatch only warns — inference and stamping happen regardless.
+With `assert_unit_consistency=True`, this is enforced: each budget that declares
+`units` must have its `lhs`/`rhs` roots infer units convertible to the declared
+ones, or `collect_budgets` raises `xbudget.UnitError`. Left at its default of
+`False`, a mismatch only warns — inference and stamping happen regardless.
 
 ```python
-xbudget.collect_budgets(grid, recipe, strict=True)   # raises UnitError on a clash
+xbudget.collect_budgets(grid, recipe, assert_unit_consistency=True)  # UnitError on a clash
 q = xbudget.BudgetQuery(grid, recipe)
 q.budget_units("mass")            # -> "kg s-1"     (the declared target)
 q.units(("mass", "rhs"))          # -> "kg.s-1"     (what the run inferred)
@@ -169,15 +198,23 @@ A real recipe nests deeply, and printed as raw JSON it is hard to read. Instead,
 arrow, the way an `xarray.Dataset` repr expands:
 
 ```python
-xbudget.show_recipe(recipe)          # every budget
-xbudget.show_recipe(recipe, "heat")  # just one
+xbudget.show_recipe(recipe)                # every budget
+xbudget.show_recipe(recipe, "heat")        # just one
+xbudget.show_recipe(recipe, "heat", grid)  # ...annotated with units
 ```
 
 In a Jupyter/VSCode notebook it displays as an interactive HTML tree (operator
 badges, diagnostics, and constants distinguished); printed at a terminal it
 falls back to an indented ASCII tree. It reads a recipe directly, so you can
 inspect one before ever running it.
 
+Each node also carries its units in square brackets — `density: 1035. [kg m-3]`,
+`lateral [×] [W]` — composed up the tree by the same arithmetic
+`assert_unit_consistency` checks. Units enter at the leaves, so pass the data you
+intend to collect against (as the third argument above) and the whole tree
+resolves; reading it top-down then shows exactly *where* a unit stops being
+inferable. Without it, only the recipe's declared constants are annotated.
+
 Once you *have* run a budget, displaying the query does the same but annotated
 with each term's resolved variable name — with terms whose diagnostics were
 missing greyed out, so the tree doubles as a map of what the run materialized: