xwmb v0.7.0: topology-aware stack, xbudget 0.8.0 recipe API, derived-variable metadata, honest budget closure

[hdrake]

Supersedes #41.

Why #41 has to be replaced rather than updated

#41 was written against an unreleased intermediate xbudget commit. Its whole
strategy — xbudget.aggregate() plus collect_budgets(..., name_scheme="legacy") —
targets an API that never shipped: v0.7.0 removed the dict-walking engine
outright, with no deprecation cycle, and 0.8.0 confirms it. Against the released
package:

xbudget.aggregate(recipe)                          -> AttributeError
collect_budgets(ds, recipe, name_scheme='legacy')  -> TypeError

So this PR goes straight to the v1-native API instead, and picks up the rest of the
stack while it is there.

The pinned stack

package	pin	why
xgcm	>= 0.10.1	boundary → padding, periodic removed, cumsum(reverse=True), the north-fold boundary, multi-tile padding fixes
xbudget	>= 0.8.0	the recipe API; also the UDUNITS units this PR's metadata is composed from
sectionate	>= 0.4.0rc1	topology-driven section finding; f_c-aware transports
regionate	>= 0.6.0rc1	boundaries are GriddedSections carrying i_c/j_c/f_c
xwmt	>= 0.3.0rc1	tile-aware integration, xeos EOS, xwmt.units
xeos	>= 0.2.2	now that xwmb's API exposes eos=

Naming the rc explicitly is load-bearing: PEP 440 admits a pre-release only for
specifiers that mention one, so >= 0.4.0 would silently resolve to 0.3.4. Bump each
to the final release as it lands. (xwmt itself requires xeos[teos10] >= 0.2.3, so
the resolver picks that up regardless of xwmb's >= 0.2.2 floor.)

What's in it

0. No backwards-compatibility shims. v0.7.0 breaks compatibility across the
whole stack — xgcm renamed boundary, xbudget deleted its engine outright — so a
shim in xwmb would only buy a caller the impression that they had migrated when they
had not. Every old spelling raises TypeError: xbudget_dict= → recipe (now a
required positional), full_xbudget_dict → full_recipe, teos10= → eos=.

default_bins also goes, but it needed replacing rather than deleting: bins
could not express what default_bins=True did (build a default target grid for this
lambda), so removing it outright would have taken a capability with it — the ECCO
example depends on it. It is now bins="default", alongside bins=<array> and
bins=None.

1. The boundary → padding migration, and arbitrary grid topology. This PR
carries modernize-topology-stack (99934ee), rebased onto current main: the
padding sweep, the split of the monolithic budget.py into single-responsibility
modules, f_c threaded through sectionate, horizontal reductions over
_horizontal_dims so they broadcast across tiles, and the ECCO LLC90 example. Three
conflicts resolved in main's favour — its git-tag versioning, and its
mass_budget(bins=...) API (which that branch predated and would have reverted; it is
now pushed down into coordinates.resolve_target_coords).

2. The xbudget 0.8.0 recipe API. WaterMassBudget builds a
BudgetQuery(grid, recipe), keeps it as self.query, and feeds
query.aggregate(decompose=…) to xwmt. Every variable name is resolved through that
query rather than hardcoded or dict-walked — hand-walking is no longer safe now that
var: null placeholders are gone and a string operand may reference the recipe's
constants: table. The argument is renamed xbudget_dict → recipe, following
xbudget; the old spelling still works and warns, via a shim mirroring
xwmt.wmt._resolve_recipe.

This fixed a live bug rather than just a name: "mass_rhs_sum_surface_exchange_flux"
is the 0.6.x spelling, matches nothing under 0.8.0, and — because the lookup was
guarded by if … in grid._ds — was silently dropping the mass source from every
budget, which in turn silently suppressed the whole closure.

3. Metadata on every derived variable. New xwmb/attrs.py. Units are composed
from the inputs with xwmt.units (a cf_units wrapper — no duplicated algebra):

mass_density  = rho_ref [kg m-3] * h [m]      -> kg m-2
layer_mass    = mass_density * areacello [m2] -> kg
mass_tendency = mass_bounds / dt              -> kg s-1

plus long_name, CF cell_methods, and the xbudget provenance of the inputs.
Following xwmt: UDUNITS-2 spellings only, and when an input's units are unknown the
units key is omitted rather than guessed, with xwmb_units_source recording
which authority answered. WaterMassBudget also warns when a grid metric is
unlabelled — xbudget multiplies areacello into nearly every term and infers units
from operands, so an unlabelled area costs the units of the entire budget, and the
published MOM6 example file ships it that way.

Two bugs surfaced here: the greater_than sign flip ran under xarray's default
keep_attrs=False and was discarding everything xwmt had just stamped; and an
identically-zero convergent transport had no units, which poisoned every sum it took
part in.

4. Not calling a residual "spurious numerical mixing" until it is one. New
xwmb/completeness.py. Anything that belongs in the budget and is missing lands in
the residual wearing the name of something it is not — silently, and with a plausible
magnitude. The audit checks the recipe's right-hand side (BudgetQuery.missing(),
incomplete_terms(), restricted to the budgets that actually feed this lambda) and
whether dM/dt, Ψ and S are present or legitimately zero (a full-domain region has no
boundary for Ψ to cross; a recipe declaring no surface mass exchange has no S to be
missing — both differ from a declared term that failed to materialize).

close_budget now always computes realized_transformation and residual,
recording absent terms in xwmb_assumed_zero, and emits spurious_numerical_mixing
only when the audit is clean. Otherwise it warns, naming each gap and the diagnostic
to go and find, and stamps xwmb_unaccounted_terms:

The tracer budget is not closed, so its residual is not attributed to spurious
numerical mixing: S (surface mass source); mass/rhs/surface_exchange_flux (missing
input(s): wfo). `residual` is still computed and carries `xwmb_unaccounted_terms`,
but it is the budget imbalance -- read it as a diagnostic of what is missing, not as
an estimate of mixing.

On a closed budget spurious_numerical_mixing is present exactly as before, so
nothing downstream changes. The report is available as wmb.completeness.

Previously this path computed nothing at all, silently, if any one of the three
terms was missing — no residual, and no explanation for its absence.

Tests

modernize-topology-stack had replaced the data-free synthetic test with MOM6-netCDF
fixtures that skip when the file is absent, which in CI meant zero budget tests
ran. Topology support is the headline of this release; it should not be exercised only
on a machine that happens to have a downloaded netCDF lying around.

xwmb/tests/synthetic.py now builds three tiny grids — plain single-tile, a bipolar
north fold (padding={"Y": {"fold": "corner"}}, with the seam vector-sign constraint
V[Ny,i] = -V[Ny,Nx-1-i] imposed so the flow is one the grid could have produced), and
two tiles joined by face_connections — following regionate's own fixtures.

The load-bearing assertion is the discrete divergence theorem: the transport xwmb
integrates along the traced boundary must equal the flux convergence summed over the
region's cells. That is exactly what a topology bug breaks, and it breaks it quietly —
a loop that misses a fold or tile seam still returns a plausible number. Both new
topologies assert it to 1e-10, as does the agreement between the along-boundary and
grid-cell-divergence methods.

Plus test_recipe_api.py, test_attrs.py, test_close.py.

Validation

• pytest xwmb/tests -q → 57 passed with the MOM6 example downloaded (54 passed,
  3 skipped without it). Env: xgcm 0.10.1, xbudget 0.8.0, xwmt 0.3.0rc1, xeos 0.2.3,
  sectionate 0.4.0rc1, regionate 0.6.0rc1.
• sphinx -b html -W on a clean tree → build succeeded. Read the Docs runs with
  fail_on_warning: true, and this branch had introduced a duplicate-object warning
  (a class re-exported at three names) plus an API page per test module.
• The global MOM6 σ₂ budget audits as complete; spurious mixing 1.95e10 kg s⁻¹
  against a 1.94e10 kg s⁻¹ storage term. Every derived variable carries units.
• examples/MOM6_water_mass_budgets.ipynb re-executed end to end, zero failing
  cells (Compatibility with xgcm 0.10 + xbudget 0.7.0 #41 had to ship its notebooks un-executed).
• examples/ECCO_AABW_watermass_budget.ipynb (the 13-tile LLC90 AABW budget)
  re-executed end to end against the 2.7 GB Zenodo dataset, zero failing cells; that
  budget also audits as fully accounted for.

Upstream bugs found along the way

Two are in sectionate, and both are being fixed there rather than papered over here:

• layer.replace("l", "i") == interface is how convergent_transport decides
  whether a layer/interface coordinate pair belongs to the same axis. It replaces
  every l, so any vertical coordinate whose stem contains one is rejected
  (lam_l → iam_i), and coincidental matches are silently accepted. The test
  fixtures here are named sigma_l/sigma_i purely to route around it.
  → convergent_transport rejects valid layer/interface pairs whose name contains an 'l' MOM6-community/sectionate#51, fix in Take the vertical coordinates from the data and the grid; drop layer/interface sectionate#7.

• A grid that registers a vertical axis cannot be traced at all.
  gridutils pads its corner-index arrays over every axis the grid registers
  (padding_width = {ax: (1, 1) for ax in grid.axes}, at three sites — including
  _OuterTopology, the main multi-tile path), so a Z axis those arrays do not
  span makes xgcm's pad raise KeyError (or, when the vertical axis has no
  padding set, ValueError: No boundary condition was specified for axis 'Z').
  This is the real reason xwmb.horizontal_grid exists — an earlier revision of
  this PR described it as something regionate/sectionate "need", which was wrong:
  it is a workaround for a bug, not an inherent requirement. Verified: with the
  padding restricted to the axes each array is actually defined on,
  GriddedRegion and MaskRegions work directly on a full X/Y/Z budget grid and
  this PR's whole suite still passes. horizontal_grid stays for now because it is
  the only thing that works against released sectionate 0.4.0rc1; retire it once
  the fix lands. → A grid that registers a vertical axis cannot be traced: padding requested over every axis MOM6-community/sectionate#52, fix in Pad only over the axes the array actually spans sectionate#8.

And one in xarray:

• xarray 2026.7.0 raises from DataArray.chunk when the array carries a
  dask-backed coordinate of object dtype — which the published MOM6 file's cftime
  time_bounds_since_init is. Every cumulative sum failed on real model output.
  coordinates.rechunk_full rechunks only the data, which is what was wanted anyway.

🤖 Generated with Claude Code

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB