From a8008ae7dfbad95fdbfc8e40c5b81abd31f3c8ca Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 28 Jul 2026 07:55:28 -0700 Subject: [PATCH 1/6] Add automatic units handling with UDUNITS inference and strict mode Every derived variable now carries an inferred `units` attribute, composed through the typed tree with real UDUNITS arithmetic via cf-units: products multiply units, sums require compatible summands, difference preserves, reciprocal inverts, lateral_divergence carries the flux unit. Constants may declare units via a {value, units} mapping; a bare number is dimensionless. `collect_budgets(strict=True)` enforces that each budget declaring a top-level `units` reconciles its lhs/rhs roots to it (else UnitError); non-strict warns. BudgetQuery gains .units() and .budget_units(). All shipped recipes are annotated and declare budget units; verified strict passes on the MOM6 example. psu is treated as dimensionless (PSS-78); degC is dimensionalized in products so heat-content terms reconcile to W. Closes #28 (mechanism; ECCO pipeline units follow). Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 26 +++ ci/environment.yml | 1 + docs/environment.yml | 1 + docs/recipes.md | 60 +++++++ pyproject.toml | 1 + xbudget/__init__.py | 2 +- xbudget/collect.py | 16 +- xbudget/evaluate.py | 155 +++++++++++++++++- xbudget/nodes.py | 10 +- xbudget/parse.py | 28 ++++ xbudget/query.py | 28 ++++ xbudget/recipes/ECCOV4r4_native.yaml | 103 +++++++++--- xbudget/recipes/MOM6.yaml | 47 ++++-- xbudget/recipes/MOM6_3Donly.yaml | 39 +++-- xbudget/recipes/MOM6_drift.yaml | 3 + xbudget/recipes/MOM6_surface.yaml | 11 +- xbudget/tests/test_units.py | 229 +++++++++++++++++++++++++++ xbudget/units.py | 160 +++++++++++++++++++ 18 files changed, 862 insertions(+), 58 deletions(-) create mode 100644 xbudget/tests/test_units.py create mode 100644 xbudget/units.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6ee4df2..85c24f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,32 @@ ## 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`). +- `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. +- `BudgetQuery` gained `.units(term)` (the stamped units of a materialized + term) and `.budget_units(budget)` (the declared target). +- 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. + ### Packaging - `pyyaml` is now declared as a dependency. It is imported at package import 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/recipes.md b/docs/recipes.md index 00f864a..6524e15 100644 --- a/docs/recipes.md +++ b/docs/recipes.md @@ -73,6 +73,25 @@ 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" +``` + +Only genuinely dimensional constants need this (`density`, `specific_heat_capacity`); +`sign`, salinity-to-mass conversion factors, and the like are dimensionless and +stay bare. See the *Units* section below. + +[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 +121,47 @@ 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 `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. + +```python +xbudget.collect_budgets(grid, recipe, strict=True) # raises 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, diff --git a/pyproject.toml b/pyproject.toml index 97cbdda..ba74c8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ + "cf-units", "numpy", "pyyaml", "xarray", diff --git a/xbudget/__init__.py b/xbudget/__init__.py index 666eef0..d3aa049 100644 --- a/xbudget/__init__.py +++ b/xbudget/__init__.py @@ -2,7 +2,7 @@ from .presets import * from .collect import * from .parse import parse_budgets, BudgetParseError -from .evaluate import evaluate_budgets, MissingDiagnosticError +from .evaluate import evaluate_budgets, MissingDiagnosticError, UnitError from .query import BudgetQuery from .display import show_recipe from .version import __version__ diff --git a/xbudget/collect.py b/xbudget/collect.py index df636ec..ee2bbe9 100644 --- a/xbudget/collect.py +++ b/xbudget/collect.py @@ -67,7 +67,9 @@ def lateral_divergence(grid, Fx, Fy): return dFx + dFy -def collect_budgets(data, recipe, allow_rechunk=True, on_missing="warn"): +def collect_budgets( + data, recipe, allow_rechunk=True, on_missing="warn", strict=False +): """Materialize every budget term described by ``recipe`` into ``data``. The recipe dict is parsed into a typed expression tree @@ -102,6 +104,15 @@ def collect_budgets(data, recipe, allow_rechunk=True, on_missing="warn"): ``xbudget_missing`` attributes (query them with :meth:`xbudget.BudgetQuery.missing`), and a term declared ``optional`` in the recipe is exempt — its absence is expected and never alarms. + strict : bool (default: False) + Enforce units. Every derived variable is stamped with an inferred + ``units`` attribute (from its inputs' units via UDUNITS) regardless of + this flag. With ``strict=True``, each budget that declares a top-level + ``units`` must additionally have its side roots infer units convertible + to the declared ones, raising + :class:`~xbudget.evaluate.UnitError` otherwise; with ``strict=False`` a + reconciliation failure only warns. Query stamped units with + :meth:`xbudget.BudgetQuery.units`. Returns ------- @@ -125,6 +136,7 @@ def collect_budgets(data, recipe, allow_rechunk=True, on_missing="warn"): budgets = parse_budgets(recipe) evaluate_budgets( - data, budgets, allow_rechunk=allow_rechunk, on_missing=on_missing + data, budgets, allow_rechunk=allow_rechunk, on_missing=on_missing, + strict=strict, ) return data diff --git a/xbudget/evaluate.py b/xbudget/evaluate.py index 9aa0ab5..e6ab328 100644 --- a/xbudget/evaluate.py +++ b/xbudget/evaluate.py @@ -57,6 +57,17 @@ _warn_if_summands_broadcast, lateral_divergence, ) +from .units import ( + convertible, + difference_units, + divergence_units, + format_units, + parse_units, + product_units, + reciprocal_units, + sum_compatible, + sum_units, +) ON_MISSING = ("warn", "raise", "ignore") @@ -73,13 +84,25 @@ def __init__(self, message, missing): self.missing = missing +class UnitError(ValueError): + """A budget's units did not reconcile under ``strict=True``. + + Carries the list of ``(budget, side, inferred, declared, kind)`` issues, so + the caller sees every mismatch/unknown at once rather than only the first. + """ + + def __init__(self, message, issues): + super().__init__(message) + self.issues = issues + + def _new_name(path): """Output variable name: the term path joined, without operator infixes.""" return "_".join(path) class _Evaluator: - def __init__(self, data, allow_rechunk=True, on_missing="warn"): + def __init__(self, data, allow_rechunk=True, on_missing="warn", strict=False): if on_missing not in ON_MISSING: raise ValueError( f"Unknown on_missing {on_missing!r}; expected one of {ON_MISSING}." @@ -92,17 +115,21 @@ def __init__(self, data, allow_rechunk=True, on_missing="warn"): self.ds = data self.allow_rechunk = allow_rechunk self.on_missing = on_missing + self.strict = strict # new variable name -> {"path", "op"} self.records = {} # Missing-diagnostic bookkeeping, drained once at the end of run(). # Only *required* (non-optional) misses land here. self._missing = [] # list of (diagnostic_name, term_path tuple) self._incomplete = [] # names of emitted variables flagged incomplete + # (budget, side, inferred, declared, kind) unit reconciliation issues. + self._unit_issues = [] def run(self, budgets): for budget in budgets.values(): for term in budget.sides.values(): self._eval_term(term) + self._check_units(budgets) self._finish() return self.records @@ -151,6 +178,77 @@ def _finish(self): ) # on_missing == "ignore": stay silent (the attrs were still stamped). + # -- unit reconciliation ------------------------------------------------ + + def _check_units(self, budgets): + """Reconcile each declared budget's root sides against its stated units. + + A budget that declares a top-level ``units`` (e.g. ``mass`` -> ``kg s-1``) + opts into the check: every materialized side root must infer units that + are dimensionally convertible to the declared ones. Budgets that declare + no ``units`` are only stamped, never checked. Issues found here join any + sum-incompatibilities recorded during evaluation and are reported once by + :meth:`_report_units`. + """ + for name, budget in budgets.items(): + declared_str = budget.metadata.get("units") + if declared_str is None: + continue + declared = parse_units(declared_str) + if declared is None: + self._unit_issues.append({ + "location": name, + "kind": "declared-unparseable", + "detail": ( + f"budget declares units {declared_str!r}, which UDUNITS " + f"cannot parse" + ), + }) + continue + for side, term in budget.sides.items(): + var = _new_name(term.path) + if var not in self.ds: + continue + inferred = parse_units(self.ds[var].attrs.get("units")) + if inferred is None: + self._unit_issues.append({ + "location": f"{name}/{side}", + "kind": "unknown", + "detail": ( + f"units could not be inferred; expected {declared_str!r}" + ), + }) + elif not convertible(inferred, declared): + self._unit_issues.append({ + "location": f"{name}/{side}", + "kind": "mismatch", + "detail": ( + f"inferred {format_units(inferred)!r} not convertible " + f"to declared {declared_str!r}" + ), + }) + self._report_units() + + def _report_units(self): + """Raise (strict) or warn (default) about accumulated unit issues, once.""" + if not self._unit_issues: + return + detail = "\n".join( + f" - {i['location']}: {i['detail']}" for i in self._unit_issues + ) + if self.strict: + raise UnitError( + "collect_budgets(strict=True): budget units did not reconcile:\n" + + detail, + issues=list(self._unit_issues), + ) + warnings.warn( + "xbudget: budget units did not fully reconcile:\n" + detail + + "\nInspect stamped units with BudgetQuery.units(...); pass " + "strict=True to fail instead.", + UserWarning, + ) + # -- term --------------------------------------------------------------- def _eval_term(self, term, optional_ctx=False): @@ -217,6 +315,9 @@ def _eval_term(self, term, optional_ctx=False): "xbudget_path": list(term.path), "xbudget_op": op.kind, } + units = format_units(opmeta.get("units") if opmeta else None) + if units is not None: + out.attrs["units"] = units self._stamp_incompleteness(out, opmeta, opt) self.ds[new_name] = out self.records[new_name] = {"path": list(term.path), "op": op.kind} @@ -264,7 +365,8 @@ def _eval_reciprocal(self, op, term, opt): self._record_missing(op.source, term.path, alarming=not opt) return None, None, None var = 1.0 / xr.where(ds[op.source] == 0, np.inf, ds[op.source]) - return var, op.source, {"missing": [], "incomplete": False} + units = reciprocal_units(parse_units(ds[op.source].attrs.get("units"))) + return var, op.source, {"missing": [], "incomplete": False, "units": units} def _eval_lateral_divergence(self, op, term, opt): fx = self._eval_term(op.fx, optional_ctx=opt) @@ -277,36 +379,47 @@ def _eval_lateral_divergence(self, op, term, opt): incomplete = bool(fx.attrs.get("xbudget_incomplete")) or bool( fy.attrs.get("xbudget_incomplete") ) - return var, [fx.name, fy.name], {"missing": [], "incomplete": incomplete} + units = divergence_units( + parse_units(fx.attrs.get("units")), parse_units(fy.attrs.get("units")) + ) + return var, [fx.name, fy.name], { + "missing": [], "incomplete": incomplete, "units": units, + } def _eval_nary(self, op, term, opt): ds = self.ds op_list = [] + unit_list = [] # units of each surviving operand, in lockstep with op_list missing = [] # labels/names of operands dropped at this node incomplete = False for name, operand in op.terms: value = None + units = None operand_optional = opt if isinstance(operand, Term): operand_optional = opt or operand.optional child = self._eval_term(operand, optional_ctx=operand_optional) if child is not None: value = child + units = parse_units(child.attrs.get("units")) if child.attrs.get("xbudget_incomplete"): incomplete = True # else: the child already recorded its own missing diagnostic. dropped_label = name elif isinstance(operand, Constant): value = operand.value + units = parse_units(operand.units) elif isinstance(operand, VarRef): if operand.name in ds: value = ds[operand.name] + units = parse_units(ds[operand.name].attrs.get("units")) else: self._record_missing(operand.name, term.path, alarming=not opt) dropped_label = operand.name if value is not None: op_list.append(value) + unit_list.append(units) continue # This operand did not resolve. @@ -323,13 +436,29 @@ def _eval_nary(self, op, term, opt): return None, None, None if op.kind == "sum": _warn_if_summands_broadcast(op_list, _new_name(term.path)) + if not sum_compatible(unit_list): + # A genuine unit clash between summands (not merely an unknown one) + # is always a recipe bug; record it for the end-of-run report. + self._unit_issues.append({ + "location": _new_name(term.path), + "kind": "sum-incompatible", + "detail": ( + "summands have incompatible units: " + + ", ".join(sorted({format_units(u) for u in unit_list if u})) + ), + }) + units = sum_units(unit_list) + else: + units = product_units(unit_list) var = sum(op_list) if op.kind == "sum" else reduce(mul, op_list, 1) if not isinstance(var, xr.DataArray): # Reduced to a pure scalar (e.g. all variable operands missing); no # variable is emitted in this case. return None, None, None provenance = [o.name if isinstance(o, xr.DataArray) else o for o in op_list] - return var, provenance, {"missing": missing, "incomplete": incomplete} + return var, provenance, { + "missing": missing, "incomplete": incomplete, "units": units, + } def _eval_difference(self, op, term, opt): if self.grid is None: @@ -352,6 +481,7 @@ def _eval_difference(self, op, term, opt): return None, None, None provenance = source.name incomplete = bool(source.attrs.get("xbudget_incomplete")) + units = difference_units(parse_units(source.attrs.get("units"))) staggered_axes = { axn: c @@ -407,10 +537,14 @@ def _eval_difference(self, op, term, opt): else: var = self.grid.diff(source.fillna(0.0), axis) - return var, provenance, {"missing": [], "incomplete": incomplete} + return var, provenance, { + "missing": [], "incomplete": incomplete, "units": units, + } -def evaluate_budgets(data, budgets, allow_rechunk=True, on_missing="warn"): +def evaluate_budgets( + data, budgets, allow_rechunk=True, on_missing="warn", strict=False +): """Evaluate parsed budgets into ``data``; return ``records``. Parameters @@ -429,6 +563,13 @@ def evaluate_budgets(data, budgets, allow_rechunk=True, on_missing="warn"): case the affected variables are stamped with ``xbudget_incomplete`` / ``xbudget_missing`` attributes, and terms declared ``optional`` in the recipe are exempt. + strict : bool, default False + Enforce budget units. Each output variable is stamped with an inferred + ``units`` attribute regardless; ``strict=True`` additionally requires + every budget that declares a top-level ``units`` to have its side roots + infer units convertible to the declared ones (and no unit clashes), + raising :class:`UnitError` otherwise. When ``False`` a reconciliation + failure only warns. Returns ------- @@ -453,5 +594,5 @@ def evaluate_budgets(data, budgets, allow_rechunk=True, on_missing="warn"): tree or want to control the two steps separately. """ return _Evaluator( - data, allow_rechunk=allow_rechunk, on_missing=on_missing + data, allow_rechunk=allow_rechunk, on_missing=on_missing, strict=strict ).run(budgets) diff --git a/xbudget/nodes.py b/xbudget/nodes.py index 90e4692..c457324 100644 --- a/xbudget/nodes.py +++ b/xbudget/nodes.py @@ -23,8 +23,16 @@ @dataclass(frozen=True) class Constant: - """A scalar factor/addend in a sum or product (e.g. a density, a sign).""" + """A scalar factor/addend in a sum or product (e.g. a density, a sign). + + ``units`` is a UDUNITS string carried alongside the value so unit inference + can include hard-coded constants (a density is ``kg m-3``, a sign is + dimensionless). A bare number in the recipe (``sign: -1.``) has no declared + units and is treated as dimensionless (``"1"``); the ``{value:, units:}`` + dict form pins a dimensional one (``density: {value: 1035., units: "kg m-3"}``). + """ value: float + units: str = "1" @dataclass(frozen=True) diff --git a/xbudget/parse.py b/xbudget/parse.py index 86c928a..a3d3ad8 100644 --- a/xbudget/parse.py +++ b/xbudget/parse.py @@ -146,6 +146,11 @@ def _parse_operand(value, path, name): # Tolerated: a placeholder operand with no content is skipped. return None if isinstance(value, dict): + # A ``{value:, units:}`` mapping is a constant with declared units, not a + # sub-term. Distinguish by the ``value`` key (and the absence of any + # operation key): a real sub-term names operations, never ``value``. + if "value" in value and not (value.keys() & OPERATION_KEYS): + return _parse_constant(value, path, name) return _parse_term(value, path, name) if isinstance(value, bool): raise BudgetParseError( @@ -162,6 +167,29 @@ def _parse_operand(value, path, name): ) +def _parse_constant(body, path, name): + """Parse a ``{value:, units:}`` constant (units default to dimensionless).""" + value = body["value"] + if isinstance(value, bool) or not isinstance(value, numbers.Number): + raise BudgetParseError( + f"Constant '{name}' at {_fmt(path)} needs a numeric 'value', got " + f"{type(value).__name__}." + ) + units = body.get("units", "1") + if not isinstance(units, str): + raise BudgetParseError( + f"Constant '{name}' at {_fmt(path)} 'units' must be a UDUNITS string, " + f"got {type(units).__name__}." + ) + extra = set(body) - {"value", "units"} + if extra: + raise BudgetParseError( + f"Constant '{name}' at {_fmt(path)} has unexpected key(s) " + f"{sorted(extra)}; only 'value' and 'units' are allowed." + ) + return Constant(float(value), units=units) + + def _single_operand(kind, body, path): """Return the single non-``var`` (name, value) of a unary op body, or None. diff --git a/xbudget/query.py b/xbudget/query.py index c2a3868..6dba6bd 100644 --- a/xbudget/query.py +++ b/xbudget/query.py @@ -329,6 +329,34 @@ def bolus_transports(self, budget="mass"): """ return dict(self.metadata(budget).get("bolus", {})) + def budget_units(self, budget): + """The top-level ``units`` a budget declares, or ``None``. + + The UDUNITS string a budget states as the units its closed sides should + reconcile to (e.g. ``"kg s-1"`` for ``mass``). Read from budget metadata; + raises ``KeyError`` for an unknown budget, as :meth:`metadata` does. This + is the *declared* target; :meth:`units` returns what a run actually + inferred for a given term. + """ + return self.metadata(budget).get("units") + + def units(self, address): + """The inferred ``units`` attribute stamped on a term's variable, or ``None``. + + Reads the ``units`` attribute the evaluator stamped onto the materialized + variable (so it also works on a dataset reopened from disk). Returns + ``None`` when the query has no dataset, the term did not materialize, or + its units could not be inferred (an input's units were unknown). Accepts + the same addresses as :meth:`var` — a term name/path, an operation-suffixed + name, or a raw diagnostic. + """ + if self._ds is None: + return None + name = self.var(address) + if name is None or name not in self._ds: + return None + return self._ds[name].attrs.get("units") + def terms(self): """Map every term path to its variable name (``None`` if not materialized).""" return {path: self._resolve_var(term) for path, term in self._by_path.items()} diff --git a/xbudget/recipes/ECCOV4r4_native.yaml b/xbudget/recipes/ECCOV4r4_native.yaml index 290a4eb..b43aa38 100755 --- a/xbudget/recipes/ECCOV4r4_native.yaml +++ b/xbudget/recipes/ECCOV4r4_native.yaml @@ -1,5 +1,6 @@ --- mass: # finite-volume mass budget in units of kg/s + units: "kg s-1" lambda: "density" thickness: "thkcello" # GM eddy-bolus mass transports (kg/s), exposed for downstream water-mass @@ -27,7 +28,9 @@ mass: # finite-volume mass budget in units of kg/s volume: "volcello" deptho_inv: reciprocal: {deptho: {var: "Depth"}} - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" rhs: sum: advection: @@ -37,7 +40,9 @@ mass: # finite-volume mass budget in units of kg/s lateral: # Eulerian lateral mass-flux convergence, div(-Fx,-Fy) = -div(Fx,Fy) product: sign: -1.0 - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" volume_flux_divergence: lateral_divergence: Fx: @@ -48,7 +53,9 @@ mass: # finite-volume mass budget in units of kg/s product: volume_flux_convergence: difference: {z_velocity: "WVELMASS_interior"} # WVELMASS zeroed at surface to avoid double-counting oceFWflx - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" area: "rA" surface_exchange_flux: product: @@ -57,13 +64,18 @@ mass: # finite-volume mass budget in units of kg/s area: "rA" heat: # finite-volume heat budget in s* in units of J/s (W) + units: "W" lambda: "THETA" lhs: sum: Eulerian_tendency: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" volume: "volcello" heatcontenttend: product: @@ -92,8 +104,12 @@ heat: # finite-volume heat budget in s* in units of J/s (W) sum: lateral: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" heat_flux_convergence: product: heat_flux_divergence: @@ -105,16 +121,24 @@ heat: # finite-volume heat budget in s* in units of J/s (W) sign: -1.0 interfacial: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" heat_flux_convergence: difference: {advr_th: "ADVr_TH"} diffusion: sum: lateral: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" explicit_heat_flux_convergence: product: heat_flux_divergence: @@ -128,37 +152,56 @@ heat: # finite-volume heat budget in s* in units of J/s (W) sign: -1.0 explicit_vertical_diffusion: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" heat_flux_convergence: difference: {dfrE_th: "DFrE_TH"} implicit_vertical_diffusion: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" heat_flux_convergence: difference: {dfrI_th: "DFrI_TH"} surface_exchange_flux: # sum: #TFLUX, QFLUX, etc. includes freshwater forcing product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" boundary_heat_flux: "boundary_forcing_heat_tendency" # deg C /s bottom_flux: product: - density: 1029.0 - specific_heat_capacity: 3994.0 + density: + value: 1029.0 + units: "kg m-3" + specific_heat_capacity: + value: 3994.0 + units: "J kg-1 K-1" surface_heat_flux: "geothermal_heat_flux_convergence" # deg C m^3 /s salt: # finite-volume salt-content budget in units of kg/s (salt mass, not salinity) + units: "kg s-1" # Salt content per cell: ρ * V * (s* SALT) * 1e-3 since SALT is in psu ~ g/kg lambda: "SALT" lhs: sum: Eulerian_tendency: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" volume: "volcello" unit_conversion: 0.001 saltcontenttend: @@ -185,7 +228,9 @@ salt: # finite-volume salt-content budget in units of kg/s (salt mass, not sali lateral: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" unit_conversion: 0.001 sign: -1.0 salt_flux_divergence: @@ -195,7 +240,9 @@ salt: # finite-volume salt-content budget in units of kg/s (salt mass, not sali interfacial: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" unit_conversion: 0.001 salt_flux_convergence: difference: {advr_slt: "ADVr_SLT"} @@ -204,7 +251,9 @@ salt: # finite-volume salt-content budget in units of kg/s (salt mass, not sali sum: lateral: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" unit_conversion: 0.001 sign: -1.0 salt_flux_divergence: @@ -214,13 +263,17 @@ salt: # finite-volume salt-content budget in units of kg/s (salt mass, not sali explicit_vertical_diffusion: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" unit_conversion: 0.001 salt_flux_convergence: difference: {dfrE_slt: "DFrE_SLT"} implicit_vertical_diffusion: product: - density: 1029.0 + density: + value: 1029.0 + units: "kg m-3" unit_conversion: 0.001 salt_flux_convergence: difference: {dfrI_slt: "DFrI_SLT"} diff --git a/xbudget/recipes/MOM6.yaml b/xbudget/recipes/MOM6.yaml index 2db0abd..30cba48 100644 --- a/xbudget/recipes/MOM6.yaml +++ b/xbudget/recipes/MOM6.yaml @@ -1,5 +1,6 @@ --- mass: # finite-volume mass budget in units of kg/s + units: "kg s-1" lambda: "density" thickness: "thkcello" lhs: @@ -7,7 +8,9 @@ mass: # finite-volume mass budget in units of kg/s Eulerian_tendency: product: thickness_tendency: "dhdt" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: sum: @@ -16,7 +19,9 @@ mass: # finite-volume mass budget in units of kg/s lateral: product: thickness_tendency: "dynamics_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" sum: zonal_convergence: @@ -34,12 +39,16 @@ mass: # finite-volume mass budget in units of kg/s interfacial: product: thickness_tendency: "vert_remap_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" surface_exchange_flux: product: thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" sum: rain_and_ice: @@ -72,6 +81,7 @@ mass: # finite-volume mass budget in units of kg/s area: "areacello" heat: # finite-volume heat budget in units of J/s + units: "W" lambda: "thetao" surface_lambda: "tos" lhs: @@ -95,10 +105,14 @@ heat: # finite-volume heat budget in units of J/s surface_ocean_flux_advective_negative_lhs: product: sign: -1. - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: sum: @@ -157,10 +171,14 @@ heat: # finite-volume heat budget in units of J/s surface_ocean_flux_advective_negative_rhs: product: sign: -1. - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" bottom_flux: product: @@ -172,6 +190,7 @@ heat: # finite-volume heat budget in units of J/s area: "areacello" salt: # finite-volume salt budget in units of kg/s + units: "kg s-1" lambda: "so" surface_lambda: "sos" lhs: @@ -198,7 +217,9 @@ salt: # finite-volume salt budget in units of kg/s unit_conversion: 0.001 lambda_mass: "sos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: @@ -229,7 +250,9 @@ salt: # finite-volume salt budget in units of kg/s unit_conversion: 0.001 lambda_mass: 0. thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" surface_ocean_flux_advective_negative_rhs: product: @@ -237,5 +260,7 @@ salt: # finite-volume salt budget in units of kg/s unit_conversion: 0.001 lambda_mass: "sos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" diff --git a/xbudget/recipes/MOM6_3Donly.yaml b/xbudget/recipes/MOM6_3Donly.yaml index 99edd23..add476b 100644 --- a/xbudget/recipes/MOM6_3Donly.yaml +++ b/xbudget/recipes/MOM6_3Donly.yaml @@ -1,5 +1,6 @@ --- mass: # finite-volume mass budget in units of kg/s + units: "kg s-1" lambda: "density" thickness: "thkcello" lhs: @@ -7,7 +8,9 @@ mass: # finite-volume mass budget in units of kg/s Eulerian_tendency: product: thickness_tendency: null - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: sum: @@ -29,15 +32,20 @@ mass: # finite-volume mass budget in units of kg/s interfacial: product: thickness_tendency: "vert_remap_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" surface_exchange_flux: product: thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" heat: # finite-volume heat budget in units of J/s + units: "W" lambda: "thetao" surface_lambda: "tos" lhs: @@ -61,10 +69,14 @@ heat: # finite-volume heat budget in units of J/s surface_ocean_flux_advective_negative_lhs: product: sign: -1. - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: sum: @@ -86,10 +98,14 @@ heat: # finite-volume heat budget in units of J/s surface_ocean_flux_advective_negative_rhs: product: sign: -1. - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" bottom_flux: product: @@ -101,6 +117,7 @@ heat: # finite-volume heat budget in units of J/s area: "areacello" salt: # finite-volume salt budget in units of kg/s + units: "kg s-1" lambda: "so" surface_lambda: "sos" lhs: @@ -127,7 +144,9 @@ salt: # finite-volume salt budget in units of kg/s unit_conversion: 0.001 lambda_mass: "sos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" rhs: @@ -152,5 +171,7 @@ salt: # finite-volume salt budget in units of kg/s unit_conversion: 0.001 lambda_mass: "sos" thickness_tendency: "boundary_forcing_h_tendency" - density: 1035. + density: + value: 1035. + units: "kg m-3" area: "areacello" diff --git a/xbudget/recipes/MOM6_drift.yaml b/xbudget/recipes/MOM6_drift.yaml index bb140bb..7a440b0 100644 --- a/xbudget/recipes/MOM6_drift.yaml +++ b/xbudget/recipes/MOM6_drift.yaml @@ -1,5 +1,6 @@ --- mass: # finite-volume mass budget in units of kg/s + units: "kg s-1" lambda: "density" thickness: "thkcello" lhs: @@ -14,6 +15,7 @@ mass: # finite-volume mass budget in units of kg/s var: null heat: # finite-volume heat budget in units of J/s + units: "W" lambda: "thetao" surface_lambda: "tos" lhs: @@ -36,6 +38,7 @@ heat: # finite-volume heat budget in units of J/s var: null salt: # finite-volume salt budget in units of kg/s + units: "kg s-1" lambda: "so" surface_lambda: "sos" lhs: diff --git a/xbudget/recipes/MOM6_surface.yaml b/xbudget/recipes/MOM6_surface.yaml index 298567f..ff0d22a 100644 --- a/xbudget/recipes/MOM6_surface.yaml +++ b/xbudget/recipes/MOM6_surface.yaml @@ -1,5 +1,6 @@ --- mass: # finite-volume mass budget in units of kg/s + units: "kg s-1" lambda: "density" thickness: "thkcello" lhs: @@ -38,6 +39,7 @@ mass: # finite-volume mass budget in units of kg/s area: "areacello" heat: # finite-volume heat budget in units of J/s + units: "W" lambda: "thetao" surface_lambda: "tos" lhs: @@ -66,19 +68,24 @@ heat: # finite-volume heat budget in units of J/s area: "areacello" advective: product: - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" mass_tendency_per_unit_area: "wfo" area: "areacello" surface_ocean_flux_advective_negative_rhs: product: sign: -1. - specific_heat_capacity: 3992. + specific_heat_capacity: + value: 3992. + units: "J kg-1 K-1" lambda_mass: "tos" mass_tendency_per_unit_area: "wfo" area: "areacello" salt: # finite-volume salt budget in units of kg/s + units: "kg s-1" lambda: "so" surface_lambda: "sos" lhs: diff --git a/xbudget/tests/test_units.py b/xbudget/tests/test_units.py new file mode 100644 index 0000000..430fa28 --- /dev/null +++ b/xbudget/tests/test_units.py @@ -0,0 +1,229 @@ +"""Units: the arithmetic helpers, unit inference, stamping, and strict mode. + +All CI-safe: the unit algebra is pure, and the inference/strict tests run on the +synthetic grid (with units attached here rather than in the shared fixture, so +the many tests that use the unit-less synthetic grid are undisturbed). +""" +import glob +import warnings + +import pytest +import yaml + +import xbudget +from xbudget import units as U +from xbudget.nodes import Constant +from xbudget.parse import BudgetParseError, parse_budgets +from xbudget.tests.conftest import build_synthetic_grid + + +# -- pure unit arithmetic --------------------------------------------------- + +def _u(s): + return U.parse_units(s) + + +class TestUnitArithmetic: + def test_parse_unknown_is_none(self): + assert U.parse_units(None) is None + assert U.parse_units("") is None + assert U.parse_units(" ") is None + assert U.parse_units("not-a-unit") is None + + def test_parse_practical_salinity_is_dimensionless(self): + # UDUNITS has no `psu`; PSS-78 salinity is dimensionless. + assert U.convertible(U.parse_units("psu"), U.DIMENSIONLESS) + assert U.convertible(U.parse_units("PSU"), _u("1")) + + def test_product_multiplies(self): + got = U.product_units([_u("m s-1"), _u("kg m-3"), _u("m2")]) + assert U.convertible(got, _u("kg s-1")) + + def test_product_with_unknown_is_unknown(self): + assert U.product_units([_u("m"), None]) is None + + def test_product_dimensionalizes_offset_temperature(self): + # A heat-content product carries degC; cf-units treats it as kelvin in a + # product, so theta * c_p * rho * dh/dt * area reconciles to W. + got = U.product_units( + [_u("degC"), _u("J kg-1 K-1"), _u("m s-1"), _u("kg m-3"), _u("m2")] + ) + assert U.convertible(got, _u("W")) + + def test_sum_requires_compatibility(self): + assert U.convertible(U.sum_units([_u("W"), _u("W")]), _u("W")) + assert U.sum_units([_u("W"), _u("kg s-1")]) is None + assert U.sum_units([_u("W"), None]) is None + + def test_sum_compatible_flags_only_real_clashes(self): + assert U.sum_compatible([_u("W"), _u("J s-1")]) # convertible + assert U.sum_compatible([_u("W"), None]) # unknown != clash + assert not U.sum_compatible([_u("W"), _u("kg s-1")]) + + def test_reciprocal_and_difference(self): + assert U.convertible(U.reciprocal_units(_u("s")), _u("s-1")) + assert U.reciprocal_units(None) is None + assert U.difference_units(_u("kg s-1")) == _u("kg s-1") + assert U.difference_units(None) is None + + def test_divergence_of_matching_fluxes(self): + assert U.convertible(U.divergence_units(_u("kg s-1"), _u("kg s-1")), _u("kg s-1")) + assert U.divergence_units(_u("kg s-1"), _u("W")) is None + + +# -- parsing constants and budget units ------------------------------------- + +class TestParseUnits: + def test_bare_number_is_dimensionless(self): + recipe = {"b": {"rhs": {"product": {"sign": -1.0, "d": "x"}}}} + prod = parse_budgets(recipe)["b"].sides["rhs"].operations[0] + const = dict(prod.terms)["sign"] + assert isinstance(const, Constant) + assert const.value == -1.0 + assert U.convertible(U.parse_units(const.units), U.DIMENSIONLESS) + + def test_dict_form_constant_carries_units(self): + recipe = { + "b": {"rhs": {"product": { + "rho": {"value": 1035.0, "units": "kg m-3"}, "d": "x", + }}} + } + prod = parse_budgets(recipe)["b"].sides["rhs"].operations[0] + rho = dict(prod.terms)["rho"] + assert rho == Constant(1035.0, units="kg m-3") + + def test_constant_value_must_be_numeric(self): + recipe = {"b": {"rhs": {"product": {"c": {"value": "oops"}}}}} + with pytest.raises(BudgetParseError, match="numeric 'value'"): + parse_budgets(recipe) + + def test_constant_rejects_extra_keys(self): + recipe = {"b": {"rhs": {"product": { + "c": {"value": 1.0, "units": "1", "typo": 3}, + }}}} + with pytest.raises(BudgetParseError, match="unexpected key"): + parse_budgets(recipe) + + def test_dict_with_operation_is_a_subterm_not_constant(self): + # A mapping that also names an operation is a sub-term even if it has a + # `value`-named child term (only a units-constant has a bare `value`). + recipe = {"b": {"rhs": {"sum": { + "t": {"product": {"value": "x"}}, # `value` here is a term label + }}}} + sum_op = parse_budgets(recipe)["b"].sides["rhs"].operations[0] + t = dict(sum_op.terms)["t"] + assert not isinstance(t, Constant) + + def test_budget_units_metadata(self): + recipe = {"mass": {"units": "kg s-1", "rhs": {"product": {"d": "x"}}}} + assert parse_budgets(recipe)["mass"].metadata["units"] == "kg s-1" + + def test_shipped_recipes_declare_budget_units(self): + for path in glob.glob("xbudget/recipes/*.yaml"): + budgets = parse_budgets(yaml.safe_load(open(path))) + for name, budget in budgets.items(): + declared = budget.metadata.get("units") + assert declared is not None, f"{path}:{name} has no units" + assert U.parse_units(declared) is not None, f"{path}:{name} {declared!r}" + + +# -- inference, stamping, and strict mode ----------------------------------- + +def _units_grid(): + """The synthetic grid with UDUNITS attributes so inference has something.""" + grid = build_synthetic_grid() + grid._ds["diag_a"].attrs["units"] = "kg m-2 s-1" + grid._ds["diag_b"].attrs["units"] = "kg m-2 s-1" + grid._ds["area"].attrs["units"] = "m2" + grid._ds["flux"].attrs["units"] = "kg s-1" + return grid + + +# A units-consistent tracer budget: every rhs summand reconciles to kg s-1. +CONSISTENT = { + "tracer": { + "units": "kg s-1", + "rhs": {"sum": { + "diffusion": {"product": {"sign": -1.0, "d": "diag_a", "area": "area"}}, + "boundary": {"product": {"d": "diag_b", "area": "area"}}, + "convergence": {"difference": {"transport": "flux"}}, + }}, + } +} + + +class TestInferenceAndStrict: + def test_units_are_stamped(self): + grid = _units_grid() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xbudget.collect_budgets(grid, CONSISTENT) + q = xbudget.BudgetQuery(grid, CONSISTENT) + for addr in [("tracer", "rhs"), ("tracer", "rhs", "diffusion"), + ("tracer", "rhs", "convergence")]: + assert U.convertible(U.parse_units(q.units(addr)), _u("kg s-1")) + + def test_constant_units_flow_into_product(self): + grid = build_synthetic_grid() + grid._ds["diag_a"].attrs["units"] = "m s-1" # a thickness tendency + grid._ds["area"].attrs["units"] = "m2" + recipe = {"mass": {"units": "kg s-1", "rhs": {"product": { + "dhdt": "diag_a", + "density": {"value": 1035.0, "units": "kg m-3"}, + "area": "area", + }}}} + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xbudget.collect_budgets(grid, recipe, strict=True) + q = xbudget.BudgetQuery(grid, recipe) + assert U.convertible(U.parse_units(q.units(("mass", "rhs"))), _u("kg s-1")) + + def test_strict_passes_when_consistent(self): + grid = _units_grid() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + # Should not raise. + xbudget.collect_budgets(grid, CONSISTENT, strict=True) + + def test_strict_raises_on_mismatch(self): + grid = _units_grid() + mism = {"tracer": {**CONSISTENT["tracer"], "units": "W"}} # declare wrong units + with pytest.raises(xbudget.UnitError) as exc: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xbudget.collect_budgets(grid, mism, strict=True) + assert exc.value.issues # carries the offending issue list + + def test_nonstrict_warns_on_mismatch(self): + grid = _units_grid() + mism = {"tracer": {**CONSISTENT["tracer"], "units": "W"}} + with pytest.warns(UserWarning, match="did not fully reconcile"): + xbudget.collect_budgets(grid, mism, strict=False) + + def test_strict_raises_on_unknown_units(self): + grid = build_synthetic_grid() # no units on any diagnostic + with pytest.raises(xbudget.UnitError): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xbudget.collect_budgets(grid, CONSISTENT, strict=True) + + def test_undeclared_budget_is_not_checked(self): + grid = _units_grid() + no_units = {"tracer": {"rhs": CONSISTENT["tracer"]["rhs"]}} + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) # no unit warning expected + warnings.filterwarnings("ignore", message=".*mismatched dimensions.*") + xbudget.collect_budgets(grid, no_units, strict=True) # must not raise + + def test_query_units_and_budget_units(self): + grid = _units_grid() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + xbudget.collect_budgets(grid, CONSISTENT) + q = xbudget.BudgetQuery(grid, CONSISTENT) + assert q.budget_units("tracer") == "kg s-1" + assert q.units(("tracer", "rhs")) is not None + # A query with no dataset knows the declared units but not inferred ones. + q_planned = xbudget.BudgetQuery(None, CONSISTENT) + assert q_planned.budget_units("tracer") == "kg s-1" + assert q_planned.units(("tracer", "rhs")) is None diff --git a/xbudget/units.py b/xbudget/units.py new file mode 100644 index 0000000..6bec6b5 --- /dev/null +++ b/xbudget/units.py @@ -0,0 +1,160 @@ +"""Unit inference for xbudget derived variables, backed by UDUNITS (cf-units). + +The evaluator materializes one variable per operation; this module gives each +one an inferred ``units`` attribute by composing the units of its inputs, using +the same arithmetic the physics does: + +- ``product`` -> multiply the operand units; +- ``sum`` -> the common unit of the summands (they must be dimensionally + compatible), else *unknown*; +- ``difference`` -> unchanged (a finite difference preserves units); +- ``reciprocal`` -> ``1 / unit``; +- ``lateral_divergence`` -> the flux unit (Fx and Fy must be compatible). + +Units are ``cf_units.Unit`` objects (real UDUNITS parsing/algebra) or ``None``. +``None`` means *unknown* — a missing, blank, or unparseable ``units`` attribute — +and is deliberately distinct from the dimensionless :data:`DIMENSIONLESS` +(``Unit("1")``). Unknown propagates: any operation with an unknown operand yields +unknown rather than a fabricated dimension. + +Every function here is pure and total (it never raises and never warns): it just +answers "what are the units?". Whether an unknown or a mismatch is *tolerable* +is policy, and lives in :mod:`xbudget.evaluate` (the ``strict`` flag) rather than +here. + +Two conventions worth knowing: + +- **Practical salinity.** UDUNITS does not define ``psu``; PSS-78 salinity is + dimensionless, so :func:`parse_units` aliases ``psu`` (and spellings of it) to + ``"1"``. A salt term is then ``salinity x mass-flux``, which reconciles to + ``kg s-1`` (the residual ``1e-3`` from a ``unit_conversion`` factor is a scale, + not a dimension, and does not affect convertibility). +- **Temperature.** ``degC`` is an offset unit, but cf-units dimensionalizes it to + kelvin inside a product automatically (``degC * kg -> kg.K``), so a heat-content + term ``theta x c_p x rho x dh/dt x area`` reconciles to ``W`` with no special + handling here. +""" +from functools import reduce +from operator import mul + +from cf_units import Unit + +__all__ = [ + "DIMENSIONLESS", + "parse_units", + "product_units", + "sum_units", + "sum_compatible", + "reciprocal_units", + "difference_units", + "divergence_units", + "convertible", + "format_units", +] + +DIMENSIONLESS = Unit("1") + +# UDUNITS has no practical-salinity unit; PSS-78 salinity is dimensionless. +_ALIASES = { + "psu": "1", + "PSU": "1", + "practical_salinity_unit": "1", + "practical_salinity_units": "1", + "pss-78": "1", + "PSS-78": "1", +} + + +def parse_units(spec): + """Parse a units string (or ``Unit``/attribute) into a ``Unit``, or ``None``. + + ``None`` (unknown) is returned for ``None``, a blank string, an unparseable + string, or a cf-units ``unknown``/``no-unit`` sentinel — anything that is not + a real dimension. This is distinct from :data:`DIMENSIONLESS`. + """ + if spec is None: + return None + if isinstance(spec, Unit): + return spec if _is_real(spec) else None + text = str(spec).strip() + if not text: + return None + text = _ALIASES.get(text, text) + try: + unit = Unit(text) + except ValueError: + return None + return unit if _is_real(unit) else None + + +def _is_real(unit): + """True unless ``unit`` is cf-units' ``unknown``/``no-unit`` sentinel.""" + return not (unit.is_unknown() or unit.is_no_unit()) + + +def product_units(units): + """Units of a product of operands; ``None`` if any operand is unknown.""" + units = list(units) + if any(u is None for u in units): + return None + if not units: + return DIMENSIONLESS + # Seed the reduction with the first real unit rather than DIMENSIONLESS so an + # offset operand (e.g. degC) is never the right operand of ``1 * degC``. + return reduce(mul, units[1:], units[0]) + + +def reciprocal_units(unit): + """Units of ``1 / x``; ``None`` if ``x``'s units are unknown.""" + if unit is None: + return None + return DIMENSIONLESS / unit + + +def difference_units(unit): + """Units of a finite difference: unchanged (may be ``None``).""" + return unit + + +def sum_compatible(units): + """Whether every *known* summand shares a dimension. + + ``True`` when fewer than two units are known (nothing to contradict); ``False`` + only when two known units are dimensionally incompatible. Unknown operands do + not make a sum *incompatible* — they make its result unknown (see + :func:`sum_units`) — so this is the signal the evaluator uses to distinguish a + real unit clash (worth a warning/error) from a merely-absent unit. + """ + known = [u for u in units if u is not None] + return all(known[0].is_convertible(u) for u in known[1:]) + + +def sum_units(units): + """Common units of the summands, or ``None``. + + ``None`` if any summand is unknown or if two summands are dimensionally + incompatible. Otherwise the units of the first summand (all are compatible). + """ + units = list(units) + if not units or any(u is None for u in units): + return None + if not sum_compatible(units): + return None + return units[0] + + +def divergence_units(fx, fy): + """Units of a horizontal flux divergence: the (compatible) flux units.""" + return sum_units([fx, fy]) + + +def convertible(unit, other): + """Whether ``unit`` is dimensionally convertible to ``other`` (both known).""" + if unit is None or other is None: + return False + return unit.is_convertible(other) + + +def format_units(unit): + """A stable string for a ``Unit`` (for stamping as an attribute), or ``None``.""" + return None if unit is None else str(unit) From 13079fdb09f16e6bad09bbdb0bb59efef49044b5 Mon Sep 17 00:00:00 2001 From: Henri Drake Date: Tue, 28 Jul 2026 08:38:25 -0700 Subject: [PATCH 2/6] Re-execute example notebooks under strict units mode Attach UDUNITS units to the derived budget inputs the ECCO notebooks build (dt, volcello, WVELMASS_interior, boundary_forcing_*) so unit inference has a complete picture, and switch the three budget notebooks (MOM6 + both ECCO) to collect_budgets(..., strict=True). Verified: strict passes on both the MOM6 and ECCO V4r4 (LLC90) datasets; derived variables now display their inferred units. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../MOM6_budget_examples_mass_heat_salt.ipynb | 256 +++++++++--------- ...ov4r4_budget_examples_mass_heat_salt.ipynb | 204 +++++++++----- .../eccov4r4_heat_budget_decomposition.ipynb | 171 ++++++++---- examples/handling_missing_diagnostics.ipynb | 80 +++--- 4 files changed, 424 insertions(+), 287 deletions(-) diff --git a/examples/MOM6_budget_examples_mass_heat_salt.ipynb b/examples/MOM6_budget_examples_mass_heat_salt.ipynb index 7531894..7f8e1cb 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-28T15:33:01.766664Z", + "iopub.status.busy": "2026-07-28T15:33:01.766388Z", + "iopub.status.idle": "2026-07-28T15:33:01.879345Z", + "shell.execute_reply": "2026-07-28T15:33:01.878869Z" } }, "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-28T15:33:01.880733Z", + "iopub.status.busy": "2026-07-28T15:33:01.880650Z", + "iopub.status.idle": "2026-07-28T15:33:04.997027Z", + "shell.execute_reply": "2026-07-28T15:33:04.996451Z" } }, "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-28T15:33:04.998864Z", + "iopub.status.busy": "2026-07-28T15:33:04.998668Z", + "iopub.status.idle": "2026-07-28T15:33:06.360523Z", + "shell.execute_reply": "2026-07-28T15:33:06.359568Z" } }, "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-28T15:33:06.364794Z", + "iopub.status.busy": "2026-07-28T15:33:06.364533Z", + "iopub.status.idle": "2026-07-28T15:33:06.436880Z", + "shell.execute_reply": "2026-07-28T15:33:06.436226Z" } }, "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-28T15:33:06.438563Z", + "iopub.status.busy": "2026-07-28T15:33:06.438455Z", + "iopub.status.idle": "2026-07-28T15:33:06.459336Z", + "shell.execute_reply": "2026-07-28T15:33:06.458379Z" } }, "outputs": [ @@ -207,11 +207,11 @@ " 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Σ
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
" ], "text/plain": [ "xbudget recipe · heat\n", - "heat lambda=thetao, surface_lambda=tos\n", + "heat units=W, lambda=thetao, surface_lambda=tos\n", "├─ lhs [Σ]\n", "│ ├─ Eulerian_tendency [×]\n", "│ │ ├─ tracer_content_tendency_per_unit_area: opottemptend\n", @@ -299,15 +299,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-28T15:33:06.460982Z", + "iopub.status.busy": "2026-07-28T15:33:06.460867Z", + "iopub.status.idle": "2026-07-28T15:33:06.883460Z", + "shell.execute_reply": "2026-07-28T15:33:06.882933Z" } }, "outputs": [], "source": [ - "xbudget.collect_budgets(grid, recipe)\n", + "# strict=True asserts each budget reconciles to its declared units (mass/salt -> kg s-1, heat -> W).\n", + "xbudget.collect_budgets(grid, recipe, strict=True)\n", "q = xbudget.BudgetQuery(grid, recipe)" ] }, @@ -327,17 +328,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-28T15:33:06.884997Z", + "iopub.status.busy": "2026-07-28T15:33:06.884889Z", + "iopub.status.idle": "2026-07-28T15:33:06.905854Z", + "shell.execute_reply": "2026-07-28T15:33:06.905432Z" } }, "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 +374,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-28T15:33:06.907231Z", + "iopub.status.busy": "2026-07-28T15:33:06.907135Z", + "iopub.status.idle": "2026-07-28T15:33:06.960491Z", + "shell.execute_reply": "2026-07-28T15:33:06.960013Z" } }, "outputs": [ @@ -842,7 +844,8 @@ "Attributes:\n", " provenance: ['heat_lhs_advection_lateral', 'heat_lhs_advection_interfa...\n", " xbudget_path: ['heat', 'lhs', 'advection']\n", - " xbudget_op: sum