diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md
new file mode 100644
index 0000000..e1a4360
--- /dev/null
+++ b/.claude/CLAUDE.md
@@ -0,0 +1,5 @@
+
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+@../AGENTS.md
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 0000000..32156e2
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,110 @@
+
+# AGENTS.md
+
+
+This file provides guidance to work with code in this repository.
+
+## Project Overview
+
+zppy_interfaces is a Python package providing extra functionality for E3SM climate model analysis. It processes log files, generates time series plots, and runs PCMDI diagnostics. The package is designed to be called by zppy or used standalone for climate model analysis.
+
+## Development Environment Setup
+
+Set up the development environment using conda:
+```bash
+conda clean --all --y
+conda env create -f conda/dev.yml -n zppy-interfaces-dev
+conda activate zppy-interfaces-dev
+pip install .
+pre-commit install
+```
+
+## Common Development Commands
+
+### Installation and Setup
+- `pip install .` - Install package in development mode
+- `pip install -e .[testing]` - Install with testing dependencies
+- `pip install -e .[qa]` - Install with quality assurance tools
+
+### Code Quality and Testing
+- `pytest` - Run all tests
+- `pytest tests/unit/budget_analysis/` - Run specific module tests
+- `pytest tests/unit/budget_analysis/test_atm_parser.py` - Run single test file
+- `pre-commit run --all-files` - Run all pre-commit hooks
+- `black .` - Format code with Black
+- `isort .` - Sort imports
+- `flake8` - Check code style
+- `mypy zppy_interfaces/` - Type checking
+
+### CLI Tools Testing
+Test the main CLI applications:
+- `zi-budget-analysis --help` - Budget analysis tool
+- `zi-global-time-series --help` - Global time series plots
+- `zi-pcmdi-link-observation --help` - PCMDI observation linking
+- `zi-pcmdi-mean-climate --help` - PCMDI mean climate diagnostics
+- `zi-pcmdi-variability-modes --help` - PCMDI variability modes
+- `zi-pcmdi-enso --help` - PCMDI ENSO diagnostics
+- `zi-pcmdi-synthetic-plots --help` - PCMDI synthetic plots
+
+## Architecture Overview
+
+### Main Components
+
+**budget_analysis/** - E3SM water and energy budget analysis
+- `__main__.py` - CLI entry point with legacy and whole-model modes
+- `parser.py` - Core budget parsing logic for coupler logs
+- `ingestion/` - Component-specific log parsers (atm, ocn, ice, lnd, cpl)
+- `plotting.py` - HTML plot generation for legacy mode
+- `viz.py` - Visualization for whole-model mode
+- `checks.py` - Budget conservation checks
+- `normalization.py` - Data normalization utilities
+
+**global_time_series/** - Global time series plot generation
+- `__main__.py` - CLI with viewer vs PDF output modes
+- `coupled_global/` - Core time series generation logic
+- `create_ocean_ts.py` - Ocean-specific time series processing
+- `utils.py` - Parameter handling utilities
+
+**pcmdi_diags/** - PCMDI diagnostics suite
+- Multiple CLI tools for different diagnostic types
+- `viewer.py` - HTML viewer generation
+- `synthetic_plots/` - Synthetic plot utilities
+
+**multi_utils/** - Shared utilities
+- `logger.py` - Logging setup for child processes
+- `viewer.py` - Common viewer functionality
+
+### Data Flow Patterns
+
+**Budget Analysis (whole-model mode):**
+1. Ingest: Parse multiple log file types (cpl, atm, ocn, ice, lnd)
+2. Normalize: Standardize data formats and units
+3. Check: Run conservation checks and compute residuals
+4. Visualize: Generate HTML reports with interactive plots
+
+**Global Time Series:**
+1. Ocean processing: Extract time series from MPAS-Analysis results (optional)
+2. Coupled analysis: Generate regional and global plots
+3. Output: HTML viewer (interactive) or PDF (static) based on make_viewer setting
+
+### Key Configuration Files
+
+- `pyproject.toml` - Package configuration, dependencies, CLI entry points
+- `conda/dev.yml` - Development environment specification
+- `.pre-commit-config.yaml` - Code quality hooks (black, isort, flake8, mypy)
+- `.flake8` - Flake8 configuration (line length 119, specific ignores)
+
+### Testing Strategy
+
+- Unit tests in `tests/unit/` organized by module
+- Example scripts in `examples/` showing realistic usage
+- Integration with pre-commit hooks for quality assurance
+- pytest with coverage reporting capabilities
+
+## Important Notes
+
+- The package handles both compressed (.gz) and uncompressed log files
+- Budget analysis supports both "legacy" (coupler-only) and "whole-model" modes
+- Time series generation can produce either interactive HTML viewers or static PDFs
+- All CLI tools use argparse with comprehensive help documentation
+- The codebase follows strict code quality standards with Black formatting and comprehensive linting
diff --git a/pyproject.toml b/pyproject.toml
index 6a04641..5dc0de0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,10 +19,12 @@ classifiers = [
dependencies = [
"beautifulsoup4",
+ "bokeh",
"lxml",
"matplotlib",
"netcdf4",
"numpy >=2.0,<3.0",
+ "pandas",
"pcmdi_metrics>=3.9.3",
"xarray >=2023.02.0",
"xcdat >=0.7.3,<1.0",
@@ -117,6 +119,7 @@ version = { attr = "zppy_interfaces.version.__version__" }
# evolution of options.entry-points
[project.scripts]
+zi-budget-analysis = "zppy_interfaces.budget_analysis.__main__:main"
zi-global-time-series = "zppy_interfaces.global_time_series.__main__:main"
zi-pcmdi-link-observation = "zppy_interfaces.pcmdi_diags.link_observation:main"
zi-pcmdi-mean-climate = "zppy_interfaces.pcmdi_diags.pcmdi_mean_cimate:main"
diff --git a/tests/unit/budget_analysis/diagnose_ocn_closure.py b/tests/unit/budget_analysis/diagnose_ocn_closure.py
new file mode 100644
index 0000000..7cbe968
--- /dev/null
+++ b/tests/unit/budget_analysis/diagnose_ocn_closure.py
@@ -0,0 +1,96 @@
+"""Diagnose ocean closure: verify log-native term names are used correctly."""
+
+import glob
+import sys
+
+sys.path.insert(0, ".")
+
+from zppy_interfaces.budget_analysis.checks import OcnClosure # noqa: E402
+from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser # noqa: E402
+from zppy_interfaces.budget_analysis.normalization import normalize # noqa: E402
+
+LOG_PATH = "/pscratch/sd/c/chengzhu/zstash/archive/logs"
+START_YEAR = 1
+END_YEAR = 50
+
+ocn = OcnParser().parse_files(
+ sorted(glob.glob(f"{LOG_PATH}/ocn.log.*.gz")), START_YEAR, END_YEAR
+)
+
+print(f"Total ocean rows: {len(ocn)}")
+print()
+
+# --- Water ---
+water = ocn[ocn["quantity"] == "water"]
+print("=== Water term counts ===")
+print(water["term"].value_counts().to_string())
+print()
+
+mc = water[water["term"] == "Mass change"]
+print(f"=== 'Mass change' rows: {len(mc)} ===")
+if not mc.empty:
+ print(
+ mc[["time", "value", "units", "table_type", "period"]]
+ .head(12)
+ .to_string(index=False)
+ )
+else:
+ print(" *** NOT FOUND ***")
+print()
+
+svf = water[water["term"] == "SUM VOLUME FLUXES"]
+print(f"=== 'SUM VOLUME FLUXES' rows: {len(svf)} ===")
+if not svf.empty:
+ print(
+ svf[["time", "value", "units", "table_type", "period"]]
+ .head(12)
+ .to_string(index=False)
+ )
+else:
+ print(" *** NOT FOUND ***")
+print()
+
+# --- Heat ---
+heat = ocn[ocn["quantity"] == "heat"]
+print("=== Heat term counts ===")
+print(heat["term"].value_counts().to_string())
+print()
+
+ec = heat[heat["term"] == "Energy change"]
+print(f"=== 'Energy change' rows: {len(ec)} ===")
+if not ec.empty:
+ print(
+ ec[["time", "value", "units", "table_type", "period"]]
+ .head(12)
+ .to_string(index=False)
+ )
+else:
+ print(" *** NOT FOUND ***")
+print()
+
+shf = heat[heat["term"] == "SUM IMP+EXP HEAT FLUXES"]
+print(f"=== 'SUM IMP+EXP HEAT FLUXES' rows: {len(shf)} ===")
+if not shf.empty:
+ print(
+ shf[["time", "value", "units", "table_type", "period"]]
+ .head(12)
+ .to_string(index=False)
+ )
+else:
+ print(" *** NOT FOUND ***")
+print()
+
+# --- Test closure checks ---
+print("=== Testing OcnClosure checks ===")
+df = normalize(ocn)
+
+for q in ["water", "heat"]:
+ check = OcnClosure(quantity=q)
+ result = check.evaluate(df)
+ if result is not None:
+ print(
+ f" {q} closure: {len(result.years)} years,"
+ f" max |residual| = {abs(result.residual).max():.2e}"
+ )
+ else:
+ print(f" {q} closure: SKIPPED (missing data)")
diff --git a/tests/unit/budget_analysis/print_monthly_interface.py b/tests/unit/budget_analysis/print_monthly_interface.py
new file mode 100644
index 0000000..d5c9804
--- /dev/null
+++ b/tests/unit/budget_analysis/print_monthly_interface.py
@@ -0,0 +1,43 @@
+"""Print monthly normalized values for cpl and ocn interface comparison."""
+
+import glob
+import sys
+
+import pandas as pd
+
+sys.path.insert(0, ".")
+
+from zppy_interfaces.budget_analysis.ingestion.cpl_parser import CplParser # noqa: E402
+from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser # noqa: E402
+from zppy_interfaces.budget_analysis.normalization import normalize # noqa: E402
+
+LOG_PATH = "/pscratch/sd/c/chengzhu/zstash/archive/logs"
+START_YEAR = 1
+END_YEAR = 50
+
+cpl = CplParser().parse_files(
+ sorted(glob.glob(f"{LOG_PATH}/cpl.log.*.gz")), START_YEAR, END_YEAR
+)
+ocn = OcnParser().parse_files(
+ sorted(glob.glob(f"{LOG_PATH}/ocn.log.*.gz")), START_YEAR, END_YEAR
+)
+df = normalize(pd.concat([cpl, ocn], ignore_index=True))
+
+cpl_m = df[
+ (df.source == "cpl")
+ & (df.term == "*SUM*")
+ & (df.component == "ocn")
+ & (df.period == "monthly")
+].sort_values("time")
+
+ocn_m = df[
+ (df.source == "ocn")
+ & (df.term == "SUM VOLUME FLUXES")
+ & (df.table_type == "flux")
+ & (df.period == "monthly")
+].sort_values("time")
+
+print("=== CPL monthly ocn *SUM* ===")
+print(cpl_m[["time", "normalized_value"]].to_string(index=False))
+print(f"\n=== OCN monthly SUM VOLUME FLUXES ({len(ocn_m)} rows) ===")
+print(ocn_m[["time", "normalized_value"]].to_string(index=False))
diff --git a/tests/unit/budget_analysis/test_atm_parser.py b/tests/unit/budget_analysis/test_atm_parser.py
new file mode 100644
index 0000000..c8192a9
--- /dev/null
+++ b/tests/unit/budget_analysis/test_atm_parser.py
@@ -0,0 +1,84 @@
+"""Quick test script for AtmParser on a sample atm.log file."""
+
+import sys
+
+from zppy_interfaces.budget_analysis.ingestion.atm_parser import AtmParser
+
+LOG_FILE = (
+ "/pscratch/sd/e/e3smtest/e3sm_scratch/pm-cpu/"
+ "SMS.ne4pg2_oQU480.F2010.pm-cpu_intel.eam-thetahy_ftype2_energy"
+ ".C.JNextIntegration20260210_205258/run/"
+ "atm.log.48753126.260210-224733.gz"
+)
+
+
+def main():
+ parser = AtmParser()
+ log_files = [LOG_FILE]
+
+ # --- Raw per-step data ---
+ nstep_te, flux_diag = parser.parse_raw(log_files)
+
+ print("=== nstep_te (energy fixer) ===")
+ print(f"Shape: {nstep_te.shape}")
+ print(f"Columns: {list(nstep_te.columns)}")
+ print(nstep_te.head(3))
+ print("...")
+ print(nstep_te.tail(3))
+ print()
+
+ print("=== flux_diag (water/energy diagnostics) ===")
+ print(f"Shape: {flux_diag.shape}")
+ print(f"Columns: {list(flux_diag.columns)}")
+ print(flux_diag.head(3))
+ print("...")
+ print(flux_diag.tail(3))
+ print()
+
+ # --- Validation ---
+ print("=== Validation ===")
+ assert (
+ nstep_te.shape[0] == 122
+ ), f"Expected 122 nstep_te rows, got {nstep_te.shape[0]}"
+ print(f" nstep_te rows: {nstep_te.shape[0]} (expected 122)")
+
+ assert (
+ flux_diag.shape[0] == 120
+ ), f"Expected 120 flux_diag rows, got {flux_diag.shape[0]}"
+ print(f" flux_diag rows: {flux_diag.shape[0]} (expected 120)")
+
+ tw_first = flux_diag["tw"].iloc[0]
+ tw_last = flux_diag["tw"].iloc[-1]
+ print(f" W(n=1) = {tw_first:.6f} kg/m2 (expect ~25.303)")
+ print(f" W(n=120)= {tw_last:.6f} kg/m2 (expect ~24.949)")
+
+ e_diff_max = flux_diag["e_diff"].abs().max()
+ print(f" max |E difference| = {e_diff_max:.3e} W/m2")
+
+ # Check date tracking
+ print(
+ f" Date range: year {nstep_te['year'].min()}-{nstep_te['year'].max()}, "
+ f"month {nstep_te['month'].min()}-{nstep_te['month'].max()}, "
+ f"day {nstep_te['day'].min()}-{nstep_te['day'].max()}"
+ )
+ print()
+
+ # --- Tidy event table (annual) ---
+ events = parser.parse_files(log_files, 1, 1)
+ print("=== Tidy event table (annual, default) ===")
+ print(f"Shape: {events.shape}")
+ print(events.to_string())
+ print()
+
+ # --- Tidy event table (monthly) ---
+ monthly_parser = AtmParser(frequency="monthly")
+ events_m = monthly_parser.parse_files(log_files, 1, 1)
+ print("=== Tidy event table (monthly) ===")
+ print(f"Shape: {events_m.shape}")
+ print(events_m.to_string())
+
+ print("\nAll checks passed.")
+
+
+if __name__ == "__main__":
+ sys.exit(main() or 0)
diff --git a/zppy_interfaces/budget_analysis/README.md b/zppy_interfaces/budget_analysis/README.md
new file mode 100644
index 0000000..458c3d1
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/README.md
@@ -0,0 +1,197 @@
+# E3SM Budget Analysis Module
+
+The `budget_analysis` module provides comprehensive water and energy budget analysis for E3SM climate model simulations. It analyzes coupler log files to extract budget data, performs conservation checks, and generates interactive visualizations.
+
+## Overview
+
+This module processes E3SM component log files (coupler, atmosphere, ocean, ice, land) to:
+- Extract water, heat, and carbon budget terms
+- Perform conservation checks across components
+- Calculate residuals and identify budget imbalances
+- Generate interactive HTML reports with time series plots
+- Support both legacy (coupler-only) and whole-model analysis modes
+
+## Modes of Operation
+
+### Legacy Mode (Default)
+- **Purpose**: Original coupler-only cumulative budget analysis
+- **Input**: Coupler log files (`cpl.log.*.gz`)
+- **Output**: Interactive HTML plots and ASCII summary tables
+- **Use case**: Quick budget overview using only coupler data
+
+### Whole-Model Mode
+- **Purpose**: Comprehensive multi-component budget analysis with conservation checks
+- **Input**: Log files from all available components (cpl, atm, ocn, ice, lnd)
+- **Output**: Detailed HTML reports with conservation diagnostics
+- **Use case**: Complete budget closure analysis across all model components
+
+## CLI Usage
+
+### Basic Legacy Analysis
+```bash
+zi-budget-analysis \
+ --log_path /path/to/case/archive/logs \
+ --start_year 114 \
+ --end_year 206
+```
+
+### Whole-Model Analysis
+```bash
+zi-budget-analysis \
+ --log_path /path/to/case/archive/logs \
+ --start_year 114 \
+ --end_year 150 \
+ --mode whole-model \
+ --budget_types water,heat \
+ --frequency monthly \
+ --output_dir ./results
+```
+
+### Command Line Options
+
+- `--log_path` (required): Directory containing log files
+- `--start_year` (required): Starting year for analysis
+- `--end_year` (required): Ending year for analysis
+- `--budget_types`: Comma-separated list (water,heat,carbon) [default: water,heat]
+- `--mode`: Analysis mode (legacy,whole-model) [default: legacy]
+- `--frequency`: Temporal frequency (monthly,annual) [default: annual]
+- `--output_dir`: Output directory [default: current directory]
+- `--output_html`: Generate HTML plots [default: True]
+
+## Architecture
+
+### Core Components
+
+**Data Ingestion** (`ingestion/`)
+- `base.py` - Abstract base parser class
+- `cpl_parser.py` - Coupler log parser
+- `atm_parser.py` - Atmosphere log parser
+- `ocn_parser.py` - Ocean log parser
+- `ice_parser.py` - Sea ice log parser
+- `lnd_parser.py` - Land log parser
+
+**Data Processing**
+- `parser.py` - Legacy mode parsing logic
+- `schema.py` - Standardized data table schema
+- `normalization.py` - Unit conversion and data standardization
+
+**Analysis**
+- `checks.py` - Budget conservation checks and residual calculations
+- `plotting.py` - Legacy mode HTML plot generation
+- `viz.py` - Whole-model mode report generation
+
+### Data Schema
+
+The module uses a standardized "tidy" DataFrame schema with columns:
+- `time`: Year (annual) or year + fractional month (monthly)
+- `component`: Model component (atm, lnd, ocn, ice_nh, ice_sh, etc.)
+- `quantity`: Budget type (water, heat, carbon)
+- `term`: Specific flux or state variable name
+- `value`: Numerical value in original units
+- `units`: Original unit string
+- `source`: Source log file type (cpl, atm, etc.)
+- `period`: Temporal frequency (annual, monthly)
+- `table_type`: Data type (flux, flux_integrated, state)
+
+## Budget Conservation Checks
+
+### Water Budget Checks
+- **Component Fluxes**: Verify coupler component flux balances
+- **Interface Matching**: Compare fluxes at component boundaries
+- **Component Closure**: Check internal conservation within each component
+- **Residual Analysis**: Calculate and track budget imbalances over time
+
+### Heat Budget Checks
+- **Component Fluxes**: Energy flux balance verification
+- **Interface Matching**: Energy flux consistency across boundaries
+- **Component Closure**: Internal energy conservation (limited for atmosphere and land)
+- **Note**: Atmospheric and land energy checks are limited due to incomplete flux data in ATM and LND logs
+
+### Carbon Budget Checks
+- **Component Fluxes**: Carbon flux balance verification (basic implementation)
+
+## Output Formats
+
+### Legacy Mode Output
+- Interactive HTML plots with time series of cumulative budgets
+- ASCII summary tables written to stdout
+- Plots saved as both PNG and HTML files
+
+### Whole-Model Mode Output
+- Comprehensive HTML reports for each budget type
+- Interactive time series plots with residual analysis
+- Component-by-component breakdown
+- Landing page linking all generated reports
+- Conservation check results with pass/fail indicators
+
+## File Structure
+
+```
+budget_analysis/
+├── __init__.py # Module initialization
+├── __main__.py # CLI entry point
+├── parser.py # Legacy parsing logic
+├── plotting.py # Legacy HTML generation
+├── schema.py # Data schema definitions
+├── normalization.py # Data normalization utilities
+├── checks.py # Conservation check definitions
+├── viz.py # Whole-model visualization
+└── ingestion/ # Component-specific parsers
+ ├── __init__.py
+ ├── base.py # Abstract parser base class
+ ├── cpl_parser.py # Coupler log parser
+ ├── atm_parser.py # Atmosphere parser
+ ├── ocn_parser.py # Ocean parser
+ ├── ice_parser.py # Sea ice parser
+ └── lnd_parser.py # Land parser
+```
+
+## Example Workflows
+
+### Quick Budget Check (Legacy)
+```python
+from zppy_interfaces.budget_analysis import main
+import sys
+
+# Simulate CLI arguments
+sys.argv = [
+ 'zi-budget-analysis',
+ '--log_path', '/path/to/logs',
+ '--start_year', '100',
+ '--end_year', '150',
+ '--budget_types', 'water'
+]
+
+main()
+```
+
+### Comprehensive Analysis (Whole-Model)
+```python
+from zppy_interfaces.budget_analysis import main
+import sys
+
+sys.argv = [
+ 'zi-budget-analysis',
+ '--log_path', '/path/to/logs',
+ '--start_year', '100',
+ '--end_year', '150',
+ '--mode', 'whole-model',
+ '--budget_types', 'water,heat',
+ '--frequency', 'monthly',
+ '--output_dir', './budget_results'
+]
+
+main()
+```
+
+## Integration with zppy
+
+This module is designed to be called by zppy for automated budget analysis in E3SM post-processing workflows. The CLI interface allows seamless integration with zppy configuration files.
+
+## Limitations
+
+- Atmospheric energy budget analysis is limited due to incomplete energy flux information in ATM log files
+- Land component energy budget analysis is not available due to missing energy flux data in LND log files
+- Carbon budget checks are basic and may need expansion for comprehensive carbon cycle analysis
+- Monthly frequency analysis requires sufficient log file temporal resolution
+- Large log files may require significant memory for processing
diff --git a/zppy_interfaces/budget_analysis/__init__.py b/zppy_interfaces/budget_analysis/__init__.py
new file mode 100644
index 0000000..c5f0aac
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/__init__.py
@@ -0,0 +1 @@
+# Budget analysis package for E3SM water and energy budget visualization
diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py
new file mode 100644
index 0000000..d1a67f9
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/__main__.py
@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+
+"""
+E3SM Water and Energy Budget Analysis CLI Tool
+
+Analyzes E3SM coupler log files to extract water and energy budget data
+and generates interactive HTML plots and ASCII summary tables.
+
+This tool is designed to be called by zppy or used standalone for
+budget conservation analysis of E3SM simulations.
+"""
+
+import argparse
+import glob
+import os
+import sys
+
+import numpy as np
+
+from .parser import initialize_budgets, parse_budget_types, process_log_files
+from .plotting import generate_html_plots
+
+
+def main() -> int:
+ """Main CLI entry point."""
+ parser = argparse.ArgumentParser(
+ description="Analyze E3SM water and energy budgets from coupler log files",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ %(prog)s --log_path /path/to/case/archive/logs --start_year 114 --end_year 206
+ %(prog)s --log_path /path/to/case/archive/logs --start_year 114 --end_year 150 --budget_types water,heat
+ """,
+ )
+
+ parser.add_argument(
+ "--log_path",
+ required=True,
+ help="Path to directory containing coupler log files (cpl.log.*.gz)",
+ )
+ parser.add_argument(
+ "--start_year", type=int, required=True, help="Starting year for analysis"
+ )
+ parser.add_argument(
+ "--end_year", type=int, required=True, help="Ending year for analysis"
+ )
+ parser.add_argument(
+ "--budget_types",
+ default="water,heat",
+ help="Comma-separated list of budget types to analyze (water,heat,carbon)",
+ )
+ parser.add_argument(
+ "--output_dir",
+ default=".",
+ help="Output directory for results (default: current directory)",
+ )
+ parser.add_argument(
+ "--output_html",
+ action="store_true",
+ default=True,
+ help="Generate HTML plots (default: True)",
+ )
+ parser.add_argument(
+ "--mode",
+ choices=["legacy", "whole-model"],
+ default="legacy",
+ help="'legacy' (coupler-only cumulative) or 'whole-model' (multi-source budget checks)",
+ )
+ parser.add_argument(
+ "--frequency",
+ choices=["monthly", "annual"],
+ default="annual",
+ help="Temporal frequency for budget data: 'monthly' or 'annual' (default: annual)",
+ )
+
+ args = parser.parse_args()
+
+ # Validate inputs
+ if args.start_year > args.end_year:
+ print("ERROR: start_year must be <= end_year")
+ return 1
+
+ log_path = args.log_path
+ if not os.path.exists(log_path):
+ print(f"ERROR: Log path does not exist: {log_path}")
+ return 1
+
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ if args.mode == "whole-model":
+ return _run_whole_model(args)
+
+ return _run_legacy(args)
+
+
+def _run_legacy(args) -> int:
+ """Original coupler-only cumulative budget pipeline."""
+ budget_types = parse_budget_types(args.budget_types)
+ valid_types = ["area", "water", "heat", "carbon"]
+ for bt in budget_types:
+ if bt not in valid_types:
+ print(f"ERROR: Invalid budget type '{bt}'. Valid types: {valid_types}")
+ return 1
+
+ print("E3SM Budget Analysis Tool (legacy mode)")
+ print("========================================")
+ print(f"Years: {args.start_year} to {args.end_year}")
+ print(f"Budget types: {budget_types}")
+ print(f"Log path: {args.log_path}")
+
+ years = np.arange(args.start_year, args.end_year + 1)
+ budgets = initialize_budgets(budget_types, years)
+
+ log_files = sorted(glob.glob(os.path.join(args.log_path, "cpl.log.*.gz")))
+ if not log_files:
+ print("ERROR: No coupler log files found")
+ return 1
+
+ print(f"Found {len(log_files)} coupler log files")
+ process_log_files(log_files, budgets)
+
+ if args.output_html:
+ generate_html_plots(budgets, budget_types, args.output_dir)
+
+ print("\nBudget analysis completed successfully!")
+ return 0
+
+
+def _run_whole_model(args) -> int:
+ """Whole-model budget pipeline: ingest -> normalize -> check -> visualize."""
+ import pandas as pd
+
+ from .checks import (
+ DEFAULT_CARBON_CHECKS,
+ DEFAULT_HEAT_CHECKS,
+ DEFAULT_WATER_CHECKS,
+ run_checks,
+ )
+ from .ingestion.atm_parser import AtmParser
+ from .ingestion.cpl_parser import CplParser
+ from .ingestion.ice_parser import IceParser
+ from .ingestion.lnd_parser import LndParser
+ from .ingestion.ocn_parser import OcnParser
+ from .normalization import normalize
+ from .viz import generate_budget_report, generate_landing_page
+
+ budget_types = parse_budget_types(args.budget_types)
+
+ print("E3SM Budget Analysis Tool (whole-model mode)")
+ print("=============================================")
+ print(f"Years: {args.start_year} to {args.end_year}")
+ print(f"Budget types: {budget_types}")
+ print(f"Frequency: {args.frequency}")
+ print(f"Log path: {args.log_path}")
+
+ # Ingest
+ print("\nIngesting log files...")
+ cpl_files = sorted(glob.glob(os.path.join(args.log_path, "cpl.log.*.gz")))
+ lnd_files = sorted(glob.glob(os.path.join(args.log_path, "lnd.log.*.gz")))
+ ocn_files = sorted(glob.glob(os.path.join(args.log_path, "ocn.log.*.gz")))
+ ice_files = sorted(glob.glob(os.path.join(args.log_path, "ice.log.*.gz")))
+ atm_files = sorted(glob.glob(os.path.join(args.log_path, "atm.log.*")))
+
+ if not cpl_files:
+ print("ERROR: No coupler log files found")
+ return 1
+ print(f" {len(cpl_files)} coupler log files")
+ print(f" {len(lnd_files)} land log files")
+ print(f" {len(ocn_files)} ocean log files")
+ print(f" {len(ice_files)} ice log files")
+ print(f" {len(atm_files)} atmosphere log files")
+
+ freq = args.frequency
+ frames = []
+ frames.append(
+ CplParser(quantities=budget_types, frequency=freq).parse_files(
+ cpl_files, args.start_year, args.end_year
+ )
+ )
+ if lnd_files:
+ frames.append(
+ LndParser(frequency=freq).parse_files(
+ lnd_files, args.start_year, args.end_year
+ )
+ )
+ if ocn_files:
+ frames.append(
+ OcnParser(frequency=freq).parse_files(
+ ocn_files, args.start_year, args.end_year
+ )
+ )
+ if ice_files:
+ frames.append(
+ IceParser(frequency=freq).parse_files(
+ ice_files, args.start_year, args.end_year
+ )
+ )
+ if atm_files:
+ frames.append(
+ AtmParser(frequency=freq).parse_files(
+ atm_files, args.start_year, args.end_year
+ )
+ )
+ events = pd.concat(frames, ignore_index=True)
+ print(f" {len(events)} total event rows")
+
+ # Normalize
+ print("\nNormalizing...")
+ events = normalize(events)
+
+ # Check and visualize per quantity
+ checks_map = {
+ "water": DEFAULT_WATER_CHECKS,
+ "heat": DEFAULT_HEAT_CHECKS,
+ "carbon": DEFAULT_CARBON_CHECKS,
+ }
+
+ report_paths = {}
+ for quantity in budget_types:
+ checks = checks_map.get(quantity)
+ if checks is None:
+ print(f"\n WARNING: No checks defined for '{quantity}', skipping")
+ continue
+
+ print(f"\nRunning {quantity} budget checks...")
+ results = run_checks(events, checks)
+
+ print(f"\nGenerating {quantity} report...")
+ html_path = generate_budget_report(
+ results, events, args.output_dir, quantity=quantity
+ )
+ if html_path:
+ report_paths[quantity] = html_path
+
+ # Landing page
+ if len(report_paths) > 0:
+ index_path = generate_landing_page(args.output_dir, report_paths)
+ print(f"\nBudget analysis completed! Landing page: {index_path}")
+ else:
+ print("\nNo reports generated.")
+
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py
new file mode 100644
index 0000000..f4fc2a4
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/checks.py
@@ -0,0 +1,788 @@
+"""Budget checks: definitions and evaluation.
+
+Note on atmospheric energy checks:
+- ATM logs provide limited energy flux information (E d(TE)/dt, E RR)
+- ATM logs lack complete energy flux breakdown needed for full closure
+- AtmClosure(quantity="heat") and AtmInterfaceMatch(quantity="heat") are
+ therefore disabled in DEFAULT_HEAT_CHECKS
+- Water budget checks work normally for ATM
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+
+from .schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+)
+
+# Days per month for 365-day no-leap calendar
+DAYS_PER_MONTH = {
+ 1: 31,
+ 2: 28,
+ 3: 31,
+ 4: 30,
+ 5: 31,
+ 6: 30,
+ 7: 31,
+ 8: 31,
+ 9: 30,
+ 10: 31,
+ 11: 30,
+ 12: 31,
+}
+
+# Seconds per year (365-day calendar)
+SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0
+
+
+@dataclass
+class CheckResult:
+ """Result of a single budget check across all time steps."""
+
+ name: str
+ description: str
+ years: np.ndarray
+ lhs: np.ndarray
+ rhs: np.ndarray
+ residual: np.ndarray # lhs - rhs
+ cumulative_residual: np.ndarray
+ lhs_label: str = "LHS"
+ rhs_label: str = "RHS"
+ components: Optional[Dict[str, np.ndarray]] = field(default=None)
+
+
+def _select(
+ df: pd.DataFrame,
+ period: Optional[str] = None,
+ **filters,
+) -> pd.DataFrame:
+ """Filter df by column=value filters, optionally by period.
+
+ If period is None, prefer 'annual' if available, else use 'monthly'.
+ Monthly data is kept at monthly resolution (COL_TIME encodes
+ year + fractional month).
+ """
+ mask = pd.Series(True, index=df.index)
+ for col, val in filters.items():
+ mask = mask & (df[col] == val)
+ subset = df[mask]
+
+ if subset.empty:
+ return subset
+
+ # Determine period to use
+ available = subset[COL_PERIOD].unique()
+ if period is not None:
+ subset = subset[subset[COL_PERIOD] == period]
+ elif "annual" in available:
+ subset = subset[subset[COL_PERIOD] == "annual"]
+ else:
+ # Monthly — keep at monthly resolution.
+ # COL_TIME already encodes year + fractional month,
+ # so each month has a unique time value.
+ subset = subset[subset[COL_PERIOD] == "monthly"]
+
+ return subset.sort_values(COL_TIME)
+
+
+class BudgetCheck:
+ """Base class for budget checks."""
+
+ def __init__(self, name: str, description: str):
+ self.name = name
+ self.description = description
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ raise NotImplementedError
+
+
+class CplComponentFluxes(BudgetCheck):
+ """Per-component net flux + global residual (*SUM*).
+
+ Components dict includes each component's cumulative net flux
+ plus a '*SUM*' entry for the global residual.
+ Supports both water and heat quantities.
+ """
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"cpl_{quantity}_component_fluxes",
+ f"Coupler cumulative net {quantity} flux per component + residual",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ rows = _select(
+ df,
+ **{COL_SOURCE: "cpl", COL_TERM: "*SUM*", COL_QUANTITY: self.quantity},
+ )
+ if rows.empty:
+ return None
+ pivot = rows.pivot_table(
+ index=COL_TIME, columns=COL_COMPONENT, values="normalized_value"
+ ).sort_index()
+ years = pivot.index.values
+
+ components = {}
+ for col in pivot.columns:
+ components[col] = np.cumsum(pivot[col].values)
+
+ residual = (
+ pivot["*SUM*"].values
+ if "*SUM*" in pivot.columns
+ else pivot.sum(axis=1).values
+ )
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ np.zeros_like(residual),
+ residual,
+ -residual,
+ np.cumsum(-residual),
+ components=components,
+ )
+
+
+class InterfaceMatch(BudgetCheck):
+ """Do the coupler and component model agree on net flux?
+
+ Compares coupler *SUM* in the component column vs component's *SUM* flux.
+ Works for any component (lnd, ocn, etc.) and quantity (water, heat).
+ """
+
+ def __init__(
+ self,
+ component: str,
+ source: str,
+ quantity: str = "water",
+ comp_sum_term: str = "*SUM*",
+ ) -> None:
+ super().__init__(
+ f"{component}_{quantity}_interface_match",
+ f"{component} net {quantity} flux: coupler vs {source} model",
+ )
+ self.component = component
+ self.source = source
+ self.quantity = quantity
+ self.comp_sum_term = comp_sum_term
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ cpl = _select(
+ df,
+ **{
+ COL_SOURCE: "cpl",
+ COL_TERM: "*SUM*",
+ COL_COMPONENT: self.component,
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ comp = _select(
+ df,
+ **{
+ COL_SOURCE: self.source,
+ COL_TERM: self.comp_sum_term,
+ COL_TABLE_TYPE: "flux",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if cpl.empty or comp.empty:
+ return None
+ merged = cpl.join(comp, lsuffix="_cpl", rsuffix="_comp", how="inner")
+ if merged.empty:
+ return None
+ years = merged.index.values
+ c = merged["normalized_value_cpl"].values
+ m = merged["normalized_value_comp"].values
+ r = c - m
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ c,
+ m,
+ r,
+ np.cumsum(r),
+ lhs_label=f"cpl ({self.component})",
+ rhs_label=f"{self.source} ({self.comp_sum_term})",
+ )
+
+
+class LndClosure(BudgetCheck):
+ """Does land storage change equal the integrated flux?
+
+ Compares *NET CHANGE* TOTAL (state table) vs *SUM* (integrated flux table).
+ """
+
+ def __init__(self) -> None:
+ super().__init__(
+ "lnd_closure",
+ "Land water closure: ΔStorage vs ∫Flux dt",
+ )
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ storage = _select(
+ df,
+ **{
+ COL_SOURCE: "lnd",
+ COL_TERM: "*NET CHANGE*_TOTAL",
+ COL_TABLE_TYPE: "state",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ flux = _select(
+ df,
+ **{COL_SOURCE: "lnd", COL_TERM: "*SUM*", COL_TABLE_TYPE: "flux_integrated"},
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if storage.empty or flux.empty:
+ return None
+ merged = storage.join(flux, lsuffix="_stor", rsuffix="_flux", how="inner")
+ if merged.empty:
+ return None
+ years = merged.index.values
+ ds = merged["normalized_value_stor"].values
+ fi = merged["normalized_value_flux"].values
+ r = ds - fi
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ ds,
+ fi,
+ r,
+ np.cumsum(r),
+ lhs_label="ΔStorage (*NET CHANGE*)",
+ rhs_label="∫Flux dt (*SUM*)",
+ )
+
+
+class OcnClosure(BudgetCheck):
+ """Does ocean mass/energy change equal the net flux?
+
+ Ocean logs are monthly. _select auto-aggregates to annual.
+ For water: compares mass_change vs *SUM* flux.
+ For heat: compares energy_change vs *SUM* flux.
+ """
+
+ CHANGE_TERM: Dict[str, str] = {
+ "water": "Mass change",
+ "heat": "Energy change",
+ }
+
+ SUM_TERM: Dict[str, str] = {
+ "water": "SUM VOLUME FLUXES",
+ "heat": "SUM IMP+EXP HEAT FLUXES",
+ }
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"ocn_{quantity}_closure",
+ f"Ocean {quantity} closure: Δ{'Mass' if quantity == 'water' else 'Energy'} vs net flux",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ change_term = self.CHANGE_TERM[self.quantity]
+ sum_term = self.SUM_TERM[self.quantity]
+
+ mass = _select(
+ df,
+ **{
+ COL_SOURCE: "ocn",
+ COL_TERM: change_term,
+ COL_TABLE_TYPE: "flux",
+ COL_QUANTITY: self.quantity,
+ },
+ )
+ if mass.empty:
+ return None
+ mass_ts = mass[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ flux = _select(
+ df,
+ **{
+ COL_SOURCE: "ocn",
+ COL_TERM: sum_term,
+ COL_TABLE_TYPE: "flux",
+ COL_QUANTITY: self.quantity,
+ },
+ )
+ if flux.empty:
+ return None
+ flux_ts = flux[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ merged = mass_ts.join(flux_ts, lsuffix="_mass", rsuffix="_flux", how="inner")
+ if merged.empty:
+ return None
+ years = merged.index.values
+ ds = merged["normalized_value_mass"].values
+ fi = merged["normalized_value_flux"].values
+ r = ds - fi
+ change_label = "ΔMass" if self.quantity == "water" else "ΔEnergy"
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ ds,
+ fi,
+ r,
+ np.cumsum(r),
+ lhs_label=f"{change_label} ({change_term})",
+ rhs_label=f"Net Flux ({sum_term})",
+ )
+
+
+class IceClosure(BudgetCheck):
+ """Ice mass/energy closure: storage change flux vs net flux.
+
+ For water: compares Mass change flux vs Net mass flux (both kg/m2s -> mm/yr).
+ For heat: compares Energy change flux vs Net energy flux (both W/m2).
+ """
+
+ CHANGE_TERM: Dict[str, str] = {
+ "water": "Mass change flux",
+ "heat": "Energy change flux",
+ }
+
+ FLUX_TERM: Dict[str, str] = {
+ "water": "Net mass flux",
+ "heat": "Net energy flux",
+ }
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"ice_{quantity}_closure",
+ f"Ice {quantity} closure: {'Mass' if quantity == 'water' else 'Energy'} change flux vs net flux",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ change_term = self.CHANGE_TERM[self.quantity]
+ flux_term = self.FLUX_TERM[self.quantity]
+
+ # Get ice storage change flux
+ storage_flux = _select(
+ df,
+ **{
+ COL_SOURCE: "ice",
+ COL_TERM: change_term,
+ COL_QUANTITY: self.quantity,
+ COL_TABLE_TYPE: "flux",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ # Get ice net flux
+ net_flux = _select(
+ df,
+ **{
+ COL_SOURCE: "ice",
+ COL_TERM: flux_term,
+ COL_QUANTITY: self.quantity,
+ COL_TABLE_TYPE: "flux",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if storage_flux.empty or net_flux.empty:
+ return None
+
+ merged = storage_flux.join(
+ net_flux, lsuffix="_change", rsuffix="_net", how="inner"
+ )
+ if merged.empty:
+ return None
+
+ years = merged.index.values
+ change = merged["normalized_value_change"].values
+ net = merged["normalized_value_net"].values
+ residual = change - net
+
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ change,
+ net,
+ residual,
+ np.cumsum(residual),
+ lhs_label=f"Storage Change ({change_term})",
+ rhs_label=f"Net Flux ({flux_term})",
+ )
+
+
+class IceInterfaceMatch(BudgetCheck):
+ """Do the coupler and ice model agree on net flux?
+
+ Special case for ice: coupler splits ice into 'ice_nh' and 'ice_sh' components,
+ so we sum both hemispheres to compare with ice model's total.
+ """
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"ice_{quantity}_interface_match",
+ f"Ice net {quantity} flux: coupler (nh+sh) vs ice model",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ # Get coupler data for both ice hemispheres
+ cpl_nh = _select(
+ df,
+ **{
+ COL_SOURCE: "cpl",
+ COL_TERM: "*SUM*",
+ COL_COMPONENT: "ice_nh",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ cpl_sh = _select(
+ df,
+ **{
+ COL_SOURCE: "cpl",
+ COL_TERM: "*SUM*",
+ COL_COMPONENT: "ice_sh",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if cpl_nh.empty or cpl_sh.empty:
+ print(f"DEBUG: Missing coupler ice hemisphere data for {self.quantity}")
+ if cpl_nh.empty:
+ print(" Missing ice_nh data")
+ if cpl_sh.empty:
+ print(" Missing ice_sh data")
+ return None
+
+ # Sum both hemispheres
+ cpl_combined = cpl_nh.join(cpl_sh, lsuffix="_nh", rsuffix="_sh", how="inner")
+ if cpl_combined.empty:
+ print("DEBUG: No overlapping time periods between ice_nh and ice_sh")
+ return None
+
+ cpl_combined["normalized_value"] = (
+ cpl_combined["normalized_value_nh"] + cpl_combined["normalized_value_sh"]
+ )
+ cpl_total = cpl_combined[["normalized_value"]]
+
+ # Get ice model data - need to determine the right term
+ # For water, try common mass flux terms, for heat try energy flux terms
+ ice_terms = {
+ "water": ["Net mass flux", "*SUM*", "Mass change flux"],
+ "heat": ["Net energy flux", "*SUM*", "Energy change flux"],
+ }
+
+ ice_model = None
+ used_term = None
+ for term in ice_terms[self.quantity]:
+ ice_candidate = _select(
+ df,
+ **{
+ COL_SOURCE: "ice",
+ COL_TERM: term,
+ COL_TABLE_TYPE: "flux",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if not ice_candidate.empty:
+ ice_model = ice_candidate
+ used_term = term
+ break
+
+ if ice_model is None:
+ return None
+
+ merged = cpl_total.join(ice_model, lsuffix="_cpl", rsuffix="_ice", how="inner")
+ if merged.empty:
+ return None
+
+ years = merged.index.values
+ c = merged["normalized_value_cpl"].values
+ m = merged["normalized_value_ice"].values
+ r = c - m
+
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ c,
+ m,
+ r,
+ np.cumsum(r),
+ lhs_label="cpl (ice_nh + ice_sh)",
+ rhs_label=f"ice ({used_term})",
+ )
+
+
+class AtmInterfaceMatch(BudgetCheck):
+ """Do the coupler and atmosphere model agree on net flux?"""
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"atm_{quantity}_interface_match",
+ f"Atmosphere {quantity} interface match: coupler vs atm model",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ # Get coupler data for atmosphere component
+ cpl_atm = _select(
+ df,
+ **{
+ COL_SOURCE: "cpl",
+ COL_TERM: "*SUM*",
+ COL_COMPONENT: "atm",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if cpl_atm.empty:
+ return None
+
+ # Get atmosphere model data - determine the right term
+ atm_terms = {
+ "water": ["W flux", "dWater", "*SUM*"],
+ "heat": ["E d(TE)/dt", "E RR", "*SUM*"],
+ }
+
+ atm_model = None
+ used_term = None
+ for term in atm_terms[self.quantity]:
+ atm_candidate = _select(
+ df,
+ **{
+ COL_SOURCE: "atm",
+ COL_TERM: term,
+ COL_TABLE_TYPE: "flux",
+ COL_QUANTITY: self.quantity,
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if not atm_candidate.empty:
+ atm_model = atm_candidate
+ used_term = term
+ break
+
+ if atm_model is None:
+ return None
+
+ merged = cpl_atm.join(atm_model, lsuffix="_cpl", rsuffix="_atm", how="inner")
+ if merged.empty:
+ return None
+
+ years = merged.index.values
+ c = merged["normalized_value_cpl"].values
+ a = merged["normalized_value_atm"].values
+ r = c - a
+
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ c,
+ a,
+ r,
+ np.cumsum(r),
+ lhs_label="Coupler Flux",
+ rhs_label=f"Atm Flux ({used_term})",
+ )
+
+
+class AtmClosure(BudgetCheck):
+ """Atmosphere mass/energy closure: storage change vs net flux.
+
+ For water: compares storage change vs water flux.
+ For heat: compares storage change vs energy flux.
+ """
+
+ STORAGE_BEG_TERM: Dict[str, str] = {
+ "water": "W tot mass beg",
+ "heat": "TE beg",
+ }
+
+ STORAGE_END_TERM: Dict[str, str] = {
+ "water": "W tot mass end",
+ "heat": "TE end",
+ }
+
+ FLUX_TERM: Dict[str, str] = {
+ "water": "W flux",
+ "heat": "E d(TE)/dt",
+ }
+
+ def __init__(self, quantity: str = "water") -> None:
+ super().__init__(
+ f"atm_{quantity}_closure",
+ f"Atmosphere {quantity} closure: storage change vs net flux",
+ )
+ self.quantity = quantity
+
+ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]:
+ beg_term = self.STORAGE_BEG_TERM[self.quantity]
+ end_term = self.STORAGE_END_TERM[self.quantity]
+ flux_term = self.FLUX_TERM[self.quantity]
+
+ # Get storage begin/end
+ storage_beg = _select(
+ df,
+ **{
+ COL_SOURCE: "atm",
+ COL_TERM: beg_term,
+ COL_QUANTITY: self.quantity,
+ COL_TABLE_TYPE: "state",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ storage_end = _select(
+ df,
+ **{
+ COL_SOURCE: "atm",
+ COL_TERM: end_term,
+ COL_QUANTITY: self.quantity,
+ COL_TABLE_TYPE: "state",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ # Get net flux
+ net_flux = _select(
+ df,
+ **{
+ COL_SOURCE: "atm",
+ COL_TERM: flux_term,
+ COL_QUANTITY: self.quantity,
+ COL_TABLE_TYPE: "flux",
+ },
+ )[[COL_TIME, "normalized_value"]].set_index(COL_TIME)
+
+ if storage_beg.empty or storage_end.empty or net_flux.empty:
+ return None
+
+ # Calculate storage change rate
+ storage_merged = storage_beg.join(
+ storage_end, lsuffix="_beg", rsuffix="_end", how="inner"
+ )
+ if storage_merged.empty:
+ return None
+
+ # Determine time period and calculate proper storage change rate
+ storage_change_rates = []
+ time_indices = storage_merged.index.values
+
+ for time_val in time_indices:
+ # Extract month from time value for monthly data
+ # Time format: year + (month - 0.5)/12 for monthly, year for annual
+ if time_val != int(time_val): # Monthly data (has fractional part)
+ year = int(time_val)
+ month_fraction = time_val - year
+ month = int(month_fraction * 12 + 0.5)
+ month = max(1, min(12, month)) # Ensure valid month range
+ time_period_seconds = float(DAYS_PER_MONTH[month] * 24 * 60 * 60)
+ else: # Annual data
+ time_period_seconds = SECONDS_PER_YEAR
+
+ # Calculate storage change in mm (already normalized)
+ storage_change_mm = (
+ storage_merged.loc[time_val, "normalized_value_end"]
+ - storage_merged.loc[time_val, "normalized_value_beg"]
+ )
+
+ # Convert to rate: mm per time_period -> kg/m2/s -> mm/yr
+ storage_change_rate_kg_m2_s = storage_change_mm / time_period_seconds
+ storage_change_rate_mm_yr = storage_change_rate_kg_m2_s * SECONDS_PER_YEAR
+
+ storage_change_rates.append(storage_change_rate_mm_yr)
+
+ storage_df = pd.DataFrame(
+ {"normalized_value": storage_change_rates}, index=storage_merged.index
+ )
+
+ # Compare storage change vs net flux
+ merged = storage_df.join(
+ net_flux, lsuffix="_change", rsuffix="_flux", how="inner"
+ )
+ if merged.empty:
+ return None
+
+ years = merged.index.values
+ change = merged["normalized_value_change"].values
+ flux = merged["normalized_value_flux"].values
+ residual = change - flux
+
+ change_label = "ΔMass" if self.quantity == "water" else "ΔEnergy"
+
+ return CheckResult(
+ self.name,
+ self.description,
+ years,
+ change,
+ flux,
+ residual,
+ np.cumsum(residual),
+ lhs_label=f"Storage Change ({change_label})",
+ rhs_label=f"Net Flux ({flux_term})",
+ )
+
+
+DEFAULT_WATER_CHECKS: List[BudgetCheck] = [
+ CplComponentFluxes(quantity="water"),
+ InterfaceMatch("lnd", "lnd", quantity="water"),
+ InterfaceMatch("ocn", "ocn", quantity="water", comp_sum_term="SUM VOLUME FLUXES"),
+ IceInterfaceMatch(quantity="water"),
+ AtmInterfaceMatch(quantity="water"),
+ LndClosure(),
+ OcnClosure(quantity="water"),
+ IceClosure(quantity="water"),
+ AtmClosure(quantity="water"),
+]
+
+DEFAULT_HEAT_CHECKS: List[BudgetCheck] = [
+ CplComponentFluxes(quantity="heat"),
+ InterfaceMatch(
+ "ocn", "ocn", quantity="heat", comp_sum_term="SUM IMP+EXP HEAT FLUXES"
+ ),
+ IceInterfaceMatch(quantity="heat"),
+ # AtmInterfaceMatch(quantity="heat"), # DISABLED: ATM logs lack complete energy flux data
+ OcnClosure(quantity="heat"),
+ IceClosure(quantity="heat"),
+ # AtmClosure(quantity="heat"), # DISABLED: ATM logs lack complete energy flux data
+]
+
+DEFAULT_CARBON_CHECKS: List[BudgetCheck] = [
+ CplComponentFluxes(quantity="carbon"),
+]
+
+
+def run_checks(
+ df: pd.DataFrame,
+ checks: Optional[List[BudgetCheck]] = None,
+) -> List[CheckResult]:
+ """Run budget checks against the normalized event table."""
+ if checks is None:
+ checks = DEFAULT_WATER_CHECKS
+ # Determine time unit label from data period
+ periods = df[COL_PERIOD].unique() if not df.empty else []
+ if "monthly" in periods:
+ time_label = "months"
+ else:
+ time_label = "years"
+ results = []
+ for check in checks:
+ result = check.evaluate(df)
+ if result is not None:
+ results.append(result)
+ print(f" Check '{check.name}': {len(result.years)} {time_label}")
+ else:
+ print(f" WARNING: Check '{check.name}' skipped (missing data)")
+ return results
diff --git a/zppy_interfaces/budget_analysis/ingestion/__init__.py b/zppy_interfaces/budget_analysis/ingestion/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/zppy_interfaces/budget_analysis/ingestion/atm_parser.py b/zppy_interfaces/budget_analysis/ingestion/atm_parser.py
new file mode 100644
index 0000000..a87602c
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/atm_parser.py
@@ -0,0 +1,426 @@
+"""Atmosphere log parser — extracts energy/water conservation diagnostics.
+
+ATM log per-step diagnostic format:
+
+ nstep, te 1 TE_before TE_after fixer_frac ps
+ n, dt, W tot mass [kg/m2] 1 3600.0 25.303...
+ n, W flux, dWater [kg/m2] 1 flux_val dwater_val
+ n, W flux-dWater [kg/m2] 1 residual
+ n, W cflx*dt loss [kg/m2] 1 loss
+ n, E d(TE)/dt, RR [W/m2] 1 dTE_dt RR
+ n, E difference [W/m2] 1 diff
+ n, E shf loss [W/m2] 1 loss
+
+Date info from ``chem_surfvals_set: ncdate=YYYMMDD`` (once per day).
+
+Note: ``nstep, te`` runs from 0..N+1 while flux diagnostics run n=1..N
+(the atm receives and applies fluxes in two consecutive steps).
+"""
+
+import gzip
+from typing import Dict, List, Tuple
+
+import pandas as pd
+
+from ..schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+ COL_UNITS,
+ COL_VALUE,
+ COLUMNS,
+)
+from .base import BaseParser
+
+# Line prefixes for each diagnostic type
+_L_NSTEP_TE = "nstep, te "
+_L_NCDATE = "chem_surfvals_set: ncdate="
+_L_W_TOT = "n, dt, W tot mass [kg/m2]"
+_L_W_FLUX = "n, W flux, dWater [kg/m2]"
+_L_W_RESID = "n, W flux-dWater [kg/m2]"
+_L_W_CFLX = "n, W cflx*dt loss [kg/m2]"
+_L_E_DTEDT = "n, E d(TE)/dt, RR [W/m2]"
+_L_E_DIFF = "n, E difference [W/m2]"
+_L_E_SHF = "n, E shf loss [W/m2]"
+
+
+def _parse_ncdate(ncdate_int: int) -> Tuple[int, int, int]:
+ """Parse ncdate integer YYYMMDD -> (year, month, day).
+
+ Examples: 10101 -> (1, 1, 1), 501231 -> (50, 12, 31).
+ """
+ day = ncdate_int % 100
+ ncdate_int //= 100
+ month = ncdate_int % 100
+ year = ncdate_int // 100
+ return year, month, day
+
+
+def _open_log(filename: str):
+ """Open a log file, handling gzip compression."""
+ if filename.endswith(".gz"):
+ return gzip.open(filename, "rt")
+ return open(filename, "r")
+
+
+def _gather_energy_data(filename: str) -> Tuple[List[Dict], List[Dict]]:
+ """Extract per-step diagnostic data from a single atm log file.
+
+ Returns (nstep_te_rows, flux_diag_rows).
+ """
+ nstep_te_rows: List[Dict] = []
+ flux_diag_rows: List[Dict] = []
+
+ # Current date context (from ncdate lines)
+ cur_year, cur_month, cur_day = -1, -1, -1
+
+ # Accumulate flux diagnostic fields for the current step
+ cur_flux: Dict = {}
+
+ with _open_log(filename) as f:
+ for line in f:
+ # --- Date tracking ---
+ if _L_NCDATE in line:
+ # chem_surfvals_set: ncdate= 10101 co2vmr=...
+ parts = line.split("ncdate=")[1].split()
+ cur_year, cur_month, cur_day = _parse_ncdate(int(parts[0]))
+ continue
+
+ # --- Energy fixer (nstep, te) ---
+ if _L_NSTEP_TE in line:
+ tokens = line.split()
+ # tokens: ['nstep,', 'te', N, te_before, te_after, fixer_frac, ps]
+ nstep_te_rows.append(
+ {
+ "nstep": int(tokens[2]),
+ "te_before": float(tokens[3]),
+ "te_after": float(tokens[4]),
+ "fixer_frac": float(tokens[5]),
+ "ps": float(tokens[6]),
+ "year": cur_year,
+ "month": cur_month,
+ "day": cur_day,
+ }
+ )
+ continue
+
+ # --- Water total mass (starts a new flux diagnostic group) ---
+ if _L_W_TOT in line:
+ # Flush previous step if any
+ if cur_flux:
+ flux_diag_rows.append(cur_flux)
+ tokens = line.split()
+ cur_flux = {
+ "nstep": int(tokens[6]),
+ "dt": float(tokens[7]),
+ "tw": float(tokens[8]),
+ "year": cur_year,
+ "month": cur_month,
+ "day": cur_day,
+ }
+ continue
+
+ # --- Water flux, dWater ---
+ if _L_W_FLUX in line:
+ tokens = line.split()
+ cur_flux["w_flux"] = float(tokens[6])
+ cur_flux["w_dwater"] = float(tokens[7])
+ continue
+
+ # --- Water residual ---
+ if _L_W_RESID in line:
+ tokens = line.split()
+ cur_flux["w_residual"] = float(tokens[5])
+ continue
+
+ # --- Water coupling flux loss ---
+ if _L_W_CFLX in line:
+ tokens = line.split()
+ cur_flux["w_cflx_loss"] = float(tokens[6])
+ continue
+
+ # --- Energy d(TE)/dt and RR ---
+ if _L_E_DTEDT in line:
+ tokens = line.split()
+ cur_flux["e_dtedt"] = float(tokens[6])
+ cur_flux["e_rr"] = float(tokens[7])
+ continue
+
+ # --- Energy difference ---
+ if _L_E_DIFF in line:
+ tokens = line.split()
+ cur_flux["e_diff"] = float(tokens[5])
+ continue
+
+ # --- Energy shf loss ---
+ if _L_E_SHF in line:
+ tokens = line.split()
+ cur_flux["e_shf_loss"] = float(tokens[6])
+ continue
+
+ # Flush last step
+ if cur_flux:
+ flux_diag_rows.append(cur_flux)
+
+ return nstep_te_rows, flux_diag_rows
+
+
+class AtmParser(BaseParser):
+ """Parse atmosphere log files for energy/water conservation diagnostics."""
+
+ def __init__(self, frequency: str = "annual") -> None:
+ super().__init__(frequency=frequency)
+
+ def parse_raw(self, log_files: List[str]) -> Tuple[pd.DataFrame, pd.DataFrame]:
+ """Parse raw per-step data from atm log files.
+
+ Returns
+ -------
+ nstep_te_df : pd.DataFrame
+ Energy fixer data per step.
+ Columns: nstep, te_before, te_after, fixer_frac, ps,
+ year, month, day.
+ flux_diag_df : pd.DataFrame
+ Flux diagnostics per step.
+ Columns: nstep, dt, tw, w_flux, w_dwater, w_residual,
+ w_cflx_loss, e_dtedt, e_rr, e_diff, e_shf_loss,
+ year, month, day.
+ """
+ all_nstep: List[Dict] = []
+ all_flux: List[Dict] = []
+
+ for fname in sorted(log_files):
+ try:
+ nstep_rows, flux_rows = _gather_energy_data(fname)
+ all_nstep.extend(nstep_rows)
+ all_flux.extend(flux_rows)
+ except Exception as e:
+ print(f"WARNING: Error processing {fname}: {e}")
+ continue
+
+ nstep_te_df = pd.DataFrame(all_nstep)
+ flux_diag_df = pd.DataFrame(all_flux)
+ return nstep_te_df, flux_diag_df
+
+ def _groupby_keys(self) -> List[str]:
+ """Return groupby columns based on frequency."""
+ if self.frequency == "monthly":
+ return ["year", "month"]
+ else: # annual
+ return ["year"]
+
+ @staticmethod
+ def _make_time(frequency: str, year: int, month: int) -> float:
+ """Encode (year, month) as a float time value."""
+ if frequency == "monthly":
+ return year + (month - 0.5) / 12.0
+ return float(year)
+
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ """Parse atm logs and return a tidy event table.
+
+ Aggregates per-step data to the configured frequency (monthly
+ or annual) for fluxes and beginning/end-of-period states.
+ """
+ nstep_te_df, flux_diag_df = self.parse_raw(log_files)
+
+ if nstep_te_df.empty and flux_diag_df.empty:
+ return pd.DataFrame(columns=COLUMNS)
+
+ group_keys = self._groupby_keys()
+ rows: List[Dict] = []
+
+ # --- Flux diagnostics ---
+ if not flux_diag_df.empty:
+ flux_df = flux_diag_df[
+ (flux_diag_df["year"] >= start_year)
+ & (flux_diag_df["year"] <= end_year)
+ ]
+ if not flux_df.empty:
+ for _key, grp in flux_df.groupby(group_keys):
+ time = self._make_time(
+ self.frequency,
+ int(grp["year"].iloc[0]),
+ int(grp["month"].iloc[0]),
+ )
+ base_water = {
+ COL_TIME: time,
+ COL_COMPONENT: "atm",
+ COL_QUANTITY: "water",
+ COL_SOURCE: "atm",
+ COL_PERIOD: self.frequency,
+ }
+ base_heat = {
+ COL_TIME: time,
+ COL_COMPONENT: "atm",
+ COL_QUANTITY: "heat",
+ COL_SOURCE: "atm",
+ COL_PERIOD: self.frequency,
+ }
+
+ dt = grp["dt"].iloc[0]
+
+ # Water flux terms (mean rate, kg/m2/s)
+ if "w_flux" in grp.columns:
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "W flux",
+ COL_VALUE: (grp["w_flux"] / dt).mean(),
+ COL_UNITS: "kg/m2/s",
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+ if "w_dwater" in grp.columns:
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "dWater",
+ COL_VALUE: (grp["w_dwater"] / dt).mean(),
+ COL_UNITS: "kg/m2/s",
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+
+ # Water state (begin/end of period)
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "W tot mass beg",
+ COL_VALUE: grp["tw"].iloc[0],
+ COL_UNITS: "kg/m2",
+ COL_TABLE_TYPE: "state",
+ }
+ )
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "W tot mass end",
+ COL_VALUE: grp["tw"].iloc[-1],
+ COL_UNITS: "kg/m2",
+ COL_TABLE_TYPE: "state",
+ }
+ )
+
+ # Water diagnostics (mean per step)
+ if "w_residual" in grp.columns:
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "W residual",
+ COL_VALUE: grp["w_residual"].mean(),
+ COL_UNITS: "kg/m2",
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+ if "w_cflx_loss" in grp.columns:
+ rows.append(
+ {
+ **base_water,
+ COL_TERM: "W cflx loss",
+ COL_VALUE: grp["w_cflx_loss"].mean(),
+ COL_UNITS: "kg/m2",
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+
+ # Energy flux terms (mean, W/m2)
+ if "e_dtedt" in grp.columns:
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "E d(TE)/dt",
+ COL_VALUE: grp["e_dtedt"].mean(),
+ COL_UNITS: "W/m2",
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+ if "e_rr" in grp.columns:
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "E RR",
+ COL_VALUE: grp["e_rr"].mean(),
+ COL_UNITS: "W/m2",
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+
+ # Energy diagnostics (mean, W/m2)
+ if "e_diff" in grp.columns:
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "E difference",
+ COL_VALUE: grp["e_diff"].mean(),
+ COL_UNITS: "W/m2",
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+ if "e_shf_loss" in grp.columns:
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "E shf loss",
+ COL_VALUE: grp["e_shf_loss"].mean(),
+ COL_UNITS: "W/m2",
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+
+ # --- Energy fixer state (TE begin/end of period) ---
+ if not nstep_te_df.empty:
+ te_df = nstep_te_df[
+ (nstep_te_df["year"] >= start_year) & (nstep_te_df["year"] <= end_year)
+ ]
+ if not te_df.empty:
+ for _key, grp in te_df.groupby(group_keys):
+ time = self._make_time(
+ self.frequency,
+ int(grp["year"].iloc[0]),
+ int(grp["month"].iloc[0]),
+ )
+ base_heat = {
+ COL_TIME: time,
+ COL_COMPONENT: "atm",
+ COL_QUANTITY: "heat",
+ COL_SOURCE: "atm",
+ COL_PERIOD: self.frequency,
+ }
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "TE beg",
+ COL_VALUE: grp["te_before"].iloc[0],
+ COL_UNITS: "J/m2",
+ COL_TABLE_TYPE: "state",
+ }
+ )
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "TE end",
+ COL_VALUE: grp["te_after"].iloc[-1],
+ COL_UNITS: "J/m2",
+ COL_TABLE_TYPE: "state",
+ }
+ )
+ rows.append(
+ {
+ **base_heat,
+ COL_TERM: "E fixer frac",
+ COL_VALUE: grp["fixer_frac"].mean(),
+ COL_UNITS: "1",
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+
+ if not rows:
+ return pd.DataFrame(columns=COLUMNS)
+ return pd.DataFrame(rows, columns=COLUMNS)
diff --git a/zppy_interfaces/budget_analysis/ingestion/base.py b/zppy_interfaces/budget_analysis/ingestion/base.py
new file mode 100644
index 0000000..73c1db3
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/base.py
@@ -0,0 +1,36 @@
+"""Base class for log file parsers."""
+
+from abc import ABC, abstractmethod
+from typing import List
+
+import pandas as pd
+
+VALID_FREQUENCIES = ("monthly", "annual")
+
+
+class BaseParser(ABC):
+ """All parsers return a tidy event table DataFrame.
+
+ Parameters
+ ----------
+ frequency : str
+ Temporal granularity of the output: ``"monthly"`` (default)
+ or ``"daily"``. For parsers whose log data is already at a
+ fixed period (e.g. coupler annual/monthly), the frequency
+ selects which records to keep. For the atm parser it controls
+ the groupby aggregation window.
+ """
+
+ def __init__(self, frequency: str = "annual") -> None:
+ if frequency not in VALID_FREQUENCIES:
+ raise ValueError(
+ f"frequency must be one of {VALID_FREQUENCIES}, got {frequency!r}"
+ )
+ self.frequency = frequency
+
+ @abstractmethod
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ """Parse log files and return a tidy event table."""
+ ...
diff --git a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py
new file mode 100644
index 0000000..9ca5ae1
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py
@@ -0,0 +1,183 @@
+"""Coupler log parser — extracts budget tables into tidy DataFrames."""
+
+import gzip
+import re
+from typing import Dict, List, Optional, TextIO, Tuple
+
+import pandas as pd
+
+from ..schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+ COL_UNITS,
+ COL_VALUE,
+ COLUMNS,
+)
+from .base import BaseParser
+
+# Header patterns for each budget quantity.
+HEADER_PATTERNS: Dict[str, str] = {
+ "water": "(seq_diag_print_mct) NET WATER BUDGET (kg/m2s*1e6):",
+ "heat": "(seq_diag_print_mct) NET HEAT BUDGET (W/m2):",
+ "carbon": "(seq_diagBGC_print_mct) NET CARBON BUDGET (kg-C/m2s*1e10):",
+}
+
+UNITS: Dict[str, str] = {
+ "water": "kg/m2s*1e6",
+ "heat": "W/m2",
+ "carbon": "kg-C/m2s*1e10",
+}
+
+
+def _normalize_component_name(name: str) -> str:
+ """Normalize component names: 'ice nh' -> 'ice_nh'."""
+ return name.strip().replace(" ", "_")
+
+
+def _parse_datestamp(datestamp: str, period: str = "monthly") -> Tuple[int, int]:
+ """Convert coupler datestamp to (year, month).
+
+ For monthly data: '10201' -> year 1, month 1 (roll back one month)
+ For annual data: '20101' -> year 1 (annual summary for year 1, output at start of year 2)
+ """
+ mmdd = datestamp[-4:]
+ month = int(mmdd[:2])
+ year = int(datestamp[:-4])
+
+ if period == "annual":
+ # Annual data: date represents start of year after the summary year
+ # e.g., '20101' = annual summary for year 1, output at start of year 2
+ return year - 1, 12 # Return summary year with month 12 for annual data
+ else:
+ # Monthly data: roll back one month (date is start of *next* period)
+ if month == 1:
+ month = 12
+ year -= 1 # Roll back year when going from Jan to Dec
+ else:
+ month -= 1
+ return year, month
+
+
+def _make_time(period: str, year: int, month: int) -> float:
+ """Encode (year, month) as a float time value.
+
+ Annual: integer year. Monthly: year + (month - 0.5) / 12.
+ """
+ if period == "monthly":
+ return year + (month - 0.5) / 12.0
+ return float(year)
+
+
+def _parse_header_line(line: str, pattern: str) -> Optional[Tuple[str, float]]:
+ """Extract period and time from a budget header line.
+
+ Returns (period, time) or None on failure.
+ """
+ if not line.startswith(pattern):
+ return None
+
+ remainder = line[len(pattern) :]
+ period_match = re.search(r"period\s*=\s*(\w+)", remainder)
+ date_match = re.search(r"date\s*=\s*(\d+)", remainder)
+ if not period_match or not date_match:
+ return None
+
+ period = period_match.group(1)
+ year, month = _parse_datestamp(date_match.group(1), period)
+ return period, _make_time(period, year, month)
+
+
+def _parse_table(f: TextIO, year: float, quantity: str, period: str) -> List[Dict]:
+ """Parse one budget table after the header line was consumed."""
+ rows: List[Dict] = []
+ units = UNITS[quantity]
+
+ # First line after header: column names separated by 2+ spaces
+ col_line = f.readline().strip()
+ col_names = [
+ _normalize_component_name(c) for c in re.split(r"\s{2,}", col_line) if c
+ ]
+
+ # Data rows until blank line
+ line = f.readline()
+ while line and line.strip():
+ parts = line.split()
+ # Find where numeric data starts (handles multi-word term names)
+ term_parts: List[str] = []
+ data_start = 0
+ for j, part in enumerate(parts):
+ try:
+ float(part)
+ data_start = j
+ break
+ except ValueError:
+ term_parts.append(part)
+ term = " ".join(term_parts)
+ values = parts[data_start:]
+
+ for i, val_str in enumerate(values):
+ if i < len(col_names):
+ rows.append(
+ {
+ COL_TIME: year,
+ COL_COMPONENT: col_names[i],
+ COL_QUANTITY: quantity,
+ COL_TERM: term,
+ COL_VALUE: float(val_str),
+ COL_UNITS: units,
+ COL_SOURCE: "cpl",
+ COL_PERIOD: period,
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+
+ line = f.readline()
+
+ return rows
+
+
+class CplParser(BaseParser):
+ """Parse coupler log budget tables into a tidy event table."""
+
+ def __init__(
+ self,
+ quantities: Optional[List[str]] = None,
+ frequency: str = "annual",
+ ):
+ super().__init__(frequency=frequency)
+ self.quantities = quantities or ["water", "heat"]
+
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ rows: List[Dict] = []
+ for fname in sorted(log_files):
+ try:
+ with gzip.open(fname, "rt") as f:
+ for line in f:
+ for quantity in self.quantities:
+ pattern = HEADER_PATTERNS.get(quantity)
+ if not pattern:
+ continue
+ result = _parse_header_line(line, pattern)
+ if result is None:
+ continue
+ period, time = result
+ if period != self.frequency:
+ continue
+ # Extract year from time for consistent filtering
+ year = int(time)
+ if start_year <= year <= end_year:
+ rows.extend(_parse_table(f, time, quantity, period))
+ except Exception as e:
+ print(f"WARNING: Error processing {fname}: {e}")
+ continue
+
+ if not rows:
+ return pd.DataFrame(columns=COLUMNS)
+ return pd.DataFrame(rows, columns=COLUMNS)
diff --git a/zppy_interfaces/budget_analysis/ingestion/ice_parser.py b/zppy_interfaces/budget_analysis/ingestion/ice_parser.py
new file mode 100644
index 0000000..1b1e4d7
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/ice_parser.py
@@ -0,0 +1,277 @@
+"""Ice log parser — extracts mass and energy conservation checks into tidy DataFrames.
+
+Ice log format (monthly, inside Conservation checks blocks):
+
+ ===================================================================================
+ Conservation checks: 0001-02-01_00:00:00
+ -----------------------------------------------------------------------------------
+ Area analysis Global NH SH
+
+ Earth radius (m) = 6.371229E+06
+ Earth area (m2) = 5.101011E+14
+ Domain area (m2) = 3.574961E+14 1.501228E+14 2.073733E+14
+ Sea-ice area (m2) = 3.038731E+13 1.207167E+13 1.831565E+13
+ -----------------------------------------------------------------------------------
+ Energy conservation check
+
+ Initial energy ice (J) = -2.695578E+22 -8.884723E+21 -1.807106E+22
+ Final energy ice (J) = -1.949620E+22 -9.304888E+21 -1.019131E+22
+ Energy change (J) = 7.459584E+21 -4.201653E+20 7.879749E+21
+ Energy change flux (W/m2) = 5.459877E+00 -3.075307E-01 5.767408E+00
+
+ Surface heat flux (W/m2) = 5.153267E-01 -7.817582E-01 1.297085E+00
+ Absorbed shortwave flux (W/m2) = 2.892564E+00 2.707821E-02 2.865486E+00
+ Ocean Shortwave flux (W/m2) = -1.960086E-01 -5.682405E-03 -1.903262E-01
+ ...
+ Net energy change (J) = 7.449107E+21 -4.186389E+20 7.867746E+21
+ Net energy flux (W/m2) = 5.452209E+00 -3.064135E-01 5.758622E+00
+ -----------------------------------------------------------------------------------
+ Mass conservation check
+
+ Initial mass ice (kg) = 7.684570E+16 2.532862E+16 5.151708E+16
+ Final mass ice (kg) = 5.825212E+16 2.706170E+16 3.119042E+16
+ Ice mass change (kg) = -1.859359E+16 1.733078E+15 -2.032666E+16
+ Ice mass change flux (kg/m2s) = -1.360916E-05 1.268488E-06 -1.487765E-05
+ ...
+"""
+
+import gzip
+import re
+from typing import Dict, List, TextIO, Tuple
+
+import pandas as pd
+
+from ..schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+ COL_UNITS,
+ COL_VALUE,
+ COLUMNS,
+)
+from .base import BaseParser
+
+
+def _parse_date(date_str: str) -> Tuple[int, int]:
+ """Parse 'YYYY-MM-DD_HH:MM:SS' -> (year, month).
+
+ The date is printed at the START of the next month,
+ so date 0002-01-01 means the check covers Dec of year 1.
+ We return the year and month of the COVERED period.
+ Same logic as ocean parser.
+ """
+ match = re.match(r"(\d+)-(\d+)-(\d+)", date_str.strip())
+ if not match:
+ return -1, -1
+ y, m = int(match.group(1)), int(match.group(2))
+ # Roll back one month
+ if m == 1:
+ return y - 1, 12
+ return y, m - 1
+
+
+def _parse_energy_section(f: TextIO, time: float) -> List[Dict]:
+ """Parse Energy conservation check section."""
+ rows: List[Dict] = []
+ base = {
+ COL_TIME: time,
+ COL_COMPONENT: "ice",
+ COL_QUANTITY: "heat",
+ COL_UNITS: "W/m2",
+ COL_SOURCE: "ice",
+ COL_PERIOD: "monthly",
+ COL_TABLE_TYPE: "flux",
+ }
+
+ # Read until we find mass section or next conservation block
+ while True:
+ line = f.readline()
+ if not line: # EOF
+ break
+
+ stripped = line.strip()
+ if not stripped:
+ continue
+
+ # Stop at next section
+ if "Mass conservation check" in line:
+ # Put this line back by seeking backwards (approximate)
+ # We'll let _parse_mass_section handle this line
+ break
+ if "Conservation checks:" in line:
+ break
+ if "---" in line and len([c for c in line if c == "-"]) > 10:
+ break
+
+ # Parse energy flux terms
+ if "(W/m2)" in line:
+ # Extract term name and values
+ parts = line.split("(W/m2)")
+ if len(parts) >= 2:
+ term_name = parts[0].strip()
+ values_part = parts[1].split("=")
+ if len(values_part) >= 2:
+ try:
+ values = values_part[1].strip().split()
+ # Use global value (first column)
+ if values:
+ val = float(values[0])
+ rows.append(
+ {
+ **base,
+ COL_TERM: term_name,
+ COL_VALUE: val,
+ }
+ )
+ except (ValueError, IndexError):
+ pass
+
+ return rows
+
+
+def _parse_mass_section(f: TextIO, time: float) -> List[Dict]:
+ """Parse Mass conservation check section."""
+ rows: List[Dict] = []
+ base = {
+ COL_TIME: time,
+ COL_COMPONENT: "ice",
+ COL_QUANTITY: "water",
+ COL_UNITS: "kg/m2s",
+ COL_SOURCE: "ice",
+ COL_PERIOD: "monthly",
+ COL_TABLE_TYPE: "flux",
+ }
+
+ # Read until we find end of mass section
+ while True:
+ line = f.readline()
+ if not line: # EOF
+ break
+
+ stripped = line.strip()
+ if not stripped:
+ continue
+
+ # Stop at end of section
+ if "Conservation checks:" in line:
+ break
+ if "---" in line and len([c for c in line if c == "-"]) > 10:
+ break
+
+ # Parse mass flux terms
+ if "(kg/m2s)" in line:
+ # Extract term name and values
+ parts = line.split("(kg/m2s)")
+ if len(parts) >= 2:
+ term_name = parts[0].strip()
+ values_part = parts[1].split("=")
+ if len(values_part) >= 2:
+ try:
+ values = values_part[1].strip().split()
+ # Use global value (first column)
+ if values:
+ val = float(values[0])
+ rows.append(
+ {
+ **base,
+ COL_TERM: term_name,
+ COL_VALUE: val,
+ }
+ )
+ except (ValueError, IndexError):
+ pass
+
+ return rows
+
+
+class IceParser(BaseParser):
+ """Parse ice log files for mass and energy conservation checks."""
+
+ def __init__(self, frequency: str = "annual") -> None:
+ super().__init__(frequency=frequency)
+
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ rows: List[Dict] = []
+ for fname in sorted(log_files):
+ try:
+ with gzip.open(fname, "rt") as f:
+ for line in f:
+ if "Conservation checks:" in line:
+ # Extract date from same line
+ date_match = re.search(
+ r"(\d{4}-\d{2}-\d{2}_\d{2}:\d{2}:\d{2})", line
+ )
+ if not date_match:
+ continue
+ date_str = date_match.group(1)
+ year, month = _parse_date(date_str)
+ if year < start_year or year > end_year:
+ continue
+ if self.frequency == "monthly":
+ time = year + (month - 0.5) / 12.0
+ else:
+ time = float(year)
+ rows.extend(self._parse_block(f, time))
+ except Exception as e:
+ print(f"WARNING: Error processing {fname}: {e}")
+ continue
+
+ if not rows:
+ return pd.DataFrame(columns=COLUMNS)
+ df = pd.DataFrame(rows, columns=COLUMNS)
+
+ # For annual frequency, aggregate monthly rows to annual means
+ if self.frequency == "annual":
+ group_keys = [
+ COL_TIME,
+ COL_COMPONENT,
+ COL_QUANTITY,
+ COL_TERM,
+ COL_UNITS,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ ]
+ df[COL_PERIOD] = "annual"
+ df = df.groupby(group_keys, as_index=False).agg(
+ {COL_VALUE: "mean", COL_PERIOD: "first"}
+ )
+
+ return df
+
+ def _parse_block(self, f: TextIO, time: float) -> List[Dict]:
+ """Parse one Conservation checks block for mass and energy data."""
+ rows: List[Dict] = []
+ found_mass = False
+ found_energy = False
+
+ line = f.readline()
+ while line:
+ # Stop when both sections are processed
+ if found_mass and found_energy:
+ break
+
+ # Stop at next conservation block
+ if "Conservation checks:" in line:
+ break
+
+ # Parse energy section
+ if "Energy conservation check" in line:
+ found_energy = True
+ energy_rows = _parse_energy_section(f, time)
+ rows.extend(energy_rows)
+
+ # Parse mass section
+ elif "Mass conservation check" in line:
+ found_mass = True
+ mass_rows = _parse_mass_section(f, time)
+ rows.extend(mass_rows)
+
+ line = f.readline()
+
+ return rows
diff --git a/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py
new file mode 100644
index 0000000..bbcda72
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py
@@ -0,0 +1,344 @@
+"""Land log parser — extracts water flux and state tables into tidy DataFrames.
+
+Land log format examples:
+
+NET WATER FLUXES : period annual: date = 20101 0
+ Time | Time
+ averaged | integrated
+ kg/m2s*1e6 | kg/m2*1e6
+--------------------------------|--------------------
+ rain 7.168... | 226051777.94
+ ...
+ *SUM* -0.005... | -174290.03
+
+WATER STATES (kg/m2*1e6): period annual: date = 20101 0
+ Canopy Snow SFC Soil Liq Soil Ice Aquifer Grid-level Err | TOTAL
+------...
+ beg ...
+ end ...
+ *NET CHANGE* ...
+------...
+ *SUM* ...
+"""
+
+import gzip
+import re
+from typing import Dict, List, Optional, TextIO, Tuple
+
+import pandas as pd
+
+from ..schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+ COL_UNITS,
+ COL_VALUE,
+ COLUMNS,
+)
+from .base import BaseParser
+
+FLUX_HEADER = "NET WATER FLUXES : period"
+STATE_HEADER = "WATER STATES (kg/m2*1e6): period"
+
+
+def _parse_datestamp(datestamp: str, period: str = "monthly") -> Tuple[int, int]:
+ """Convert datestamp to (year, month). Same convention as coupler.
+
+ For monthly data: '10201' -> year 1, month 1 (roll back one month)
+ For annual data: '20101' -> year 1 (annual summary for year 1, output at start of year 2)
+ """
+ mmdd = datestamp[-4:]
+ month = int(mmdd[:2])
+ year = int(datestamp[:-4])
+
+ if period == "annual":
+ # Annual data: date represents start of year after the summary year
+ return year - 1, 12 # Return summary year with month 12 for annual data
+ else:
+ # Monthly data: roll back one month (date is start of *next* period)
+ if month == 1:
+ month = 12
+ year -= 1 # Roll back year when going from Jan to Dec
+ else:
+ month -= 1
+ return year, month
+
+
+def _make_time(period: str, year: int, month: int) -> float:
+ """Encode (year, month) as a float time value."""
+ if period == "monthly":
+ return year + (month - 0.5) / 12.0
+ return float(year)
+
+
+def _parse_period_and_time(line: str) -> Optional[Tuple[str, float]]:
+ """Extract period and time from a header line."""
+ period_match = re.search(r"period\s+(\w+):", line)
+ date_match = re.search(r"date\s*=\s*(\d+)", line)
+ if not period_match or not date_match:
+ return None
+ period = period_match.group(1)
+ year, month = _parse_datestamp(date_match.group(1), period)
+ return period, _make_time(period, year, month)
+
+
+def _parse_flux_table(f: TextIO, year: float, period: str) -> List[Dict]:
+ """Parse a NET WATER FLUXES table.
+
+ Format:
+ [2 header lines: Time averaged | Time integrated]
+ [1 units line: kg/m2s*1e6 | kg/m2*1e6]
+ [separator line: ---...|---...]
+ term rate_value | integrated_value
+ ...
+ [separator line]
+ *SUM* rate_value | integrated_value
+ """
+ rows: List[Dict] = []
+
+ # Skip 3 header lines (Time averaged/integrated labels + units)
+ for _ in range(3):
+ f.readline()
+
+ # Skip separator
+ f.readline()
+
+ # Parse data rows until we hit an empty line or non-data line
+ line = f.readline()
+ while line and line.strip():
+ stripped = line.strip()
+ # Skip separator lines
+ if stripped.startswith("---"):
+ line = f.readline()
+ continue
+
+ # Split on '|' to separate rate and integrated columns
+ parts = stripped.split("|")
+ if len(parts) < 2:
+ line = f.readline()
+ continue
+
+ rate_part = parts[0].strip()
+ integrated_part = parts[1].strip()
+
+ # Parse term name and rate value
+ rate_tokens = rate_part.split()
+ if not rate_tokens:
+ line = f.readline()
+ continue
+
+ # Find where numeric data starts
+ term_parts: List[str] = []
+ rate_val = None
+ for token in rate_tokens:
+ try:
+ rate_val = float(token)
+ break
+ except ValueError:
+ term_parts.append(token)
+
+ term = " ".join(term_parts)
+ if not term or rate_val is None:
+ line = f.readline()
+ continue
+
+ # Parse integrated value
+ try:
+ integrated_val = float(integrated_part)
+ except ValueError:
+ integrated_val = None
+
+ # Emit flux rate row
+ rows.append(
+ {
+ COL_TIME: year,
+ COL_COMPONENT: "lnd",
+ COL_QUANTITY: "water",
+ COL_TERM: term,
+ COL_VALUE: rate_val,
+ COL_UNITS: "kg/m2s*1e6",
+ COL_SOURCE: "lnd",
+ COL_PERIOD: period,
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+
+ # Emit flux integrated row
+ if integrated_val is not None:
+ rows.append(
+ {
+ COL_TIME: year,
+ COL_COMPONENT: "lnd",
+ COL_QUANTITY: "water",
+ COL_TERM: term,
+ COL_VALUE: integrated_val,
+ COL_UNITS: "kg/m2*1e6",
+ COL_SOURCE: "lnd",
+ COL_PERIOD: period,
+ COL_TABLE_TYPE: "flux_integrated",
+ }
+ )
+
+ line = f.readline()
+
+ return rows
+
+
+def _parse_state_table(f: TextIO, year: float, period: str) -> List[Dict]:
+ """Parse a WATER STATES table.
+
+ Format:
+ [column header line: Canopy Snow SFC Soil Liq Soil Ice Aquifer Grid-level Err | TOTAL]
+ [separator line]
+ beg val1 val2 ... | total
+ end val1 val2 ... | total
+ *NET CHANGE* val1 val2 ... | total
+ [separator line]
+ *SUM* ... | total
+ """
+ rows: List[Dict] = []
+
+ # Column header line
+ col_line = f.readline()
+ # Split on '|' — left side has pool names, right has TOTAL
+ col_parts = col_line.split("|")
+ left_header = col_parts[0].strip() if col_parts else ""
+ # Parse pool names from left header (separated by 2+ spaces)
+ pool_names = [p.strip() for p in re.split(r"\s{2,}", left_header) if p.strip()]
+
+ # Skip separator
+ f.readline()
+
+ # Parse data rows
+ line = f.readline()
+ while line and line.strip():
+ stripped = line.strip()
+ if stripped.startswith("---"):
+ line = f.readline()
+ continue
+
+ # Split on '|'
+ parts = stripped.split("|")
+ left_part = parts[0].strip()
+ right_part = parts[1].strip() if len(parts) > 1 else ""
+
+ tokens = left_part.split()
+ if not tokens:
+ line = f.readline()
+ continue
+
+ # Find row label and values
+ label_parts: List[str] = []
+ data_start = 0
+ for j, token in enumerate(tokens):
+ try:
+ float(token)
+ data_start = j
+ break
+ except ValueError:
+ label_parts.append(token)
+
+ row_label = " ".join(label_parts)
+ values = tokens[data_start:]
+
+ # Parse TOTAL from right side of '|'
+ total_val = None
+ if right_part:
+ try:
+ total_val = float(right_part)
+ except ValueError:
+ pass
+
+ # Emit per-pool values
+ for i, val_str in enumerate(values):
+ if i < len(pool_names):
+ try:
+ val = float(val_str)
+ except ValueError:
+ continue
+ rows.append(
+ {
+ COL_TIME: year,
+ COL_COMPONENT: "lnd",
+ COL_QUANTITY: "water",
+ COL_TERM: f"{row_label}_{pool_names[i]}",
+ COL_VALUE: val,
+ COL_UNITS: "kg/m2*1e6",
+ COL_SOURCE: "lnd",
+ COL_PERIOD: period,
+ COL_TABLE_TYPE: "state",
+ }
+ )
+
+ # Emit TOTAL
+ if total_val is not None:
+ rows.append(
+ {
+ COL_TIME: year,
+ COL_COMPONENT: "lnd",
+ COL_QUANTITY: "water",
+ COL_TERM: f"{row_label}_TOTAL",
+ COL_VALUE: total_val,
+ COL_UNITS: "kg/m2*1e6",
+ COL_SOURCE: "lnd",
+ COL_PERIOD: period,
+ COL_TABLE_TYPE: "state",
+ }
+ )
+
+ line = f.readline()
+
+ return rows
+
+
+class LndParser(BaseParser):
+ """Parse land log files for water flux and state tables."""
+
+ def __init__(self, frequency: str = "annual") -> None:
+ super().__init__(frequency=frequency)
+
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ rows: List[Dict] = []
+ for fname in sorted(log_files):
+ try:
+ with gzip.open(fname, "rt") as f:
+ for line in f:
+ stripped = line.strip()
+
+ if stripped.startswith(FLUX_HEADER):
+ result = _parse_period_and_time(stripped)
+ if result is None:
+ continue
+ period, time = result
+ if period != self.frequency:
+ continue
+ # Extract year from time for consistent filtering
+ year = int(time)
+ if start_year <= year <= end_year:
+ rows.extend(_parse_flux_table(f, time, period))
+
+ elif stripped.startswith(STATE_HEADER):
+ result = _parse_period_and_time(stripped)
+ if result is None:
+ continue
+ period, time = result
+ if period != self.frequency:
+ continue
+ # Extract year from time for consistent filtering
+ year = int(time)
+ if start_year <= year <= end_year:
+ rows.extend(_parse_state_table(f, time, period))
+
+ except Exception as e:
+ print(f"WARNING: Error processing {fname}: {e}")
+ continue
+
+ if not rows:
+ return pd.DataFrame(columns=COLUMNS)
+ return pd.DataFrame(rows, columns=COLUMNS)
diff --git a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py
new file mode 100644
index 0000000..356bc19
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py
@@ -0,0 +1,398 @@
+"""Ocean log parser — extracts mass conservation checks into tidy DataFrames.
+
+Ocean log format (monthly, inside CONSERVATION CHECKS blocks):
+
+ date: 0001-02-01_00:00:00
+ Conversion factors:
+ Earth area in E3SM: A = 5.1010114020779156E+14 m^2
+ Averaging time interval: dt = 2676600.0000 s, 30.9792 days
+ ...
+ MASS CONSERVATION CHECK
+
+ MASS FLUXES
+ MPAS-Ocean name kg/s (F) coupler name short name kg/m^2/s*1e6 (F/A)
+ frazilFreshwaterFlux -1.28194649E+08 o2x_Fioo_frazil wfreeze -0.25131222
+ ...
+ SUM VOLUME FLUXES -7.63382509E+07 -0.14965317
+
+ CHANGE IN MASS: computed from ocean domain
+ kg
+ Initial mass 1.36527973E+21
+ Final mass 1.36527952E+21
+ Mass change -2.04326962E+14
+
+ MASS CONSERVATION SUMMARY
+ kg kg/s kg/m^2/s*1e6
+ Mass change ...
+ Net mass flux ...
+ Absolute mass error ...
+
+ RELATIVE MASS ERROR = ...
+"""
+
+import gzip
+import re
+from typing import Dict, List, Optional, TextIO, Tuple
+
+import pandas as pd
+
+from ..schema import (
+ COL_COMPONENT,
+ COL_PERIOD,
+ COL_QUANTITY,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ COL_TERM,
+ COL_TIME,
+ COL_UNITS,
+ COL_VALUE,
+ COLUMNS,
+)
+from .base import BaseParser
+
+
+def _parse_date(date_str: str) -> Tuple[int, int]:
+ """Parse 'YYYY-MM-DD_HH:MM:SS' -> (year, month).
+
+ The date is printed at the START of the next month,
+ so date 0002-01-01 means the check covers Dec of year 1.
+ We return the year and month of the COVERED period.
+ """
+ match = re.match(r"(\d+)-(\d+)-(\d+)", date_str.strip())
+ if not match:
+ return -1, -1
+ y, m = int(match.group(1)), int(match.group(2))
+ # Roll back one month
+ if m == 1:
+ return y - 1, 12
+ return y, m - 1
+
+
+def _parse_mass_fluxes(f: TextIO) -> List[Tuple[str, float]]:
+ """Parse MASS FLUXES table, return list of (short_name, kg/m2s*1e6 value)."""
+ fluxes: List[Tuple[str, float]] = []
+
+ # Skip blank line + "MASS FLUXES" header + column header line
+ line = f.readline() # blank
+ line = f.readline() # "MASS FLUXES"
+ line = f.readline() # column headers
+
+ line = f.readline()
+ while line and line.strip():
+ parts = line.split()
+ if not parts:
+ break
+ # Last token is the kg/m^2/s*1e6 value, second-to-last is short_name
+ # Format: MPAS_name kg/s_val coupler_name short_name kg/m2s*1e6_val
+ # OR: SUM VOLUME FLUXES kg/s_val [empty coupler/short] kg/m2s*1e6_val
+ try:
+ val = float(parts[-1])
+ except ValueError:
+ line = f.readline()
+ continue
+
+ # Determine term name
+ if "SUM VOLUME FLUXES" in line:
+ term = "SUM VOLUME FLUXES"
+ else:
+ # short_name is the second-to-last token
+ term = parts[-2]
+
+ fluxes.append((term, val))
+ line = f.readline()
+
+ return fluxes
+
+
+def _parse_mass_change(f: TextIO) -> Optional[Dict[str, float]]:
+ """Parse CHANGE IN MASS block, return dict of values in kg."""
+ result: Dict[str, float] = {}
+ # Expect lines like:
+ # Initial mass 1.36527973E+21
+ # Final mass 1.36527952E+21
+ # Mass change -2.04326962E+14
+ for _ in range(4): # header line + 3 data lines
+ line = f.readline()
+ if "Initial mass" in line:
+ result["initial_mass_kg"] = float(line.split()[-1])
+ elif "Final mass" in line:
+ result["final_mass_kg"] = float(line.split()[-1])
+ elif "Mass change" in line:
+ result["mass_change_kg"] = float(line.split()[-1])
+ return result if result else None
+
+
+def _parse_mass_summary(f: TextIO) -> Optional[Dict[str, float]]:
+ """Parse MASS CONSERVATION SUMMARY, return kg/m2s*1e6 values.
+
+ Format (immediately after the MASS CONSERVATION SUMMARY header line):
+ kg kg/s kg/m^2/s*1e6
+ Mass change ... ... -0.14965317
+ Net mass flux ... ... -0.14965317
+ Absolute mass error ... ... 0.00000000
+ """
+ result: Dict[str, float] = {}
+ line = f.readline() # column headers: kg kg/s kg/m^2/s*1e6
+ line = f.readline() # Mass change
+ if "Mass change" in line:
+ result["mass_change"] = float(line.split()[-1])
+ line = f.readline() # Net mass flux
+ if "Net mass flux" in line:
+ result["net_mass_flux"] = float(line.split()[-1])
+ line = f.readline() # Absolute mass error
+ if "Absolute mass error" in line:
+ result["absolute_mass_error"] = float(line.split()[-1])
+ return result if result else None
+
+
+def _parse_heat_fluxes(f: TextIO) -> List[Tuple[str, float]]:
+ """Parse HEAT FLUXES sections (explicit + implicit), return (term, W/m2) pairs.
+
+ Reads through explicit and implicit heat flux tables until
+ 'SUM IMP+EXP HEAT FLUXES' is found.
+ """
+ fluxes: List[Tuple[str, float]] = []
+
+ line = f.readline()
+ while line:
+ stripped = line.strip()
+ if not stripped:
+ line = f.readline()
+ continue
+
+ # Stop after we've captured the combined sum
+ if "SUM IMP+EXP" in line and "HEAT FLUXES" in line:
+ parts = line.split()
+ try:
+ val = float(parts[-1])
+ fluxes.append(("SUM IMP+EXP HEAT FLUXES", val))
+ except ValueError:
+ pass
+ break
+
+ # Skip header/label lines
+ if "HEAT FLUXES" in line or "MPAS-Ocean name" in line:
+ line = f.readline()
+ continue
+
+ # Parse explicit/implicit individual flux rows and SUM lines
+ parts = line.split()
+ if not parts:
+ line = f.readline()
+ continue
+
+ try:
+ val = float(parts[-1])
+ except ValueError:
+ line = f.readline()
+ continue
+
+ if "SUM EXPLICIT" in line:
+ term = "SUM EXPLICIT HEAT FLUXES"
+ elif "SUM IMPLICIT" in line:
+ term = "SUM IMPLICIT HEAT FLUXES"
+ else:
+ # Use short_name (second-to-last) if available, else MPAS name
+ term = parts[-2] if len(parts) >= 3 else parts[0]
+
+ fluxes.append((term, val))
+ line = f.readline()
+
+ return fluxes
+
+
+def _parse_energy_summary(f: TextIO) -> Optional[Dict[str, float]]:
+ """Parse ENERGY CONSERVATION SUMMARY, return W/m^2 values.
+
+ Format:
+ J W (J/dt) W/m^2 (J/dt/A)
+ Energy change ... ... 7.89999048
+ Net energy flux ... ... 7.89999048
+ Absolute energy error ... ... 0.00000000
+ """
+ result: Dict[str, float] = {}
+ line = f.readline() # column headers
+ line = f.readline() # Energy change
+ if "Energy change" in line:
+ result["energy_change"] = float(line.split()[-1])
+ line = f.readline() # Net energy flux
+ if "Net energy flux" in line:
+ result["net_energy_flux"] = float(line.split()[-1])
+ line = f.readline() # Absolute energy error
+ if "Absolute energy error" in line:
+ result["absolute_energy_error"] = float(line.split()[-1])
+ return result if result else None
+
+
+class OcnParser(BaseParser):
+ """Parse ocean log files for mass and energy conservation checks."""
+
+ def __init__(self, frequency: str = "annual") -> None:
+ super().__init__(frequency=frequency)
+
+ def parse_files(
+ self, log_files: List[str], start_year: int, end_year: int
+ ) -> pd.DataFrame:
+ rows: List[Dict] = []
+ for fname in sorted(log_files):
+ try:
+ with gzip.open(fname, "rt") as f:
+ for line in f:
+ if "CONSERVATION CHECKS" in line and "date:" not in line:
+ # Next line has the date
+ date_line = f.readline()
+ if "date:" not in date_line:
+ continue
+ date_str = date_line.split("date:")[1].strip()
+ year, month = _parse_date(date_str)
+ if year < start_year or year > end_year:
+ continue
+ if self.frequency == "monthly":
+ time = year + (month - 0.5) / 12.0
+ else:
+ time = float(year)
+ rows.extend(self._parse_block(f, time))
+ except Exception as e:
+ print(f"WARNING: Error processing {fname}: {e}")
+ continue
+
+ if not rows:
+ return pd.DataFrame(columns=COLUMNS)
+ df = pd.DataFrame(rows, columns=COLUMNS)
+
+ # For annual frequency, aggregate monthly rows to annual means
+ if self.frequency == "annual":
+ group_keys = [
+ COL_TIME,
+ COL_COMPONENT,
+ COL_QUANTITY,
+ COL_TERM,
+ COL_UNITS,
+ COL_SOURCE,
+ COL_TABLE_TYPE,
+ ]
+ df[COL_PERIOD] = "annual"
+ df = df.groupby(group_keys, as_index=False).agg(
+ {COL_VALUE: "mean", COL_PERIOD: "first"}
+ )
+
+ return df
+
+ def _parse_block(self, f: TextIO, time: float) -> List[Dict]:
+ """Parse one CONSERVATION CHECKS block for mass and energy data."""
+ rows: List[Dict] = []
+ found_mass = False
+ found_energy = False
+
+ line = f.readline()
+ while line:
+ if found_mass and found_energy:
+ break
+ if "CONSERVATION CHECKS" in line and "date:" not in line:
+ break
+
+ if "MASS CONSERVATION CHECK" in line and "SUMMARY" not in line:
+ found_mass = True
+ rows.extend(self._parse_mass_section(f, time))
+
+ elif "ENERGY CONSERVATION CHECK" in line and "SUMMARY" not in line:
+ found_energy = True
+ rows.extend(self._parse_energy_section(f, time))
+
+ line = f.readline()
+
+ return rows
+
+ def _parse_mass_section(self, f: TextIO, time: float) -> List[Dict]:
+ """Parse MASS CONSERVATION CHECK: fluxes + summary."""
+ rows: List[Dict] = []
+ base = {
+ COL_TIME: time,
+ COL_COMPONENT: "ocn",
+ COL_QUANTITY: "water",
+ COL_UNITS: "kg/m2s*1e6",
+ COL_SOURCE: "ocn",
+ COL_PERIOD: "monthly",
+ }
+
+ for term, val in _parse_mass_fluxes(f):
+ rows.append(
+ {**base, COL_TERM: term, COL_VALUE: val, COL_TABLE_TYPE: "flux"}
+ )
+
+ line = f.readline()
+ while line:
+ if "MASS CONSERVATION SUMMARY" in line:
+ summary = _parse_mass_summary(f)
+ if summary:
+ if "mass_change" in summary:
+ rows.append(
+ {
+ **base,
+ COL_TERM: "Mass change",
+ COL_VALUE: summary["mass_change"],
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+ if "absolute_mass_error" in summary:
+ rows.append(
+ {
+ **base,
+ COL_TERM: "Absolute mass error",
+ COL_VALUE: summary["absolute_mass_error"],
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+ break
+ elif "SALT CONSERVATION" in line:
+ break
+ line = f.readline()
+
+ return rows
+
+ def _parse_energy_section(self, f: TextIO, time: float) -> List[Dict]:
+ """Parse ENERGY CONSERVATION CHECK: fluxes + summary."""
+ rows: List[Dict] = []
+ base = {
+ COL_TIME: time,
+ COL_COMPONENT: "ocn",
+ COL_QUANTITY: "heat",
+ COL_UNITS: "W/m2",
+ COL_SOURCE: "ocn",
+ COL_PERIOD: "monthly",
+ }
+
+ for term, val in _parse_heat_fluxes(f):
+ rows.append(
+ {**base, COL_TERM: term, COL_VALUE: val, COL_TABLE_TYPE: "flux"}
+ )
+
+ line = f.readline()
+ while line:
+ if "ENERGY CONSERVATION SUMMARY" in line:
+ summary = _parse_energy_summary(f)
+ if summary:
+ if "energy_change" in summary:
+ rows.append(
+ {
+ **base,
+ COL_TERM: "Energy change",
+ COL_VALUE: summary["energy_change"],
+ COL_TABLE_TYPE: "flux",
+ }
+ )
+ if "absolute_energy_error" in summary:
+ rows.append(
+ {
+ **base,
+ COL_TERM: "Absolute energy error",
+ COL_VALUE: summary["absolute_energy_error"],
+ COL_TABLE_TYPE: "diagnostic",
+ }
+ )
+ break
+ elif "RELATIVE ENERGY" in line:
+ break
+ line = f.readline()
+
+ return rows
diff --git a/zppy_interfaces/budget_analysis/normalization.py b/zppy_interfaces/budget_analysis/normalization.py
new file mode 100644
index 0000000..8601823
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/normalization.py
@@ -0,0 +1,79 @@
+"""Normalize the tidy event table for budget analysis."""
+
+import pandas as pd
+
+from .schema import COL_QUANTITY, COL_SOURCE, COL_TABLE_TYPE, COL_UNITS, COL_VALUE
+
+# Seconds per year (365-day calendar)
+SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0
+
+
+def normalize(df: pd.DataFrame) -> pd.DataFrame:
+ """Apply all normalizations. Returns a new DataFrame with added columns.
+
+ Adds 'normalized_value' and 'normalized_units' columns:
+
+ Water:
+ - Flux rates (kg/m2s*1e6) -> mm/yr
+ - Flux integrated (kg/m2*1e6) -> mm
+ - State values (kg/m2*1e6) -> mm
+
+ Heat:
+ - Flux rates kept in W/m2 (cumulative scaling to J/m2 at plot time)
+
+ Carbon:
+ - Flux rates (kg-C/m2s*1e10) -> kg-C/m2*1e10/yr
+ """
+ df = df.copy()
+ df["normalized_value"] = df[COL_VALUE].copy()
+ df["normalized_units"] = df[COL_UNITS].copy()
+
+ # --- Water ---
+ water_mask = df[COL_QUANTITY] == "water"
+
+ # Flux rates: kg/m2s * 1e6 -> mm/yr (ocean), kg/m2s -> mm/yr (ice, atm)
+ water_flux = water_mask & (df[COL_TABLE_TYPE] == "flux")
+
+ # Ice and atm data: kg/m2s -> mm/yr (no 1e6 factor)
+ raw_flux = water_flux & df[COL_SOURCE].isin(["ice", "atm"])
+ df.loc[raw_flux, "normalized_value"] = (
+ df.loc[raw_flux, COL_VALUE] * SECONDS_PER_YEAR
+ )
+
+ # Ocean/other data: kg/m2s*1e6 -> mm/yr (with 1e6 factor)
+ scaled_flux = water_flux & ~df[COL_SOURCE].isin(["ice", "atm"])
+ df.loc[scaled_flux, "normalized_value"] = (
+ df.loc[scaled_flux, COL_VALUE] * SECONDS_PER_YEAR / 1e6
+ )
+ df.loc[water_flux, "normalized_units"] = "mm/yr"
+
+ # Integrated fluxes and states: kg/m2 * 1e6 -> mm (ocean), kg/m2 -> mm (ice, atm)
+ water_integrated = water_mask & df[COL_TABLE_TYPE].isin(
+ ["flux_integrated", "state"]
+ )
+
+ # Ice and atm states: kg/m2 -> mm (no 1e6 factor)
+ raw_integrated = water_integrated & df[COL_SOURCE].isin(["ice", "atm"])
+ df.loc[raw_integrated, "normalized_value"] = df.loc[raw_integrated, COL_VALUE]
+
+ # Ocean/other states: kg/m2 * 1e6 -> mm (with 1e6 factor)
+ scaled_integrated = water_integrated & ~df[COL_SOURCE].isin(["ice", "atm"])
+ df.loc[scaled_integrated, "normalized_value"] = (
+ df.loc[scaled_integrated, COL_VALUE] / 1e6
+ )
+
+ df.loc[water_integrated, "normalized_units"] = "mm"
+
+ # --- Heat ---
+ # Keep W/m2 as-is (no conversion needed)
+ # Cumulative residual plots scale to J/m2 at plot time
+
+ # --- Carbon ---
+ # Flux rates: kg-C/m2s * 1e10 -> kg-C/m2*1e10 /yr
+ carbon_flux = (df[COL_QUANTITY] == "carbon") & (df[COL_TABLE_TYPE] == "flux")
+ df.loc[carbon_flux, "normalized_value"] = (
+ df.loc[carbon_flux, COL_VALUE] * SECONDS_PER_YEAR
+ )
+ df.loc[carbon_flux, "normalized_units"] = "kg-C/m2*1e10/yr"
+
+ return df
diff --git a/zppy_interfaces/budget_analysis/parser.py b/zppy_interfaces/budget_analysis/parser.py
new file mode 100644
index 0000000..7cb42b4
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/parser.py
@@ -0,0 +1,154 @@
+"""E3SM coupler log budget data parsing."""
+
+import gzip
+import re
+from typing import Dict, List, Optional, TextIO
+
+import numpy as np
+import numpy.ma as ma
+
+
+class Budget:
+ """
+ Class to parse and store E3SM coupler budget data from log files.
+
+ Parses budget tables from E3SM coupler log files and stores data
+ in structured arrays for analysis and visualization.
+ """
+
+ def __init__(self, header: str, years: np.ndarray):
+ """
+ Initialize Budget object.
+
+ Args:
+ header: Budget header string to search for in log files
+ years: Array of years to process
+ """
+ # Identifying header
+ self.header = header
+
+ # Years to save
+ self.years = years
+ # Dictionary to look up year indices
+ self.iyear = {key: i for i, key in enumerate(self.years)}
+
+ # Extract units
+ self.units = re.findall(r"\(([^)]+)", self.header)[1]
+
+ # To be defined later
+ self.cols: Optional[List[str]] = None
+ self.rows: Optional[List[str]] = None
+ self.icol: Optional[Dict[str, int]] = None
+ self.irow: Optional[Dict[str, int]] = None
+ self.data: Optional[ma.MaskedArray] = None
+
+ def parse(self, f: TextIO, datestamp: str) -> None:
+ """
+ Parse budget data from file for given datestamp.
+
+ Args:
+ f: Open file object
+ datestamp: Date stamp string from log file
+ """
+ # Check if year is within range
+ year = int(datestamp[:-4]) - 1
+ if year not in self.iyear:
+ return
+
+ # Read table header, extract column names
+ tmp = f.readline().strip()
+ cols = re.split(r"\s{2,}", tmp)
+
+ # Store or check column names
+ if self.cols is None:
+ self.cols = cols
+ self.icol = {key: i for i, key in enumerate(cols)}
+ elif self.cols != cols:
+ print("ERROR: cols mismatched")
+
+ # Read table rows
+ lines = []
+ tmp = f.readline()
+ while tmp.strip():
+ parts = tmp.split()
+ # Handle multi-word row names (e.g., "surface co2", "black carbon")
+ # Find where the numeric data starts
+ row_name_parts = []
+ data_start = 0
+ for j, part in enumerate(parts):
+ try:
+ float(part)
+ data_start = j
+ break
+ except ValueError:
+ row_name_parts.append(part)
+ row_name = " ".join(row_name_parts)
+ data_values = parts[data_start:]
+ lines.append((row_name, data_values))
+ tmp = f.readline()
+
+ # Store or check row names
+ rows = [line[0] for line in lines]
+ if self.rows is None:
+ self.rows = rows
+ self.irow = {key: i for i, key in enumerate(rows)}
+ elif self.rows != rows:
+ print("ERROR: rows mismatched")
+
+ # Store in 3d array
+ if self.data is None:
+ self.data = ma.masked_all((len(self.years), len(self.rows), len(self.cols)))
+ iyear = self.iyear[year]
+
+ for i, (row_name, data_values) in enumerate(lines):
+ try:
+ # Convert string values to float
+ numeric_values = [float(v) for v in data_values]
+ self.data[iyear, i, :] = numeric_values
+ except (ValueError, TypeError) as e:
+ print(f"ERROR converting row '{row_name}' values {data_values}: {e}")
+ # Keep as masked values if conversion fails
+
+ return
+
+
+def parse_budget_types(budget_types_str: str) -> list[str]:
+ """Parse comma-separated budget types string."""
+ return [bt.strip() for bt in budget_types_str.split(",") if bt.strip()]
+
+
+def initialize_budgets(budget_types: list[str], years: np.ndarray) -> dict[str, Budget]:
+ """Initialize budget objects for specified types and years."""
+ budget_headers = {
+ "area": "(seq_diag_print_mct) NET AREA BUDGET (m2/m2): period = annual: date =",
+ "water": "(seq_diag_print_mct) NET WATER BUDGET (kg/m2s*1e6): period = annual: date =",
+ "heat": "(seq_diag_print_mct) NET HEAT BUDGET (W/m2): period = annual: date =",
+ "carbon": "(seq_diagBGC_print_mct) NET CARBON BUDGET (kg-C/m2s*1e10): period = annual: date =",
+ }
+
+ budgets = {}
+ for budget_type in budget_types:
+ if budget_type in budget_headers:
+ budgets[budget_type] = Budget(budget_headers[budget_type], years)
+ else:
+ print(f"WARNING: Unknown budget type {budget_type}, skipping")
+
+ return budgets
+
+
+def process_log_files(log_files: list[str], budgets: dict[str, Budget]) -> None:
+ """Process coupler log files and extract budget data."""
+ for fname in log_files:
+ print(f"Processing {fname}")
+ try:
+ with gzip.open(fname, "rt") as f:
+ line = f.readline()
+ while line != "":
+ for budget_type, budget_obj in budgets.items():
+ if line.startswith(budget_obj.header):
+ datestamp = line.replace(budget_obj.header, "").split()[0]
+ budget_obj.parse(f, datestamp)
+ line = f.readline()
+ except Exception as e:
+ print(f"ERROR processing {fname}: {e}")
+ continue
diff --git a/zppy_interfaces/budget_analysis/plotting.py b/zppy_interfaces/budget_analysis/plotting.py
new file mode 100644
index 0000000..753564a
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/plotting.py
@@ -0,0 +1,158 @@
+"""Budget visualization functions using Bokeh and matplotlib."""
+
+import os
+from typing import Dict, List
+
+import numpy as np
+
+from .parser import Budget
+
+# Seconds per year (365 days)
+DT_SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0
+
+# Unit conversion factors
+# Water: kg/m2s*1e6 -> mm (since 1 kg/m2 = 1 mm, multiply by dt/1e6)
+WATER_CONVERSION = DT_SECONDS_PER_YEAR / 1e6
+# Energy: W/m2 -> J/m2 *1e9 (since W = J/s, multiply by dt/1e9)
+ENERGY_CONVERSION = DT_SECONDS_PER_YEAR / 1e9
+# Carbon: kg-C/m2s*1e10 -> kg-C/m2 (multiply by dt/1e10)
+CARBON_CONVERSION = DT_SECONDS_PER_YEAR / 1e10
+
+
+def generate_html_plots( # noqa: C901
+ budgets: Dict[str, Budget], budget_names: List[str], output_dir: str
+) -> None:
+ """Generate interactive HTML plots using bokeh."""
+ try:
+ import itertools
+
+ from bokeh.layouts import column
+ from bokeh.models import ColumnDataSource, Legend
+ from bokeh.palettes import Category10
+ from bokeh.plotting import figure, output_file, save
+ except ImportError as e:
+ print(f"ERROR: bokeh package not available: {e}")
+ return
+
+ # Also try to import matplotlib for PNG fallback
+ try:
+ import matplotlib.pyplot as plt
+
+ matplotlib_available = True
+ except ImportError:
+ matplotlib_available = False
+
+ # Create budget plots
+ plots = []
+ for budget_name in budget_names:
+ if budget_name not in budgets:
+ continue
+
+ if budgets[budget_name].data is None:
+ continue
+
+ b = budgets[budget_name]
+
+ # List of colors for plots
+ colors = itertools.cycle(Category10[10])
+
+ # Determine unit conversion factor based on budget type
+ if budget_name == "water":
+ conversion_factor = WATER_CONVERSION
+ converted_units = "mm"
+ elif budget_name == "heat":
+ conversion_factor = ENERGY_CONVERSION
+ converted_units = "J/m2 *1e9"
+ elif budget_name == "carbon":
+ conversion_factor = CARBON_CONVERSION
+ converted_units = "kg-C/m2"
+ else:
+ conversion_factor = 1.0
+ converted_units = b.units
+
+ # Create ColumnDataSource
+ data = {}
+ data["years"] = b.years
+
+ for krow, vrow in b.irow.items(): # type: ignore
+ for kcol, vcol in b.icol.items(): # type: ignore
+ raw_data = b.data[:, vrow, vcol] # type: ignore
+ # Apply unit conversion and compute cumulative sum
+ converted_data = raw_data * conversion_factor
+ cumsum_data = np.cumsum(converted_data) - converted_data[0]
+ data[krow + "_" + kcol] = cumsum_data
+ source = ColumnDataSource(data=data)
+
+ # Determine which row to plot
+ if "*SUM*" in b.irow: # type: ignore
+ row_name = "*SUM*"
+ plot_title = f"*SUM* annual cumulative {budget_name} budget"
+ else:
+ row_name = list(b.irow.keys())[0] # type: ignore
+ plot_title = f"{row_name} annual cumulative {budget_name} budget"
+
+ # Create Bokeh plot
+ p = figure(
+ title=plot_title,
+ height=400,
+ width=1200,
+ x_axis_label="year",
+ y_axis_label=f"{budget_name} budget ({converted_units})",
+ )
+ p.add_layout(Legend(), "right")
+
+ # Add lines to Bokeh plot
+ for k, v in b.icol.items(): # type: ignore
+ line_name = f"{row_name}_{k}"
+ if line_name in data:
+ p.line(
+ x="years",
+ y=line_name,
+ legend_label=k,
+ line_width=2,
+ color=next(colors),
+ source=source,
+ )
+ p.legend.click_policy = "hide"
+ plots.append(p)
+
+ # Create matplotlib PNG fallback
+ if matplotlib_available:
+ plt.figure(figsize=(12, 6))
+
+ colors_mpl = plt.cm.tab10(np.linspace(0, 1, len(b.icol))) # type: ignore
+ for i, (k, v) in enumerate(b.icol.items()): # type: ignore
+ line_name = f"{row_name}_{k}"
+ if line_name in data:
+ plt.plot(
+ data["years"],
+ data[line_name],
+ label=k,
+ linewidth=2,
+ color=colors_mpl[i],
+ )
+
+ plt.title(plot_title)
+ plt.xlabel("year")
+ plt.ylabel(f"{budget_name} budget ({converted_units})")
+ plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
+ plt.grid(True, alpha=0.3)
+ plt.tight_layout()
+
+ png_file = os.path.join(output_dir, f"budget_{budget_name}.png")
+ plt.savefig(png_file, dpi=150, bbox_inches="tight")
+ plt.close()
+
+ if plots:
+ # Save Bokeh HTML
+ html_file = os.path.join(output_dir, "budgets.html")
+
+ c = column(children=plots, sizing_mode="stretch_width")
+ output_file(html_file)
+ save(c)
+
+ print(f"Interactive HTML plots written to {html_file}")
+ if matplotlib_available:
+ print(f"PNG fallback plots also created in {output_dir}")
+ else:
+ print("No plots generated - no valid budget data found")
diff --git a/zppy_interfaces/budget_analysis/schema.py b/zppy_interfaces/budget_analysis/schema.py
new file mode 100644
index 0000000..068883d
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/schema.py
@@ -0,0 +1,35 @@
+"""Column schema for the tidy budget event table."""
+
+from typing import List
+
+import pandas as pd
+
+# Column name constants
+COL_TIME = "time" # float: year (annual) or year + (month-0.5)/12 (monthly)
+COL_COMPONENT = (
+ "component" # str: "atm", "lnd", "rof", "ocn", "ice_nh", "ice_sh", "glc", "*SUM*"
+)
+COL_QUANTITY = "quantity" # str: "water" (later "heat", "carbon")
+COL_TERM = "term" # str: flux or state term name
+COL_VALUE = "value" # float64: raw value in original units
+COL_UNITS = "units" # str: original units string
+COL_SOURCE = "source" # str: "cpl", "lnd", etc.
+COL_PERIOD = "period" # str: "annual" or "monthly"
+COL_TABLE_TYPE = "table_type" # str: "flux", "flux_integrated", "state"
+
+COLUMNS: List[str] = [
+ COL_TIME,
+ COL_COMPONENT,
+ COL_QUANTITY,
+ COL_TERM,
+ COL_VALUE,
+ COL_UNITS,
+ COL_SOURCE,
+ COL_PERIOD,
+ COL_TABLE_TYPE,
+]
+
+
+def empty_event_table() -> pd.DataFrame:
+ """Return an empty DataFrame with the correct schema."""
+ return pd.DataFrame(columns=COLUMNS)
diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py
new file mode 100644
index 0000000..a22354d
--- /dev/null
+++ b/zppy_interfaces/budget_analysis/viz.py
@@ -0,0 +1,258 @@
+"""Visualization for whole-model budget analysis.
+
+Generates an HTML report with Bokeh plots for each budget check result.
+"""
+
+import os
+from typing import Dict, List
+
+import numpy as np
+import pandas as pd
+from bokeh.layouts import column
+from bokeh.models import Div, FixedTicker
+from bokeh.palettes import Category10
+from bokeh.plotting import figure, output_file, save
+
+from .checks import CheckResult
+from .normalization import SECONDS_PER_YEAR
+
+# Units per quantity for plot labels
+QUANTITY_UNITS: Dict[str, Dict[str, str]] = {
+ "water": {"flux": "mm/yr", "cumulative": "mm"},
+ "heat": {"flux": "W/m2", "cumulative": "J/m2*1e9"},
+ "carbon": {"flux": "kg-C/m2s*1e10", "cumulative": "kg-C/m2*1e10"},
+}
+
+
+def _cum_scale(quantity: str) -> float:
+ """Scale factor for cumulative plots.
+
+ Heat: W/m2 cumsum over years -> J/m2 (multiply by seconds_per_year).
+ Water: mm/yr cumsum over years -> mm (no scaling needed).
+ """
+ if quantity == "heat":
+ return SECONDS_PER_YEAR / 1e9
+ return 1.0
+
+
+def generate_budget_report(
+ results: List[CheckResult],
+ df: pd.DataFrame,
+ output_dir: str,
+ quantity: str = "water",
+) -> str:
+ """Generate an HTML report with budget check plots.
+
+ Returns path to the HTML file.
+ """
+ units = QUANTITY_UNITS.get(quantity, {"flux": "", "cumulative": ""})
+ flux_units = units["flux"]
+ cum_units = units["cumulative"]
+ scale = _cum_scale(quantity)
+ plots: list = []
+
+ for r in results:
+ if r.name == f"cpl_{quantity}_component_fluxes":
+ plots.append(Div(text=f"{quantity.title()} Budget Overview
"))
+ plots.append(_plot_cumulative_components(r, quantity, cum_units, scale))
+
+ elif r.name.endswith("_interface_match"):
+ # name format: {component}_{quantity}_interface_match
+ comp = r.name.replace(f"_{quantity}_interface_match", "")
+ plots.append(Div(text=f"Interface Match: {comp}
"))
+ plots.append(
+ _plot_comparison(
+ r,
+ f"{comp} {quantity.title()} Flux (cpl vs {comp})",
+ flux_units,
+ )
+ )
+ plots.append(
+ _plot_residual(r, f"Interface Residual (cpl - {comp})", flux_units)
+ )
+ plots.append(
+ _plot_cumulative(
+ r,
+ f"Interface Cumulative Residual (cpl - {comp})",
+ cum_units,
+ scale,
+ )
+ )
+
+ elif r.name == "lnd_closure":
+ plots.append(Div(text="Land Water Closure
"))
+ plots.append(_plot_comparison(r, "Land ΔStorage vs ∫Flux dt", cum_units))
+ plots.append(_plot_residual(r, "Closure Residual", cum_units))
+ plots.append(_plot_cumulative(r, "Closure Cumulative Residual", cum_units))
+
+ elif r.name.startswith("ocn_") and r.name.endswith("_closure"):
+ label = "Water" if "water" in r.name else "Heat"
+ change_label = "ΔMass" if label == "Water" else "ΔEnergy"
+ plots.append(Div(text=f"Ocean {label} Closure
"))
+ plots.append(
+ _plot_comparison(r, f"Ocean {change_label} vs Net Flux", flux_units)
+ )
+ plots.append(_plot_residual(r, "Ocean Closure Residual", flux_units))
+ plots.append(
+ _plot_cumulative(
+ r,
+ "Ocean Closure Cumulative Residual",
+ cum_units,
+ scale,
+ )
+ )
+
+ elif r.name.startswith("ice_") and r.name.endswith("_closure"):
+ label = "Water" if "water" in r.name else "Heat"
+ change_label = "ΔMass" if label == "Water" else "ΔEnergy"
+ plots.append(Div(text=f"Ice {label} Closure
"))
+ plots.append(
+ _plot_comparison(r, f"Ice {change_label} vs Net Flux", flux_units)
+ )
+ plots.append(_plot_residual(r, "Ice Closure Residual", flux_units))
+ plots.append(
+ _plot_cumulative(
+ r,
+ "Ice Closure Cumulative Residual",
+ cum_units,
+ scale,
+ )
+ )
+
+ elif r.name.startswith("atm_") and r.name.endswith("_closure"):
+ label = "Water" if "water" in r.name else "Heat"
+ change_label = "ΔMass" if label == "Water" else "ΔEnergy"
+ plots.append(Div(text=f"Atmosphere {label} Closure
"))
+ plots.append(
+ _plot_comparison(r, f"Atm {change_label} vs Net Flux", flux_units)
+ )
+ plots.append(_plot_residual(r, "Atm Closure Residual", flux_units))
+ plots.append(
+ _plot_cumulative(
+ r,
+ "Atm Closure Cumulative Residual",
+ cum_units,
+ scale,
+ )
+ )
+
+ if not plots:
+ print(f"No {quantity} plots generated — no check results available")
+ return ""
+
+ html_path = os.path.join(output_dir, f"{quantity}_budget_report.html")
+ output_file(html_path, title=f"E3SM {quantity.title()} Budget Analysis")
+ save(column(plots, sizing_mode="stretch_width"))
+ print(f"Report written to {html_path}")
+ return html_path
+
+
+def generate_landing_page(output_dir: str, report_paths: Dict[str, str]) -> str:
+ """Generate an index.html landing page linking to individual budget reports.
+
+ Returns path to the landing page.
+ """
+ links = []
+ for quantity, path in sorted(report_paths.items()):
+ filename = os.path.basename(path)
+ links.append(
+ f'{quantity.title()} Budget Report'
+ )
+ html = f"""
+
+E3SM Budget Analysis
+
+E3SM Budget Analysis
+
+
+"""
+ index_path = os.path.join(output_dir, "index.html")
+ with open(index_path, "w") as f:
+ f.write(html)
+ print(f"Landing page written to {index_path}")
+ return index_path
+
+
+def _make_figure(title: str, y_label: str, years: np.ndarray) -> figure:
+ p = figure(
+ title=title,
+ height=350,
+ width=1200,
+ x_axis_label="time (year)",
+ y_axis_label=y_label,
+ )
+ # For monthly data, force ticks at integer years so Bokeh doesn't
+ # pick fractional intervals (e.g. 0.2) that don't align with months.
+ # For annual data, let Bokeh auto-tick (handles long runs like 500 yrs).
+ if np.any(years != np.floor(years)):
+ int_years = sorted(set(int(y) for y in years))
+ p.xaxis.ticker = FixedTicker(ticks=int_years)
+ return p
+
+
+def _plot_residual(r: CheckResult, title: str, units: str) -> figure:
+ """Plot residual time series with a zero reference line."""
+ p = _make_figure(title, f"residual ({units})", r.years)
+ p.line(r.years, r.residual, line_width=2, color="red")
+ p.line(
+ r.years,
+ np.zeros_like(r.years),
+ line_width=1,
+ color="gray",
+ line_dash="dashed",
+ )
+ return p
+
+
+def _plot_cumulative(
+ r: CheckResult, title: str, units: str, scale: float = 1.0
+) -> figure:
+ """Plot cumulative residual, optionally scaled for unit conversion."""
+ p = _make_figure(title, f"cumulative residual ({units})", r.years)
+ p.line(r.years, r.cumulative_residual * scale, line_width=2, color="darkred")
+ return p
+
+
+def _plot_comparison(r: CheckResult, title: str, units: str) -> figure:
+ """Plot LHS and RHS on the same axes."""
+ p = _make_figure(title, units, r.years)
+ p.line(r.years, r.lhs, line_width=2, color="blue", legend_label=r.lhs_label)
+ p.line(r.years, r.rhs, line_width=2, color="orange", legend_label=r.rhs_label)
+ p.legend.click_policy = "hide"
+ return p
+
+
+def _plot_cumulative_components(
+ r: CheckResult, quantity: str, units: str, scale: float = 1.0
+) -> figure:
+ """Cumulative net flux per component, with *SUM* residual highlighted."""
+ p = _make_figure(
+ f"Cumulative Net {quantity.title()} Flux per Component", units, r.years
+ )
+ if r.components is None:
+ return p
+ # Plot component lines, highlight *SUM* as thick dashed red
+ other_names = sorted(k for k in r.components if k != "*SUM*")
+ colors = Category10[max(3, len(other_names) + 1)]
+ for i, name in enumerate(other_names):
+ p.line(
+ r.years,
+ r.components[name] * scale,
+ line_width=2,
+ color=colors[i % len(colors)],
+ legend_label=name,
+ )
+ if "*SUM*" in r.components:
+ p.line(
+ r.years,
+ r.components["*SUM*"] * scale,
+ line_width=3,
+ color="red",
+ line_dash="dashed",
+ legend_label="*SUM* (residual)",
+ )
+ p.legend.click_policy = "hide"
+ p.legend.location = "top_left"
+ return p