Typed engine refactor + LLC/ECCO integration (native xgcm)#26
Open
hdrake wants to merge 22 commits into
Open
Conversation
this commit incorporates an optional "allow_rechunk" option for collect_budgets and budget_fill_dict
…ow_rechunk" option
By explicitly adding chunks to the example MOM6 dataset, we should be able to better catch chunk-related problems in the future.
Pins the full set of variables collect_budgets materializes from the example MOM6 grid (108 vars) with their shape, provenance, and numerical fingerprint, as a safety net for the upcoming internals refactor. The ~600MB example dataset is not tracked, so the test skips when it is absent; the small reference fingerprint is committed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Fix misplaced grid guard in budget_fill_dict's `difference` branch:
the `else: raise(...must be xgcm.Grid...)` was attached to
`if var_pref is None` instead of a grid check, so (a) a difference op
on a plain Dataset hit an undefined `staggered_axes` (opaque
NameError) instead of a clear error, and (b) a difference term reached
after another operation in the same node spuriously raised even with a
valid grid. Now raises a clear ValueError up front when no grid.
- Collapse 4 copy-pasted 'missing variable' warnings into one helper.
- Remove mutable default arg (new_b={}) in _deep_search.
- Add regression tests covering both difference-guard failure modes.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Introduces an explicit, immutable expression tree for xbudget recipes so the engine can dispatch on node type instead of probing dict keys: - nodes.py: Budget/Term + Sum/Product/Difference/Reciprocal/Constant/ VarRef dataclasses. A Term may carry multiple operations (e.g. a bulk product and an equivalent finer sum). - parse.py: dict -> AST with schema validation (BudgetParseError names the offending path), the single source of schema truth. - evaluate.py: pure evaluator that materializes one variable per operation using simplified path-based names (operator infixes dropped, redundant 'copy' variables collapsed), with xbudget_path/provenance attrs and a legacy<->new alias map. Differential tests prove numerical equivalence to the legacy engine: on the example MOM6 grid the 108 legacy variables collapse to 57, all matching to rtol=1e-9, with every legacy name reachable via the alias map. Covered by a synthetic grid (CI) and the gated MOM6 grid. Public API is unchanged so far; rewiring follows in Phase 2b. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
collect_budgets now parses the convention into the typed tree and runs
the pure evaluator instead of the legacy dict-walker:
- Emits the simplified path-based variable names by default; no longer
mutates the input recipe dict; returns the data object for chaining.
- Adds name_scheme='legacy' to additionally write the historical
operator-suffixed names as aliases for migration.
- Handles leaf terms ({"var": "diag"}) by aliasing the referenced
diagnostic, matching legacy behavior.
- budget_fill_dict retained as the (now documented-as-legacy) reference
engine that the equivalence test pins against.
- Exports parse_budgets/evaluate_budgets/BudgetParseError.
Tests updated: collect_budgets assertions use the new names with added
coverage for legacy mode and no-recipe-mutation; the equivalence oracle
drives budget_fill_dict directly; characterization reference regenerated
for the typed engine (57 vars). Full suite: 51 passed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- name_scheme='legacy' now reproduces the old engine exactly (historical names + in-place dict filling), so get_vars/aggregate keep working; the default v1 mode stays clean (new names, no mutation). - Add CHANGELOG.md with the migration guide and old->new name table. - Notebook: set name_scheme='legacy' on the single collect_budgets call so the rest of the tutorial (old names, get_vars, aggregate) is unchanged. - Update legacy-mode test to assert faithful old-name output + dict fill. Full suite: 51 passed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Documents the released (v0.6.2) architecture: the xbudget_dict convention tree, the collect.py dict-walking engine (collect_budgets/budget_fill_dict), the query helpers, and key behaviors (in-place dict filling, missing-diagnostic handling, string-path variable naming). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rewrites the architecture section to describe the parse->evaluate engine (nodes/parse/evaluate), the v1 vs legacy name schemes, collect_budgets no longer mutating the recipe, the retained legacy reference engine, and the new test layout. Points to CHANGELOG.md for migration. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…etadata Acting on the three-agent audit of the refactor: Engine: - Remove the `reciprocal` operation entirely (unused by any convention, untested): dropped from nodes/parse/evaluate and the legacy budget_fill_dict; removed its test. - Fix latent multi-operation naming bug: the bare path name now goes to the first sum/product op (matching the legacy 'copy'), so a difference/unary op listed first no longer steals it. Only affects never-shipped difference-first multi-op terms; shipped conventions unchanged. Tests: - Add a leaf-term case to the synthetic equivalence fixture; make the equivalence count invariant robust to leaf terms (and assert injectivity). Docs/clarity: - Note that get_vars/aggregate/disaggregate require name_scheme='legacy'. - presets.load_preset_budget lists all four shipped conventions. - Fix disaggregate docstring example; clarify MOM6.yaml shortwave comments (drop the reference to a non-existent finite_difference op). - CLAUDE.md: drop reciprocal, document the xbudget_op attr. Metadata: - Bump version to 0.7.0; retitle CHANGELOG and note reciprocal removal. - Add Python 3.11 to the CI matrix and the test env (matches requires-python). - Add __all__ to collect.py so `import *` stops leaking np/xr/etc. Full suite: 50 passed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
When pushing a new commit to an open PR, update the PR body to match the latest commit — including checking off completed task-list items — via gh pr edit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Reconciles the typed-engine refactor with origin/main's LLC90/ECCO feature (#25), allow_rechunk (#24), and v0.6.3, on a clean PR branch. Native xgcm replaces the custom LLC op: - New collect.lateral_divergence() uses grid.diff({axis: F}, axis, other_component=...) + face_connections; verified bit-exact vs the old diff_2d_flux_llc90 on the real ECCO LLC90 grid (heat/salt adv+diff fluxes). - xbudget/llc90/operations.py deleted; both engines now use the native helper. Typed engine gains full ECCO support: - Restored Reciprocal node; added LateralDivergence node (Fx/Fy fluxes). - Difference now accepts a computed sub-term operand (not just a raw var), so the heat/salt LHS Eulerian content-tendency is reproduced. - Parser tolerates (warns + skips) unavailable-diagnostic placeholders and malformed-but-ignorable terms, matching the legacy engine. Validation: the typed engine reproduces the legacy engine on the full ECCO grid exactly (71 vars, 0 mismatches, all 132 legacy names mapped) and MOM6 is unchanged. Full suite: 52 passed. Also: Phase-1 difference-guard fix re-applied; _deep_search mutable-default and duplicated warnings cleaned; legacy-only notes on get_vars/aggregate/ disaggregate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- test_evaluate_equivalence: add gated ECCO LLC90 typed-vs-legacy test (71 vars, 0 mismatches; exercises reciprocal, difference-of-sub-term, and native lateral_divergence on the 13-tile face-connected grid). - Fix the existing ECCO notebook for the typed engine: drop the deleted xbudget.llc90 import and use name_scheme='legacy' (aggregate/get_vars read the filled dict). - Add eccov4r4_heat_budget_decomposition.ipynb: decomposes the ECCO heat budget via aggregate(decompose=...), mirroring the MOM6 decomposition workflow. (Needs execution against the ECCO data to populate outputs.) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- CHANGELOG: correct the (wrong) 'reciprocal removed' note — it is restored; document ECCO/LLC90 typed-engine support, native-xgcm lateral_divergence, difference-of-sub-term, the parser-tolerance behavior (and the bolus malformation it surfaces), and the unreleased-xgcm dependency. - CLAUDE.md: add reciprocal/lateral_divergence ops, ECCO convention + grid loader + notebooks, the lenient-parser and xgcm-version notes, and the ECCO equivalence test. - docs/source: add the heat-decomposition notebook to the toctree; mention the new operations. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The lateral bolus_mass_flux_convergence term in ECCOV4r4_native.yaml was missing its enclosing 'product:' wrapper, so its sign/density/ volume_flux_divergence children sat directly on the term and were ignored by both engines — the GM bolus velocity (UVELSTAR/VVELSTAR) contributed nothing to the mass budget. Restore the wrapper so the bolus convergence is materialized. This changes ECCO mass-budget results (the lateral bolus term is no longer zero: |sum| ~ 9.6e13 kg/s on the example grid). The typed engine still reproduces the legacy engine exactly (now 75 vars / 140 legacy names, 0 mismatches). Update the gated ECCO equivalence test counts and assert the bolus convergence variable is produced; correct CHANGELOG/CLAUDE accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
version.py is already 0.7.0 (the authoritative source via hatch). Sync the conda recipe version; its sha256 must be regenerated against the published 0.7.0 sdist at release time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Plotting cells made robust to terms lacking a vertical (k) or time dimension (e.g. the geothermal bottom flux is static in k,tile,j,i). Executed end-to-end against the ECCO LLC90 data: 8/8 cells, 0 errors, 2 figures (decomposed maps + global-mean time series). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces
xbudget's recursive dict-walking engine with a typed expression tree (parse → evaluate), while keeping the convention/YAML format unchanged and preserving numerical results exactly. Also integrates the LLC90/ECCO work frommainonto the new engine, switches the LLC flux divergence to native xgcm, and fixes a physics bug in the ECCO mass budget.The recipe (
xbudget_dict) is the public input format; internally it is now parsed into immutable dataclass nodes (xbudget/nodes.py) and evaluated by a pure evaluator (xbudget/evaluate.py), with a single schema-validating parser (xbudget/parse.py).budget_fill_dictis retained as the legacy reference engine and pins equivalence.Why
The old engine inferred a node's meaning from which magic dict keys it carried, built variable names by string concatenation and reverse-parsed them, and mutated both the dataset and the recipe in place. This made it hard to extend, hard to test in isolation, and (in one case) silently wrong. The typed engine dispatches on node type, is pure with respect to the recipe, and emits one clearly-named variable per operation.
What changed
Typed engine
parse_budgets(dict) -> {name: Budget},evaluate_budgets(data, budgets) -> (alias_map, records).collect_budgets(data, xbudget_dict, allow_rechunk=True, name_scheme="v1")parses then evaluates.v1(default) uses simplified names (operator infixes dropped, duplicate "copy" vars collapsed) and does not mutate the recipe;legacyreproduces the historical names and fills the recipe dict in place.xbudget_path,xbudget_op, andprovenanceattrs.LLC90 / ECCO
reciprocal,lateral_divergence, anddifference-of-a-computed-sub-term are all supported in the typed engine, so the native-grid ECCO mass/heat/salt budgets evaluate underv1.lateral_divergencenow uses native xgcm (grid.diffwithother_component+face_connections) instead of the hand-rolleddiff_2d_flux_llc90; verified bit-for-bit identical on the real ECCO LLC90 grid.xbudget/llc90is removed.Physics fix
bolus_mass_flux_convergence(lateral) term was missing itsproduct:wrapper, so the eddy bolus transport was silently dropped from the mass budget. Fixed in the convention; the bolus convergence is now materialized. This changes ECCO mass-budget results.Other
differencegrid-guard (opaqueNameError/ spurious raise → clearValueError).name_scheme="legacy"for the example notebooks that useget_vars/aggregate.eccov4r4_heat_budget_decomposition.ipynb(decomposes the ECCO heat budget).0.7.0;CHANGELOG.mdmigration guide; CLAUDE.md / docs updates.Validation
lateral_divergence, and asserts the bolus convergence is materialized.Notes for the reviewer
lateral_divergenceneeds native face-connected differencing, only on xgcmmain(post-0.9.0). Pins should tighten once a release ships it.🤖 Generated with Claude Code