Add script to call zi-budget-analysis

zppy/__main__.py

@@ -12,6 +12,7 @@
 from mache import MachineInfo
 from validate import Validator
 
+from zppy.budget import budget
 from zppy.bundle import Bundle, predefined_bundles
 from zppy.climo import climo
 from zppy.e3sm_diags import e3sm_diags
@@ -271,6 +272,9 @@ def _launch_scripts(config: ConfigObj, script_dir, job_ids_file, plugins) -> Non
     # ilamb tasks
     existing_bundles = ilamb(config, script_dir, existing_bundles, job_ids_file)
 
+    # budget tasks
+    existing_bundles = budget(config, script_dir, existing_bundles, job_ids_file)
+
     # pcmdi_diags tasks
     existing_bundles = pcmdi_diags(config, script_dir, existing_bundles, job_ids_file)
 

zppy/budget.py

@@ -0,0 +1,79 @@
+from typing import Any, Dict, List, Tuple
+
+from configobj import ConfigObj
+
+from zppy.bundle import handle_bundles
+from zppy.utils import (
+    check_status,
+    get_file_names,
+    get_tasks,
+    get_years,
+    initialize_template,
+    make_executable,
+    submit_script,
+    write_settings_file,
+)
+
+
+# -----------------------------------------------------------------------------
+def budget(config: ConfigObj, script_dir: str, existing_bundles, job_ids_file):
+
+    template, _ = initialize_template(config, "budget.bash")
+
+    # --- List of budget tasks ---
+    tasks: List[Dict[str, Any]] = get_tasks(config, "budget")
+    if len(tasks) == 0:
+        return existing_bundles
+
+    # --- Generate and submit budget scripts ---
+    for c in tasks:
+
+        dependencies: List[str] = []
+
+        # Loop over year sets
+        year_sets: List[Tuple[int, int]] = get_years(c["years"])
+        for s in year_sets:
+            c["year1"] = s[0]
+            c["year2"] = s[1]
+            if ("last_year" in c.keys()) and (c["year2"] > c["last_year"]):
+                continue  # Skip this year set
+            c["scriptDir"] = script_dir
+
+            prefix = f"budget_{c['year1']:04d}-{c['year2']:04d}"
+            print(prefix)
+            c["prefix"] = prefix
+            bash_file, settings_file, status_file = get_file_names(script_dir, prefix)
+            skip: bool = check_status(status_file)
+            if skip:
+                continue
+            # Create script
+            with open(bash_file, "w") as f:
+                f.write(template.render(**c))
+            make_executable(bash_file)
+            c["dependencies"] = dependencies
+            write_settings_file(settings_file, c, s)
+            export = "NONE"
+            existing_bundles = handle_bundles(
+                c,
+                bash_file,
+                export,
+                dependFiles=dependencies,
+                existing_bundles=existing_bundles,
+            )
+            if not c["dry_run"]:
+                if c["bundle"] == "":
+                    # Submit job
+                    submit_script(
+                        bash_file,
+                        status_file,
+                        export,
+                        job_ids_file,
+                        dependFiles=dependencies,
+                        fail_on_dependency_skip=c["fail_on_dependency_skip"],
+                    )
+                else:
+                    print(f"...adding to bundle {c['bundle']}")
+
+            print(f"   environment_commands={c['environment_commands']}")
+
+    return existing_bundles

zppy/defaults/default.ini

@@ -179,6 +179,19 @@ input_files = string(default="eam.h2")
 # If not set, zppy will attempt to infer this value from the topography file.
 res = string(default="")
 
+[budget]
+# Budget types to analyze: water, heat (comma-separated)
+budget_types = string(default="water,heat")
+# Generate HTML plots with interactive visualizations
+output_html = boolean(default=True)
+# Generate ASCII summary tables
+output_ascii = boolean(default=True)
+
+  [[__many__]]
+  budget_types = string(default=None)
+  output_html = boolean(default=None)
+  output_ascii = boolean(default=None)
+
 [e3sm_diags]
 # See https://e3sm-project.github.io/e3sm_diags/_build/html/master/available-parameters.html
 backend = string(default="mpl")

zppy/templates/budget.bash

@@ -0,0 +1,147 @@
+#!/bin/bash
+{% include 'inclusions/slurm_header.bash' %}
+{% include 'inclusions/boilerplate.bash' %}
+set -e
+{{ environment_commands }}
+set +e
+
+# Generate water and energy budget analysis plots
+################################################################################
+
+# Create temporary workdir
+hash=`mktemp --dry-run -d XXXX`
+workdir=tmp.{{ prefix }}.${id}.${hash}
+mkdir ${workdir}
+cd ${workdir}
+echo "Working in temporary directory: ${workdir}"
+
+# Execute budget analysis using external CLI tool
+echo "Running budget analysis for years {{ year1 }} to {{ year2 }}..."
+
+# Construct log path from zppy parameters
+LOG_PATH="{{ input }}/{{ case }}/archive/logs"
+
+zi-budget-analysis \
+  --log_path "${LOG_PATH}" \
+  --start_year {{ year1 }} \
+  --end_year {{ year2 }} \
+  --budget_types {{ budget_types }} \
+{%- if output_html %}
+  --output_html \
+{%- endif %}
+{%- if output_ascii %}
+  --output_ascii \
+{%- endif %}
+  --output_dir .
+
+if [ $? != 0 ]; then
+  cd {{ scriptDir }}
+  echo 'ERROR (1)' > {{ prefix }}.status
+  exit 1
+fi
+
+# Copy results to case directory
+echo
+echo "===== COPY RESULTS TO CASE DIRECTORY ====="
+echo
+
+case_dir={{ output }}
+dest_dir=${case_dir}/post/budget/{{ year1 }}-{{ year2 }}
+mkdir -p ${dest_dir}
+if [ $? != 0 ]; then
+  cd {{ scriptDir }}
+  echo 'ERROR (2)' > {{ prefix }}.status
+  exit 2
+fi
+
+# Copy all generated files
+cp *.html ${dest_dir}/ 2>/dev/null || true
+cp *.txt ${dest_dir}/ 2>/dev/null || true
+cp *.png ${dest_dir}/ 2>/dev/null || true
+
+echo "Results copied to ${dest_dir}/"
+
+# Copy output to web server
+echo
+echo "===== COPY FILES TO WEB SERVER ====="
+echo
+
+# Create web directory
+www={{ www }}
+case={{ case }}
+web_dir=${www}/${case}/budget/{{ year1 }}-{{ year2 }}
+mkdir -p ${web_dir}
+if [ $? != 0 ]; then
+  cd {{ scriptDir }}
+  echo 'ERROR (3)' > {{ prefix }}.status
+  exit 3
+fi
+
+{% if machine in ['pm-cpu', 'pm-gpu'] %}
+# For NERSC, make sure it is world readable
+f=`realpath ${web_dir}`
+while [[ $f != "/" ]]
+do
+  owner=`stat --format '%U' $f`
+  if [ "${owner}" = "${USER}" ]; then
+    chgrp e3sm $f
+    chmod go+rx $f
+  fi
+  f=$(dirname $f)
+done
+{% endif %}
+
+# Copy HTML files for web viewing
+cp *.html ${web_dir}/ 2>/dev/null || true
+if [ $? != 0 ]; then
+  cd {{ scriptDir }}
+  echo 'ERROR (4)' > {{ prefix }}.status
+  exit 4
+fi
+
+{% if machine in ['pm-cpu', 'pm-gpu'] %}
+# For NERSC, change permissions of new files
+pushd ${web_dir}
+chgrp -R e3sm *.html 2>/dev/null || true
+chmod -R go+rX,go-w *.html 2>/dev/null || true
+popd
+{% endif %}
+
+{% if machine in ['anvil', 'chrysalis'] %}
+# For LCRC, change permissions of new files
+pushd ${web_dir}
+chmod -R go+rX,go-w *.html 2>/dev/null || true
+popd
+{% endif %}
+
+echo "Web files copied to ${web_dir}/"
+
+# Clean up temporary workdir
+echo
+echo "===== CLEANUP TEMPORARY DIRECTORY ====="
+echo
+
+cd {{ scriptDir }}
+if [[ "${debug,,}" != "true" ]]; then
+  rm -rf ${workdir}
+  if [ $? = 0 ]; then
+    echo "Successfully cleaned up ${workdir}"
+  else
+    echo "Warning: Failed to clean up ${workdir}"
+  fi
+else
+  echo "Debug mode: keeping temporary directory ${workdir}"
+fi
+
+################################################################################
+# Update status file and exit
+{% raw %}
+ENDTIME=$(date +%s)
+ELAPSEDTIME=$(($ENDTIME - $STARTTIME))
+{% endraw %}
+echo ==============================================
+echo "Elapsed time: $ELAPSEDTIME seconds"
+echo ==============================================
+rm -f {{ prefix }}.status
+echo 'OK' > {{ prefix }}.status
+exit 0