xbudget.query.BudgetQuery

class xbudget.query.BudgetQuery(data, recipe)[source]

Bases: object

Look up the variables an xbudget run produced.

Parameters:
  • data (xgcm.Grid, xr.Dataset, or None) –

    The object passed to collect_budgets(). Used to tell which terms were actually materialized: a term whose diagnostics were missing from the dataset is skipped by the evaluator (with a warning) and is reported here as unavailable rather than as a name that does not exist.

    None inspects a recipe offline, with no dataset to check against. Results are then planned names — what the recipe would emit if every diagnostic were present — which may differ from a real run (see _resolve_var()).

  • recipe (dict) – The recipe used for the run (e.g. from load_preset_budget()).

Examples

Address a term by its path and get back the flat variable name the engine emitted — no need to hand-assemble the compound name yourself:

>>> xbudget.collect_budgets(grid, recipe)
>>> q = xbudget.BudgetQuery(grid, recipe)
>>> q.var(("mass", "rhs", "advection", "lateral", "zonal_convergence"))
'mass_rhs_advection_lateral_zonal_convergence'

Ask what a term is built from, not just its own name:

>>> q.get_vars("heat_rhs")
{'var': 'heat_rhs',
 'sum': ['heat_rhs_diffusion', 'heat_rhs_surface_exchange_flux', ...]}

Read budget state the engine does not build but downstream code depends on — here the mass-budget layer thickness that xwmt.WaterMass needs:

>>> q.thickness("mass")
'thkcello'

See also

xbudget.parse_budgets, xbudget.evaluate_budgets

Methods

aggregate([decompose])

Collapse each budget to its top-level terms and their variable names.

bolus_transports([budget])

The GM eddy-bolus mass-transport variables a budget declares, or {}.

get_vars(address)

Describe a term (or raw diagnostic): its variable and its operands.

incomplete_terms()

Names of every materialized variable flagged xbudget_incomplete.

is_complete(address)

Whether a term was built from all the inputs its recipe names.

lambda_var(budget)

The lambda coordinate variable a budget declares, or None.

metadata([budget])

Return a budget's metadata (the non-lhs/rhs keys of a recipe).

missing([address])

Report inputs a recipe named that the dataset did not supply.

surface_lambda(budget)

The surface_lambda variable a budget declares, or None.

terms()

Map every term path to its variable name (None if not materialized).

thickness([budget])

The layer-thickness variable a budget declares, or None.

var(address)

The dataset variable name for a term, or None if not materialized.

Attributes

duplicate_names

Names produced by more than one term path (should be empty).

aggregate(decompose=())[source]

Collapse each budget to its top-level terms and their variable names.

Parameters:

decompose (str or iterable of str, optional) – Term names to expand into their summed parts instead of reporting them whole. Expanded keys are joined to their parent’s with an underscore ("advection" -> "advection_lateral"), applied recursively. Matching is exact.

Returns:

{budget: {**metadata, side: {label: variable_name}}}. Terms that were not materialized are omitted.

Return type:

dict

Examples

>>> q.aggregate()["heat"]["rhs"]
{'advection': 'heat_rhs_advection', 'diffusion': 'heat_rhs_diffusion'}
>>> q.aggregate(decompose=["diffusion"])["heat"]["rhs"]
{'advection': 'heat_rhs_advection',
 'diffusion_lateral': 'heat_rhs_diffusion_lateral',
 'diffusion_vertical': 'heat_rhs_diffusion_vertical'}
bolus_transports(budget='mass')[source]

The GM eddy-bolus mass-transport variables a budget declares, or {}.

Returns the bolus metadata block — the face mass transports (kg/s) of the parameterized eddy (GM bolus) velocity, e.g. {"x_mass_transport": "bolus_x_mass_transport", "y_...": ..., "z_...": ...}.

The bolus is deliberately absent from the (Eulerian) volume budget: it is non-divergent and carries zero net volume, so including it would only break native closure. In density coordinates, though, the bolus transport across isopycnals is a real water-mass-transformation term, so it is exposed here for downstream consumers (xwmb) rather than discarded.

Returns {} if the budget declares no bolus. Raises KeyError for an unknown budget, as metadata() does.

property duplicate_names

Names produced by more than one term path (should be empty).

"_".join(path) is not injective, so two distinct terms can collide on one output name — in which case the second silently overwrites the first in the dataset. Non-empty means the recipe needs a rename.

get_vars(address)[source]

Describe a term (or raw diagnostic): its variable and its operands.

Parameters:

address (str, tuple, or list) – A variable name, a term path, or a raw diagnostic name. A list (or array) requests a batch and returns a list of results.

Returns:

{"var": name} plus, for each of the term’s operations, a "sum"/"product"/… key listing its operands (sub-term variable names, raw diagnostic names, and constants, in recipe order). For a raw diagnostic: {"var": name, "referenced_by": [paths]}.

Return type:

dict

incomplete_terms()[source]

Names of every materialized variable flagged xbudget_incomplete.

Read straight from the dataset the evaluator wrote, so it also works on a dataset reopened from disk. Includes terms flagged only because a descendant dropped an input; use missing() for the inputs that were actually absent. Empty when there is no dataset to inspect.

is_complete(address)[source]

Whether a term was built from all the inputs its recipe names.

Returns True for a fully-materialized term, False for one that is incomplete (a dropped operand somewhere beneath it) or was not materialized at all, and None when the query has no dataset to check against (BudgetQuery(None, recipe)), where completeness is unknowable.

lambda_var(budget)[source]

The lambda coordinate variable a budget declares, or None.

(Named lambda_var because lambda is a Python keyword; the recipe key itself is lambda.) Raises KeyError for an unknown budget, as metadata() does.

metadata(budget=None)[source]

Return a budget’s metadata (the non-lhs/rhs keys of a recipe).

Budget metadata declares the quantities the engine does not itself build but that describe the budget’s state — e.g. the mass budget’s layer thickness and its coordinate lambda. Exposing it here gives downstream code (xwmt/xwmb) one engine-independent contract for those names instead of reaching into raw recipe keys.

Parameters:

budget (str or None) – A budget name (e.g. "mass", "heat"). If None (default), return {budget_name: metadata_dict} for every budget.

Returns:

The metadata dict for budget (e.g. {"lambda": "density", "thickness": "thkcello"}), or a dict of them keyed by budget name. Shallow copies, so mutating the result never rewrites the parsed recipe (metadata values are scalars).

Return type:

dict

missing(address=None)[source]

Report inputs a recipe named that the dataset did not supply.

With address, return the list of inputs that term was built without ([] if it is complete). With no argument, return {term_path: [missing inputs]} for every term that dropped at least one required input — both partial sums and products that did not materialize. Terms declared optional (and everything beneath them) are omitted: their absence is expected, which is the whole point of the flag.

Raises ValueError when the query has no dataset (BudgetQuery(None, recipe)), because what is present is exactly what cannot be known then.

surface_lambda(budget)[source]

The surface_lambda variable a budget declares, or None.

Raises KeyError for an unknown budget, as metadata() does.

terms()[source]

Map every term path to its variable name (None if not materialized).

thickness(budget='mass')[source]

The layer-thickness variable a budget declares, or None.

The mass budget’s thickness (e.g. "thkcello") is its prognostic state variable — layer thickness in metres — and the core input xwmt.WaterMass needs to build its vertical (mass) metrics. Reading it through this accessor keeps that dependency explicit and independent of the recipe’s dict layout or the naming scheme used to collect it.

Returns None if the budget declares no thickness. Raises KeyError if the recipe has no such budget at all — including the defaulted "mass", so a recipe without a mass budget is reported rather than silently indistinguishable from one that omits the key.

var(address)[source]

The dataset variable name for a term, or None if not materialized.

Parameters:

address (str or tuple) – A variable name ("heat_lhs_advection") or a term path (("heat", "lhs", "advection")).

Raises:

KeyError – If the address matches no term, with a suggestion of close matches.