From a9800636ad2f87bf70f651595e211f8f55c3e308 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 29 Jan 2026 18:48:29 -0600 Subject: [PATCH 01/15] Add E3SM Budget Analysis CLI Tool --- pyproject.toml | 2 + zppy_interfaces/budget_analysis/__init__.py | 1 + zppy_interfaces/budget_analysis/__main__.py | 138 ++++++++++++++++ zppy_interfaces/budget_analysis/parser.py | 140 ++++++++++++++++ zppy_interfaces/budget_analysis/plotting.py | 168 ++++++++++++++++++++ 5 files changed, 449 insertions(+) create mode 100644 zppy_interfaces/budget_analysis/__init__.py create mode 100644 zppy_interfaces/budget_analysis/__main__.py create mode 100644 zppy_interfaces/budget_analysis/parser.py create mode 100644 zppy_interfaces/budget_analysis/plotting.py diff --git a/pyproject.toml b/pyproject.toml index 6a04641..f3a4985 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ classifiers = [ dependencies = [ "beautifulsoup4", + "bokeh", "lxml", "matplotlib", "netcdf4", @@ -117,6 +118,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/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..6c1efae --- /dev/null +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -0,0 +1,138 @@ +#!/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_ascii_output, 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)", + ) + 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( + "--output_ascii", + action="store_true", + default=True, + help="Generate ASCII summary tables (default: True)", + ) + + args = parser.parse_args() + + # Validate inputs + if args.start_year > args.end_year: + print("ERROR: start_year must be <= end_year") + return 1 + + # Use provided log path + log_path = args.log_path + + if not os.path.exists(log_path): + print(f"ERROR: Log path does not exist: {log_path}") + return 1 + + # Parse budget types + budget_types = parse_budget_types(args.budget_types) + valid_types = ["area", "water", "heat"] + for bt in budget_types: + if bt not in valid_types: + print(f"ERROR: Invalid budget type '{bt}'. Valid types: {valid_types}") + return 1 + + # Create output directory if it doesn't exist + os.makedirs(args.output_dir, exist_ok=True) + + print("E3SM Budget Analysis Tool") + print("=========================") + print(f"Years: {args.start_year} to {args.end_year}") + print(f"Budget types: {budget_types}") + print(f"Log path: {log_path}") + print(f"Output directory: {args.output_dir}") + + # Set up years array + years = np.arange(args.start_year, args.end_year + 1) + + # Initialize budget objects + budgets = initialize_budgets(budget_types, years) + + # Find and process log files + log_pattern = os.path.join(log_path, "cpl.log.*.gz") + log_files = sorted(glob.glob(log_pattern)) + + if not log_files: + print(f"ERROR: No coupler log files found at {log_pattern}") + return 1 + + print(f"Found {len(log_files)} coupler log files") + + # Process log files + process_log_files(log_files, budgets) + + # Generate outputs + print("\nGenerating output files...") + + # Generate ASCII summaries if requested + if args.output_ascii: + for budget_type, budget_obj in budgets.items(): + generate_ascii_output(budget_obj, budget_type, args.output_dir) + + # Generate HTML plots if requested + if args.output_html: + generate_html_plots(budgets, budget_types, args.output_dir) + + print("\nBudget analysis completed successfully!") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/zppy_interfaces/budget_analysis/parser.py b/zppy_interfaces/budget_analysis/parser.py new file mode 100644 index 0000000..6f9e7d8 --- /dev/null +++ b/zppy_interfaces/budget_analysis/parser.py @@ -0,0 +1,140 @@ +"""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(): + lines.append(tmp.split()) + 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, values in enumerate(lines): + data_values = values[1:] + 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 '{values[0]}' 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 =", + } + + 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..f35502b --- /dev/null +++ b/zppy_interfaces/budget_analysis/plotting.py @@ -0,0 +1,168 @@ +"""Budget visualization functions using Bokeh and matplotlib.""" + +import os +from typing import Dict, List + +import numpy as np + +from .parser import Budget + + +def generate_ascii_output( + budget_obj: Budget, budget_name: str, output_dir: str +) -> None: + """Generate ASCII summary table for a budget.""" + if budget_obj.data is None: + print(f"No data available for {budget_name} budget") + return + + # Calculate average over the period + avg = np.average(budget_obj.data[:, :, :], axis=0) + + # Generate output filename + filename = os.path.join(output_dir, f"{budget_name}_budget_summary.txt") + + with open(filename, "w") as f: + f.write(f"----- Average {budget_name} budget years {budget_obj.years[0]:04d} ") + f.write(f"to {budget_obj.years[-1]:04d} ({budget_obj.units}) -----\n") + f.write("\n") + + # Write header + ncols = len(budget_obj.cols) # type: ignore + header_line = f"{'':10s}" + "".join([f"{col:>12s} " for col in budget_obj.cols]) # type: ignore + f.write(header_line + "\n") + + # Write data rows + for row in budget_obj.rows: # type: ignore + irow = budget_obj.irow[row] # type: ignore + data_line = f"{row:10s}" + "".join( + [f"{avg[irow, i]:12.6f} " for i in range(ncols)] + ) + f.write(data_line + "\n") + + f.write("-" * 60 + "\n") + + print(f"ASCII summary written to {filename}") + + +def generate_html_plots( + 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]) + + # 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 + cumsum_data = np.cumsum(raw_data) - raw_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 ({b.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 ({b.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") From ea38d5ee58f728c98e15abac57c23d2a5296282d Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Thu, 5 Feb 2026 16:44:44 -0800 Subject: [PATCH 02/15] update units; add carbon budgets --- zppy_interfaces/budget_analysis/__main__.py | 17 +----- zppy_interfaces/budget_analysis/parser.py | 22 +++++-- zppy_interfaces/budget_analysis/plotting.py | 66 +++++++++------------ 3 files changed, 49 insertions(+), 56 deletions(-) diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index 6c1efae..4915850 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -18,7 +18,7 @@ import numpy as np from .parser import initialize_budgets, parse_budget_types, process_log_files -from .plotting import generate_ascii_output, generate_html_plots +from .plotting import generate_html_plots def main() -> int: @@ -47,7 +47,7 @@ def main() -> int: parser.add_argument( "--budget_types", default="water,heat", - help="Comma-separated list of budget types to analyze (water,heat)", + help="Comma-separated list of budget types to analyze (water,heat,carbon)", ) parser.add_argument( "--output_dir", @@ -60,12 +60,6 @@ def main() -> int: default=True, help="Generate HTML plots (default: True)", ) - parser.add_argument( - "--output_ascii", - action="store_true", - default=True, - help="Generate ASCII summary tables (default: True)", - ) args = parser.parse_args() @@ -83,7 +77,7 @@ def main() -> int: # Parse budget types budget_types = parse_budget_types(args.budget_types) - valid_types = ["area", "water", "heat"] + 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}") @@ -121,11 +115,6 @@ def main() -> int: # Generate outputs print("\nGenerating output files...") - # Generate ASCII summaries if requested - if args.output_ascii: - for budget_type, budget_obj in budgets.items(): - generate_ascii_output(budget_obj, budget_type, args.output_dir) - # Generate HTML plots if requested if args.output_html: generate_html_plots(budgets, budget_types, args.output_dir) diff --git a/zppy_interfaces/budget_analysis/parser.py b/zppy_interfaces/budget_analysis/parser.py index 6f9e7d8..7cb42b4 100644 --- a/zppy_interfaces/budget_analysis/parser.py +++ b/zppy_interfaces/budget_analysis/parser.py @@ -70,7 +70,21 @@ def parse(self, f: TextIO, datestamp: str) -> None: lines = [] tmp = f.readline() while tmp.strip(): - lines.append(tmp.split()) + 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 @@ -86,14 +100,13 @@ def parse(self, f: TextIO, datestamp: str) -> None: self.data = ma.masked_all((len(self.years), len(self.rows), len(self.cols))) iyear = self.iyear[year] - for i, values in enumerate(lines): - data_values = values[1:] + 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 '{values[0]}' values {data_values}: {e}") + print(f"ERROR converting row '{row_name}' values {data_values}: {e}") # Keep as masked values if conversion fails return @@ -110,6 +123,7 @@ def initialize_budgets(budget_types: list[str], years: np.ndarray) -> dict[str, "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 = {} diff --git a/zppy_interfaces/budget_analysis/plotting.py b/zppy_interfaces/budget_analysis/plotting.py index f35502b..8a73bae 100644 --- a/zppy_interfaces/budget_analysis/plotting.py +++ b/zppy_interfaces/budget_analysis/plotting.py @@ -7,42 +7,16 @@ from .parser import Budget +# Seconds per year (365 days) +DT_SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0 -def generate_ascii_output( - budget_obj: Budget, budget_name: str, output_dir: str -) -> None: - """Generate ASCII summary table for a budget.""" - if budget_obj.data is None: - print(f"No data available for {budget_name} budget") - return - - # Calculate average over the period - avg = np.average(budget_obj.data[:, :, :], axis=0) - - # Generate output filename - filename = os.path.join(output_dir, f"{budget_name}_budget_summary.txt") - - with open(filename, "w") as f: - f.write(f"----- Average {budget_name} budget years {budget_obj.years[0]:04d} ") - f.write(f"to {budget_obj.years[-1]:04d} ({budget_obj.units}) -----\n") - f.write("\n") - - # Write header - ncols = len(budget_obj.cols) # type: ignore - header_line = f"{'':10s}" + "".join([f"{col:>12s} " for col in budget_obj.cols]) # type: ignore - f.write(header_line + "\n") - - # Write data rows - for row in budget_obj.rows: # type: ignore - irow = budget_obj.irow[row] # type: ignore - data_line = f"{row:10s}" + "".join( - [f"{avg[irow, i]:12.6f} " for i in range(ncols)] - ) - f.write(data_line + "\n") - - f.write("-" * 60 + "\n") - - print(f"ASCII summary written to {filename}") +# 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( @@ -82,6 +56,20 @@ def generate_html_plots( # 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 @@ -89,7 +77,9 @@ def generate_html_plots( 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 - cumsum_data = np.cumsum(raw_data) - raw_data[0] + # 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) @@ -107,7 +97,7 @@ def generate_html_plots( height=400, width=1200, x_axis_label="year", - y_axis_label=f"{budget_name} budget ({b.units})", + y_axis_label=f"{budget_name} budget ({converted_units})", ) p.add_layout(Legend(), "right") @@ -144,7 +134,7 @@ def generate_html_plots( plt.title(plot_title) plt.xlabel("year") - plt.ylabel(f"{budget_name} budget ({b.units})") + 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() From 3b3d79c80610e05eba68cac7ee7ffe6f5def6ed4 Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Fri, 6 Feb 2026 21:15:50 -0800 Subject: [PATCH 03/15] add water budget for lnd ocn and cpl --- pyproject.toml | 1 + zppy_interfaces/budget_analysis/__main__.py | 103 ++++-- zppy_interfaces/budget_analysis/checks.py | 285 ++++++++++++++++ .../budget_analysis/ingestion/__init__.py | 0 .../budget_analysis/ingestion/base.py | 17 + .../budget_analysis/ingestion/cpl_parser.py | 145 +++++++++ .../budget_analysis/ingestion/lnd_parser.py | 308 ++++++++++++++++++ .../budget_analysis/ingestion/ocn_parser.py | 245 ++++++++++++++ .../budget_analysis/normalization.py | 38 +++ zppy_interfaces/budget_analysis/plotting.py | 2 +- zppy_interfaces/budget_analysis/schema.py | 35 ++ zppy_interfaces/budget_analysis/viz.py | 131 ++++++++ 12 files changed, 1285 insertions(+), 25 deletions(-) create mode 100644 zppy_interfaces/budget_analysis/checks.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/__init__.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/base.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/cpl_parser.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/lnd_parser.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/ocn_parser.py create mode 100644 zppy_interfaces/budget_analysis/normalization.py create mode 100644 zppy_interfaces/budget_analysis/schema.py create mode 100644 zppy_interfaces/budget_analysis/viz.py diff --git a/pyproject.toml b/pyproject.toml index f3a4985..5dc0de0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "matplotlib", "netcdf4", "numpy >=2.0,<3.0", + "pandas", "pcmdi_metrics>=3.9.3", "xarray >=2023.02.0", "xcdat >=0.7.3,<1.0", diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index 4915850..5d1cf34 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -60,6 +60,12 @@ def main() -> int: 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)", + ) args = parser.parse_args() @@ -68,14 +74,21 @@ def main() -> int: print("ERROR: start_year must be <= end_year") return 1 - # Use provided log path log_path = args.log_path - if not os.path.exists(log_path): print(f"ERROR: Log path does not exist: {log_path}") return 1 - # Parse budget types + 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: @@ -83,39 +96,23 @@ def main() -> int: print(f"ERROR: Invalid budget type '{bt}'. Valid types: {valid_types}") return 1 - # Create output directory if it doesn't exist - os.makedirs(args.output_dir, exist_ok=True) - - print("E3SM Budget Analysis Tool") - print("=========================") + 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: {log_path}") - print(f"Output directory: {args.output_dir}") + print(f"Log path: {args.log_path}") - # Set up years array years = np.arange(args.start_year, args.end_year + 1) - - # Initialize budget objects budgets = initialize_budgets(budget_types, years) - # Find and process log files - log_pattern = os.path.join(log_path, "cpl.log.*.gz") - log_files = sorted(glob.glob(log_pattern)) - + log_files = sorted(glob.glob(os.path.join(args.log_path, "cpl.log.*.gz"))) if not log_files: - print(f"ERROR: No coupler log files found at {log_pattern}") + print("ERROR: No coupler log files found") return 1 print(f"Found {len(log_files)} coupler log files") - - # Process log files process_log_files(log_files, budgets) - # Generate outputs - print("\nGenerating output files...") - - # Generate HTML plots if requested if args.output_html: generate_html_plots(budgets, budget_types, args.output_dir) @@ -123,5 +120,63 @@ def main() -> int: return 0 +def _run_whole_model(args) -> int: + """Whole-model budget pipeline: ingest -> normalize -> check -> visualize.""" + import pandas as pd + + from .checks import run_checks + from .ingestion.cpl_parser import CplParser + from .ingestion.lnd_parser import LndParser + from .ingestion.ocn_parser import OcnParser + from .normalization import normalize + from .viz import generate_budget_report + + print("E3SM Budget Analysis Tool (whole-model mode)") + print("=============================================") + print(f"Years: {args.start_year} to {args.end_year}") + 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"))) + + 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") + + frames = [] + frames.append(CplParser().parse_files(cpl_files, args.start_year, args.end_year)) + if lnd_files: + frames.append( + LndParser().parse_files(lnd_files, args.start_year, args.end_year) + ) + if ocn_files: + frames.append( + OcnParser().parse_files(ocn_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 + print("\nRunning budget checks...") + results = run_checks(events) + + # Visualize + print("\nGenerating report...") + html_path = generate_budget_report(results, events, args.output_dir) + + print(f"\nBudget analysis completed! Report: {html_path}") + 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..11b672b --- /dev/null +++ b/zppy_interfaces/budget_analysis/checks.py @@ -0,0 +1,285 @@ +"""Budget checks: definitions and evaluation.""" + +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, +) + + +@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 + 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'. + For monthly data, aggregate to annual by summing per year. + """ + 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 only — aggregate to annual per year + # Rates (flux) get averaged; totals (state, flux_integrated) get summed + group_keys = [ + COL_TIME, + COL_COMPONENT, + COL_QUANTITY, + COL_TERM, + COL_SOURCE, + COL_TABLE_TYPE, + ] + flux_rows = subset[subset[COL_TABLE_TYPE] == "flux"] + other_rows = subset[subset[COL_TABLE_TYPE] != "flux"] + parts = [] + if not flux_rows.empty: + parts.append( + flux_rows.groupby(group_keys, as_index=False).agg( + {"normalized_value": "mean", "normalized_units": "first"} + ) + ) + if not other_rows.empty: + parts.append( + other_rows.groupby(group_keys, as_index=False).agg( + {"normalized_value": "sum", "normalized_units": "first"} + ) + ) + if parts: + subset = pd.concat(parts, ignore_index=True) + else: + return subset.iloc[:0] + + 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 water flux + global residual (*SUM*). + + Components dict includes each component's cumulative net flux + plus a '*SUM*' entry for the global residual. + """ + + def __init__(self) -> None: + super().__init__( + "cpl_component_fluxes", + "Coupler cumulative net water flux per component + residual", + ) + + def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: + rows = _select(df, **{COL_SOURCE: "cpl", COL_TERM: "*SUM*"}) + 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 water flux? + + Compares coupler *SUM* in the component column vs component's *SUM* flux. + Works for any component (lnd, ocn, etc.). + """ + + def __init__(self, component: str, source: str) -> None: + super().__init__( + f"{component}_interface_match", + f"{component} net water flux: coupler vs {source} model", + ) + self.component = component # coupler column name (e.g. "lnd", "ocn") + self.source = source # source tag in event table (e.g. "lnd", "ocn") + + def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: + cpl = _select( + df, **{COL_SOURCE: "cpl", COL_TERM: "*SUM*", COL_COMPONENT: self.component} + )[[COL_TIME, "normalized_value"]].set_index(COL_TIME) + + comp = _select( + df, **{COL_SOURCE: self.source, COL_TERM: "*SUM*", COL_TABLE_TYPE: "flux"} + )[[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)) + + +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)) + + +class OcnClosure(BudgetCheck): + """Does ocean mass change equal the net flux? + + Ocean logs are monthly. _select auto-aggregates to annual. + Compares mass_change (state) vs *SUM* flux. + """ + + def __init__(self) -> None: + super().__init__( + "ocn_closure", + "Ocean water closure: ΔMass vs net flux", + ) + + def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: + mass = _select( + df, + **{COL_SOURCE: "ocn", COL_TERM: "mass_change", COL_TABLE_TYPE: "flux"}, + ) + 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*", COL_TABLE_TYPE: "flux"}, + ) + 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 + return CheckResult(self.name, self.description, years, ds, fi, r, np.cumsum(r)) + + +DEFAULT_WATER_CHECKS: List[BudgetCheck] = [ + CplComponentFluxes(), + InterfaceMatch("lnd", "lnd"), + InterfaceMatch("ocn", "ocn"), + LndClosure(), + OcnClosure(), +] + + +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 + 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)} years") + 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/base.py b/zppy_interfaces/budget_analysis/ingestion/base.py new file mode 100644 index 0000000..17d2447 --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/base.py @@ -0,0 +1,17 @@ +"""Base class for log file parsers.""" + +from abc import ABC, abstractmethod +from typing import List + +import pandas as pd + + +class BaseParser(ABC): + """All parsers return a tidy event table DataFrame.""" + + @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..08428d9 --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py @@ -0,0 +1,145 @@ +"""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):", +} + +UNITS: Dict[str, str] = { + "water": "kg/m2s*1e6", +} + + +def _normalize_component_name(name: str) -> str: + """Normalize component names: 'ice nh' -> 'ice_nh'.""" + return name.strip().replace(" ", "_") + + +def _parse_datestamp(datestamp: str) -> int: + """Convert coupler datestamp to year. + + The date is reported at the start of the next period. + E.g. '20101' -> strip last 4 chars -> '2' -> minus 1 -> year 1. + """ + return int(datestamp[:-4]) - 1 + + +def _parse_header_line(line: str, pattern: str) -> Optional[Tuple[str, int]]: + """Extract period and year from a budget header line. + + Returns (period, year) 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 = _parse_datestamp(date_match.group(1)) + return period, year + + +def _parse_table(f: TextIO, year: int, 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): + self.quantities = quantities or ["water"] + + 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, year = result + if start_year <= year <= end_year: + rows.extend(_parse_table(f, year, 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/lnd_parser.py b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py new file mode 100644 index 0000000..dacfcad --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py @@ -0,0 +1,308 @@ +"""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) -> int: + """Convert datestamp to year. Same convention as coupler.""" + return int(datestamp[:-4]) - 1 + + +def _parse_period_and_year(line: str) -> Optional[Tuple[str, int]]: + """Extract period and year 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 = _parse_datestamp(date_match.group(1)) + return period, year + + +def _parse_flux_table(f: TextIO, year: int, 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: int, 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 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_year(stripped) + if result is None: + continue + period, year = result + if start_year <= year <= end_year: + rows.extend(_parse_flux_table(f, year, period)) + + elif stripped.startswith(STATE_HEADER): + result = _parse_period_and_year(stripped) + if result is None: + continue + period, year = result + if start_year <= year <= end_year: + rows.extend(_parse_state_table(f, year, 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..50aef5e --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py @@ -0,0 +1,245 @@ +"""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*" + 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 + + +class OcnParser(BaseParser): + """Parse ocean log files for mass conservation checks.""" + + 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 + rows.extend(self._parse_block(f, year, month)) + 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) + + def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: + """Parse one CONSERVATION CHECKS block for mass data.""" + rows: List[Dict] = [] + + # Scan for MASS CONSERVATION CHECK within this block + line = f.readline() + while line: + if "MASS CONSERVATION CHECK" in line and "SUMMARY" not in line: + # Parse flux table + fluxes = _parse_mass_fluxes(f) + for term, val in fluxes: + rows.append( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "water", + COL_TERM: term, + COL_VALUE: val, + COL_UNITS: "kg/m2s*1e6", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "flux", + } + ) + + # Continue reading for SUMMARY (has mass change in kg/m2s*1e6) + 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( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "water", + COL_TERM: "mass_change", + COL_VALUE: summary["mass_change"], + COL_UNITS: "kg/m2s*1e6", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "flux", + } + ) + if "absolute_mass_error" in summary: + rows.append( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "water", + COL_TERM: "absolute_mass_error", + COL_VALUE: summary["absolute_mass_error"], + COL_UNITS: "kg/m2s*1e6", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "diagnostic", + } + ) + break # Done with this block + elif "SALT CONSERVATION" in line: + break # Past mass section + line = f.readline() + break # Only one MASS CONSERVATION CHECK per block + elif "===" in line or line.strip() == "": + pass + 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..240e12f --- /dev/null +++ b/zppy_interfaces/budget_analysis/normalization.py @@ -0,0 +1,38 @@ +"""Normalize the tidy event table for budget analysis.""" + +import pandas as pd + +from .schema import 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: + - Flux rates (kg/m2s*1e6) -> mm/yr + - Flux integrated (kg/m2*1e6) -> mm + - State values (kg/m2*1e6) -> mm + """ + df = df.copy() + df["normalized_value"] = df[COL_VALUE].copy() + df["normalized_units"] = df[COL_UNITS].copy() + + # Flux rates: kg/m2s * 1e6 -> mm/yr + # 1 kg/m2 = 1 mm, so (val * 1e-6) kg/m2/s * seconds_per_year = mm/yr + flux_mask = df[COL_TABLE_TYPE] == "flux" + df.loc[flux_mask, "normalized_value"] = ( + df.loc[flux_mask, COL_VALUE] * SECONDS_PER_YEAR / 1e6 + ) + df.loc[flux_mask, "normalized_units"] = "mm/yr" + + # Integrated fluxes and states: kg/m2 * 1e6 -> mm + integrated_mask = df[COL_TABLE_TYPE].isin(["flux_integrated", "state"]) + df.loc[integrated_mask, "normalized_value"] = ( + df.loc[integrated_mask, COL_VALUE] / 1e6 + ) + df.loc[integrated_mask, "normalized_units"] = "mm" + + return df diff --git a/zppy_interfaces/budget_analysis/plotting.py b/zppy_interfaces/budget_analysis/plotting.py index 8a73bae..753564a 100644 --- a/zppy_interfaces/budget_analysis/plotting.py +++ b/zppy_interfaces/budget_analysis/plotting.py @@ -19,7 +19,7 @@ CARBON_CONVERSION = DT_SECONDS_PER_YEAR / 1e10 -def generate_html_plots( +def generate_html_plots( # noqa: C901 budgets: Dict[str, Budget], budget_names: List[str], output_dir: str ) -> None: """Generate interactive HTML plots using bokeh.""" diff --git a/zppy_interfaces/budget_analysis/schema.py b/zppy_interfaces/budget_analysis/schema.py new file mode 100644 index 0000000..97bfadd --- /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" # int: year the period ends (e.g. year 1 annual → time=1) +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..04f0e54 --- /dev/null +++ b/zppy_interfaces/budget_analysis/viz.py @@ -0,0 +1,131 @@ +"""Visualization for whole-model budget analysis. + +Generates an HTML report with Bokeh plots for each budget check result. +""" + +import os +from typing import List + +import numpy as np +import pandas as pd +from bokeh.layouts import column +from bokeh.models import Div +from bokeh.palettes import Category10 +from bokeh.plotting import figure, output_file, save + +from .checks import CheckResult + + +def generate_budget_report( + results: List[CheckResult], + df: pd.DataFrame, + output_dir: str, +) -> str: + """Generate an HTML report with budget check plots. + + Returns path to the HTML file. + """ + plots: list = [] + + for r in results: + if r.name == "cpl_component_fluxes": + plots.append(Div(text="

Conservation Overview

")) + plots.append(_plot_cumulative_components(r)) + + elif r.name.endswith("_interface_match"): + comp = r.name.replace("_interface_match", "") + plots.append(Div(text=f"

Interface Match: {comp}

")) + plots.append( + _plot_comparison(r, f"{comp} Water Flux (cpl vs {comp})", "mm/yr") + ) + plots.append( + _plot_residual(r, f"Interface Residual (cpl - {comp})", "mm/yr") + ) + + elif r.name == "lnd_closure": + plots.append(Div(text="

Land Water Closure

")) + plots.append(_plot_comparison(r, "Land ΔStorage vs ∫Flux dt", "mm")) + plots.append(_plot_residual(r, "Closure Residual", "mm")) + plots.append(_plot_cumulative(r, "Closure Cumulative Residual", "mm")) + + elif r.name == "ocn_closure": + plots.append(Div(text="

Ocean Water Closure

")) + plots.append(_plot_comparison(r, "Ocean ΔMass vs Net Flux", "mm")) + plots.append(_plot_residual(r, "Ocean Closure Residual", "mm")) + plots.append(_plot_cumulative(r, "Ocean Closure Cumulative Residual", "mm")) + + if not plots: + print("No plots generated — no check results available") + return "" + + html_path = os.path.join(output_dir, "water_budget_report.html") + output_file(html_path, title="E3SM Water Budget Analysis") + save(column(plots, sizing_mode="stretch_width")) + print(f"Report written to {html_path}") + return html_path + + +def _make_figure(title: str, y_label: str) -> figure: + return figure( + title=title, + height=350, + width=1200, + x_axis_label="year", + y_axis_label=y_label, + ) + + +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})") + 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) -> figure: + """Plot cumulative residual.""" + p = _make_figure(title, f"cumulative residual ({units})") + p.line(r.years, r.cumulative_residual, 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) + p.line(r.years, r.lhs, line_width=2, color="blue", legend_label="LHS") + p.line(r.years, r.rhs, line_width=2, color="orange", legend_label="RHS") + p.legend.click_policy = "hide" + return p + + +def _plot_cumulative_components(r: CheckResult) -> figure: + """Cumulative net water flux per component, with *SUM* residual highlighted.""" + p = _make_figure("Cumulative Net Water Flux per Component", "mm") + 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], + line_width=2, + color=colors[i % len(colors)], + legend_label=name, + ) + if "*SUM*" in r.components: + p.line( + r.years, + r.components["*SUM*"], + line_width=3, + color="red", + line_dash="dashed", + legend_label="*SUM* (residual)", + ) + p.legend.click_policy = "hide" + p.legend.location = "top_left" + return p From f19219b14fd416bf9bc4a2fe27e56a84b3b1442a Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Mon, 9 Feb 2026 15:43:41 -0800 Subject: [PATCH 04/15] update figure labels --- zppy_interfaces/budget_analysis/checks.py | 38 +++++++++++++++++++++-- zppy_interfaces/budget_analysis/viz.py | 9 ++++-- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index 11b672b..81a5d22 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -28,6 +28,8 @@ class CheckResult: 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) @@ -176,7 +178,17 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: 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)) + return CheckResult( + self.name, + self.description, + years, + c, + m, + r, + np.cumsum(r), + lhs_label=f"cpl ({self.component})", + rhs_label=f"{self.source} (*SUM*)", + ) class LndClosure(BudgetCheck): @@ -215,7 +227,17 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: 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)) + 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): @@ -255,7 +277,17 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: ds = merged["normalized_value_mass"].values fi = merged["normalized_value_flux"].values r = ds - fi - return CheckResult(self.name, self.description, years, ds, fi, r, np.cumsum(r)) + return CheckResult( + self.name, + self.description, + years, + ds, + fi, + r, + np.cumsum(r), + lhs_label="ΔMass (mass_change)", + rhs_label="Net Flux (*SUM*)", + ) DEFAULT_WATER_CHECKS: List[BudgetCheck] = [ diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py index 04f0e54..c9a207e 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -41,6 +41,11 @@ def generate_budget_report( plots.append( _plot_residual(r, f"Interface Residual (cpl - {comp})", "mm/yr") ) + plots.append( + _plot_cumulative( + r, f"Interface Cumulative Residual (cpl - {comp})", "mm" + ) + ) elif r.name == "lnd_closure": plots.append(Div(text="

Land Water Closure

")) @@ -95,8 +100,8 @@ def _plot_cumulative(r: CheckResult, title: str, units: str) -> figure: def _plot_comparison(r: CheckResult, title: str, units: str) -> figure: """Plot LHS and RHS on the same axes.""" p = _make_figure(title, units) - p.line(r.years, r.lhs, line_width=2, color="blue", legend_label="LHS") - p.line(r.years, r.rhs, line_width=2, color="orange", legend_label="RHS") + 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 From c0a7666c220b33ff22ce5e1d4f9f42059bd6afe0 Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Tue, 10 Feb 2026 10:48:52 -0800 Subject: [PATCH 05/15] add heat/energy budget --- .../budget_analysis/diagnose_ocn_closure.py | 80 ++++++++ .../print_monthly_interface.py | 43 +++++ zppy_interfaces/budget_analysis/__main__.py | 48 ++++- zppy_interfaces/budget_analysis/checks.py | 113 +++++++++--- .../budget_analysis/ingestion/cpl_parser.py | 4 +- .../budget_analysis/ingestion/ocn_parser.py | 171 ++++++++++++++++-- .../budget_analysis/normalization.py | 33 +++- zppy_interfaces/budget_analysis/viz.py | 150 ++++++++++++--- 8 files changed, 553 insertions(+), 89 deletions(-) create mode 100644 tests/unit/budget_analysis/diagnose_ocn_closure.py create mode 100644 tests/unit/budget_analysis/print_monthly_interface.py 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..54934cc --- /dev/null +++ b/tests/unit/budget_analysis/diagnose_ocn_closure.py @@ -0,0 +1,80 @@ +"""Diagnose ocean closure: verify log-native term names are used correctly.""" + +import glob +import sys + +import pandas as pd + +sys.path.insert(0, ".") + +from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser +from zppy_interfaces.budget_analysis.normalization import normalize + +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) + +from zppy_interfaces.budget_analysis.checks import OcnClosure + +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, 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..48ff85e --- /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 +from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser +from zppy_interfaces.budget_analysis.normalization import normalize + +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*") + & (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* ({len(ocn_m)} rows) ===") +print(ocn_m[["time", "normalized_value"]].to_string(index=False)) diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index 5d1cf34..72c2583 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -124,16 +124,19 @@ def _run_whole_model(args) -> int: """Whole-model budget pipeline: ingest -> normalize -> check -> visualize.""" import pandas as pd - from .checks import run_checks + from .checks import DEFAULT_HEAT_CHECKS, DEFAULT_WATER_CHECKS, run_checks from .ingestion.cpl_parser import CplParser from .ingestion.lnd_parser import LndParser from .ingestion.ocn_parser import OcnParser from .normalization import normalize - from .viz import generate_budget_report + 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"Log path: {args.log_path}") # Ingest @@ -150,7 +153,11 @@ def _run_whole_model(args) -> int: print(f" {len(ocn_files)} ocean log files") frames = [] - frames.append(CplParser().parse_files(cpl_files, args.start_year, args.end_year)) + frames.append( + CplParser(quantities=budget_types).parse_files( + cpl_files, args.start_year, args.end_year + ) + ) if lnd_files: frames.append( LndParser().parse_files(lnd_files, args.start_year, args.end_year) @@ -166,15 +173,36 @@ def _run_whole_model(args) -> int: print("\nNormalizing...") events = normalize(events) - # Check - print("\nRunning budget checks...") - results = run_checks(events) + # Check and visualize per quantity + checks_map = { + "water": DEFAULT_WATER_CHECKS, + "heat": DEFAULT_HEAT_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 - # Visualize - print("\nGenerating report...") - html_path = generate_budget_report(results, events, args.output_dir) + # 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.") - print(f"\nBudget analysis completed! Report: {html_path}") return 0 diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index 81a5d22..9a2e4a8 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -103,20 +103,25 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: class CplComponentFluxes(BudgetCheck): - """Per-component net water flux + global residual (*SUM*). + """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) -> None: + def __init__(self, quantity: str = "water") -> None: super().__init__( - "cpl_component_fluxes", - "Coupler cumulative net water flux per component + residual", + 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*"}) + rows = _select( + df, + **{COL_SOURCE: "cpl", COL_TERM: "*SUM*", COL_QUANTITY: self.quantity}, + ) if rows.empty: return None pivot = rows.pivot_table( @@ -146,27 +151,47 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: class InterfaceMatch(BudgetCheck): - """Do the coupler and component model agree on net water flux? + """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.). + Works for any component (lnd, ocn, etc.) and quantity (water, heat). """ - def __init__(self, component: str, source: str) -> None: + def __init__( + self, + component: str, + source: str, + quantity: str = "water", + comp_sum_term: str = "*SUM*", + ) -> None: super().__init__( - f"{component}_interface_match", - f"{component} net water flux: coupler vs {source} model", + f"{component}_{quantity}_interface_match", + f"{component} net {quantity} flux: coupler vs {source} model", ) - self.component = component # coupler column name (e.g. "lnd", "ocn") - self.source = source # source tag in event table (e.g. "lnd", "ocn") + 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} + 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: "*SUM*", COL_TABLE_TYPE: "flux"} + 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: @@ -187,7 +212,7 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: r, np.cumsum(r), lhs_label=f"cpl ({self.component})", - rhs_label=f"{self.source} (*SUM*)", + rhs_label=f"{self.source} ({self.comp_sum_term})", ) @@ -241,22 +266,42 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: class OcnClosure(BudgetCheck): - """Does ocean mass change equal the net flux? + """Does ocean mass/energy change equal the net flux? Ocean logs are monthly. _select auto-aggregates to annual. - Compares mass_change (state) vs *SUM* flux. + For water: compares mass_change vs *SUM* flux. + For heat: compares energy_change vs *SUM* flux. """ - def __init__(self) -> None: + 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__( - "ocn_closure", - "Ocean water closure: ΔMass vs net flux", + 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: "mass_change", COL_TABLE_TYPE: "flux"}, + **{ + COL_SOURCE: "ocn", + COL_TERM: change_term, + COL_TABLE_TYPE: "flux", + COL_QUANTITY: self.quantity, + }, ) if mass.empty: return None @@ -264,7 +309,12 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: flux = _select( df, - **{COL_SOURCE: "ocn", COL_TERM: "*SUM*", COL_TABLE_TYPE: "flux"}, + **{ + COL_SOURCE: "ocn", + COL_TERM: sum_term, + COL_TABLE_TYPE: "flux", + COL_QUANTITY: self.quantity, + }, ) if flux.empty: return None @@ -277,6 +327,7 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: 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, @@ -285,17 +336,23 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: fi, r, np.cumsum(r), - lhs_label="ΔMass (mass_change)", - rhs_label="Net Flux (*SUM*)", + lhs_label=f"{change_label} ({change_term})", + rhs_label=f"Net Flux ({sum_term})", ) DEFAULT_WATER_CHECKS: List[BudgetCheck] = [ - CplComponentFluxes(), - InterfaceMatch("lnd", "lnd"), - InterfaceMatch("ocn", "ocn"), + CplComponentFluxes(quantity="water"), + InterfaceMatch("lnd", "lnd", quantity="water"), + InterfaceMatch("ocn", "ocn", quantity="water", comp_sum_term="SUM VOLUME FLUXES"), LndClosure(), - OcnClosure(), + OcnClosure(quantity="water"), +] + +DEFAULT_HEAT_CHECKS: List[BudgetCheck] = [ + CplComponentFluxes(quantity="heat"), + InterfaceMatch("ocn", "ocn", quantity="heat", comp_sum_term="SUM IMP+EXP HEAT FLUXES"), + OcnClosure(quantity="heat"), ] diff --git a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py index 08428d9..39d613d 100644 --- a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py @@ -23,10 +23,12 @@ # 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):", } UNITS: Dict[str, str] = { "water": "kg/m2s*1e6", + "heat": "W/m2", } @@ -116,7 +118,7 @@ class CplParser(BaseParser): """Parse coupler log budget tables into a tidy event table.""" def __init__(self, quantities: Optional[List[str]] = None): - self.quantities = quantities or ["water"] + self.quantities = quantities or ["water", "heat"] def parse_files( self, log_files: List[str], start_year: int, end_year: int diff --git a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py index 50aef5e..9e954d8 100644 --- a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py @@ -93,7 +93,7 @@ def _parse_mass_fluxes(f: TextIO) -> List[Tuple[str, float]]: # Determine term name if "SUM VOLUME FLUXES" in line: - term = "*SUM*" + term = "SUM VOLUME FLUXES" else: # short_name is the second-to-last token term = parts[-2] @@ -145,8 +145,87 @@ def _parse_mass_summary(f: TextIO) -> Optional[Dict[str, float]]: 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 conservation checks.""" + """Parse ocean log files for mass and energy conservation checks.""" def parse_files( self, log_files: List[str], start_year: int, end_year: int @@ -175,13 +254,21 @@ def parse_files( return pd.DataFrame(rows, columns=COLUMNS) def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: - """Parse one CONSERVATION CHECKS block for mass data.""" + """Parse one CONSERVATION CHECKS block for mass and energy data.""" rows: List[Dict] = [] + found_mass = False + found_energy = False - # Scan for MASS CONSERVATION CHECK within this block line = f.readline() while line: + # Stop if we hit the next block or have found both sections + 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 # Parse flux table fluxes = _parse_mass_fluxes(f) for term, val in fluxes: @@ -199,7 +286,7 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: } ) - # Continue reading for SUMMARY (has mass change in kg/m2s*1e6) + # Continue reading for SUMMARY line = f.readline() while line: if "MASS CONSERVATION SUMMARY" in line: @@ -211,7 +298,7 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: COL_TIME: year, COL_COMPONENT: "ocn", COL_QUANTITY: "water", - COL_TERM: "mass_change", + COL_TERM: "Mass change", COL_VALUE: summary["mass_change"], COL_UNITS: "kg/m2s*1e6", COL_SOURCE: "ocn", @@ -225,19 +312,81 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: COL_TIME: year, COL_COMPONENT: "ocn", COL_QUANTITY: "water", - COL_TERM: "absolute_mass_error", - COL_VALUE: summary["absolute_mass_error"], + COL_TERM: "Absolute mass error", + COL_VALUE: summary[ + "absolute_mass_error" + ], COL_UNITS: "kg/m2s*1e6", COL_SOURCE: "ocn", COL_PERIOD: "monthly", COL_TABLE_TYPE: "diagnostic", } ) - break # Done with this block + break elif "SALT CONSERVATION" in line: - break # Past mass section + break line = f.readline() - break # Only one MASS CONSERVATION CHECK per block + + elif "ENERGY CONSERVATION CHECK" in line and "SUMMARY" not in line: + found_energy = True + # Parse heat flux tables + heat_fluxes = _parse_heat_fluxes(f) + for term, val in heat_fluxes: + rows.append( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "heat", + COL_TERM: term, + COL_VALUE: val, + COL_UNITS: "W/m2", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "flux", + } + ) + + # Continue reading for ENERGY CONSERVATION SUMMARY + 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( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "heat", + COL_TERM: "Energy change", + COL_VALUE: summary["energy_change"], + COL_UNITS: "W/m2", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "flux", + } + ) + if "absolute_energy_error" in summary: + rows.append( + { + COL_TIME: year, + COL_COMPONENT: "ocn", + COL_QUANTITY: "heat", + COL_TERM: "Absolute energy error", + COL_VALUE: summary[ + "absolute_energy_error" + ], + COL_UNITS: "W/m2", + COL_SOURCE: "ocn", + COL_PERIOD: "monthly", + COL_TABLE_TYPE: "diagnostic", + } + ) + break + elif "RELATIVE ENERGY" in line: + break + line = f.readline() + elif "===" in line or line.strip() == "": pass line = f.readline() diff --git a/zppy_interfaces/budget_analysis/normalization.py b/zppy_interfaces/budget_analysis/normalization.py index 240e12f..d81d018 100644 --- a/zppy_interfaces/budget_analysis/normalization.py +++ b/zppy_interfaces/budget_analysis/normalization.py @@ -2,7 +2,7 @@ import pandas as pd -from .schema import COL_TABLE_TYPE, COL_UNITS, COL_VALUE +from .schema import COL_QUANTITY, COL_TABLE_TYPE, COL_UNITS, COL_VALUE # Seconds per year (365-day calendar) SECONDS_PER_YEAR = 365.0 * 24.0 * 60.0 * 60.0 @@ -12,27 +12,40 @@ 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 (W/m2) -> J/m2 (cumulative energy per year) """ 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 - # 1 kg/m2 = 1 mm, so (val * 1e-6) kg/m2/s * seconds_per_year = mm/yr - flux_mask = df[COL_TABLE_TYPE] == "flux" - df.loc[flux_mask, "normalized_value"] = ( - df.loc[flux_mask, COL_VALUE] * SECONDS_PER_YEAR / 1e6 + water_flux = water_mask & (df[COL_TABLE_TYPE] == "flux") + df.loc[water_flux, "normalized_value"] = ( + df.loc[water_flux, COL_VALUE] * SECONDS_PER_YEAR / 1e6 ) - df.loc[flux_mask, "normalized_units"] = "mm/yr" + df.loc[water_flux, "normalized_units"] = "mm/yr" # Integrated fluxes and states: kg/m2 * 1e6 -> mm - integrated_mask = df[COL_TABLE_TYPE].isin(["flux_integrated", "state"]) - df.loc[integrated_mask, "normalized_value"] = ( - df.loc[integrated_mask, COL_VALUE] / 1e6 + water_integrated = water_mask & df[COL_TABLE_TYPE].isin( + ["flux_integrated", "state"] ) - df.loc[integrated_mask, "normalized_units"] = "mm" + df.loc[water_integrated, "normalized_value"] = ( + df.loc[water_integrated, COL_VALUE] / 1e6 + ) + df.loc[water_integrated, "normalized_units"] = "mm" + + # --- Heat --- + # Keep W/m2 as-is (no conversion needed) + # Cumulative residual plots will accumulate W/m2 values over years return df diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py index c9a207e..6f66eff 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -4,7 +4,7 @@ """ import os -from typing import List +from typing import Dict, List import numpy as np import pandas as pd @@ -14,62 +14,143 @@ 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"}, +} + + +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 == "cpl_component_fluxes": - plots.append(Div(text="

Conservation Overview

")) - plots.append(_plot_cumulative_components(r)) + 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"): - comp = r.name.replace("_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} Water Flux (cpl vs {comp})", "mm/yr") + _plot_comparison( + r, f"{comp} {quantity.title()} Flux (cpl vs {comp})", + flux_units, + ) ) plots.append( - _plot_residual(r, f"Interface Residual (cpl - {comp})", "mm/yr") + _plot_residual( + r, f"Interface Residual (cpl - {comp})", flux_units + ) ) plots.append( _plot_cumulative( - r, f"Interface Cumulative Residual (cpl - {comp})", "mm" + 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", "mm")) - plots.append(_plot_residual(r, "Closure Residual", "mm")) - plots.append(_plot_cumulative(r, "Closure Cumulative Residual", "mm")) + 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 == "ocn_closure": - plots.append(Div(text="

Ocean Water Closure

")) - plots.append(_plot_comparison(r, "Ocean ΔMass vs Net Flux", "mm")) - plots.append(_plot_residual(r, "Ocean Closure Residual", "mm")) - plots.append(_plot_cumulative(r, "Ocean Closure Cumulative Residual", "mm")) + 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, + ) + ) if not plots: - print("No plots generated — no check results available") + print(f"No {quantity} plots generated — no check results available") return "" - html_path = os.path.join(output_dir, "water_budget_report.html") - output_file(html_path, title="E3SM Water Budget Analysis") + 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

    +
      +{"".join(links)} +
    + +""" + 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) -> figure: return figure( title=title, @@ -85,15 +166,20 @@ def _plot_residual(r: CheckResult, title: str, units: str) -> figure: p = _make_figure(title, f"residual ({units})") 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" + 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) -> figure: - """Plot cumulative residual.""" +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})") - p.line(r.years, r.cumulative_residual, line_width=2, color="darkred") + p.line( + r.years, r.cumulative_residual * scale, line_width=2, color="darkred" + ) return p @@ -101,14 +187,20 @@ def _plot_comparison(r: CheckResult, title: str, units: str) -> figure: """Plot LHS and RHS on the same axes.""" p = _make_figure(title, units) 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.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) -> figure: - """Cumulative net water flux per component, with *SUM* residual highlighted.""" - p = _make_figure("Cumulative Net Water Flux per Component", "mm") +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 + ) if r.components is None: return p # Plot component lines, highlight *SUM* as thick dashed red @@ -117,7 +209,7 @@ def _plot_cumulative_components(r: CheckResult) -> figure: for i, name in enumerate(other_names): p.line( r.years, - r.components[name], + r.components[name] * scale, line_width=2, color=colors[i % len(colors)], legend_label=name, @@ -125,7 +217,7 @@ def _plot_cumulative_components(r: CheckResult) -> figure: if "*SUM*" in r.components: p.line( r.years, - r.components["*SUM*"], + r.components["*SUM*"] * scale, line_width=3, color="red", line_dash="dashed", From 6944c774b284184203fe0bf96d14b1cbe0e7f771 Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Tue, 10 Feb 2026 12:57:46 -0800 Subject: [PATCH 06/15] add carbon coupler budget --- .../budget_analysis/diagnose_ocn_closure.py | 38 +++- .../print_monthly_interface.py | 10 +- zppy_interfaces/budget_analysis/__main__.py | 8 +- zppy_interfaces/budget_analysis/checks.py | 8 +- .../budget_analysis/ingestion/cpl_parser.py | 2 + .../budget_analysis/ingestion/ocn_parser.py | 215 ++++++++---------- .../budget_analysis/normalization.py | 15 +- zppy_interfaces/budget_analysis/viz.py | 65 +++--- 8 files changed, 184 insertions(+), 177 deletions(-) diff --git a/tests/unit/budget_analysis/diagnose_ocn_closure.py b/tests/unit/budget_analysis/diagnose_ocn_closure.py index 54934cc..7cbe968 100644 --- a/tests/unit/budget_analysis/diagnose_ocn_closure.py +++ b/tests/unit/budget_analysis/diagnose_ocn_closure.py @@ -3,12 +3,11 @@ import glob import sys -import pandas as pd - sys.path.insert(0, ".") -from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser -from zppy_interfaces.budget_analysis.normalization import normalize +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 @@ -30,7 +29,11 @@ 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)) + print( + mc[["time", "value", "units", "table_type", "period"]] + .head(12) + .to_string(index=False) + ) else: print(" *** NOT FOUND ***") print() @@ -38,7 +41,11 @@ 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)) + print( + svf[["time", "value", "units", "table_type", "period"]] + .head(12) + .to_string(index=False) + ) else: print(" *** NOT FOUND ***") print() @@ -52,7 +59,11 @@ 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)) + print( + ec[["time", "value", "units", "table_type", "period"]] + .head(12) + .to_string(index=False) + ) else: print(" *** NOT FOUND ***") print() @@ -60,7 +71,11 @@ 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)) + print( + shf[["time", "value", "units", "table_type", "period"]] + .head(12) + .to_string(index=False) + ) else: print(" *** NOT FOUND ***") print() @@ -69,12 +84,13 @@ print("=== Testing OcnClosure checks ===") df = normalize(ocn) -from zppy_interfaces.budget_analysis.checks import OcnClosure - 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, max |residual| = {abs(result.residual).max():.2e}") + 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 index 48ff85e..d5c9804 100644 --- a/tests/unit/budget_analysis/print_monthly_interface.py +++ b/tests/unit/budget_analysis/print_monthly_interface.py @@ -7,9 +7,9 @@ sys.path.insert(0, ".") -from zppy_interfaces.budget_analysis.ingestion.cpl_parser import CplParser -from zppy_interfaces.budget_analysis.ingestion.ocn_parser import OcnParser -from zppy_interfaces.budget_analysis.normalization import normalize +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 @@ -32,12 +32,12 @@ ocn_m = df[ (df.source == "ocn") - & (df.term == "*SUM*") + & (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* ({len(ocn_m)} rows) ===") +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/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index 72c2583..b04280a 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -124,7 +124,12 @@ def _run_whole_model(args) -> int: """Whole-model budget pipeline: ingest -> normalize -> check -> visualize.""" import pandas as pd - from .checks import DEFAULT_HEAT_CHECKS, DEFAULT_WATER_CHECKS, run_checks + from .checks import ( + DEFAULT_CARBON_CHECKS, + DEFAULT_HEAT_CHECKS, + DEFAULT_WATER_CHECKS, + run_checks, + ) from .ingestion.cpl_parser import CplParser from .ingestion.lnd_parser import LndParser from .ingestion.ocn_parser import OcnParser @@ -177,6 +182,7 @@ def _run_whole_model(args) -> int: checks_map = { "water": DEFAULT_WATER_CHECKS, "heat": DEFAULT_HEAT_CHECKS, + "carbon": DEFAULT_CARBON_CHECKS, } report_paths = {} diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index 9a2e4a8..e27dabf 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -351,10 +351,16 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: DEFAULT_HEAT_CHECKS: List[BudgetCheck] = [ CplComponentFluxes(quantity="heat"), - InterfaceMatch("ocn", "ocn", quantity="heat", comp_sum_term="SUM IMP+EXP HEAT FLUXES"), + InterfaceMatch( + "ocn", "ocn", quantity="heat", comp_sum_term="SUM IMP+EXP HEAT FLUXES" + ), OcnClosure(quantity="heat"), ] +DEFAULT_CARBON_CHECKS: List[BudgetCheck] = [ + CplComponentFluxes(quantity="carbon"), +] + def run_checks( df: pd.DataFrame, diff --git a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py index 39d613d..f97c4d4 100644 --- a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py @@ -24,11 +24,13 @@ 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", } diff --git a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py index 9e954d8..41e6a90 100644 --- a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py @@ -261,7 +261,6 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: line = f.readline() while line: - # Stop if we hit the next block or have found both sections if found_mass and found_energy: break if "CONSERVATION CHECKS" in line and "date:" not in line: @@ -269,126 +268,106 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: if "MASS CONSERVATION CHECK" in line and "SUMMARY" not in line: found_mass = True - # Parse flux table - fluxes = _parse_mass_fluxes(f) - for term, val in fluxes: - rows.append( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "water", - COL_TERM: term, - COL_VALUE: val, - COL_UNITS: "kg/m2s*1e6", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "flux", - } - ) - - # Continue reading for SUMMARY - 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( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "water", - COL_TERM: "Mass change", - COL_VALUE: summary["mass_change"], - COL_UNITS: "kg/m2s*1e6", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "flux", - } - ) - if "absolute_mass_error" in summary: - rows.append( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "water", - COL_TERM: "Absolute mass error", - COL_VALUE: summary[ - "absolute_mass_error" - ], - COL_UNITS: "kg/m2s*1e6", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "diagnostic", - } - ) - break - elif "SALT CONSERVATION" in line: - break - line = f.readline() + rows.extend(self._parse_mass_section(f, year)) elif "ENERGY CONSERVATION CHECK" in line and "SUMMARY" not in line: found_energy = True - # Parse heat flux tables - heat_fluxes = _parse_heat_fluxes(f) - for term, val in heat_fluxes: - rows.append( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "heat", - COL_TERM: term, - COL_VALUE: val, - COL_UNITS: "W/m2", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "flux", - } - ) - - # Continue reading for ENERGY CONSERVATION SUMMARY - 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( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "heat", - COL_TERM: "Energy change", - COL_VALUE: summary["energy_change"], - COL_UNITS: "W/m2", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "flux", - } - ) - if "absolute_energy_error" in summary: - rows.append( - { - COL_TIME: year, - COL_COMPONENT: "ocn", - COL_QUANTITY: "heat", - COL_TERM: "Absolute energy error", - COL_VALUE: summary[ - "absolute_energy_error" - ], - COL_UNITS: "W/m2", - COL_SOURCE: "ocn", - COL_PERIOD: "monthly", - COL_TABLE_TYPE: "diagnostic", - } - ) - break - elif "RELATIVE ENERGY" in line: - break - line = f.readline() - - elif "===" in line or line.strip() == "": - pass + rows.extend(self._parse_energy_section(f, year)) + + line = f.readline() + + return rows + + def _parse_mass_section(self, f: TextIO, year: int) -> List[Dict]: + """Parse MASS CONSERVATION CHECK: fluxes + summary.""" + rows: List[Dict] = [] + base = { + COL_TIME: year, + 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, year: int) -> List[Dict]: + """Parse ENERGY CONSERVATION CHECK: fluxes + summary.""" + rows: List[Dict] = [] + base = { + COL_TIME: year, + 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 index d81d018..0a6a0c4 100644 --- a/zppy_interfaces/budget_analysis/normalization.py +++ b/zppy_interfaces/budget_analysis/normalization.py @@ -19,7 +19,10 @@ def normalize(df: pd.DataFrame) -> pd.DataFrame: - State values (kg/m2*1e6) -> mm Heat: - - Flux rates (W/m2) -> J/m2 (cumulative energy per year) + - 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() @@ -46,6 +49,14 @@ def normalize(df: pd.DataFrame) -> pd.DataFrame: # --- Heat --- # Keep W/m2 as-is (no conversion needed) - # Cumulative residual plots will accumulate W/m2 values over years + # 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/viz.py b/zppy_interfaces/budget_analysis/viz.py index 6f66eff..946d4bb 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -20,6 +20,7 @@ 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"}, } @@ -52,12 +53,8 @@ def generate_budget_report( 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) - ) + 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 @@ -65,48 +62,43 @@ def generate_budget_report( plots.append(Div(text=f"

    Interface Match: {comp}

    ")) plots.append( _plot_comparison( - r, f"{comp} {quantity.title()} Flux (cpl vs {comp})", + r, + f"{comp} {quantity.title()} Flux (cpl vs {comp})", flux_units, ) ) plots.append( - _plot_residual( - r, f"Interface Residual (cpl - {comp})", flux_units - ) + _plot_residual(r, f"Interface Residual (cpl - {comp})", flux_units) ) plots.append( _plot_cumulative( - r, f"Interface Cumulative Residual (cpl - {comp})", - cum_units, scale, + 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_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) - ) + 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) + _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, + r, + "Ocean Closure Cumulative Residual", + cum_units, + scale, ) ) @@ -121,9 +113,7 @@ def generate_budget_report( return html_path -def generate_landing_page( - output_dir: str, report_paths: Dict[str, str] -) -> str: +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. @@ -166,7 +156,10 @@ def _plot_residual(r: CheckResult, title: str, units: str) -> figure: p = _make_figure(title, f"residual ({units})") 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", + r.years, + np.zeros_like(r.years), + line_width=1, + color="gray", line_dash="dashed", ) return p @@ -177,9 +170,7 @@ def _plot_cumulative( ) -> figure: """Plot cumulative residual, optionally scaled for unit conversion.""" p = _make_figure(title, f"cumulative residual ({units})") - p.line( - r.years, r.cumulative_residual * scale, line_width=2, color="darkred" - ) + p.line(r.years, r.cumulative_residual * scale, line_width=2, color="darkred") return p @@ -187,9 +178,7 @@ def _plot_comparison(r: CheckResult, title: str, units: str) -> figure: """Plot LHS and RHS on the same axes.""" p = _make_figure(title, units) 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.line(r.years, r.rhs, line_width=2, color="orange", legend_label=r.rhs_label) p.legend.click_policy = "hide" return p @@ -198,9 +187,7 @@ 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 - ) + p = _make_figure(f"Cumulative Net {quantity.title()} Flux per Component", units) if r.components is None: return p # Plot component lines, highlight *SUM* as thick dashed red From 37eb5419a1282040bd5f9b091248cab99a54e5d0 Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Wed, 11 Feb 2026 14:34:44 -0800 Subject: [PATCH 07/15] add atm log parser --- tests/unit/budget_analysis/test_atm_parser.py | 76 ++++ zppy_interfaces/budget_analysis/__main__.py | 7 + .../budget_analysis/ingestion/atm_parser.py | 398 ++++++++++++++++++ 3 files changed, 481 insertions(+) create mode 100644 tests/unit/budget_analysis/test_atm_parser.py create mode 100644 zppy_interfaces/budget_analysis/ingestion/atm_parser.py 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..d384950 --- /dev/null +++ b/tests/unit/budget_analysis/test_atm_parser.py @@ -0,0 +1,76 @@ +"""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 --- + events = parser.parse_files(log_files, 1, 1) + print("=== Tidy event table ===") + print(f"Shape: {events.shape}") + print(events.to_string()) + + print("\nAll checks passed.") + + +if __name__ == "__main__": + sys.exit(main() or 0) diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index b04280a..aca3359 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -130,6 +130,7 @@ def _run_whole_model(args) -> int: DEFAULT_WATER_CHECKS, run_checks, ) + from .ingestion.atm_parser import AtmParser from .ingestion.cpl_parser import CplParser from .ingestion.lnd_parser import LndParser from .ingestion.ocn_parser import OcnParser @@ -149,6 +150,7 @@ def _run_whole_model(args) -> int: 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"))) + atm_files = sorted(glob.glob(os.path.join(args.log_path, "atm.log.*"))) if not cpl_files: print("ERROR: No coupler log files found") @@ -156,6 +158,7 @@ def _run_whole_model(args) -> int: 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(atm_files)} atmosphere log files") frames = [] frames.append( @@ -171,6 +174,10 @@ def _run_whole_model(args) -> int: frames.append( OcnParser().parse_files(ocn_files, args.start_year, args.end_year) ) + if atm_files: + frames.append( + AtmParser().parse_files(atm_files, args.start_year, args.end_year) + ) events = pd.concat(frames, ignore_index=True) print(f" {len(events)} total event rows") 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..61a4d8a --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/atm_parser.py @@ -0,0 +1,398 @@ +"""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 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 parse_files( + self, log_files: List[str], start_year: int, end_year: int + ) -> pd.DataFrame: + """Parse atm logs and return a tidy monthly event table. + + Aggregates per-step data to monthly averages for fluxes and + beginning/end-of-month values for 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) + + rows: List[Dict] = [] + + # --- Flux diagnostics (monthly aggregation) --- + 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 (year, month), grp in flux_df.groupby(["year", "month"]): + base_water = { + COL_TIME: year, + COL_COMPONENT: "atm", + COL_QUANTITY: "water", + COL_SOURCE: "atm", + COL_PERIOD: "monthly", + } + base_heat = { + COL_TIME: year, + COL_COMPONENT: "atm", + COL_QUANTITY: "heat", + COL_SOURCE: "atm", + COL_PERIOD: "monthly", + } + + dt = grp["dt"].iloc[0] + + # Water flux terms (monthly 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 month) + 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 (monthly 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 (monthly 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 (monthly 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 month) --- + 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 (year, month), grp in te_df.groupby(["year", "month"]): + base_heat = { + COL_TIME: year, + COL_COMPONENT: "atm", + COL_QUANTITY: "heat", + COL_SOURCE: "atm", + COL_PERIOD: "monthly", + } + 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) From 741d002e83fa1ed8ebfa3025798034810f905f8e Mon Sep 17 00:00:00 2001 From: ChengzhuZhang Date: Wed, 11 Feb 2026 15:46:16 -0800 Subject: [PATCH 08/15] support vis monthly budgets in addition to annual --- tests/unit/budget_analysis/test_atm_parser.py | 12 +++- zppy_interfaces/budget_analysis/__main__.py | 22 +++++-- zppy_interfaces/budget_analysis/checks.py | 44 ++++--------- .../budget_analysis/ingestion/atm_parser.py | 64 +++++++++++++------ .../budget_analysis/ingestion/base.py | 21 +++++- .../budget_analysis/ingestion/cpl_parser.py | 55 ++++++++++++---- .../budget_analysis/ingestion/lnd_parser.py | 59 +++++++++++------ .../budget_analysis/ingestion/ocn_parser.py | 43 ++++++++++--- zppy_interfaces/budget_analysis/schema.py | 2 +- zppy_interfaces/budget_analysis/viz.py | 25 +++++--- 10 files changed, 240 insertions(+), 107 deletions(-) diff --git a/tests/unit/budget_analysis/test_atm_parser.py b/tests/unit/budget_analysis/test_atm_parser.py index d384950..c8192a9 100644 --- a/tests/unit/budget_analysis/test_atm_parser.py +++ b/tests/unit/budget_analysis/test_atm_parser.py @@ -63,11 +63,19 @@ def main(): ) print() - # --- Tidy event table --- + # --- Tidy event table (annual) --- events = parser.parse_files(log_files, 1, 1) - print("=== Tidy event table ===") + 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.") diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index aca3359..b551ef4 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -66,6 +66,12 @@ def main() -> int: 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() @@ -143,6 +149,7 @@ def _run_whole_model(args) -> int: 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 @@ -160,23 +167,30 @@ def _run_whole_model(args) -> int: print(f" {len(ocn_files)} ocean log files") print(f" {len(atm_files)} atmosphere log files") + freq = args.frequency frames = [] frames.append( - CplParser(quantities=budget_types).parse_files( + CplParser(quantities=budget_types, frequency=freq).parse_files( cpl_files, args.start_year, args.end_year ) ) if lnd_files: frames.append( - LndParser().parse_files(lnd_files, args.start_year, args.end_year) + LndParser(frequency=freq).parse_files( + lnd_files, args.start_year, args.end_year + ) ) if ocn_files: frames.append( - OcnParser().parse_files(ocn_files, args.start_year, args.end_year) + OcnParser(frequency=freq).parse_files( + ocn_files, args.start_year, args.end_year + ) ) if atm_files: frames.append( - AtmParser().parse_files(atm_files, args.start_year, args.end_year) + 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") diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index e27dabf..8cc783b 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -41,7 +41,8 @@ def _select( """Filter df by column=value filters, optionally by period. If period is None, prefer 'annual' if available, else use 'monthly'. - For monthly data, aggregate to annual by summing per year. + 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(): @@ -58,35 +59,10 @@ def _select( elif "annual" in available: subset = subset[subset[COL_PERIOD] == "annual"] else: - # Monthly only — aggregate to annual per year - # Rates (flux) get averaged; totals (state, flux_integrated) get summed - group_keys = [ - COL_TIME, - COL_COMPONENT, - COL_QUANTITY, - COL_TERM, - COL_SOURCE, - COL_TABLE_TYPE, - ] - flux_rows = subset[subset[COL_TABLE_TYPE] == "flux"] - other_rows = subset[subset[COL_TABLE_TYPE] != "flux"] - parts = [] - if not flux_rows.empty: - parts.append( - flux_rows.groupby(group_keys, as_index=False).agg( - {"normalized_value": "mean", "normalized_units": "first"} - ) - ) - if not other_rows.empty: - parts.append( - other_rows.groupby(group_keys, as_index=False).agg( - {"normalized_value": "sum", "normalized_units": "first"} - ) - ) - if parts: - subset = pd.concat(parts, ignore_index=True) - else: - return subset.iloc[:0] + # 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) @@ -369,12 +345,18 @@ def run_checks( """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)} years") + 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/atm_parser.py b/zppy_interfaces/budget_analysis/ingestion/atm_parser.py index 61a4d8a..a87602c 100644 --- a/zppy_interfaces/budget_analysis/ingestion/atm_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/atm_parser.py @@ -172,6 +172,9 @@ def _gather_energy_data(filename: str) -> Tuple[List[Dict], List[Dict]]: 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. @@ -203,47 +206,67 @@ def parse_raw(self, log_files: List[str]) -> Tuple[pd.DataFrame, pd.DataFrame]: 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 monthly event table. + """Parse atm logs and return a tidy event table. - Aggregates per-step data to monthly averages for fluxes and - beginning/end-of-month values for states. + 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 (monthly aggregation) --- + # --- 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 (year, month), grp in flux_df.groupby(["year", "month"]): + 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: year, + COL_TIME: time, COL_COMPONENT: "atm", COL_QUANTITY: "water", COL_SOURCE: "atm", - COL_PERIOD: "monthly", + COL_PERIOD: self.frequency, } base_heat = { - COL_TIME: year, + COL_TIME: time, COL_COMPONENT: "atm", COL_QUANTITY: "heat", COL_SOURCE: "atm", - COL_PERIOD: "monthly", + COL_PERIOD: self.frequency, } dt = grp["dt"].iloc[0] - # Water flux terms (monthly mean rate, kg/m2/s) + # Water flux terms (mean rate, kg/m2/s) if "w_flux" in grp.columns: rows.append( { @@ -265,7 +288,7 @@ def parse_files( } ) - # Water state (begin/end of month) + # Water state (begin/end of period) rows.append( { **base_water, @@ -285,7 +308,7 @@ def parse_files( } ) - # Water diagnostics (monthly mean per step) + # Water diagnostics (mean per step) if "w_residual" in grp.columns: rows.append( { @@ -307,7 +330,7 @@ def parse_files( } ) - # Energy flux terms (monthly mean, W/m2) + # Energy flux terms (mean, W/m2) if "e_dtedt" in grp.columns: rows.append( { @@ -329,7 +352,7 @@ def parse_files( } ) - # Energy diagnostics (monthly mean, W/m2) + # Energy diagnostics (mean, W/m2) if "e_diff" in grp.columns: rows.append( { @@ -351,19 +374,24 @@ def parse_files( } ) - # --- Energy fixer state (TE begin/end of month) --- + # --- 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 (year, month), grp in te_df.groupby(["year", "month"]): + 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: year, + COL_TIME: time, COL_COMPONENT: "atm", COL_QUANTITY: "heat", COL_SOURCE: "atm", - COL_PERIOD: "monthly", + COL_PERIOD: self.frequency, } rows.append( { diff --git a/zppy_interfaces/budget_analysis/ingestion/base.py b/zppy_interfaces/budget_analysis/ingestion/base.py index 17d2447..73c1db3 100644 --- a/zppy_interfaces/budget_analysis/ingestion/base.py +++ b/zppy_interfaces/budget_analysis/ingestion/base.py @@ -5,9 +5,28 @@ import pandas as pd +VALID_FREQUENCIES = ("monthly", "annual") + class BaseParser(ABC): - """All parsers return a tidy event table DataFrame.""" + """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( diff --git a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py index f97c4d4..05f16fb 100644 --- a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py @@ -39,19 +39,39 @@ def _normalize_component_name(name: str) -> str: return name.strip().replace(" ", "_") -def _parse_datestamp(datestamp: str) -> int: - """Convert coupler datestamp to year. +def _parse_datestamp(datestamp: str) -> Tuple[int, int]: + """Convert coupler datestamp to (year, month). The date is reported at the start of the next period. - E.g. '20101' -> strip last 4 chars -> '2' -> minus 1 -> year 1. + E.g. '20101' -> MMDD='0101', year_part='2', year=2-1=1, month=01. + For annual period the month is ignored downstream. """ - return int(datestamp[:-4]) - 1 + mmdd = datestamp[-4:] + month = int(mmdd[:2]) + year = int(datestamp[:-4]) - 1 + # Roll back one month (date is start of *next* period) + if month == 1: + month = 12 + # year already decremented above + 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, int]]: - """Extract period and year from a budget header line. +def _parse_header_line(line: str, pattern: str) -> Optional[Tuple[str, float]]: + """Extract period and time from a budget header line. - Returns (period, year) or None on failure. + Returns (period, time) or None on failure. """ if not line.startswith(pattern): return None @@ -63,11 +83,11 @@ def _parse_header_line(line: str, pattern: str) -> Optional[Tuple[str, int]]: return None period = period_match.group(1) - year = _parse_datestamp(date_match.group(1)) - return period, year + year, month = _parse_datestamp(date_match.group(1)) + return period, _make_time(period, year, month) -def _parse_table(f: TextIO, year: int, quantity: str, period: str) -> List[Dict]: +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] @@ -119,7 +139,12 @@ def _parse_table(f: TextIO, year: int, quantity: str, period: str) -> List[Dict] class CplParser(BaseParser): """Parse coupler log budget tables into a tidy event table.""" - def __init__(self, quantities: Optional[List[str]] = None): + def __init__( + self, + quantities: Optional[List[str]] = None, + frequency: str = "annual", + ): + super().__init__(frequency=frequency) self.quantities = quantities or ["water", "heat"] def parse_files( @@ -137,9 +162,11 @@ def parse_files( result = _parse_header_line(line, pattern) if result is None: continue - period, year = result - if start_year <= year <= end_year: - rows.extend(_parse_table(f, year, quantity, period)) + period, time = result + if period != self.frequency: + continue + if start_year <= time <= end_year + 1: + rows.extend(_parse_table(f, time, quantity, period)) except Exception as e: print(f"WARNING: Error processing {fname}: {e}") continue diff --git a/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py index dacfcad..4abb781 100644 --- a/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py @@ -45,23 +45,37 @@ STATE_HEADER = "WATER STATES (kg/m2*1e6): period" -def _parse_datestamp(datestamp: str) -> int: - """Convert datestamp to year. Same convention as coupler.""" - return int(datestamp[:-4]) - 1 - - -def _parse_period_and_year(line: str) -> Optional[Tuple[str, int]]: - """Extract period and year from a header line.""" +def _parse_datestamp(datestamp: str) -> Tuple[int, int]: + """Convert datestamp to (year, month). Same convention as coupler.""" + mmdd = datestamp[-4:] + month = int(mmdd[:2]) + year = int(datestamp[:-4]) - 1 + if month == 1: + month = 12 + 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 = _parse_datestamp(date_match.group(1)) - return period, year + year, month = _parse_datestamp(date_match.group(1)) + return period, _make_time(period, year, month) -def _parse_flux_table(f: TextIO, year: int, period: str) -> List[Dict]: +def _parse_flux_table(f: TextIO, year: float, period: str) -> List[Dict]: """Parse a NET WATER FLUXES table. Format: @@ -163,7 +177,7 @@ def _parse_flux_table(f: TextIO, year: int, period: str) -> List[Dict]: return rows -def _parse_state_table(f: TextIO, year: int, period: str) -> List[Dict]: +def _parse_state_table(f: TextIO, year: float, period: str) -> List[Dict]: """Parse a WATER STATES table. Format: @@ -273,6 +287,9 @@ def _parse_state_table(f: TextIO, year: int, period: str) -> List[Dict]: 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: @@ -284,20 +301,24 @@ def parse_files( stripped = line.strip() if stripped.startswith(FLUX_HEADER): - result = _parse_period_and_year(stripped) + result = _parse_period_and_time(stripped) if result is None: continue - period, year = result - if start_year <= year <= end_year: - rows.extend(_parse_flux_table(f, year, period)) + period, time = result + if period != self.frequency: + continue + if start_year <= time <= end_year + 1: + rows.extend(_parse_flux_table(f, time, period)) elif stripped.startswith(STATE_HEADER): - result = _parse_period_and_year(stripped) + result = _parse_period_and_time(stripped) if result is None: continue - period, year = result - if start_year <= year <= end_year: - rows.extend(_parse_state_table(f, year, period)) + period, time = result + if period != self.frequency: + continue + if start_year <= time <= end_year + 1: + rows.extend(_parse_state_table(f, time, period)) except Exception as e: print(f"WARNING: Error processing {fname}: {e}") diff --git a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py index 41e6a90..356bc19 100644 --- a/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/ocn_parser.py @@ -227,6 +227,9 @@ def _parse_energy_summary(f: TextIO) -> Optional[Dict[str, float]]: 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: @@ -244,16 +247,38 @@ def parse_files( year, month = _parse_date(date_str) if year < start_year or year > end_year: continue - rows.extend(self._parse_block(f, year, month)) + 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) - return pd.DataFrame(rows, 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, year: int, month: int) -> List[Dict]: + 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 @@ -268,21 +293,21 @@ def _parse_block(self, f: TextIO, year: int, month: int) -> List[Dict]: if "MASS CONSERVATION CHECK" in line and "SUMMARY" not in line: found_mass = True - rows.extend(self._parse_mass_section(f, year)) + 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, year)) + rows.extend(self._parse_energy_section(f, time)) line = f.readline() return rows - def _parse_mass_section(self, f: TextIO, year: int) -> List[Dict]: + def _parse_mass_section(self, f: TextIO, time: float) -> List[Dict]: """Parse MASS CONSERVATION CHECK: fluxes + summary.""" rows: List[Dict] = [] base = { - COL_TIME: year, + COL_TIME: time, COL_COMPONENT: "ocn", COL_QUANTITY: "water", COL_UNITS: "kg/m2s*1e6", @@ -325,11 +350,11 @@ def _parse_mass_section(self, f: TextIO, year: int) -> List[Dict]: return rows - def _parse_energy_section(self, f: TextIO, year: int) -> List[Dict]: + def _parse_energy_section(self, f: TextIO, time: float) -> List[Dict]: """Parse ENERGY CONSERVATION CHECK: fluxes + summary.""" rows: List[Dict] = [] base = { - COL_TIME: year, + COL_TIME: time, COL_COMPONENT: "ocn", COL_QUANTITY: "heat", COL_UNITS: "W/m2", diff --git a/zppy_interfaces/budget_analysis/schema.py b/zppy_interfaces/budget_analysis/schema.py index 97bfadd..068883d 100644 --- a/zppy_interfaces/budget_analysis/schema.py +++ b/zppy_interfaces/budget_analysis/schema.py @@ -5,7 +5,7 @@ import pandas as pd # Column name constants -COL_TIME = "time" # int: year the period ends (e.g. year 1 annual → time=1) +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*" ) diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py index 946d4bb..dffef78 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -9,7 +9,7 @@ import numpy as np import pandas as pd from bokeh.layouts import column -from bokeh.models import Div +from bokeh.models import Div, FixedTicker from bokeh.palettes import Category10 from bokeh.plotting import figure, output_file, save @@ -141,19 +141,26 @@ def generate_landing_page(output_dir: str, report_paths: Dict[str, str]) -> str: return index_path -def _make_figure(title: str, y_label: str) -> figure: - return figure( +def _make_figure(title: str, y_label: str, years: np.ndarray) -> figure: + p = figure( title=title, height=350, width=1200, - x_axis_label="year", + 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})") + p = _make_figure(title, f"residual ({units})", r.years) p.line(r.years, r.residual, line_width=2, color="red") p.line( r.years, @@ -169,14 +176,14 @@ 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})") + 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) + 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" @@ -187,7 +194,9 @@ 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) + 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 From aa66c307a7697aa52f7058bba50c30be3293f959 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 12 Feb 2026 21:59:49 -0600 Subject: [PATCH 09/15] fix time misalignment between cpl and components --- .../budget_analysis/ingestion/cpl_parser.py | 33 ++++++++++------- .../budget_analysis/ingestion/lnd_parser.py | 35 +++++++++++++------ 2 files changed, 45 insertions(+), 23 deletions(-) diff --git a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py index 05f16fb..9ca5ae1 100644 --- a/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/cpl_parser.py @@ -39,23 +39,28 @@ def _normalize_component_name(name: str) -> str: return name.strip().replace(" ", "_") -def _parse_datestamp(datestamp: str) -> Tuple[int, int]: +def _parse_datestamp(datestamp: str, period: str = "monthly") -> Tuple[int, int]: """Convert coupler datestamp to (year, month). - The date is reported at the start of the next period. - E.g. '20101' -> MMDD='0101', year_part='2', year=2-1=1, month=01. - For annual period the month is ignored downstream. + 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]) - 1 - # Roll back one month (date is start of *next* period) - if month == 1: - month = 12 - # year already decremented above + 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: - month -= 1 - return year, month + # 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: @@ -83,7 +88,7 @@ def _parse_header_line(line: str, pattern: str) -> Optional[Tuple[str, float]]: return None period = period_match.group(1) - year, month = _parse_datestamp(date_match.group(1)) + year, month = _parse_datestamp(date_match.group(1), period) return period, _make_time(period, year, month) @@ -165,7 +170,9 @@ def parse_files( period, time = result if period != self.frequency: continue - if start_year <= time <= end_year + 1: + # 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}") diff --git a/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py index 4abb781..bbcda72 100644 --- a/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/lnd_parser.py @@ -45,16 +45,27 @@ STATE_HEADER = "WATER STATES (kg/m2*1e6): period" -def _parse_datestamp(datestamp: str) -> Tuple[int, int]: - """Convert datestamp to (year, month). Same convention as coupler.""" +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]) - 1 - if month == 1: - month = 12 + 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: - month -= 1 - return year, month + # 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: @@ -71,7 +82,7 @@ def _parse_period_and_time(line: str) -> Optional[Tuple[str, float]]: if not period_match or not date_match: return None period = period_match.group(1) - year, month = _parse_datestamp(date_match.group(1)) + year, month = _parse_datestamp(date_match.group(1), period) return period, _make_time(period, year, month) @@ -307,7 +318,9 @@ def parse_files( period, time = result if period != self.frequency: continue - if start_year <= time <= end_year + 1: + # 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): @@ -317,7 +330,9 @@ def parse_files( period, time = result if period != self.frequency: continue - if start_year <= time <= end_year + 1: + # 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: From 87ba0051095dffb00a316de1e0304feefe0da5c7 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 12 Feb 2026 22:45:32 -0600 Subject: [PATCH 10/15] add ice log parser --- zppy_interfaces/budget_analysis/__main__.py | 9 + .../budget_analysis/ingestion/ice_parser.py | 268 ++++++++++++++++++ 2 files changed, 277 insertions(+) create mode 100644 zppy_interfaces/budget_analysis/ingestion/ice_parser.py diff --git a/zppy_interfaces/budget_analysis/__main__.py b/zppy_interfaces/budget_analysis/__main__.py index b551ef4..d1a67f9 100644 --- a/zppy_interfaces/budget_analysis/__main__.py +++ b/zppy_interfaces/budget_analysis/__main__.py @@ -138,6 +138,7 @@ def _run_whole_model(args) -> int: ) 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 @@ -157,6 +158,7 @@ def _run_whole_model(args) -> int: 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: @@ -165,6 +167,7 @@ def _run_whole_model(args) -> int: 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 @@ -186,6 +189,12 @@ def _run_whole_model(args) -> int: 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( 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..1b2053e --- /dev/null +++ b/zppy_interfaces/budget_analysis/ingestion/ice_parser.py @@ -0,0 +1,268 @@ +"""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, 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. + 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 + + # 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: + 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] = [] + + # Read through the conservation block + while True: + line = f.readline() + if not line: # EOF + break + + # Stop at next conservation block + if "Conservation checks:" in line: + break + + # Parse energy section + if "Energy conservation check" in line: + energy_rows = _parse_energy_section(f, time) + rows.extend(energy_rows) + # Note: _parse_energy_section will stop when it hits "Mass conservation check" + + # Parse mass section + elif "Mass conservation check" in line: + mass_rows = _parse_mass_section(f, time) + rows.extend(mass_rows) + + return rows \ No newline at end of file From 4ec5156d90f45e6f9c69b0b3998a9a86395c04d8 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Fri, 13 Feb 2026 13:46:58 -0600 Subject: [PATCH 11/15] handle ice log variation;complete ice budget analysis --- zppy_interfaces/budget_analysis/checks.py | 186 ++++++++++++++++++ .../budget_analysis/ingestion/ice_parser.py | 27 ++- .../budget_analysis/normalization.py | 17 +- zppy_interfaces/budget_analysis/viz.py | 17 ++ 4 files changed, 234 insertions(+), 13 deletions(-) diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index 8cc783b..ee3984c 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -317,12 +317,196 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: ) +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})", + ) + + 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"), LndClosure(), OcnClosure(quantity="water"), + IceClosure(quantity="water"), ] DEFAULT_HEAT_CHECKS: List[BudgetCheck] = [ @@ -330,7 +514,9 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: InterfaceMatch( "ocn", "ocn", quantity="heat", comp_sum_term="SUM IMP+EXP HEAT FLUXES" ), + IceInterfaceMatch(quantity="heat"), OcnClosure(quantity="heat"), + IceClosure(quantity="heat"), ] DEFAULT_CARBON_CHECKS: List[BudgetCheck] = [ diff --git a/zppy_interfaces/budget_analysis/ingestion/ice_parser.py b/zppy_interfaces/budget_analysis/ingestion/ice_parser.py index 1b2053e..1b1e4d7 100644 --- a/zppy_interfaces/budget_analysis/ingestion/ice_parser.py +++ b/zppy_interfaces/budget_analysis/ingestion/ice_parser.py @@ -37,7 +37,7 @@ import gzip import re -from typing import Dict, List, Optional, TextIO, Tuple +from typing import Dict, List, TextIO, Tuple import pandas as pd @@ -104,6 +104,8 @@ def _parse_energy_section(f: TextIO, time: float) -> List[Dict]: 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: @@ -157,7 +159,7 @@ def _parse_mass_section(f: TextIO, time: float) -> List[Dict]: # Stop at end of section if "Conservation checks:" in line: break - if "===================" in line: + if "---" in line and len([c for c in line if c == "-"]) > 10: break # Parse mass flux terms @@ -202,7 +204,9 @@ def parse_files( 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) + 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) @@ -243,11 +247,13 @@ def parse_files( 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 - # Read through the conservation block - while True: - line = f.readline() - if not line: # EOF + line = f.readline() + while line: + # Stop when both sections are processed + if found_mass and found_energy: break # Stop at next conservation block @@ -256,13 +262,16 @@ def _parse_block(self, f: TextIO, time: float) -> List[Dict]: # Parse energy section if "Energy conservation check" in line: + found_energy = True energy_rows = _parse_energy_section(f, time) rows.extend(energy_rows) - # Note: _parse_energy_section will stop when it hits "Mass conservation check" # Parse mass section elif "Mass conservation check" in line: + found_mass = True mass_rows = _parse_mass_section(f, time) rows.extend(mass_rows) - return rows \ No newline at end of file + line = f.readline() + + return rows diff --git a/zppy_interfaces/budget_analysis/normalization.py b/zppy_interfaces/budget_analysis/normalization.py index 0a6a0c4..5ca59bb 100644 --- a/zppy_interfaces/budget_analysis/normalization.py +++ b/zppy_interfaces/budget_analysis/normalization.py @@ -2,7 +2,7 @@ import pandas as pd -from .schema import COL_QUANTITY, COL_TABLE_TYPE, COL_UNITS, COL_VALUE +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 @@ -31,10 +31,19 @@ def normalize(df: pd.DataFrame) -> pd.DataFrame: # --- Water --- water_mask = df[COL_QUANTITY] == "water" - # Flux rates: kg/m2s * 1e6 -> mm/yr + # Flux rates: kg/m2s * 1e6 -> mm/yr (ocean), kg/m2s -> mm/yr (ice) water_flux = water_mask & (df[COL_TABLE_TYPE] == "flux") - df.loc[water_flux, "normalized_value"] = ( - df.loc[water_flux, COL_VALUE] * SECONDS_PER_YEAR / 1e6 + + # Ice data: kg/m2s -> mm/yr (no 1e6 factor) + ice_flux = water_flux & (df[COL_SOURCE] == "ice") + df.loc[ice_flux, "normalized_value"] = ( + df.loc[ice_flux, COL_VALUE] * SECONDS_PER_YEAR + ) + + # Ocean/other data: kg/m2s*1e6 -> mm/yr (with 1e6 factor) + other_flux = water_flux & (df[COL_SOURCE] != "ice") + df.loc[other_flux, "normalized_value"] = ( + df.loc[other_flux, COL_VALUE] * SECONDS_PER_YEAR / 1e6 ) df.loc[water_flux, "normalized_units"] = "mm/yr" diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py index dffef78..afbf958 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -102,6 +102,23 @@ def generate_budget_report( ) ) + 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, + ) + ) + if not plots: print(f"No {quantity} plots generated — no check results available") return "" From a3d7e3514aa6a60e1929a92d50b8e8d5d3158608 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Fri, 13 Feb 2026 18:13:09 -0600 Subject: [PATCH 12/15] complete atm closure check(water only), energy budget missing terms --- zppy_interfaces/budget_analysis/checks.py | 242 +++++++++++++++++- .../budget_analysis/normalization.py | 30 ++- zppy_interfaces/budget_analysis/viz.py | 17 ++ 3 files changed, 277 insertions(+), 12 deletions(-) diff --git a/zppy_interfaces/budget_analysis/checks.py b/zppy_interfaces/budget_analysis/checks.py index ee3984c..f4fc2a4 100644 --- a/zppy_interfaces/budget_analysis/checks.py +++ b/zppy_interfaces/budget_analysis/checks.py @@ -1,4 +1,12 @@ -"""Budget checks: definitions and evaluation.""" +"""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 @@ -16,6 +24,25 @@ 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: @@ -499,14 +526,225 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: ) +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] = [ @@ -515,8 +753,10 @@ def evaluate(self, df: pd.DataFrame) -> Optional[CheckResult]: "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] = [ diff --git a/zppy_interfaces/budget_analysis/normalization.py b/zppy_interfaces/budget_analysis/normalization.py index 5ca59bb..8601823 100644 --- a/zppy_interfaces/budget_analysis/normalization.py +++ b/zppy_interfaces/budget_analysis/normalization.py @@ -31,29 +31,37 @@ def normalize(df: pd.DataFrame) -> pd.DataFrame: # --- Water --- water_mask = df[COL_QUANTITY] == "water" - # Flux rates: kg/m2s * 1e6 -> mm/yr (ocean), kg/m2s -> mm/yr (ice) + # Flux rates: kg/m2s * 1e6 -> mm/yr (ocean), kg/m2s -> mm/yr (ice, atm) water_flux = water_mask & (df[COL_TABLE_TYPE] == "flux") - # Ice data: kg/m2s -> mm/yr (no 1e6 factor) - ice_flux = water_flux & (df[COL_SOURCE] == "ice") - df.loc[ice_flux, "normalized_value"] = ( - df.loc[ice_flux, COL_VALUE] * SECONDS_PER_YEAR + # 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) - other_flux = water_flux & (df[COL_SOURCE] != "ice") - df.loc[other_flux, "normalized_value"] = ( - df.loc[other_flux, COL_VALUE] * SECONDS_PER_YEAR / 1e6 + 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 + # 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"] ) - df.loc[water_integrated, "normalized_value"] = ( - df.loc[water_integrated, COL_VALUE] / 1e6 + + # 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 --- diff --git a/zppy_interfaces/budget_analysis/viz.py b/zppy_interfaces/budget_analysis/viz.py index afbf958..a22354d 100644 --- a/zppy_interfaces/budget_analysis/viz.py +++ b/zppy_interfaces/budget_analysis/viz.py @@ -119,6 +119,23 @@ def generate_budget_report( ) ) + 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 "" From 4e6c024c1f46b6d6a118859db25d7b4e2e4db32c Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Wed, 18 Feb 2026 13:05:16 -0600 Subject: [PATCH 13/15] add CLAUDE.md and README.md --- CLAUDE.md | 108 ++++++++++++ zppy_interfaces/budget_analysis/README.md | 197 ++++++++++++++++++++++ 2 files changed, 305 insertions(+) create mode 100644 CLAUDE.md create mode 100644 zppy_interfaces/budget_analysis/README.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..f5c503d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,108 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working 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 \ No newline at end of file diff --git a/zppy_interfaces/budget_analysis/README.md b/zppy_interfaces/budget_analysis/README.md new file mode 100644 index 0000000..19112a6 --- /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 \ No newline at end of file From af2a8a9c0f80c1cfb69ec667387590b3f821dad9 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Wed, 18 Feb 2026 13:26:19 -0600 Subject: [PATCH 14/15] fix pre-commit error --- CLAUDE.md | 2 +- zppy_interfaces/budget_analysis/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index f5c503d..0ac4144 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -105,4 +105,4 @@ Test the main CLI applications: - 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 \ No newline at end of file +- The codebase follows strict code quality standards with Black formatting and comprehensive linting diff --git a/zppy_interfaces/budget_analysis/README.md b/zppy_interfaces/budget_analysis/README.md index 19112a6..458c3d1 100644 --- a/zppy_interfaces/budget_analysis/README.md +++ b/zppy_interfaces/budget_analysis/README.md @@ -194,4 +194,4 @@ This module is designed to be called by zppy for automated budget analysis in E3 - 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 \ No newline at end of file +- Large log files may require significant memory for processing From 4d9becddd008b60c072fe51a78ff51bed032ce13 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Wed, 18 Feb 2026 14:15:11 -0600 Subject: [PATCH 15/15] update AGENTS.md and CLAUDE.md follow project standard --- .claude/CLAUDE.md | 5 +++++ CLAUDE.md => AGENTS.md | 6 ++++-- 2 files changed, 9 insertions(+), 2 deletions(-) create mode 100644 .claude/CLAUDE.md rename CLAUDE.md => AGENTS.md (97%) 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/CLAUDE.md b/AGENTS.md similarity index 97% rename from CLAUDE.md rename to AGENTS.md index 0ac4144..32156e2 100644 --- a/CLAUDE.md +++ b/AGENTS.md @@ -1,6 +1,8 @@ -# CLAUDE.md + +# AGENTS.md -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +This file provides guidance to work with code in this repository. ## Project Overview