Handling missing diagnostics¶
One recipe is meant to serve many datasets, and no two models write out exactly the same diagnostics. So when a term references a variable your dataset does not have, xbudget skips that term rather than raising — but it is careful never to do so in a way that misleads you into thinking a term is there when it isn’t.
This notebook shows what that looks like in practice. It needs no downloads: the data is made up on the spot, so you can run it anywhere xbudget is installed.
import warnings
import numpy as np
import xarray as xr
import xbudget
A budget that closes¶
We start from a heat budget that balances exactly — the same shape as the
Quickstart. The right-hand side is advection + diffusion + a surface flux, where
the surface flux is a product of a per-area flux and the cell area.
rng = np.random.default_rng(0)
dims, coords = ("x", "y"), {"x": [0, 1, 2], "y": [0, 1, 2]}
advection = rng.random((3, 3))
diffusion = rng.random((3, 3))
surface_flux = rng.random((3, 3)) # W m-2
cell_area = np.full((3, 3), 1.0e6) # m2
def make_dataset(include_surface=True):
data = {
"heat_tendency": (dims, advection + diffusion + surface_flux * cell_area),
"advective_flux_convergence": (dims, advection),
"diffusive_flux_convergence": (dims, diffusion),
"cell_area": (dims, cell_area),
}
if include_surface:
data["surface_heat_flux"] = (dims, surface_flux)
return xr.Dataset(data, coords=coords)
recipe = {
"heat": {
"lhs": {"sum": {"tendency": {"var": "heat_tendency"}}},
"rhs": {
"sum": {
"advection": {"var": "advective_flux_convergence"},
"diffusion": {"var": "diffusive_flux_convergence"},
"surface_forcing": {
"product": {"flux": "surface_heat_flux", "area": "cell_area"},
},
},
},
}
}
ds = make_dataset(include_surface=True)
xbudget.collect_budgets(ds, recipe)
q = xbudget.BudgetQuery(ds, recipe)
residual = ds[q.var("heat_lhs")] - ds[q.var("heat_rhs")]
print("closes to:", float(abs(residual).max()))
q.aggregate()["heat"]["rhs"]
closes to: 0.0
{'advection': 'heat_rhs_advection',
'diffusion': 'heat_rhs_diffusion',
'surface_forcing': 'heat_rhs_surface_forcing'}
Everything materialized, and is_complete confirms nothing was dropped:
print("is_complete(heat_rhs):", q.is_complete("heat_rhs"))
print("incomplete_terms():", q.incomplete_terms())
is_complete(heat_rhs): True
incomplete_terms(): []
Now break it: a missing diagnostic¶
Suppose this model did not write surface_heat_flux. Collecting the same recipe
against that dataset does not error — but it warns, once, with a summary that
names both the missing diagnostic and the budget term it left incomplete.
ds_missing = make_dataset(include_surface=False)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
xbudget.collect_budgets(ds_missing, recipe) # default on_missing="warn"
for w in caught:
print(w.category.__name__, "->", w.message)
UserWarning -> xbudget: missing diagnostic(s) ['surface_heat_flux']; 1 budget term(s) are now incomplete: ['heat_rhs']. The affected variables carry `xbudget_incomplete`/`xbudget_missing` attrs; query them with BudgetQuery.missing(). Mark expected-absent terms `optional` in the recipe to silence this, or use on_missing='raise' to fail instead.
The surface_forcing term is a product, and one of its factors is now absent.
A product needs every factor, so the term is not built — and crucially it is
not fabricated as an array of zeros (which would quietly enter the sum as a
real, null contribution). It is simply gone, and the sum above it is flagged
incomplete.
q_missing = xbudget.BudgetQuery(ds_missing, recipe)
print("var(heat_rhs_surface_forcing):", q_missing.var("heat_rhs_surface_forcing"))
print("is_complete(heat_rhs): ", q_missing.is_complete("heat_rhs"))
print("incomplete_terms(): ", q_missing.incomplete_terms())
print()
print("missing():")
for path, gaps in q_missing.missing().items():
print(" ", path, "->", gaps)
var(heat_rhs_surface_forcing): None
is_complete(heat_rhs): False
incomplete_terms(): ['heat_rhs']
missing():
('heat', 'rhs') -> ['surface_forcing']
('heat', 'rhs', 'surface_forcing') -> ['surface_heat_flux']
The incompleteness is not just a warning that scrolls past — it is stamped onto the variable itself, so it survives being written to a netCDF file and reopened later:
dict(ds_missing["heat_rhs"].attrs)
{'provenance': ['heat_rhs_advection', 'heat_rhs_diffusion'],
'xbudget_path': ['heat', 'rhs'],
'xbudget_op': 'sum',
'xbudget_incomplete': 1,
'xbudget_missing': ['surface_forcing']}
provenance records what went in; xbudget_missing records what was supposed
to and didn’t. get_vars shows the same, alongside the operands:
q_missing.get_vars("heat_rhs")
{'var': 'heat_rhs',
'sum': ['heat_rhs_advection', 'heat_rhs_diffusion'],
'missing': ['surface_forcing']}
Note that aggregate() lists only what actually materialized, so the dropped
surface_forcing simply isn’t in it — which is exactly the trap missing()
exists to close. When a real budget fails to close, reach for is_complete() and
missing() first.
q_missing.aggregate()["heat"]["rhs"]
{'advection': 'heat_rhs_advection', 'diffusion': 'heat_rhs_diffusion'}
Choosing how loud to be: on_missing¶
The default "warn" is right when one recipe legitimately serves datasets with
different diagnostics. When you instead want a guarantee that the recipe matched
the data, pass on_missing="raise": any required gap becomes a
MissingDiagnosticError, whose .missing attribute lists every gap at once.
try:
xbudget.collect_budgets(make_dataset(include_surface=False), recipe,
on_missing="raise")
except xbudget.MissingDiagnosticError as err:
print("raised:", err)
print()
print(".missing:", err.missing)
raised: collect_budgets(on_missing='raise'): missing diagnostic(s) ['surface_heat_flux']; 1 budget term(s) are now incomplete: ['heat_rhs']. Provide the diagnostics, mark the term(s) `optional` in the recipe, or use on_missing='warn'/'ignore'.
.missing: [('surface_heat_flux', ('heat', 'rhs', 'surface_forcing'))]
on_missing="ignore" is the opposite — no warning at all — but it still stamps
the xbudget_incomplete / xbudget_missing attributes, so the durable record is
there even when the warning isn’t.
Declaring an expected absence: optional¶
Sometimes a diagnostic is legitimately absent on some datasets, and you don’t
want a warning every single run. Deleting the term from the recipe would work, but
it also erases the fact that the term was ever expected. Instead, mark it
optional: true. Its whole subtree is then exempt: dropped silently, with no
incomplete flag on its parent, and no error even under on_missing="raise".
optional_recipe = {
"heat": {
"lhs": {"sum": {"tendency": {"var": "heat_tendency"}}},
"rhs": {
"sum": {
"advection": {"var": "advective_flux_convergence"},
"diffusion": {"var": "diffusive_flux_convergence"},
"surface_forcing": {
"optional": True, # <- expected-absent here
"product": {"flux": "surface_heat_flux", "area": "cell_area"},
},
},
},
}
}
ds_opt = make_dataset(include_surface=False)
# on_missing="raise" would fire for a *required* gap; an optional one is exempt.
xbudget.collect_budgets(ds_opt, optional_recipe, on_missing="raise")
q_opt = xbudget.BudgetQuery(ds_opt, optional_recipe)
print("is_complete(heat_rhs):", q_opt.is_complete("heat_rhs"))
print("incomplete_terms(): ", q_opt.incomplete_terms())
print("missing(): ", q_opt.missing())
is_complete(heat_rhs): True
incomplete_terms(): []
missing(): {}
heat_rhs is reported complete, because with surface_forcing declared
optional its absence is exactly what the recipe said to expect — the honest,
self-documenting alternative to quietly deleting the term.
In short¶
A missing diagnostic skips its term; it never silently pretends the term is there.
A
productmissing a factor is dropped, not fabricated as zero.Incompleteness is stamped on the data (
xbudget_incomplete/xbudget_missing) and queryable throughis_complete(),missing(), andincomplete_terms()— even on a reopened dataset.on_missingpickswarn/raise/ignore;optional: truedeclares an expected absence.