Workflow for ZPPY- PCMDI for E3SM diagnostics

conda/chrisalis.yml

@@ -0,0 +1,34 @@
+name: zppy_dev
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  # Base
+  # =================
+  - python=3.9.13
+  - pip=22.2.2
+  - configobj=5.0.6
+  - jinja2=3.1.2
+  - mache>=1.5.0
+  - mpas_tools>=0.15.0
+  - pillow=9.2.0
+  # Developer Tools
+  # =================
+  # If versions are updated, also update 'rev' in `.pre-commit-config.yaml`
+  - black=22.8.0        # version from https://anaconda.org/conda-forge/black
+  - flake8=5.0.4        # version from https://anaconda.org/conda-forge/flake8
+  # This line also implicitly installs isort
+  - flake8-isort=4.2.0  # version from https://anaconda.org/conda-forge/flake8-isort
+  - mypy=0.982          # version from https://anaconda.org/conda-forge/mypy
+  - pre-commit=2.20.0   # version from https://anaconda.org/conda-forge/pre-commit
+  - tbump=6.9.0
+  # Documentation
+  # If versions are updated, also update in `.github/workflows/build_workflow.yml`
+  # =================
+  - sphinx=5.2.3
+  - sphinx-multiversion=0.2.4
+  - sphinx_rtd_theme=1.0.0
+  # Need to pin docutils because 0.17 has a bug with unordered lists
+  # https://github.com/readthedocs/sphinx_rtd_theme/issues/1115
+  - docutils=0.16
+prefix: /home/ac.szhang/.conda/envs/zppy_dev

zppy/__main__.py

@@ -15,6 +15,7 @@
 from zppy.global_time_series import global_time_series
 from zppy.ilamb import ilamb
 from zppy.mpas_analysis import mpas_analysis
+from zppy.pcmdi_diags import pcmdi_diags
 from zppy.tc_analysis import tc_analysis
 from zppy.ts import ts
 from zppy.utils import checkStatus, submitScript
@@ -209,6 +210,9 @@ def main():  # noqa: C901
     # ilamb tasks
     existing_bundles = ilamb(config, scriptDir, existing_bundles, job_ids_file)
 
+    # pcmdi_diags tasks
+    existing_bundles = pcmdi_diags(config, scriptDir, existing_bundles, job_ids_file)
+
     # zppy external plugins
     for plugin in plugins:
         # Get plugin module function

zppy/pcmdi_diags.py

@@ -0,0 +1,158 @@
+import os
+import pprint
+from typing import List
+
+import jinja2
+
+from zppy.bundle import handle_bundles
+from zppy.utils import (
+    add_dependencies,
+    checkStatus,
+    getTasks,
+    getYears,
+    makeExecutable,
+    print_url,
+    submitScript,
+)
+
+
+# -----------------------------------------------------------------------------
+def pcmdi_diags(config, scriptDir, existing_bundles, job_ids_file):
+
+    # Initialize jinja2 template engine
+    templateLoader = jinja2.FileSystemLoader(
+        searchpath=config["default"]["templateDir"]
+    )
+    templateEnv = jinja2.Environment(loader=templateLoader)
+    template = templateEnv.get_template("pcmdi_diags.bash")
+
+    # --- List of pcmdi_diags tasks ---
+    tasks = getTasks(config, "pcmdi_diags")
+    if len(tasks) == 0:
+        return existing_bundles
+
+    # --- Generate and submit pcmdi_diags scripts ---
+    dependencies: List[str] = []
+
+    for c in tasks:
+
+        c["scriptDir"] = scriptDir
+
+        if "ts_num_years" in c.keys():
+            c["ts_num_years"] = int(c["ts_num_years"])
+
+        # procedure type for e3sm_to_cmip
+        c["cmor_tables_prefix"] = c["diagnostics_base_path"]
+
+        # Loop over year sets
+        year_sets = getYears(c["ts_years"])
+        if ("ref_years" in c.keys()) and (c["ref_years"] != [""]):
+            ref_year_sets = getYears(c["ref_years"])
+        else:
+            ref_year_sets = year_sets
+        for s, rs in zip(year_sets, ref_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["ref_year1"] = rs[0]
+            c["ref_year2"] = rs[1]
+            if c["subsection"]:
+                c["sub"] = c["subsection"]
+            else:
+                c["sub"] = c["grid"]
+            # Make a guess for observation paths, if need be
+            if ("ts_num_years" in c.keys()) and (c["obs_ts"] == ""):
+                c[
+                    "obs_ts"
+                ] = f"{c['diagnostics_base_path']}/observations/Atm/time-series/"
+            if c["run_type"] == "model_vs_obs":
+                prefix = "pcmdi_diags_%s_%s_%04d-%04d" % (
+                    c["sub"],
+                    c["tag"],
+                    c["year1"],
+                    c["year2"],
+                )
+            elif c["run_type"] == "model_vs_model":
+                prefix = "pcmdi_diags_%s_%s_%04d-%04d_vs_%04d-%04d" % (
+                    c["sub"],
+                    c["tag"],
+                    c["year1"],
+                    c["year2"],
+                    c["ref_year1"],
+                    c["ref_year2"],
+                )
+                reference_data_path = (
+                    c["reference_data_path"].split("/post")[0] + "/post"
+                )
+                if ("ts_num_years" in c.keys()) and (c["reference_data_path_ts"] == ""):
+                    c[
+                        "reference_data_path_ts"
+                    ] = f"{reference_data_path}/atm/{c['grid']}/cmip_ts/monthly"
+            else:
+                raise ValueError("Invalid run_type={}".format(c["run_type"]))
+            print(prefix)
+            c["prefix"] = prefix
+            scriptFile = os.path.join(scriptDir, "%s.bash" % (prefix))
+            statusFile = os.path.join(scriptDir, "%s.status" % (prefix))
+            settingsFile = os.path.join(scriptDir, "%s.settings" % (prefix))
+            skip = checkStatus(statusFile)
+            if skip:
+                continue
+
+            # Create script
+            with open(scriptFile, "w") as f:
+                f.write(template.render(**c))
+            makeExecutable(scriptFile)
+
+            # Iterate from year1 to year2 incrementing by the number of years per time series file.
+            if "ts_num_years" in c.keys():
+                for yr in range(c["year1"], c["year2"], c["ts_num_years"]):
+                    start_yr = yr
+                    end_yr = yr + c["ts_num_years"] - 1
+                    if (
+                        ("mean_climate" in c["sets"])
+                        or ("variability_mode_atm" in c["sets"])
+                        or ("variability_mode_cpl" in c["sets"])
+                        or ("enso" in c["sets"])
+                    ):
+                        add_dependencies(
+                            dependencies,
+                            scriptDir,
+                            "ts",
+                            "atm_monthly_180x360_aave",
+                            start_yr,
+                            end_yr,
+                            c["ts_num_years"],
+                        )
+            with open(settingsFile, "w") as sf:
+                p = pprint.PrettyPrinter(indent=2, stream=sf)
+                p.pprint(c)
+                p.pprint(s)
+
+            export = "ALL"
+            existing_bundles = handle_bundles(
+                c,
+                scriptFile,
+                export,
+                dependFiles=dependencies,
+                existing_bundles=existing_bundles,
+            )
+            if not c["dry_run"]:
+                if c["bundle"] == "":
+                    # Submit job
+                    submitScript(
+                        scriptFile,
+                        statusFile,
+                        export,
+                        job_ids_file,
+                        dependFiles=dependencies,
+                    )
+
+                else:
+                    print("...adding to bundle '%s'" % (c["bundle"]))
+
+            print(f"   environment_commands={c['environment_commands']}")
+            print_url(c, "pcmdi_diags")
+
+    return existing_bundles