diff --git a/CHANGELOG.md b/CHANGELOG.md index d581319..c1631df 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CLAUDE.md b/CLAUDE.md index 99b2de0..440c9a6 100644 --- a/CLAUDE.md +++ b/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_` — 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,7 +71,9 @@ 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"`. @@ -71,12 +81,13 @@ recipe ──parse_budgets──▶ typed tree (nodes.py) ──evaluate_budgets - `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 diff --git a/ci/environment.yml b/ci/environment.yml index 6091719..6c66ccf 100644 --- a/ci/environment.yml +++ b/ci/environment.yml @@ -4,6 +4,7 @@ channels: - nodefaults dependencies: - python>=3.11 + - cf-units - cftime - dask - netcdf4 diff --git a/docs/environment.yml b/docs/environment.yml index c0e8f1e..8e3e836 100644 --- a/docs/environment.yml +++ b/docs/environment.yml @@ -3,6 +3,7 @@ channels: - conda-forge dependencies: - python=3.12 + - cf-units - cftime - dask - ipython diff --git a/docs/quickstart.md b/docs/quickstart.md index 46e745e..be22d74 100644 --- a/docs/quickstart.md +++ b/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 diff --git a/docs/recipes.md b/docs/recipes.md index 00f864a..f6154a1 100644 --- a/docs/recipes.md +++ b/docs/recipes.md @@ -73,6 +73,78 @@ lateral: area: "areacello" ``` +**Constants can carry units.** A bare number is treated as *dimensionless*. To +give a constant a physical unit — so it participates in unit inference — write it +as a `{value:, units:}` mapping with a [UDUNITS][udunits] string: + +```yaml +product: + thickness_tendency: "dhdt" # its own units come from the dataset + density: + value: 1035. + units: "kg m-3" + area: "areacello" +``` + +A genuinely dimensionless factor — a `sign`, a fraction — stays bare. + +````{warning} +A **unit-conversion factor is the one thing that looks dimensionless and is +not.** If it exists to cancel a scale that an input carries in its own `units` +string, it has to declare the reciprocal scale, or the scale survives into the +result and the stamped `units` misreport the values by exactly that factor. + +CF writes practical salinity as a scale: ECCO's salt fluxes are `1e-3 m3 s-1`. +A recipe multiplying one by `0.001` to reach kg/s must say so: + +```yaml +constants: + salinity_scale: + value: 0.001 + units: "1e3" # value x units == 1: numerically the identity, + # dimensionally the thing that cancels the input's 1e-3 +``` + +Written as a bare `0.001` instead, the values are still right but the term +stamps itself `0.001 kg s-1` — a thousand-fold lie to anyone who reads the +attribute. `assert_unit_consistency=True` catches this (see *Units* 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, so a budget cannot be named that. + +References are resolved for the operands of a `sum` or a `product`: a string +there is looked up in the table first, and otherwise treated as a diagnostic +name. (The operand of a `difference` or a `reciprocal` is always a diagnostic — +differencing a constant is not a thing you want.) A constant sharing a name with +a dataset variable therefore *shadows* it — name them distinctly. All the shipped +recipes use the table 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 untouched — `lambda` (the tracer the budget is written in), `thickness`, `surface_lambda`: @@ -102,6 +174,58 @@ q.surface_lambda("heat") # -> "tos" q.metadata("mass") # -> {"lambda": "density", "thickness": "thkcello"} ``` +## Units + +Every variable xbudget materializes is stamped with an inferred `units` +attribute. Units compose through the tree with real [UDUNITS][udunits] arithmetic +(via [`cf-units`][cf-units]): a `product` multiplies its operands' units, a `sum` +takes the common unit of its summands (which must be dimensionally compatible), a +`difference` preserves units, a `reciprocal` inverts them, and a +`lateral_divergence` carries the flux unit. Each input's units come from its own +`units` attribute in the dataset; constants supply theirs from the recipe (see +above). If any input's units are unknown (a missing or unparseable `units` +attribute), the result's units are left unset rather than guessed. + +A budget can declare the units its closed sides should reconcile to, as a +top-level `units` key: + +```yaml +mass: + units: "kg s-1" + lambda: "density" + ... +``` + +With `assert_unit_consistency=True`, this is enforced: each budget that declares +`units` must have its `lhs`/`rhs` roots infer *the same* units — same dimension +and same magnitude — or `collect_budgets` raises `xbudget.UnitError`. Left at its +default of `False`, a mismatch only warns — inference and stamping happen +regardless. + +Magnitude counts because a leading numeric factor is part of a UDUNITS spec, and +therefore 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 the conversion-factor trap above, and this is what catches it. + +One thing it does *not* check: a side root that never materialized, because a +diagnostic it needs was absent, is skipped rather than failed. Absent inputs are +`on_missing`'s business — reporting the same problem twice under the name "units" +would only obscure it. + +```python +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) +``` + +Two conventions the shipped recipes rely on: practical salinity (`psu`), which +UDUNITS does not define, is treated as dimensionless (PSS-78), and offset +temperature (`degC`) is dimensionalized to kelvin inside a product, so +heat-content terms (θ·cₚ·ρ·…) reconcile to `W`. + +[cf-units]: https://cf-units.readthedocs.io/ + ## Seeing a recipe A real recipe nests deeply, and printed as raw JSON it is hard to read. Instead, @@ -109,8 +233,9 @@ 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 @@ -118,6 +243,19 @@ 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 (`xbudget.units.infer_units`, if you want to +call it yourself). 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. + +Given the data, the annotation answers *what a run will build*, not *what the +recipe describes*: a term whose diagnostics are absent is left blank, and its +parents still show their units, because that is what the run will produce from +the operands that remain. + 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: @@ -140,6 +278,45 @@ xbudget.BudgetQuery(grid, recipe) # same tree, annotated with run state `sum` and `product` take any number of operands, each of which may itself be a sub-term, so budgets nest as deeply as you need. +(the-grid)= +### The grid + +`sum`, `product` and `reciprocal` work on a plain `xarray.Dataset`. The two +operations that cross cell faces — `difference` and `lateral_divergence` — need +an `xgcm.Grid`, because only the grid knows which dimension is staggered against +which. Pass one in place of the dataset and everything else is unchanged: + +```python +import xgcm + +grid = xgcm.Grid( + ds, + coords={ # your staggered dimension names + "X": {"center": "xh", "outer": "xq"}, + "Y": {"center": "yh", "outer": "yq"}, + }, + metrics={("X", "Y"): ["areacello"]}, # optional + padding="fill", # boundary condition + autoparse_metadata=False, +) +xbudget.collect_budgets(grid, recipe) # derived vars land in grid._ds +``` + +Position names are xgcm's (`center` plus one of `left`/`right`/`inner`/`outer`). +Without a grid, a recipe using either operation fails immediately with +`Input 'data' must be an 'xgcm.Grid' instance…` rather than producing something +subtly wrong. + +```{note} +xbudget requires **xgcm ≥ 0.10**, which renamed `boundary=` to `padding=` and +removed `periodic=` — both now raise `ValueError`. Older tutorials and answers +still show the previous spelling; `periodic=False` becomes `padding="fill"`. +``` + +For a face-connected topology (the ECCO LLC90 tiles, say), also pass +`face_connections`; `lateral_divergence` then stitches the seams correctly. See +`examples/load_example_ecco_grid.py` for a complete 13-tile grid. + **`difference`** takes a single operand and differences it across whichever axis it is staggered on — the axis is inferred from the operand's dimensions, so you do not name it: diff --git a/examples/MOM6_budget_examples_mass_heat_salt.ipynb b/examples/MOM6_budget_examples_mass_heat_salt.ipynb index 7531894..cebbc7e 100644 --- a/examples/MOM6_budget_examples_mass_heat_salt.ipynb +++ b/examples/MOM6_budget_examples_mass_heat_salt.ipynb @@ -6,10 +6,10 @@ "id": "182a8d3a-09f8-448e-a029-d03db69bbe44", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:40:59.237337Z", - "iopub.status.busy": "2026-07-22T16:40:59.237116Z", - "iopub.status.idle": "2026-07-22T16:40:59.382976Z", - "shell.execute_reply": "2026-07-22T16:40:59.382641Z" + "iopub.execute_input": "2026-07-28T20:48:59.697600Z", + "iopub.status.busy": "2026-07-28T20:48:59.697309Z", + "iopub.status.idle": "2026-07-28T20:48:59.858677Z", + "shell.execute_reply": "2026-07-28T20:48:59.858334Z" } }, "outputs": [], @@ -24,10 +24,10 @@ "id": "a5f0eef4-7368-4f12-bc35-e09387046539", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:40:59.384742Z", - "iopub.status.busy": "2026-07-22T16:40:59.384641Z", - "iopub.status.idle": "2026-07-22T16:41:01.032907Z", - "shell.execute_reply": "2026-07-22T16:41:01.032576Z" + "iopub.execute_input": "2026-07-28T20:48:59.860799Z", + "iopub.status.busy": "2026-07-28T20:48:59.860664Z", + "iopub.status.idle": "2026-07-28T20:49:01.297665Z", + "shell.execute_reply": "2026-07-28T20:49:01.297334Z" } }, "outputs": [], @@ -61,10 +61,10 @@ "id": "5aa1b5c9-60d4-4b5e-9c7d-d0c25efb5093", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:01.034610Z", - "iopub.status.busy": "2026-07-22T16:41:01.034441Z", - "iopub.status.idle": "2026-07-22T16:41:02.318088Z", - "shell.execute_reply": "2026-07-22T16:41:02.317747Z" + "iopub.execute_input": "2026-07-28T20:49:01.299719Z", + "iopub.status.busy": "2026-07-28T20:49:01.299517Z", + "iopub.status.idle": "2026-07-28T20:49:02.608288Z", + "shell.execute_reply": "2026-07-28T20:49:02.607930Z" } }, "outputs": [ @@ -95,10 +95,10 @@ "id": "449ee6d3-57ca-4e3d-94dc-076887309454", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:02.320650Z", - "iopub.status.busy": "2026-07-22T16:41:02.320528Z", - "iopub.status.idle": "2026-07-22T16:41:02.350279Z", - "shell.execute_reply": "2026-07-22T16:41:02.350054Z" + "iopub.execute_input": "2026-07-28T20:49:02.610725Z", + "iopub.status.busy": "2026-07-28T20:49:02.610614Z", + "iopub.status.idle": "2026-07-28T20:49:02.660840Z", + "shell.execute_reply": "2026-07-28T20:49:02.660417Z" } }, "outputs": [], @@ -123,10 +123,10 @@ "id": "2d158238-b8a5-4fe5-8a59-7198d9e759ce", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:02.351654Z", - "iopub.status.busy": "2026-07-22T16:41:02.351570Z", - "iopub.status.idle": "2026-07-22T16:41:02.367544Z", - "shell.execute_reply": "2026-07-22T16:41:02.367307Z" + "iopub.execute_input": "2026-07-28T20:49:02.662529Z", + "iopub.status.busy": "2026-07-28T20:49:02.662441Z", + "iopub.status.idle": "2026-07-28T20:49:02.689895Z", + "shell.execute_reply": "2026-07-28T20:49:02.689609Z" } }, "outputs": [ @@ -196,6 +196,12 @@ " font-style: italic;\n", " margin-left: 0.5em;\n", "}\n", + ".xbdg-wrap .xbdg-units {\n", + " font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;\n", + " font-size: 0.9em;\n", + " color: var(--xbdg-muted);\n", + " margin-left: 0.4em;\n", + "}\n", ".xbdg-wrap .xbdg-leaf { padding: 1px 0; }\n", ".xbdg-wrap .xbdg-label { color: var(--xbdg-muted); }\n", ".xbdg-wrap .xbdg-diag {\n", @@ -207,71 +213,71 @@ " color: var(--xbdg-const);\n", "}\n", ".xbdg-wrap .xbdg-missing { opacity: 0.45; }\n", - "
xbudget recipe · heat
heatlambda=thetaosurface_lambda=tos
lhsΣ
Eulerian_tendency×
tracer_content_tendency_per_unit_area: opottemptend
area: areacello
advectionΣ
lateral×
sign: -1.
tracer_content_tendency_per_unit_area: T_advection_xy
area: areacello
interfacial×
sign: -1.
tracer_content_tendency_per_unit_area: Th_tendency_vert_remap
area: areacello
surface_ocean_flux_advective_negative_lhs×
sign: -1.
specific_heat_capacity: 3992.
lambda_mass: tos
thickness_tendency: boundary_forcing_h_tendency
density: 1035.
area: areacello
rhsΣ
diffusionΣ
lateral×
tracer_content_tendency_per_unit_area: opottemppmdiff
area: areacello
interfacial×
tracer_content_tendency_per_unit_area: opottempdiff
area: areacello
surface_exchange_flux×Σ
×product
tracer_content_tendency_per_unit_area: boundary_forcing_heat_tendency
area: areacello
Σsum
nonadvectiveΣ
latent×
tracer_content_tendency_per_unit_area: hflso
area: areacello
sensible×
tracer_content_tendency_per_unit_area: hfsso
area: areacello
longwave×
tracer_content_tendency_per_unit_area: rlntds
area: areacello
shortwave×
tracer_content_tendency_per_unit_area: rsdoabsorb
area: areacello
advective×
tracer_content_tendency_per_unit_area: heat_content_surfwater
area: areacello
surface_ocean_flux_advective_negative_rhs×
sign: -1.
specific_heat_capacity: 3992.
lambda_mass: tos
thickness_tendency: boundary_forcing_h_tendency
density: 1035.
area: areacello
bottom_flux×
tracer_content_tendency_per_unit_area: internal_heat_heat_tendency
area: areacello
frazil_ice×
tracer_content_tendency_per_unit_area: frazil_heat_tendency
area: areacello
" + "
xbudget recipe · heat
heatunits=Wlambda=thetaosurface_lambda=tos
lhsΣ[W]
Eulerian_tendency×[W]
tracer_content_tendency_per_unit_area: opottemptend[W m-2]
area: areacello[m2]
advectionΣ[W]
lateral×[W]
sign: -1.
tracer_content_tendency_per_unit_area: T_advection_xy[W m-2]
area: areacello[m2]
interfacial×[W]
sign: -1.
tracer_content_tendency_per_unit_area: Th_tendency_vert_remap[W m-2]
area: areacello[m2]
surface_ocean_flux_advective_negative_lhs×[W]
sign: -1.
specific_heat_capacity: 3992.[J kg-1 K-1]
lambda_mass: tos[degC]
thickness_tendency: boundary_forcing_h_tendency[m s-1]
density: 1035.[kg m-3]
area: areacello[m2]
rhsΣ[W]
diffusionΣ[W]
lateral×[W]
tracer_content_tendency_per_unit_area: opottemppmdiff[W m-2]
area: areacello[m2]
interfacial×[W]
tracer_content_tendency_per_unit_area: opottempdiff[W m-2]
area: areacello[m2]
surface_exchange_flux×Σ[W]
×product
tracer_content_tendency_per_unit_area: boundary_forcing_heat_tendency[W m-2]
area: areacello[m2]
Σsum
nonadvectiveΣ[W]
latent×[W]
tracer_content_tendency_per_unit_area: hflso[W m-2]
area: areacello[m2]
sensible×[W]
tracer_content_tendency_per_unit_area: hfsso[W m-2]
area: areacello[m2]
longwave×[W]
tracer_content_tendency_per_unit_area: rlntds[W m-2]
area: areacello[m2]
shortwave×[W]
tracer_content_tendency_per_unit_area: rsdoabsorb[W m-2]
area: areacello[m2]
advective×[W]
tracer_content_tendency_per_unit_area: heat_content_surfwater[W m-2]
area: areacello[m2]
surface_ocean_flux_advective_negative_rhs×[W]
sign: -1.
specific_heat_capacity: 3992.[J kg-1 K-1]
lambda_mass: tos[degC]
thickness_tendency: boundary_forcing_h_tendency[m s-1]
density: 1035.[kg m-3]
area: areacello[m2]
bottom_flux×[W]
tracer_content_tendency_per_unit_area: internal_heat_heat_tendency[W m-2]
area: areacello[m2]
frazil_ice×[W]
tracer_content_tendency_per_unit_area: frazil_heat_tendency[W m-2]
area: areacello[m2]
" ], "text/plain": [ "xbudget recipe · heat\n", - "heat lambda=thetao, surface_lambda=tos\n", - "├─ lhs [Σ]\n", - "│ ├─ Eulerian_tendency [×]\n", - "│ │ ├─ tracer_content_tendency_per_unit_area: opottemptend\n", - "│ │ └─ area: areacello\n", - "│ ├─ advection [Σ]\n", - "│ │ ├─ lateral [×]\n", + "heat units=W, lambda=thetao, surface_lambda=tos\n", + "├─ lhs [Σ] [W]\n", + "│ ├─ Eulerian_tendency [×] [W]\n", + "│ │ ├─ tracer_content_tendency_per_unit_area: opottemptend [W m-2]\n", + "│ │ └─ area: areacello [m2]\n", + "│ ├─ advection [Σ] [W]\n", + "│ │ ├─ lateral [×] [W]\n", "│ │ │ ├─ sign: -1.\n", - "│ │ │ ├─ tracer_content_tendency_per_unit_area: T_advection_xy\n", - "│ │ │ └─ area: areacello\n", - "│ │ └─ interfacial [×]\n", + "│ │ │ ├─ tracer_content_tendency_per_unit_area: T_advection_xy [W m-2]\n", + "│ │ │ └─ area: areacello [m2]\n", + "│ │ └─ interfacial [×] [W]\n", "│ │ ├─ sign: -1.\n", - "│ │ ├─ tracer_content_tendency_per_unit_area: Th_tendency_vert_remap\n", - "│ │ └─ area: areacello\n", - "│ └─ surface_ocean_flux_advective_negative_lhs [×]\n", + "│ │ ├─ tracer_content_tendency_per_unit_area: Th_tendency_vert_remap [W m-2]\n", + "│ │ └─ area: areacello [m2]\n", + "│ └─ surface_ocean_flux_advective_negative_lhs [×] [W]\n", "│ ├─ sign: -1.\n", - "│ ├─ specific_heat_capacity: 3992.\n", - "│ ├─ lambda_mass: tos\n", - "│ ├─ thickness_tendency: boundary_forcing_h_tendency\n", - "│ ├─ density: 1035.\n", - "│ └─ area: areacello\n", - "└─ rhs [Σ]\n", - " ├─ diffusion [Σ]\n", - " │ ├─ lateral [×]\n", - " │ │ ├─ tracer_content_tendency_per_unit_area: opottemppmdiff\n", - " │ │ └─ area: areacello\n", - " │ └─ interfacial [×]\n", - " │ ├─ tracer_content_tendency_per_unit_area: opottempdiff\n", - " │ └─ area: areacello\n", - " ├─ surface_exchange_flux [× Σ]\n", - " │ ├─ tracer_content_tendency_per_unit_area: boundary_forcing_heat_tendency\n", - " │ ├─ area: areacello\n", - " │ ├─ nonadvective [Σ]\n", - " │ │ ├─ latent [×]\n", - " │ │ │ ├─ tracer_content_tendency_per_unit_area: hflso\n", - " │ │ │ └─ area: areacello\n", - " │ │ ├─ sensible [×]\n", - " │ │ │ ├─ tracer_content_tendency_per_unit_area: hfsso\n", - " │ │ │ └─ area: areacello\n", - " │ │ ├─ longwave [×]\n", - " │ │ │ ├─ tracer_content_tendency_per_unit_area: rlntds\n", - " │ │ │ └─ area: areacello\n", - " │ │ └─ shortwave [×]\n", - " │ │ ├─ tracer_content_tendency_per_unit_area: rsdoabsorb\n", - " │ │ └─ area: areacello\n", - " │ └─ advective [×]\n", - " │ ├─ tracer_content_tendency_per_unit_area: heat_content_surfwater\n", - " │ └─ area: areacello\n", - " ├─ surface_ocean_flux_advective_negative_rhs [×]\n", + "│ ├─ specific_heat_capacity: 3992. [J kg-1 K-1]\n", + "│ ├─ lambda_mass: tos [degC]\n", + "│ ├─ thickness_tendency: boundary_forcing_h_tendency [m s-1]\n", + "│ ├─ density: 1035. [kg m-3]\n", + "│ └─ area: areacello [m2]\n", + "└─ rhs [Σ] [W]\n", + " ├─ diffusion [Σ] [W]\n", + " │ ├─ lateral [×] [W]\n", + " │ │ ├─ tracer_content_tendency_per_unit_area: opottemppmdiff [W m-2]\n", + " │ │ └─ area: areacello [m2]\n", + " │ └─ interfacial [×] [W]\n", + " │ ├─ tracer_content_tendency_per_unit_area: opottempdiff [W m-2]\n", + " │ └─ area: areacello [m2]\n", + " ├─ surface_exchange_flux [× Σ] [W]\n", + " │ ├─ tracer_content_tendency_per_unit_area: boundary_forcing_heat_tendency [W m-2]\n", + " │ ├─ area: areacello [m2]\n", + " │ ├─ nonadvective [Σ] [W]\n", + " │ │ ├─ latent [×] [W]\n", + " │ │ │ ├─ tracer_content_tendency_per_unit_area: hflso [W m-2]\n", + " │ │ │ └─ area: areacello [m2]\n", + " │ │ ├─ sensible [×] [W]\n", + " │ │ │ ├─ tracer_content_tendency_per_unit_area: hfsso [W m-2]\n", + " │ │ │ └─ area: areacello [m2]\n", + " │ │ ├─ longwave [×] [W]\n", + " │ │ │ ├─ tracer_content_tendency_per_unit_area: rlntds [W m-2]\n", + " │ │ │ └─ area: areacello [m2]\n", + " │ │ └─ shortwave [×] [W]\n", + " │ │ ├─ tracer_content_tendency_per_unit_area: rsdoabsorb [W m-2]\n", + " │ │ └─ area: areacello [m2]\n", + " │ └─ advective [×] [W]\n", + " │ ├─ tracer_content_tendency_per_unit_area: heat_content_surfwater [W m-2]\n", + " │ └─ area: areacello [m2]\n", + " ├─ surface_ocean_flux_advective_negative_rhs [×] [W]\n", " │ ├─ sign: -1.\n", - " │ ├─ specific_heat_capacity: 3992.\n", - " │ ├─ lambda_mass: tos\n", - " │ ├─ thickness_tendency: boundary_forcing_h_tendency\n", - " │ ├─ density: 1035.\n", - " │ └─ area: areacello\n", - " ├─ bottom_flux [×]\n", - " │ ├─ tracer_content_tendency_per_unit_area: internal_heat_heat_tendency\n", - " │ └─ area: areacello\n", - " └─ frazil_ice [×]\n", - " ├─ tracer_content_tendency_per_unit_area: frazil_heat_tendency\n", - " └─ area: areacello" + " │ ├─ specific_heat_capacity: 3992. [J kg-1 K-1]\n", + " │ ├─ lambda_mass: tos [degC]\n", + " │ ├─ thickness_tendency: boundary_forcing_h_tendency [m s-1]\n", + " │ ├─ density: 1035. [kg m-3]\n", + " │ └─ area: areacello [m2]\n", + " ├─ bottom_flux [×] [W]\n", + " │ ├─ tracer_content_tendency_per_unit_area: internal_heat_heat_tendency [W m-2]\n", + " │ └─ area: areacello [m2]\n", + " └─ frazil_ice [×] [W]\n", + " ├─ tracer_content_tendency_per_unit_area: frazil_heat_tendency [W m-2]\n", + " └─ area: areacello [m2]" ] }, "execution_count": 5, @@ -280,7 +286,7 @@ } ], "source": [ - "xbudget.show_recipe(recipe, \"heat\")" + "xbudget.show_recipe(recipe, \"heat\", grid)" ] }, { @@ -288,9 +294,9 @@ "id": "fbbcd656-22be-44b3-9a79-412f4bd03d86", "metadata": {}, "source": [ - "The `var: null` entries in the dictionary denote where diagnostics are not directly available. In many cases, however, they can be reconstructed by taking the sum or product of available diagnostics.\n", + "Each term is annotated with the operation that builds it (`Σ` sum, `×` product, `Δ` difference) and, in square brackets, the units `xbudget` infers for it — read from the diagnostics in `grid` at the leaves and composed upwards, so both sides of the heat budget come out in `W`. Passing `grid` is what makes that possible; without it only the recipe's declared constants carry units.\n", "\n", - "We can use the `xbudget.collect_budgets` function to leverage the information in this metadata dictionary to fill in all of these gaps in the budget." + "A term with no `var` of its own is one no single diagnostic provides: it is *reconstructed* from the sum or product of the diagnostics beneath it. We can use the `xbudget.collect_budgets` function to leverage the information in this metadata dictionary to fill in all of these gaps in the budget." ] }, { @@ -299,15 +305,16 @@ "id": "527f1b10", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:02.368898Z", - "iopub.status.busy": "2026-07-22T16:41:02.368818Z", - "iopub.status.idle": "2026-07-22T16:41:02.516061Z", - "shell.execute_reply": "2026-07-22T16:41:02.515703Z" + "iopub.execute_input": "2026-07-28T20:49:02.691270Z", + "iopub.status.busy": "2026-07-28T20:49:02.691161Z", + "iopub.status.idle": "2026-07-28T20:49:02.850291Z", + "shell.execute_reply": "2026-07-28T20:49:02.850022Z" } }, "outputs": [], "source": [ - "xbudget.collect_budgets(grid, recipe)\n", + "# assert_unit_consistency=True asserts each budget reconciles to its declared units (mass/salt -> kg s-1, heat -> W).\n", + "xbudget.collect_budgets(grid, recipe, assert_unit_consistency=True)\n", "q = xbudget.BudgetQuery(grid, recipe)" ] }, @@ -327,17 +334,18 @@ "id": "aa8797da-0f49-40aa-be98-b03ba49b380a", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:02.517700Z", - "iopub.status.busy": "2026-07-22T16:41:02.517602Z", - "iopub.status.idle": "2026-07-22T16:41:02.528471Z", - "shell.execute_reply": "2026-07-22T16:41:02.528173Z" + "iopub.execute_input": "2026-07-28T20:49:02.851801Z", + "iopub.status.busy": "2026-07-28T20:49:02.851715Z", + "iopub.status.idle": "2026-07-28T20:49:02.863201Z", + "shell.execute_reply": "2026-07-28T20:49:02.862959Z" } }, "outputs": [ { "data": { "text/plain": [ - "{'lambda': 'thetao',\n", + "{'units': 'W',\n", + " 'lambda': 'thetao',\n", " 'surface_lambda': 'tos',\n", " 'lhs': {'Eulerian_tendency': 'heat_lhs_Eulerian_tendency',\n", " 'advection': 'heat_lhs_advection',\n", @@ -372,10 +380,10 @@ "id": "7c9bfd4c-aaff-44ca-b360-d3b43f3addf1", "metadata": { "execution": { - "iopub.execute_input": "2026-07-22T16:41:02.529815Z", - "iopub.status.busy": "2026-07-22T16:41:02.529734Z", - "iopub.status.idle": "2026-07-22T16:41:02.548238Z", - "shell.execute_reply": "2026-07-22T16:41:02.548034Z" + "iopub.execute_input": "2026-07-28T20:49:02.864542Z", + "iopub.status.busy": "2026-07-28T20:49:02.864464Z", + "iopub.status.idle": "2026-07-28T20:49:02.886341Z", + "shell.execute_reply": "2026-07-28T20:49:02.886097Z" } }, "outputs": [ @@ -842,7 +850,8 @@ "Attributes:\n", " provenance: ['heat_lhs_advection_lateral', 'heat_lhs_advection_interfa...\n", " xbudget_path: ['heat', 'lhs', 'advection']\n", - " xbudget_op: sum