diff --git a/.github/workflows/build_workflow.yml b/.github/workflows/build_workflow.yml index 0c83b65..ace0d1d 100644 --- a/.github/workflows/build_workflow.yml +++ b/.github/workflows/build_workflow.yml @@ -54,6 +54,7 @@ jobs: needs: check-jobs-to-skip if: ${{ needs.check-jobs-to-skip.outputs.should_skip != 'true' }} runs-on: ubuntu-latest + timeout-minutes: 15 strategy: matrix: python-version: ["3.11", "3.12", "3.13"] @@ -79,7 +80,7 @@ jobs: miniforge-variant: Miniforge3 miniforge-version: latest environment-file: conda/dev.yml - channel-priority: flexible # Changed from strict to flexible + channel-priority: strict auto-update-conda: true python-version: ${{ matrix.python-version }} channels: conda-forge @@ -142,7 +143,7 @@ jobs: # miniforge-variant: Miniforge3 # miniforge-version: latest # environment-file: conda/dev.yml - # channel-priority: flexible # Changed from strict to flexible + # channel-priority: strict # auto-update-conda: true # python-version: "3.13" # Use stable Python version for docs diff --git a/conda/dev.yml b/conda/dev.yml index b0d25fc..d48b043 100644 --- a/conda/dev.yml +++ b/conda/dev.yml @@ -18,10 +18,11 @@ dependencies: - beautifulsoup4 - lxml - matplotlib-base >=3.8.2,<3.10 - - mpas_tools >=0.21.0 + - mpas_tools - netcdf4 - numpy >=2.0,<3.0 - output_viewer=1.3.3 + - pcmdi_metrics>=3.9.3 - xarray >=2023.02.0 - xcdat >=0.7.3,<1.0 # Testing @@ -40,10 +41,13 @@ dependencies: - black ==25.1.0 - flake8 ==7.3.0 - isort ==6.0.1 - - mypy ==1.18.2 + #- mypy ==1.18.2 - pre-commit ==4.3.0 - types-PyYAML >=6.0.0 # Developer Tools # ======================= - tbump=6.9.0 - ipykernel + # pip dependencies + - pip: + - mypy==1.18.2 diff --git a/pyproject.toml b/pyproject.toml index 30f3c32..9c85245 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "matplotlib >=3.8.2,<3.10", "netcdf4", "numpy >=2.0,<3.0", + "pcmdi_metrics>=3.9.3", "xarray >=2023.02.0", "xcdat >=0.7.3,<1.0", ] @@ -116,6 +117,11 @@ version = { attr = "zppy_interfaces.version.__version__" } # evolution of options.entry-points [project.scripts] 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" +zi-pcmdi-variability-modes = "zppy_interfaces.pcmdi_diags.pcmdi_variability_modes:main" +zi-pcmdi-enso = "zppy_interfaces.pcmdi_diags.pcmdi_enso:main" +zi-pcmdi-synthetic-plots = "zppy_interfaces.pcmdi_diags.pcmdi_synthetic_plots:main" [project.urls] Documentation = "https://docs.e3sm.org/zppy-interfaces" diff --git a/tests/unit/pcmdi_diags/test_pcmdi_mean_climate.py b/tests/unit/pcmdi_diags/test_pcmdi_mean_climate.py new file mode 100644 index 0000000..ae2c6d6 --- /dev/null +++ b/tests/unit/pcmdi_diags/test_pcmdi_mean_climate.py @@ -0,0 +1,27 @@ +from typing import List + +from zppy_interfaces.pcmdi_diags.pcmdi_mean_cimate import generate_mean_clim_cmds + + +def test_generate_mean_clim_cmds(): + # Sample of mean_climate default vars in zppy/defaults/default.ini + variables: List[str] = ["tauu", "tauv", "ta-200"] + # Example: /lcrc/group/e3sm/ac.forsyth2/zppy_pr719_output/unique_id_48/v3.LR.amip_0101/post/scripts/tmp.pcmdi_diags_mean_climate_model_vs_obs_2005-2014.915900.07Jq/pcmdi_diags/climo_ref_mean_climate_catalogue.json + obs_dic = { + "tauu": {"set": "default"}, + "tauv": {"set": "default"}, + "ta": {"set": "default"}, # Not in the example json file above + } + # Example: Appears after "AC" in .nc files in /lcrc/group/e3sm/ac.forsyth2/zppy_pr719_output/unique_id_48/v3.LR.amip_0101/post/scripts/tmp.pcmdi_diags_mean_climate_model_vs_obs_2005-2014.915900.07Jq/climo + case_id: str = "v20250923" + actual = generate_mean_clim_cmds( + variables=variables, + obs_dic=obs_dic, + case_id=case_id, + ) + expected: List[str] = [ + "mean_climate_driver.py -p parameterfile.py --vars tauu -r default --case_id v20250923", + "mean_climate_driver.py -p parameterfile.py --vars tauv -r default --case_id v20250923", + "mean_climate_driver.py -p parameterfile.py --vars ta-200 -r default --case_id v20250923", + ] + assert actual == expected diff --git a/tests/unit/pcmdi_diags/test_pcmdi_setup.py b/tests/unit/pcmdi_diags/test_pcmdi_setup.py new file mode 100644 index 0000000..a27a1a5 --- /dev/null +++ b/tests/unit/pcmdi_diags/test_pcmdi_setup.py @@ -0,0 +1,23 @@ +from zppy_interfaces.pcmdi_diags.pcmdi_setup import ( + DataCatalogueBuilder, + LandSeaMaskGenerator, +) + + +def test_DataCatalogueBuilder(): + dcb = DataCatalogueBuilder("", [], "", [], [], "", "") + + assert dcb._get_base_varname("ta-200") == "ta" + assert dcb._get_base_varname("ta_200") == "ta" + assert dcb._get_base_varname("pr") == "pr" + + +def test_LandSeaMaskGenerator(): + lsmg = LandSeaMaskGenerator("", "", "", "") + assert lsmg._parse_flag("True") + assert lsmg._parse_flag("Y") + assert lsmg._parse_flag("Yes") + assert lsmg._parse_flag("true") + assert lsmg._parse_flag("y") + assert lsmg._parse_flag("yes") + assert not lsmg._parse_flag("False") diff --git a/tests/unit/pcmdi_diags/test_pcmdi_variability_modes.py b/tests/unit/pcmdi_diags/test_pcmdi_variability_modes.py new file mode 100644 index 0000000..b3a2d81 --- /dev/null +++ b/tests/unit/pcmdi_diags/test_pcmdi_variability_modes.py @@ -0,0 +1,66 @@ +from zppy_interfaces.pcmdi_diags.pcmdi_variability_modes import ( + VariabilityMetricsCollector, + generate_varmode_cmds, +) + + +def test_VariabilityMetricsCollector(): + vmc = VariabilityMetricsCollector( + ["mode"], + "png", + "mip", + "exp", + "model", + "relm", + "v20250923", + "dir_%(metric_type)", + "", + ) + assert vmc.input_dir == "dir_variability_modes" + assert vmc.model_name == "mip.exp.model_relm" + assert vmc.seasons == ["DJF", "MAM", "JJA", "SON", "yearly", "monthly"] + assert vmc.fig_sets["MOV_eoftest"] == ["diagnostic_results", "EG_Spec*"] + assert vmc.fig_sets["MOV_compose"] == ["graphics", "*compare_obs"] + assert vmc.fig_sets["MOV_telecon"] == ["graphics", "*teleconnection"] + assert vmc.fig_sets["MOV_pattern"] == ["graphics", "*"] + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "invalid.txt") + == "graphics_mode_DJF_unknown.png" + ) + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "North_test.txt") + == "graphics_mode_DJF_EG_Spec.png" + ) + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "_cbf_.txt") + == "graphics_mode_DJF_cbf.png" + ) + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "EOF1.txt") + == "graphics_mode_DJF_eof1.png" + ) + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "EOF2.txt") + == "graphics_mode_DJF_eof2.png" + ) + assert ( + vmc._classify_output_name("graphics", "mode", "DJF", "EOF3.txt") + == "graphics_mode_DJF_eof3.png" + ) + + +def test_generate_varmode_cmds(): + actual = generate_varmode_cmds( + ["mode1", "mode2"], + "varOBS", + "reftyrs", + "reftyre", + "refname", + "refpath", + "v20250923", + ) + expected = [ + "variability_modes_driver.py -p parameterfile.py --variability_mode mode1 --eofn_mod 1 --eofn_obs 1 --varOBS varOBS --osyear reftyrs --oeyear reftyre --reference_data_name refname --reference_data_path refpath --case_id v20250923", + "variability_modes_driver.py -p parameterfile.py --variability_mode mode2 --eofn_mod 1 --eofn_obs 1 --varOBS varOBS --osyear reftyrs --oeyear reftyre --reference_data_name refname --reference_data_path refpath --case_id v20250923", + ] + assert actual == expected diff --git a/tests/unit/pcmdi_diags/test_synthetic_plots_utils.py b/tests/unit/pcmdi_diags/test_synthetic_plots_utils.py new file mode 100644 index 0000000..7ab4661 --- /dev/null +++ b/tests/unit/pcmdi_diags/test_synthetic_plots_utils.py @@ -0,0 +1,22 @@ +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import get_highlight_models + + +def test_get_highlight_models(): + all_models = [ + "CESM2-FV2", + "CESM2-WACCM", + "CESM2-WACCM-FV2", + "GFDL-AM4", + "GFDL-CM4", + "GFDL-ESM4", + "E3SM-1-0", + "E3SM-2-0", + ] + model_name = ["CESM2-FV2", "E3SM"] + actual = get_highlight_models(all_models, model_name) + expected = [ + "E3SM-1-0", + "E3SM-2-0", + "CESM2-FV2", + ] + assert actual == expected diff --git a/tests/unit/pcmdi_diags/test_viewer.py b/tests/unit/pcmdi_diags/test_viewer.py new file mode 100644 index 0000000..0e72222 --- /dev/null +++ b/tests/unit/pcmdi_diags/test_viewer.py @@ -0,0 +1,6 @@ +from zppy_interfaces.pcmdi_diags.viewer import safe_join + + +def test_safe_join(): + assert safe_join("a", "b") == "a/b" + assert safe_join("a/", "b") == "a/b" diff --git a/zppy_interfaces/pcmdi_diags/__init__.py b/zppy_interfaces/pcmdi_diags/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/zppy_interfaces/pcmdi_diags/link_observation.py b/zppy_interfaces/pcmdi_diags/link_observation.py new file mode 100644 index 0000000..897ff34 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/link_observation.py @@ -0,0 +1,198 @@ +import argparse +import glob +import json +import os +import re +import shutil +import sys +from typing import Dict, List + +from pcmdi_metrics.io import xcdat_open + +from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger +from zppy_interfaces.pcmdi_diags.utils import ALT_OBS_MAP + +# Set up the root logger and module level logger. The module level logger is +# a child of the root logger. +_setup_root_logger() +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class LinkObservationParameters(object): + def __init__(self, args: Dict[str, str]): + self.model_name: str = f"{args['model_name_ref']}.{args['tableID_ref']}" + self.variables: List[str] = args["vars"].split(",") + self.obs_sets: List[str] = args["obs_sets"].split(",") + self.obs_ts: str = args["obs_ts"] + self.obstmp_dir: str = args["obstmp_dir"] + + +class ObservationLinker: + def __init__( + self, + model_name, + variables, + obs_sets, + ts_dir_ref_source, + obstmp_dir, + obs_alias_file, + altobs_dic, + ): + self.model_name = model_name + self.variables = variables + self.obs_sets = obs_sets + self.ts_dir_ref_source = ts_dir_ref_source + self.obstmp_dir = obstmp_dir + self.obs_dic = json.load(open(obs_alias_file)) + self.altobs_dic = altobs_dic + + def _resolve_obs_file(self, varin, obsid): + if varin not in self.obs_dic or obsid not in self.obs_dic[varin]: + print(f"[Warning] No alias found for variable '{varin}' in obsid '{obsid}'") + return None, None + + obsname = self.obs_dic[varin][obsid] + obsstr = ( + obsname.replace("_", "*").replace("-", "*") + if "ceres_ebaf" in obsname + else obsname + ) + pattern = os.path.join(self.ts_dir_ref_source, obsstr, f"{varin}_*.nc") + fpaths = sorted(glob.glob(pattern)) + + if fpaths and os.path.exists(fpaths[0]): + return fpaths[0], varin + + # Try altobs mapping + if varin in self.altobs_dic: + alt_var = self.altobs_dic[varin] + pattern_alt = os.path.join( + self.ts_dir_ref_source, obsstr, f"{alt_var}_*.nc" + ) + fpaths = sorted(glob.glob(pattern_alt)) + if fpaths and os.path.exists(fpaths[0]): + return fpaths[0], alt_var + + print(f"[Warning] Observation file not found for {varin} ({obsid})") + return None, None + + def link_obs_data(self): + for i, vv in enumerate(self.variables): + varin = re.split(r"_|-", vv)[0] if "_" in vv or "-" in vv else vv + if len(self.obs_sets) > 1 and len(self.obs_sets) == len(self.variables): + obsid = self.obs_sets[i] + else: + obsid = self.obs_sets[0] + + filepath, resolved_var = self._resolve_obs_file(varin, obsid) + if filepath: + template = os.path.basename(filepath) + parts = template.replace(".nc", "").split("_") + if len(parts) < 3: + print(f"[Error] Unexpected filename format: {template}") + continue + yms, yme = parts[-2][:6], parts[-1][:6] + obsname = self.obs_dic[varin][obsid].replace(".", "_") + out = os.path.join( + self.obstmp_dir, + f"{self.model_name.replace('put_model_here', obsname)}.{varin}.{yms}-{yme}.nc", + ) + + if not os.path.exists(out): + os.makedirs(os.path.dirname(out), exist_ok=True) + if resolved_var == varin: + os.symlink(filepath, out) + print(f"[Info] Linked {resolved_var} → {out}") + else: + ds = xcdat_open(filepath) + ds = ds.rename({resolved_var: varin}) + ds.to_netcdf(out) + print( + f"[Info] Renamed and saved {resolved_var} as {varin} → {out}" + ) + else: + print(f"[Info] Skipping existing file: {out}") + + def derive_var(self, vout, var_dic): + template = None + out = None + ds_out = None + + for i, (var, scale) in enumerate(var_dic.items()): + fpaths = sorted(glob.glob(os.path.join(self.obstmp_dir, f"*.{var}.*.nc"))) + if not fpaths: + print( + f"[Warning] No file found for base variable '{var}' needed to derive '{vout}'" + ) + continue + + ds = xcdat_open(fpaths[0]) + if i == 0: + template = os.path.basename(fpaths[0]) + out = os.path.join( + self.obstmp_dir, template.replace(f".{var}.", f".{vout}.") + ) + shutil.copy(fpaths[0], out) + ds_out = ds.rename_vars({var: vout}) + ds_out[vout] = ds_out[vout] * scale + else: + ds_other = xcdat_open(fpaths[0]) + if ds_out: + ds_out[vout] = ds_out[vout] + ds_other[var] * scale + else: + raise ValueError("ds_out is None") + + if template and ds_out: + ds_out.to_netcdf(out) + print(f"[Info] Derived variable '{vout}' written to {out}") + + def process_derived_variables(self): + for vv in self.variables: + if vv in ["rltcre", "rstcre"]: + fpaths = sorted(glob.glob(os.path.join(self.obstmp_dir, f"*{vv}_*.nc"))) + if not fpaths: + if vv == "rstcre": + self.derive_var("rstcre", {"rsutcs": 1, "rsut": -1}) + elif vv == "rltcre": + self.derive_var("rltcre", {"rlutcs": 1, "rlut": -1}) + + +# Functions ################################################################### +def main(): + parameters: LinkObservationParameters = _get_args() + linker = ObservationLinker( + model_name=parameters.model_name, + variables=parameters.variables, + obs_sets=parameters.obs_sets, + ts_dir_ref_source=parameters.obs_ts, + obstmp_dir=parameters.obstmp_dir, + altobs_dic=ALT_OBS_MAP, + obs_alias_file="reference_alias.json", + ) + linker.link_obs_data() + linker.process_derived_variables() + + +def _get_args() -> LinkObservationParameters: + # Parser + parser: argparse.ArgumentParser = argparse.ArgumentParser( + usage="zi-pcmdi-link-observation " + ) + parser.add_argument("--model_name_ref", type=str) + parser.add_argument("--tableID_ref", type=str) + parser.add_argument("--vars", type=str) + parser.add_argument("--obs_sets", type=str) + parser.add_argument("--obs_ts", type=str) + parser.add_argument("--obstmp_dir", type=str) + parser.add_argument("--debug", type=str) + + # Ignore the first arg + # (zi-pcmdi-link-observation) + args: argparse.Namespace = parser.parse_args(sys.argv[1:]) + + if args.debug and args.debug.lower() == "true": + logger.setLevel("DEBUG") + logger.debug("Debug logging enabled") + + return LinkObservationParameters(vars(args)) diff --git a/zppy_interfaces/pcmdi_diags/pcmdi_enso.py b/zppy_interfaces/pcmdi_diags/pcmdi_enso.py new file mode 100644 index 0000000..741077e --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/pcmdi_enso.py @@ -0,0 +1,503 @@ +import argparse +import glob +import json +import os +import re +import sys +import time +from collections import OrderedDict +from typing import Dict, List + +from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger +from zppy_interfaces.pcmdi_diags.pcmdi_setup import CoreOutput, CoreParameters, set_up +from zppy_interfaces.pcmdi_diags.utils import ( + ALT_OBS_MAP, + run_parallel_jobs, + run_serial_jobs, +) + +# Set up the root logger and module level logger. The module level logger is +# a child of the root logger. +_setup_root_logger() +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class ENSOParameters(object): + def __init__(self, args: Dict[str, str]): + self.enso_groups: str = args["enso_groups"] + + +class EnsoDiagnosticsCollector: + def __init__( + self, fig_format, refname, model_name_parts, case_id, input_dir, output_dir + ): + self.fig_format = fig_format + self.refname = refname + self.mip, self.exp, self.model, self.relm = model_name_parts + self.case_id = case_id + self.model_name = f"{self.mip}.{self.exp}.{self.model}_{self.relm}" + self.input_dir = input_dir.replace("%(metric_type)", "enso_metric") + self.output_dir = output_dir + self.diag_metric = "enso_metric" + self.fig_sets = OrderedDict([("ENSO_metric", ["graphics", "*"])]) + + def collect_figures(self, groups) -> bool: + logger.info("Entering EnsoDiagnosticsCollector.collect_figures") + success: bool = True + for fset, (subdir, pattern) in self.fig_sets.items(): + logger.info(f"Processing {fset}, ({subdir}, {pattern})") + fdir = self.input_dir.replace("%(output_type)", subdir) + logger.info(f"Processing fdir={fdir}") + found_groups: List[str] = os.listdir(fdir) + if sorted(groups) != sorted(found_groups): + logger.error( + f"Groups mismatch: expected {sorted(groups)}, found {sorted(found_groups)} in {fdir}" + ) + success = False + continue + for group in groups: + logger.info(f"Processing group={group}") + template = os.path.join(fdir, group, f"{pattern}.{self.fig_format}") + logger.info( + f"template={template}, pattern.fig_format={pattern}.{self.fig_format}" + ) + fpaths = sorted(glob.glob(template)) + + if not fpaths: + logger.error( + f"fpaths={fpaths}, self.input_dir={self.input_dir}, template={os.path.abspath(template)}, files in template={os.listdir(os.path.join(fdir, group))}" + ) + success = False + for fpath in fpaths: + logger.info(f"Processing fpath={fpath}") + tail = fpath.split("/")[-1].split(f"{self.model}_{self.relm}")[-1] + outpath = os.path.join( + self.output_dir.replace("%(group_type)", fset), group + ) + logger.info(f"outpath={outpath}") + os.makedirs(outpath, exist_ok=True) + outfile = f"{group}{tail}" + os.rename(fpath, os.path.join(outpath, outfile)) + return success + + def collect_metrics(self) -> bool: + logger.info("Entering EnsoDiagnosticsCollector.collect_metrics") + success: bool = True + inpath = self.input_dir.replace("%(output_type)", "metrics_results") + fpaths = sorted(glob.glob(os.path.join(inpath, "*/*.json"))) + + if not fpaths: + logger.error( + f"fpaths={fpaths}, self.input_dir={self.input_dir}, inpath={os.path.abspath(inpath)}, files in inpath={os.listdir(inpath)}" + ) + success = False + for fpath in fpaths: + logger.info(f"Processing fpath={fpath}") + refmode = fpath.split("/")[-2] + reffile = fpath.split("/")[-1] + outpath = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), + self.diag_metric, + refmode, + ) + logger.info(f"outpath={outpath}") + os.makedirs(outpath, exist_ok=True) + + base_filename = ( + f"{refmode}.{self.model_name}.vs.{self.refname}.{self.case_id}.json" + ) + outfile = ( + base_filename.replace(".json", ".diveDown.json") + if "diveDown" in reffile + else base_filename + ) + os.rename(fpath, os.path.join(outpath, outfile)) + return success + + def collect_diags(self) -> bool: + logger.info("Entering EnsoDiagnosticsCollector.collect_diags") + success: bool = True + inpath = self.input_dir.replace("%(output_type)", "diagnostic_results") + fpaths = sorted(glob.glob(os.path.join(inpath, "*/*.nc"))) + + if not fpaths: + logger.error( + f"fpaths={fpaths}, self.input_dir={self.input_dir}, inpath={os.path.abspath(inpath)}, files in inpath={os.listdir(inpath)}" + ) + success = False + for fpath in fpaths: + logger.info(f"Processing fpath={fpath}") + refmode = fpath.split("/")[-2] + reffile = fpath.split("/")[-1] + outpath = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), + self.diag_metric, + refmode, + ) + logger.info(f"outpath={outpath}") + os.makedirs(outpath, exist_ok=True) + + os.rename(fpath, os.path.join(outpath, reffile)) + return success + + def run(self, groups): + logger.info("Entering EnsoDiagnosticsCollector.run") + figures_success: bool = self.collect_figures(groups) + metrics_success: bool = self.collect_metrics() + diags_success: bool = self.collect_diags() + if figures_success and metrics_success and diags_success: + logger.info("Completing EnsoDiagnosticsCollector.run") + else: + raise RuntimeError( + "EnsoDiagnosticsCollector.run failed: " + f"figures_success={figures_success}, metrics_success={metrics_success}, diags_success={diags_success}" + ) + + +# Functions ################################################################### +def main(): + logger.error("zi-pcmdi-enso is not yet supported. Exiting.") + sys.exit(1) + args: Dict[str, str] = _get_args() + core_parameters = CoreParameters(args) + enso_parameters = ENSOParameters(args) + core_output: CoreOutput = set_up(core_parameters) + + ############################################# + # call enso_driver.py to process diagnostics + ############################################# + build_enso_obsvar_catalog(core_output.obs_dic, core_parameters.variables) + build_enso_obsvar_landmask(core_output.obs_dic, core_parameters.variables) + # now start enso driver + check_enso_input() + lstcmd = generate_enso_cmds(enso_parameters.enso_groups, core_parameters.case_id) + logger.info( + f"input_template={core_output.input_template}; if the directories based on this template are empty, lstcmd={lstcmd} failed to produce output." + ) + if (len(lstcmd) > 0) and core_parameters.multiprocessing: + logger.info(f"Running parallel jobs for {lstcmd}") + try: + results = run_parallel_jobs(lstcmd, core_parameters.num_workers) + check_enso_output(results) + except RuntimeError as e: + logger.error(f"Execution failed: {e}") + raise e + elif (len(lstcmd) > 0) and not core_parameters.multiprocessing: + logger.info(f"Running serial jobs for {lstcmd}") + try: + results = run_serial_jobs(lstcmd) + check_enso_output(results) + except RuntimeError as e: + logger.error(f"Execution failed: {e}") + raise e + else: + logger.info("no jobs to run...") + logger.info("successfully finish all jobs....") + # time delay to ensure process completely finished + time.sleep(5) + # Initialize and run collector + obs_dict = json.load(open("obs_catalogue.json")) + obs_name = list(obs_dict.keys())[0] + collector = EnsoDiagnosticsCollector( + fig_format=core_parameters.figure_format, + refname=obs_name, + model_name_parts=core_parameters.model_name.split("."), + case_id=core_parameters.case_id, + input_dir=core_output.input_template, + output_dir=core_output.out_path, + ) + enso_groups: List[str] = enso_parameters.enso_groups.split(",") + collector.run(enso_groups) + + +def _get_args() -> Dict[str, str]: + # Parser + parser: argparse.ArgumentParser = argparse.ArgumentParser( + usage="zi-pcmdi-enso " + ) + + # For CoreParameters + parser.add_argument("--num_workers", type=str) + parser.add_argument("--multiprocessing", type=str) + parser.add_argument("--subsection", type=str) + parser.add_argument("--climo_ts_dir_primary", type=str) # needs ts_dir_primary + parser.add_argument("--climo_ts_dir_ref", type=str) # needs ts_dir_ref + parser.add_argument("--model_name", type=str) + parser.add_argument("--model_tableID", type=str) + parser.add_argument("--figure_format", type=str) + parser.add_argument("--run_type", type=str) + parser.add_argument("--obs_sets", type=str) # run_type == "model_vs_obs" only + parser.add_argument( + "--model_name_ref", type=str + ) # run_type == "model_vs_model" only + parser.add_argument("--vars", type=str) + parser.add_argument("--tableID_ref", type=str) # run_type == "model_vs_model" only + parser.add_argument("--generate_sftlf", type=str) + parser.add_argument("--case_id", type=str) + parser.add_argument("--results_dir", type=str) + parser.add_argument("--debug", type=str) + + # For ENSOParameters + parser.add_argument("--enso_groups", type=str) + + # Ignore the first arg + # (zi-pcmdi-enso) + args: argparse.Namespace = parser.parse_args(sys.argv[1:]) + + if args.debug and args.debug.lower() == "true": + logger.setLevel("DEBUG") + logger.debug("Debug logging enabled") + + return vars(args) + + +def build_enso_obsvar_catalog( + obs_dic: Dict, variables: List[str], output_file: str = "obs_catalogue.json" +) -> None: + """ + Organize observational data for the ENSO driver based on the variable list. + + Parameters: + obs_dic (dict): Dictionary mapping variable names to their observation sets and data files. + variables (list): List of variable names to process. + output_file (str): Output JSON file path to save the observation catalogue. + """ + refr_dic: OrderedDict = OrderedDict() + + for var in variables: + vkey = re.split(r"[_-]", var)[0] if "_" in var or "-" in var else var + + if vkey not in obs_dic: + raise KeyError( + f"Variable key '{vkey}' not found in observation dictionary. Available keys are {obs_dic.keys()}" + ) + + refset = obs_dic[vkey]["set"] + refname = obs_dic[vkey].get(refset) + + if not refname: + raise KeyError( + f"Reference name not found for variable '{vkey}' and set '{refset}'." + ) + + refr_dic.setdefault(refname, {})[vkey] = obs_dic[vkey][refname] + + with open(output_file, "w") as f: + json.dump(refr_dic, f, indent=4, sort_keys=False, separators=(",", ": ")) + + logger.info(f"[INFO] Observation catalogue written to: {output_file}") + + +def build_enso_obsvar_landmask( + obs_dic: Dict, + variables: List[str], + output_file: str = "obs_landmask.json", + mask_dir: str = "fixed", +) -> None: + """ + Organize observational land/sea mask mapping for ENSO diagnostics. + + Parameters: + obs_dic (dict): Dictionary mapping variables to observation metadata. + variables (list): List of variable names used in ENSO analysis. + output_file (str): Path to output the landmask JSON. + mask_dir (str): Directory prefix where the landmask files are located. + """ + relf_dic: OrderedDict = OrderedDict() + + for var in variables: + vkey = re.split(r"[_-]", var)[0] if "_" in var or "-" in var else var + + if vkey not in obs_dic: + raise KeyError( + f"Variable key '{vkey}' not found in observation dictionary." + ) + + refset = obs_dic[vkey]["set"] + refname = obs_dic[vkey].get(refset) + + if not refname: + raise KeyError( + f"Reference name not found for variable '{vkey}' and set '{refset}'." + ) + + relf_dic.setdefault(refname, os.path.join(mask_dir, f"sftlf.{refname}.nc")) + + with open(output_file, "w") as f: + json.dump(relf_dic, f, indent=4, sort_keys=False, separators=(",", ": ")) + + logger.info(f"[INFO] Landmask mapping written to: {output_file}") + + +def check_enso_input(): + current_dir: str = os.path.abspath(os.getcwd()) + ts_dir: str = os.path.join(current_dir, "ts") + if not os.path.exists(ts_dir): + raise FileNotFoundError(f"{ts_dir} (input for enso_driver) does not exist.") + if not os.listdir(ts_dir): + raise FileNotFoundError(f"{ts_dir} is empty.") + else: + for obs_var_name, cmip_var_name in ALT_OBS_MAP.items(): + logger.info( + f"Symlinking cmip-standard {cmip_var_name} to observational variable name {obs_var_name}, if present" + ) + found_nc_file = glob.glob(f"{ts_dir}/*.{cmip_var_name}.*.nc") + if found_nc_file: + source_file = found_nc_file[0] + link_name = found_nc_file[0].replace( + f".{cmip_var_name}.", f".{obs_var_name}." + ) + os.symlink(source_file, link_name) + found_txt_file = glob.glob(f"{ts_dir}/{cmip_var_name}_files.txt") + if found_txt_file: + source_file = found_txt_file[0] + link_name = f"{ts_dir}/{obs_var_name}_files.txt" + os.symlink(source_file, link_name) + + +def generate_enso_cmds( + enso_groups_str, + case_id, + param_file="parameterfile.py", + driver_script="enso_driver.py", +): + """ + Generate ENSO driver command-line strings for given metric groups. + + Parameters: + enso_groups_str: Comma-separated list of ENSO metric groups. + case_id: Case identifier. + param_file: Parameter file used by the driver script. + driver_script: ENSO driver script filename. + + Returns: + cmds: List of shell command strings to run. + """ + enso_groups = enso_groups_str.split(",") + commands = [ + "{} -p {} --metricsCollection {} --case_id {}".format( + driver_script, param_file, group, case_id + ) + for group in enso_groups + ] + current_dir: str = os.path.abspath(os.getcwd()) + logger.info(f"Commands will be run from current_dir={current_dir}") + dir_contents: List[str] = os.listdir(current_dir) + if param_file not in dir_contents: + logger.error( + f"Parameter file '{param_file}' not found in current directory: {current_dir}" + ) + raise FileNotFoundError(f"Parameter file '{param_file}' not found.") + + return commands + + +def check_enso_output(results): + logger.info("Checking ENSO output.") + success: bool = True + for i, (stdout, stderr, return_code) in enumerate(results): + logger.info(f"Command {i + 1} finished:") + logger.info(f"STDOUT: {stdout}") + logger.info(f"STDERR: {stderr}") + logger.info(f"Return code: {return_code}") + if not check_vars(stdout): + logger.error(f"Command {i + 1} failed to produce expected variables.") + success = False + if not check_output_dirs(stdout): + logger.error( + f"Command {i + 1} failed to produce expected output directories." + ) + success = False + if not success: + raise RuntimeError("ENSO output check failed.") + logger.info("ENSO output check passed.") + + +def check_vars(stdout: str) -> bool: + """ + Check if the output from an enso_driver.py command contains expected variables. + + Parameters: + stdout (str): Standard output from the command execution. + + Returns: + bool: True if expected variables are found, False otherwise. + """ + success: bool = True + match_object = re.search(r"list_variables:\s*\[(.*?)\]", stdout) + # Special-case "optional" missing variables (these quantities may not be frequently output) + optional_missing = {"ssh", "thf"} + if match_object: + variables_content = match_object.group(1) + # Split by comma and clean up each variable name + requested_variables = [] + for var in variables_content.split(","): + # Remove quotes, whitespace, and extract just the variable name + clean_var = re.sub(r"['\"\s]", "", var.strip()) + if clean_var: # Only care about non-empty strings + requested_variables.append(clean_var) + # Now, check if we actually have data for these variables + current_dir: str = os.path.abspath(os.getcwd()) + variables_missing_data: List[str] = [] + for var in requested_variables: + found_nc_file = glob.glob(f"ts/*.{var}.*.nc") + found_txt_file = glob.glob(f"ts/{var}_files.txt") + if (not found_nc_file) or (not found_txt_file): + variables_missing_data.append(var) + # Check for references + if var in ALT_OBS_MAP: + alt_var = ALT_OBS_MAP[var] + found_nc_file_alt = glob.glob(f"ts/*.{alt_var}.*.nc") + found_txt_file_alt = glob.glob(f"ts/{alt_var}_files.txt") + if found_nc_file_alt or found_txt_file_alt: + logger.error( + f"Found alternative variable '{alt_var}' for '{var}' in {current_dir}/ts. This indicates that the variable derivation/mapping has not been applied correctly." + ) + ts_dir = os.path.join(current_dir, "ts") + if variables_missing_data: + if set(variables_missing_data) <= optional_missing: + # Only ssh/thf are missing → warn but do not fail + logger.warning( + f"Optional variables missing: {variables_missing_data} in directory{ts_dir}" + ) + success = True + else: + # Other variables missing → error and fail + logger.error( + f"Variables missing data: {variables_missing_data} in directory {ts_dir}" + ) + logger.error(f"Full contents of {ts_dir}: {os.listdir(ts_dir)}") + success = False + else: + logger.info( + f"All requested variables {requested_variables} found in directory {ts_dir}" + ) + else: + logger.error("No variable list found in stdout.") + success = False + return success + + +def check_output_dirs(stdout: str) -> bool: + current_dir: str = os.path.abspath(os.getcwd()) + success: bool = True + for output_type in ["graphics", "diagnostic_results", "metrics_results"]: + match_object = re.search(f"output directory for {output_type}:(.*)", stdout) + if match_object: + subdir = match_object.group(1).strip() + combined_dir = os.path.join(current_dir, subdir) + if not os.path.exists(combined_dir): + logger.error( + f"{output_type} output directory does not exist: {combined_dir}" + ) + success = False + else: + if not os.listdir(combined_dir): + logger.error( + f"{output_type} output directory is empty: {combined_dir}" + ) + success = False + # else: success = True + # else: success = True # Don't assume we have any particular directory + return success diff --git a/zppy_interfaces/pcmdi_diags/pcmdi_mean_cimate.py b/zppy_interfaces/pcmdi_diags/pcmdi_mean_cimate.py new file mode 100644 index 0000000..c29a713 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/pcmdi_mean_cimate.py @@ -0,0 +1,257 @@ +import argparse +import glob +import json +import os +import re +import sys +import time +from collections import OrderedDict +from typing import Dict, List + +from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger +from zppy_interfaces.pcmdi_diags.pcmdi_setup import CoreOutput, CoreParameters, set_up +from zppy_interfaces.pcmdi_diags.utils import run_parallel_jobs, run_serial_jobs + +# Set up the root logger and module level logger. The module level logger is +# a child of the root logger. +_setup_root_logger() +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class MeanClimateParameters(object): + def __init__(self, args: Dict[str, str]): + self.regions: List[str] = args["regions"].split(",") + + +class MeanClimateMetricsCollector: + def __init__( + self, + regions, + variables, + fig_format, + model_info, + case_id, + input_template, + output_dir, + ): + self.regions = regions + self.variables = variables + self.fig_format = fig_format + self.mip, self.exp, self.model, self.relm = model_info + self.case_id = case_id + self.input_template = input_template + self.output_dir = output_dir + self.diag_metric = "mean_climate" + self.seasons = ["DJF", "MAM", "JJA", "SON", "AC"] + self.model_name = f"{self.mip}.{self.exp}.{self.model}_{self.relm}" + + def collect(self): + self._collect_figures() + self._collect_metrics() + self._collect_diags() + + def _collect_figures(self): + fig_sets = OrderedDict() + fig_sets["CLIM_patttern"] = ["graphics", "*"] + + for fset, (fig_type, prefix) in fig_sets.items(): + for region in self.regions: + for season in self.seasons: + for var in self.variables: + indir = self.input_template.replace( + "%(metric_type)", self.diag_metric + ) + indir = indir.replace("%(output_type)", fig_type) + search_path = os.path.join( + indir, var, f"{prefix}{region}_{season}*.{self.fig_format}" + ) + fpaths = sorted(glob.glob(search_path)) + + for fpath in fpaths: + refname = os.path.basename(fpath).split("_")[0] + filname = f"{refname}_{region}_{season}.{self.fig_format}" + outpath = os.path.join( + self.output_dir.replace("%(group_type)", fset), + region, + season, + ) + os.makedirs(outpath, exist_ok=True) + outfile = os.path.join(outpath, filname) + os.rename(fpath, outfile) + + def _collect_diags(self): + inpath = self.input_template.replace("%(metric_type)", self.diag_metric) + inpath = inpath.replace("%(output_type)", "diagnostic_results") + outpath = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), self.diag_metric + ) + + os.makedirs(outpath, exist_ok=True) + fpaths = sorted(glob.glob(os.path.join(inpath, "*/*/*/*/*/*/*.nc"))) + + for fpath in fpaths: + filname = fpath.split("/")[-1] + outfile = os.path.join(outpath, filname) + os.rename(fpath, outfile) + + def _collect_metrics(self): + inpath = self.input_template.replace("%(metric_type)", self.diag_metric) + inpath = inpath.replace("%(output_type)", "metrics_results") + outpath = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), self.diag_metric + ) + + os.makedirs(outpath, exist_ok=True) + fpaths = sorted(glob.glob(os.path.join(inpath, "*.json"))) + + for fpath in fpaths: + refname = os.path.basename(fpath).split("_")[:2] + filname = f"{refname[0]}.{refname[1]}.{self.model_name}.{self.case_id}.json" + outfile = os.path.join(outpath, filname) + os.rename(fpath, outfile) + + +# Functions ################################################################### +def main(): + args: Dict[str, str] = _get_args() + core_parameters = CoreParameters(args) + mean_climate_parameters = MeanClimateParameters(args) + core_output: CoreOutput = set_up(core_parameters) + + # assign region to each variable + save_variable_regions(core_parameters.variables, mean_climate_parameters.regions) + # generate the command list + lstcmd = generate_mean_clim_cmds( + variables=core_parameters.variables, + obs_dic=core_output.obs_dic, + case_id=core_parameters.case_id, + ) + #################################################### + # call pcmdi mean climate diagnostics + #################################################### + if (len(lstcmd) > 0) and core_output.multiprocessing: + try: + results = run_parallel_jobs(lstcmd, core_parameters.num_workers) + for i, (stdout, stderr, return_code) in enumerate(results): + print(f"\nCommand {i + 1} finished:") + print(f"STDOUT: {stdout}") + print(f"STDERR: {stderr}") + print(f"Return code: {return_code}") + except RuntimeError as e: + print(f"Execution failed: {e}") + elif len(lstcmd) > 0: + try: + results = run_serial_jobs(lstcmd) + for i, (stdout, stderr, return_code) in enumerate(results): + print(f"\nCommand {i + 1} finished:") + print(f"STDOUT: {stdout}") + print(f"STDERR: {stderr}") + print(f"Return code: {return_code}") + except RuntimeError as e: + print(f"Execution failed: {e}") + else: + print("no jobs to run,continue....") + print("successfully finish all jobs....") + # time delay to ensure process completely finished + time.sleep(5) + # orgnize diagnostic output + model_info_str: List[str] = core_parameters.model_name.split(".") + if len(model_info_str) == 4: + # (mip, exp, model, relm) + # model_info_tuple: Tuple[str, str, str, str] = + model_info_tuple = tuple(model_info_str) + else: + raise ValueError( + f"(mip, exp, model, relm) cannot be parsed from {core_parameters.model_name}" + ) + collector = MeanClimateMetricsCollector( + regions=mean_climate_parameters.regions, + variables=core_parameters.variables, + fig_format=core_parameters.figure_format, + model_info=model_info_tuple, + case_id=core_parameters.case_id, + input_template=core_output.input_template, + output_dir=core_output.out_path, + ) + collector.collect() + + +def _get_args() -> Dict[str, str]: + # Parser + parser: argparse.ArgumentParser = argparse.ArgumentParser( + usage="zi-pcmdi-mean-climate " + ) + + # For CoreParameters + parser.add_argument("--num_workers", type=str) + parser.add_argument("--multiprocessing", type=str) + parser.add_argument("--subsection", type=str) + parser.add_argument("--climo_ts_dir_primary", type=str) # needs climo_dir_primary + parser.add_argument("--climo_ts_dir_ref", type=str) # needs climo_dir_ref + parser.add_argument("--model_name", type=str) + parser.add_argument("--model_tableID", type=str) + parser.add_argument("--figure_format", type=str) + parser.add_argument("--run_type", type=str) + parser.add_argument("--obs_sets", type=str) # run_type == "model_vs_obs" only + parser.add_argument( + "--model_name_ref", type=str + ) # run_type == "model_vs_model" only + parser.add_argument("--vars", type=str) + parser.add_argument("--tableID_ref", type=str) # run_type == "model_vs_model" only + parser.add_argument("--generate_sftlf", type=str) + parser.add_argument("--case_id", type=str) + parser.add_argument("--results_dir", type=str) + parser.add_argument("--debug", type=str) + + # For MeanClimateParameters + parser.add_argument("--regions", type=str) + + # Ignore the first arg + # (zi-pcmdi-mean-climate) + args: argparse.Namespace = parser.parse_args(sys.argv[1:]) + + if args.debug and args.debug.lower() == "true": + logger.setLevel("DEBUG") + logger.debug("Debug logging enabled") + + return vars(args) + + +def save_variable_regions(variables, regions, output_path="regions.json"): + """ + Maps each variable (simplified key) to a list of regions and saves to JSON. + """ + region_map = OrderedDict() + for var in variables: + var_key = re.split(r"[_-]", var)[0] if "_" in var or "-" in var else var + region_map[var_key] = regions + + with open(output_path, "w") as f: + json.dump(region_map, f, sort_keys=False, indent=4, separators=(",", ": ")) + return region_map + + +def generate_mean_clim_cmds(variables, obs_dic, case_id): + """ + Generates a list of shell commands for mean climate diagnostics. + """ + commands = [] + for var in variables: + var_key = re.split(r"[_-]", var)[0] if "_" in var or "-" in var else var + if var_key in obs_dic: + refset = obs_dic[var_key]["set"] + cmd = " ".join( + [ + "mean_climate_driver.py", + "-p parameterfile.py", + "--vars", + var, + "-r", + refset, + "--case_id", + case_id, + ] + ) + commands.append(cmd) + return commands diff --git a/zppy_interfaces/pcmdi_diags/pcmdi_setup.py b/zppy_interfaces/pcmdi_diags/pcmdi_setup.py new file mode 100644 index 0000000..d88981f --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/pcmdi_setup.py @@ -0,0 +1,423 @@ +import glob +import json +import os +import re +from collections import OrderedDict +from datetime import datetime +from typing import Dict, List, Tuple + +import xarray as xr +from pcmdi_metrics.io import xcdat_open +from pcmdi_metrics.utils import create_land_sea_mask + +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class CoreParameters(object): + def __init__(self, args: Dict[str, str]): + self.num_workers: int = int(args["num_workers"]) + self.multiprocessing: bool = args["multiprocessing"].lower() == "true" + self.subsection: str = args["subsection"] + self.test_data_path: str = args["climo_ts_dir_primary"] + self.reference_data_path: str = args["climo_ts_dir_ref"] + self.model_name: str = args["model_name"] + self.model_tableID: str = args["model_tableID"] + self.figure_format: str = args["figure_format"] + self.run_type: str = args["run_type"] + self.obs_sets: str = args["obs_sets"] # run_type == "model_vs_obs" only + self.model_name_ref: str = args[ + "model_name_ref" + ] # run_type == "model_vs_model" only + self.variables: List[str] = args["vars"].split(",") + self.tableID_ref: str = args["tableID_ref"] # run_type == "model_vs_model" only + # Whether to generate the land/sea mask + self.generate_flag: str = args["generate_sftlf"] + self.case_id: str = args["case_id"] + self.results_dir: str = args["results_dir"] + + +class CoreOutput(object): + def __init__(self, multiprocessing, obs_dic, input_template, out_path): + self.multiprocessing = multiprocessing + self.obs_dic = obs_dic + self.input_template = input_template + self.out_path = out_path + + +class DataCatalogueBuilder: + def __init__( + self, + test_path: str, + test_set: List[str], + ref_path: str, + ref_set: List[str], + variables: List[str], + label, + output_dir, + ): + self.test_path: str = test_path + self.test_set: List[str] = test_set + self.ref_path: str = ref_path + self.ref_set: List[str] = ref_set + self.variables: List[str] = variables + self.label = label + self.output_dir = output_dir + + self.test_info: OrderedDict = OrderedDict() + self.ref_info: OrderedDict = OrderedDict() + + def build_catalogues(self) -> Tuple[OrderedDict, OrderedDict]: + if not self.variables: + logger.info("DataCatalogueBuilder's variable list is empty") + else: + logger.info(f"DataCatalogueBuilder's vars: {self.variables}") + for idx, var in enumerate(self.variables): + logger.info(f"Building catalogue for {var}") + varin = self._get_base_varname(var) + logger.info(f"Looking for {varin}, the base var name of {var}") + logger.info(f"Finding test files in {self.test_path}") + test_files = sorted( + glob.glob(os.path.join(self.test_path, f"*.{varin}.*.nc")) + ) + logger.info(f"Finding ref files in {self.ref_path}") + ref_files = sorted( + glob.glob(os.path.join(self.ref_path, f"*.{varin}.*.nc")) + ) + + if ( + test_files + and ref_files + and os.path.exists(test_files[0]) + and os.path.exists(ref_files[0]) + ): + logger.info( + f"Extracting & assigining metadata for {varin}, the base var name of {var}" + ) + for fileset, info_dict, dataset, dataset_set in [ + (test_files[0], self.test_info, self.variables, self.test_set), + (ref_files[0], self.ref_info, self.variables, self.ref_set), + ]: + metadata = self._extract_metadata(fileset, varin, var) + self._assign_metadata( + info_dict, varin, dataset, dataset_set, idx, metadata + ) + else: + logger.info( + f"NOT extracting & assigining metadata for {varin}, the base var name of {var}." + ) + logger.info(f"test_files={test_files}") + logger.info(f"ref_files={ref_files}") + if test_files: + logger.info(f"test_files[0]={test_files[0]}") + if ref_files: + logger.info(f"ref_files[0]={ref_files[0]}") + + # `odict_keys([])` evaluates as False/None would. + if self.test_info.keys(): + self._save_catalogue(self.test_path, self.test_info) + else: + logger.info(f"test_info has no data to dump to {self.test_path}") + if self.ref_info.keys(): + self._save_catalogue(self.ref_path, self.ref_info) + else: + logger.info(f"ref_info has no data to dump to {self.ref_path}") + + return self.test_info, self.ref_info + + def _get_base_varname(self, var): + return re.split("_|-", var)[0] if ("_" in var or "-" in var) else var + + def _extract_metadata(self, filepath, varin, var): + filename = os.path.basename(filepath) + logger.info(f"Extracting metadata from {filename}, dervied from {filepath}") + parts = filename.split(".") + if len(parts) < 7: + # Example file in tmp-dir/ts: + # e3sm.amip.v3-LR.0101.Amon.ts.200501-201412.nc + logger.error( + f"Filename {filename} does not have at least 7 parts when split by '.', unexpected format." + ) + yymm_range = parts[6].split("-") + if len(yymm_range) != 2: + logger.error( + f"Filename {filename} has unexpected date range format in part '{parts[6]}'." + ) + logger.info( + f"Parsing {filename}, determined mip={parts[0]}, exp={parts[1]}, model={parts[2]}, realization={parts[3]}, tableID={parts[4]}, yymms={yymm_range[0]}, yymme={yymm_range[1]}" + ) + return { + "mip": parts[0], + "exp": parts[1], + "model": parts[2], + "realization": parts[3], + "tableID": parts[4], + "yymms": yymm_range[0], + "yymme": yymm_range[1], + "var_in_file": varin, + "var_name": var, + "file_path": filepath, + "template": filename, + } + + def _assign_metadata( + self, target_dict: OrderedDict, varin, dataset, dataset_names, idx, metadata + ): + if varin not in target_dict: + target_dict[varin] = {} + kset = ( + dataset_names[0] + if len(dataset_names) != len(dataset) + else dataset_names[idx] + ) + model = metadata["model"] + + target_dict[varin]["set"] = kset + target_dict[varin][kset] = model + target_dict[varin][model] = metadata + + def _save_catalogue(self, source_path: str, data_dict: OrderedDict): + filename = f"{source_path}_{self.label}_catalogue.json" + filepath = os.path.join(self.output_dir, filename) + logger.info( + f"Saving catalogue {filepath}, absolute path {os.path.abspath(filepath)}" + ) + with open(filepath, "w") as f: + json.dump(data_dict, f, indent=4, sort_keys=False, separators=(",", ": ")) + + +class LandSeaMaskGenerator: + def __init__(self, test_path, ref_path, subsection, fixed_dir="fixed"): + self.test_path = test_path + self.ref_path = ref_path + self.subsection = subsection + self.fixed_dir = fixed_dir + + def run(self, enable_flag): + if self._parse_flag(enable_flag): + for group_path in [self.test_path, self.ref_path]: + self._process_group(group_path) + + def _parse_flag(self, flag): + return str(flag).lower() in ["true", "y", "yes"] + + def _process_group(self, group): + catalog_path = os.path.join( + "pcmdi_diags", f"{group}_{self.subsection}_catalogue.json" + ) + + if not os.path.exists(catalog_path): + print( + f"Warning: Catalogue not found at {catalog_path}, absolute path {os.path.abspath(catalog_path)}" + ) + return + + with open(catalog_path) as f: + data_catalog = json.load(f) + + for var, meta in data_catalog.items(): + dataset = meta["set"] + model = meta[dataset] + input_file = meta[model]["file_path"] + output_file = os.path.join(self.fixed_dir, f"sftlf.{model}.nc") + + if not os.path.exists(self.fixed_dir): + os.makedirs(self.fixed_dir) + + if not os.path.exists(output_file): + self._generate_mask(input_file, output_file, model) + + def _generate_mask(self, input_path, output_path, model_name): + ds = xcdat_open(input_path, decode_times=True) + ds = ds.bounds.add_missing_bounds() + + try: + mask = create_land_sea_mask(ds, method="regionmask") + print("Land mask estimated using regionmask method.") + except Exception: + mask = create_land_sea_mask(ds, method="pcmdi") + print("Land mask estimated using PCMDI method.") + + mask = mask * 100.0 + mask.attrs.update( + {"long_name": "land_area_fraction", "units": "%", "id": "sftlf"} + ) + + mask_ds = mask.to_dataset(name="sftlf").compute() + mask_ds = mask_ds.bounds.add_missing_bounds() + mask_ds = mask_ds.fillna(1.0e20) + + mask_ds.attrs.update( + { + "model": model_name, + "associated_files": input_path, + "history": f"File processed: {datetime.now().strftime('%Y%m%d')}", + } + ) + + comp = dict(_FillValue=1.0e20) + encoding = { + v: comp for v in set(mask_ds.data_vars.keys()) | set(mask_ds.coords.keys()) + } + + mask_ds.to_netcdf(output_path, encoding=encoding) + + del ds, mask_ds, mask + + +# Functions ################################################################### + + +def set_up(parameters: CoreParameters) -> CoreOutput: + # Determine multiprocessing usage + multiprocessing: bool = ( + parameters.multiprocessing if parameters.num_workers >= 2 else False + ) + # Dataset identifiers + test_data_set: List[str] = [parameters.model_name.split(".")[1]] + reference_data_set: List[str] + if parameters.run_type == "model_vs_obs": + reference_data_set = parameters.obs_sets.split(",") + elif parameters.run_type == "model_vs_model": + reference_data_set = [parameters.model_name_ref.split(".")[1]] + else: + raise ValueError(f"Invalid run_type={parameters.run_type}") + ############################################################### + # Check and process derived quantities; these quantities are + # likely not included as default in e3sm_to_cmip module + ############################################################### + for var in parameters.variables: + varin = re.split(r"[_-]", var)[0] if "_" in var or "-" in var else var + test_fpaths = sorted( + glob.glob(os.path.join(parameters.test_data_path, f"*.{var}.*.nc")) + ) + if not test_fpaths: + derive_missing_variable( + varin, + parameters.test_data_path, + f"{parameters.model_name}.{parameters.model_tableID}", + ) + if parameters.run_type == "model_vs_model": + ref_fpaths = sorted( + glob.glob( + os.path.join(parameters.reference_data_path, f"*.{var}.*.nc") + ) + ) + if not ref_fpaths: + derive_missing_variable( + varin, + parameters.reference_data_path, + f"{parameters.model_name_ref}.{parameters.tableID_ref}", + ) + ####################################################### + # collect and document data info in a dictionary + # for convenience of pcmdi processing + ####################################################### + builder = DataCatalogueBuilder( + parameters.test_data_path, + test_data_set, + parameters.reference_data_path, + reference_data_set, + parameters.variables, + parameters.subsection, + "pcmdi_diags", + ) + _, obs_dic = builder.build_catalogues() + if not obs_dic.keys(): + raise ValueError("obs_dic has no keys!") + ########################################################## + # land/sea mask is needed in PCMDI diagnostics, check and + # generate it here as these data are not always available + # for model or observations + ########################################################## + # Instantiate and run + mask_generator = LandSeaMaskGenerator( + test_path=parameters.test_data_path, + ref_path=parameters.reference_data_path, + subsection=parameters.subsection, + fixed_dir="fixed", + ) + mask_generator.run(parameters.generate_flag) + # Diagnostic input file templates + input_template = os.path.join( + "pcmdi_diags", + "%(output_type)", + "%(metric_type)", + parameters.model_name.split(".")[0], + parameters.model_name.split(".")[1], + parameters.case_id, + ) + # Diagnostic output path templates + out_path = os.path.join(parameters.results_dir, "%(group_type)") + logger.info(f"out_path={out_path}") + return CoreOutput(multiprocessing, obs_dic, input_template, out_path) + + +def derive_missing_variable(varin, path, model_id): + """ + Derive variable with existing variables, preserving coordinates and attributes. + + Args: + varin (str): Name of the derived variable (e.g., 'rstcre'). + path (str): Directory to look for/create the file. + model_id (str): Identifier for constructing output filenames. + """ + derived_var_map = { + "rstcre": {"rsutcs": 1, "rsut": -1}, + "rltcre": {"rlutcs": 1, "rlut": -1}, + } + + if varin not in derived_var_map: + return # Nothing to derive + + var_dic = derived_var_map[varin] + derived_data = None + base_ds = None + output_file = None + + for i, (src_var, scale) in enumerate(var_dic.items()): + fpaths = sorted(glob.glob(os.path.join(path, f"*.{src_var}.*.nc"))) + if not fpaths: + raise FileNotFoundError( + f"No file found for source variable '{src_var}' in {path}" + ) + fpath = fpaths[0] + ds = xcdat_open(fpath) + data = ds[src_var] * scale + + if i == 0: + base_ds = ds.copy(deep=True) + derived_data = data.copy(deep=True) + template = os.path.basename(fpath) + output_file = os.path.join( + path, template.replace(f".{src_var}.", f".{varin}.") + ) + else: + derived_data = derived_data + data + + if base_ds is not None and derived_data is not None: + derived_da = xr.DataArray( + data=derived_data.data, + coords=derived_data.coords, + dims=derived_data.dims, + attrs=derived_data.attrs, + ) + + out_ds = base_ds.drop_vars(list(var_dic.keys()), errors="ignore") + out_ds[varin] = derived_da + + # Optional: annotate global attributes + out_ds.attrs.update( + { + "derived_variable": varin, + "derived_from": ", ".join(var_dic.keys()), + "model_id": model_id, + } + ) + + out_ds.to_netcdf(output_file) + print(f"Derived variable '{varin}' written to {output_file}") + + return diff --git a/zppy_interfaces/pcmdi_diags/pcmdi_synthetic_plots.py b/zppy_interfaces/pcmdi_diags/pcmdi_synthetic_plots.py new file mode 100644 index 0000000..aa24802 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/pcmdi_synthetic_plots.py @@ -0,0 +1,260 @@ +import argparse +import json +import os +import shutil +import sys +from typing import Dict, List + +from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.synthetic_metrics_plotter import ( + SyntheticMetricsPlotter, +) +from zppy_interfaces.pcmdi_diags.viewer import ( + collect_config, + generate_data_html, + generate_methodology_html, + generate_viewer_html, +) + +# Set up the root logger and module level logger. The module level logger is +# a child of the root logger. +_setup_root_logger() +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class SyntheticPlotsParameters(object): + def __init__(self, args: Dict[str, str]): + self.figure_format: str = args["figure_format"] + self.www: str = args["www"] + self.save_all_data: bool = str(args["save_all_data"]).lower() in ( + "true", + "1", + "yes", + ) + self.results_dir: str = args["results_dir"] + self.case: str = args["case"] + self.model_name: str = args["model_name"] + self.model_tableID: str = args["model_tableID"] + self.web_dir: str = args["web_dir"] + self.clim_viewer: bool = str(args["clim_viewer"]).lower() in ( + "true", + "1", + "yes", + ) + self.clim_vars: List[str] = args["clim_vars"].split(",") + self.clim_years: str = args["clim_years"] + self.clim_regions: List[str] = args["clim_regions"].split(",") + self.cmip_clim_dir: str = args["cmip_clim_dir"] + self.cmip_clim_set: str = args["cmip_clim_set"] + self.mova_viewer: bool = str(args["mova_viewer"]).lower() in ( + "true", + "1", + "yes", + ) + self.mova_modes: List[str] = args["mova_modes"].split(",") + self.mova_vars: List[str] = args["mova_vars"].split(",") + self.mova_years: str = args["mova_years"] + self.movc_viewer: bool = str(args["movc_viewer"]).lower() in ( + "true", + "1", + "yes", + ) + self.movc_modes: List[str] = args["movc_modes"].split(",") + self.movc_vars: List[str] = args["movc_vars"].split(",") + self.movc_years: str = args["movc_years"] + self.cmip_movs_dir: str = args["cmip_movs_dir"] + self.cmip_movs_set: str = args["cmip_movs_set"] + self.enso_viewer: bool = str(args["enso_viewer"]).lower() in ( + "true", + "1", + "yes", + ) + self.enso_vars: List[str] = args["enso_vars"].split(",") + self.enso_years: str = args["enso_years"] + self.cmip_enso_dir: str = args["cmip_enso_dir"] + self.cmip_enso_set: str = args["cmip_enso_set"] + self.pcmdi_webtitle: str = args["pcmdi_webtitle"] + self.pcmdi_version: str = args["pcmdi_version"] + self.run_type: str = args["run_type"] + self.pcmdi_external_prefix: str = args["pcmdi_external_prefix"] + self.pcmdi_viewer_template: str = args["pcmdi_viewer_template"] + + +# Functions ################################################################### +def main(): + args: Dict[str, str] = _get_args() + parameters = SyntheticPlotsParameters(args) + + ######################################### + # plot synthetic figures for pcmdi metrics + ######################################### + logger.info("generate synthetic metrics plot ...") + test_input_path = os.path.join( + parameters.www, + "put_model_here", + "pcmdi_diags", + parameters.results_dir, + "metrics_data", + "%(group_type)", + ) + metric_dict = json.load(open("synthetic_metrics_list.json")) + plotter = SyntheticMetricsPlotter( + # Core + case_name=parameters.case, + test_name=parameters.model_name, + table_id=parameters.model_tableID, + figure_format=parameters.figure_format, + metric_dict=metric_dict, + save_data=parameters.save_all_data, + base_test_input_path=test_input_path, + results_dir=os.path.join(parameters.web_dir, parameters.results_dir), + # Mean climate + clim_viewer=parameters.clim_viewer, + clim_vars=parameters.clim_vars, + clim_regions=parameters.clim_regions, + cmip_clim_dir=parameters.cmip_clim_dir, + cmip_clim_set=parameters.cmip_clim_set, + # MOVA + mova_viewer=parameters.mova_viewer, + mova_modes=parameters.mova_modes, + # MOVC + movc_viewer=parameters.movc_viewer, + movc_modes=parameters.movc_modes, + cmip_movs_dir=parameters.cmip_movs_dir, + cmip_movs_set=parameters.cmip_movs_set, + # ENSO + enso_viewer=parameters.enso_viewer, + cmip_enso_dir=parameters.cmip_enso_dir, + cmip_enso_set=parameters.cmip_enso_set, + ) + + # Generate Summary Metrics plots + # e.g., "climatology,enso,variability" + figure_sets = [] + if parameters.clim_viewer: + figure_sets.append("climatology") + if parameters.mova_viewer: + figure_sets.append("variability(ATM)") + if parameters.movc_viewer: + figure_sets.append("variability(CPL)") + if parameters.enso_viewer: + figure_sets.append("enso") + + logger.info(f"Generating groups={figure_sets}") + # This calls the _handle_{figure_set} functions + # Those call the {figure_set}_plot_driver functions + plotter.generate() + + logger.info("Generating viewer page for diagnostics...") + subtitle = parameters.run_type.replace("_", " ").capitalize() + + # Set up paths + obs_dir = os.path.join( + parameters.pcmdi_external_prefix, "observations", "Atm", "time-series" + ) + pmp_dir = os.path.join(parameters.pcmdi_external_prefix, "pcmdi_data") + out_dir = os.path.join(parameters.web_dir, parameters.results_dir, "viewer") + os.makedirs(out_dir, exist_ok=True) + # Copy logo + web_logo_src = os.path.join( + parameters.pcmdi_external_prefix, + parameters.pcmdi_viewer_template, + "e3sm_pmp_logo.png", + ) + web_logo_dst = os.path.join(out_dir, "e3sm_pmp_logo.png") + shutil.copy(web_logo_src, web_logo_dst) + # Build config + config = collect_config( + title=parameters.pcmdi_webtitle, + subtitle=subtitle, + version=parameters.pcmdi_version, + case_id=parameters.case, + diag_dir=parameters.web_dir, + obs_dir=obs_dir, + pmp_dir=pmp_dir, + out_dir=out_dir, + clim_viewer=parameters.clim_viewer, + clim_period=parameters.clim_years, + clim_regions=parameters.clim_regions, + clim_vars=parameters.clim_vars, + mova_viewer=parameters.mova_viewer, + mova_modes=parameters.mova_modes, + mova_period=parameters.mova_years, + movc_viewer=parameters.movc_viewer, + movc_modes=parameters.movc_modes, + movc_period=parameters.movc_years, + enso_viewer=parameters.enso_viewer, + enso_period=parameters.enso_years, + ) + # Render viewer + generate_methodology_html(config) + generate_data_html(config) + generate_viewer_html(config) + + +def str2bool(v): + if isinstance(v, bool): + return v + val = str(v).lower() + if val in ("yes", "true", "t", "1", "y", "on"): + return True + elif val in ("no", "false", "f", "0", "n", "off"): + return False + else: + raise argparse.ArgumentTypeError(f"Invalid boolean value: {v}") + + +def _get_args() -> Dict[str, str]: + # Parser + parser: argparse.ArgumentParser = argparse.ArgumentParser( + usage="zi-pcmdi-synthetic-plots " + ) + + # For SyntheticPlotsParameters + parser.add_argument("--synthetic_sets", type=str) + parser.add_argument("--figure_format", type=str) + parser.add_argument("--www", type=str) + parser.add_argument("--results_dir", type=str) + parser.add_argument("--case", type=str) + parser.add_argument("--model_name", type=str) + parser.add_argument("--model_tableID", type=str) + parser.add_argument("--web_dir", type=str) + parser.add_argument("--clim_viewer", type=str2bool) + parser.add_argument("--clim_vars", type=str) + parser.add_argument("--clim_years", type=str) + parser.add_argument("--clim_regions", type=str) + parser.add_argument("--cmip_clim_dir", type=str) + parser.add_argument("--cmip_clim_set", type=str) + parser.add_argument("--mova_viewer", type=str2bool) + parser.add_argument("--mova_modes", type=str) + parser.add_argument("--mova_vars", type=str) + parser.add_argument("--mova_years", type=str) + parser.add_argument("--movc_viewer", type=str2bool) + parser.add_argument("--movc_modes", type=str) + parser.add_argument("--movc_vars", type=str) + parser.add_argument("--movc_years", type=str) + parser.add_argument("--cmip_movs_dir", type=str) + parser.add_argument("--cmip_movs_set", type=str) + parser.add_argument("--enso_viewer", type=str2bool) + parser.add_argument("--enso_vars", type=str) + parser.add_argument("--enso_years", type=str) + parser.add_argument("--cmip_enso_dir", type=str) + parser.add_argument("--cmip_enso_set", type=str) + parser.add_argument("--pcmdi_webtitle", type=str) + parser.add_argument("--pcmdi_version", type=str) + parser.add_argument("--run_type", type=str) + parser.add_argument("--pcmdi_external_prefix", type=str) + parser.add_argument("--pcmdi_viewer_template", type=str) + parser.add_argument("--save_all_data", type=str2bool) + parser.add_argument("--debug", type=str) + + # Ignore the first arg + # (zi-pcmdi-synthetic-plots) + args: argparse.Namespace = parser.parse_args(sys.argv[1:]) + + if args.debug and args.debug.lower() == "true": + logger.setLevel("DEBUG") + logger.debug("Debug logging enabled") + + return vars(args) diff --git a/zppy_interfaces/pcmdi_diags/pcmdi_variability_modes.py b/zppy_interfaces/pcmdi_diags/pcmdi_variability_modes.py new file mode 100644 index 0000000..22ac7ad --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/pcmdi_variability_modes.py @@ -0,0 +1,286 @@ +import argparse +import glob +import os +import sys +import time +from collections import OrderedDict +from typing import Dict, List + +from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger +from zppy_interfaces.pcmdi_diags.pcmdi_setup import CoreOutput, CoreParameters, set_up +from zppy_interfaces.pcmdi_diags.utils import run_parallel_jobs, run_serial_jobs + +# Set up the root logger and module level logger. The module level logger is +# a child of the root logger. +_setup_root_logger() +logger = _setup_child_logger(__name__) + + +# Classes ##################################################################### +class VariabilityModesParameters(object): + def __init__(self, args: Dict[str, str]): + self.var_modes: List[str] = args["var_modes"].split(",") + # self.vars is distinct from the list version in CoreParameters + self.vars: str = args["vars"] + + +class VariabilityMetricsCollector: + def __init__( + self, modes, fig_format, mip, exp, model, relm, case_id, input_dir, output_dir + ): + self.modes = modes + self.fig_format = fig_format + self.mip = mip + self.exp = exp + self.model = model + self.relm = relm + self.case_id = case_id + self.input_dir = input_dir.replace("%(metric_type)", "variability_modes") + self.output_dir = output_dir + self.model_name = f"{mip}.{exp}.{model}_{relm}" + self.seasons = ["DJF", "MAM", "JJA", "SON", "yearly", "monthly"] + self.fig_sets = OrderedDict( + { + "MOV_eoftest": ["diagnostic_results", "EG_Spec*"], + "MOV_compose": ["graphics", "*compare_obs"], + "MOV_telecon": ["graphics", "*teleconnection"], + "MOV_pattern": ["graphics", "*"], + } + ) + + def collect(self): + self._collect_figures() + self._collect_metrics() + self._collect_diags() + + def _collect_figures(self): + for fig_set, (out_type, pattern_base) in self.fig_sets.items(): + for mode in self.modes: + for season in self.seasons: + indir = self.input_dir.replace("%(output_type)", out_type) + template = ( + f"{pattern_base}_{mode}_{season}*.{self.fig_format}" + if fig_set == "MOV_eoftest" + else f"{mode}_*_{season}_{pattern_base}.{self.fig_format}" + ) + search_path = os.path.join(indir, mode, "*", template) + matched_files = sorted(glob.glob(search_path)) + + for fpath in matched_files: + filename = os.path.basename(fpath) + outfile = self._classify_output_name( + fig_set, mode, season, filename + ) + outdir = os.path.join( + self.output_dir.replace("%(group_type)", "MOV_metric"), + fig_set, + season, + ) + os.makedirs(outdir, exist_ok=True) + os.rename(fpath, os.path.join(outdir, outfile)) + + def _classify_output_name(self, fig_set, mode, season, filename): + suffix = "unknown" + if "North_test" in filename: + suffix = "EG_Spec" + elif "_cbf_" in filename: + suffix = "cbf" + elif "EOF1" in filename: + suffix = "eof1" + elif "EOF2" in filename: + suffix = "eof2" + elif "EOF3" in filename: + suffix = "eof3" + return f"{fig_set}_{mode}_{season}_{suffix}.{self.fig_format}" + + def _collect_metrics(self): + metrics_dir = self.input_dir.replace("%(output_type)", "metrics_results") + json_files = sorted(glob.glob(os.path.join(metrics_dir, "*/*/*.json"))) + + for fpath in json_files: + refmode = fpath.split("/")[-3] + refname = fpath.split("/")[-2] + reffile = fpath.split("/")[-1] + + eof_lookup = {"PSA1": "EOF2", "NPO": "EOF2", "NPGO": "EOF2", "PSA2": "EOF3"} + refeof = eof_lookup.get(refmode, "EOF1") + + outdir = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), + "variability_modes", + refmode, + refname, + ) + os.makedirs(outdir, exist_ok=True) + + base_name = f"var_mode_{refmode}.{refeof}.{self.model_name}.vs.{refname}.{self.case_id}" + if "diveDown" in reffile: + outfile = os.path.join(outdir, f"{base_name}.diveDown.json") + else: + outfile = os.path.join(outdir, f"{base_name}.json") + + os.rename(fpath, outfile) + + def _collect_diags(self): + diags_dir = self.input_dir.replace("%(output_type)", "diagnostic_results") + json_files = sorted(glob.glob(os.path.join(diags_dir, "*/*/*.nc"))) + + for fpath in json_files: + refmode = fpath.split("/")[-3] + refname = fpath.split("/")[-2] + reffile = fpath.split("/")[-1] + + outdir = os.path.join( + self.output_dir.replace("%(group_type)", "metrics_data"), + "variability_modes", + refmode, + refname, + ) + os.makedirs(outdir, exist_ok=True) + + outfile = os.path.join(outdir, reffile) + + os.rename(fpath, outfile) + + +# Functions ################################################################### +def main(): + args: Dict[str, str] = _get_args() + core_parameters = CoreParameters(args) + variability_modes_parameters = VariabilityModesParameters(args) + core_output: CoreOutput = set_up(core_parameters) + + ########################################## + # call pcmdi mode variability diagnostics + ########################################## + # from configuration file + varOBS = variability_modes_parameters.vars + if varOBS not in core_output.obs_dic: + raise KeyError( + f"VarOBS '{varOBS}' not found in obs_dic. Available keys are {core_output.obs_dic.keys()}" + ) + refset = core_output.obs_dic[varOBS]["set"] + refname = core_output.obs_dic[varOBS][refset] + refpath = core_output.obs_dic[varOBS][refname]["file_path"] + reftyrs = int(str(core_output.obs_dic[varOBS][refname]["yymms"])[0:4]) + reftyre = int(str(core_output.obs_dic[varOBS][refname]["yymme"])[0:4]) + # Call the function + lstcmd = generate_varmode_cmds( + modes=variability_modes_parameters.var_modes, + varOBS=varOBS, + reftyrs=reftyrs, + reftyre=reftyre, + refname=refname, + refpath=refpath, + case_id=core_parameters.case_id, + ) + if (len(lstcmd) > 0) and core_parameters.multiprocessing: + try: + results = run_parallel_jobs(lstcmd, core_parameters.num_workers) + for i, (stdout, stderr, return_code) in enumerate(results): + print(f"\nCommand {i + 1} finished:") + print(f"STDOUT: {stdout}") + print(f"STDERR: {stderr}") + print(f"Return code: {return_code}") + except RuntimeError as e: + print(f"Execution failed: {e}") + elif len(lstcmd) > 0: + try: + results = run_serial_jobs(lstcmd) + for i, (stdout, stderr, return_code) in enumerate(results): + print(f"\nCommand {i + 1} finished:") + print(f"STDOUT: {stdout}") + print(f"STDERR: {stderr}") + print(f"Return code: {return_code}") + except RuntimeError as e: + print(f"Execution failed: {e}") + else: + print("no jobs to run,continue...") + print("successfully finish all jobs....") + # time delay to ensure process completely finished + time.sleep(5) + # Create the collector instance + split_name: List[str] = core_parameters.model_name.split(".") + collector = VariabilityMetricsCollector( + modes=variability_modes_parameters.var_modes, + fig_format=core_parameters.figure_format, + mip=split_name[0], + exp=split_name[1], + model=split_name[2], + relm=split_name[3], + case_id=core_parameters.case_id, + input_dir=core_output.input_template, + output_dir=core_output.out_path, + ) + # Run the collection process + collector.collect() + + +def _get_args() -> Dict[str, str]: + # Parser + parser: argparse.ArgumentParser = argparse.ArgumentParser( + usage="zi-pcmdi-variability-modes " + ) + + # For CoreParameters + parser.add_argument("--num_workers", type=str) + parser.add_argument("--multiprocessing", type=str) + parser.add_argument("--subsection", type=str) + parser.add_argument("--climo_ts_dir_primary", type=str) # needs ts_dir_primary + parser.add_argument("--climo_ts_dir_ref", type=str) # needs ts_dir_ref + parser.add_argument("--model_name", type=str) + parser.add_argument("--model_tableID", type=str) + parser.add_argument("--figure_format", type=str) + parser.add_argument("--run_type", type=str) + parser.add_argument("--obs_sets", type=str) # run_type == "model_vs_obs" only + parser.add_argument( + "--model_name_ref", type=str + ) # run_type == "model_vs_model" only + parser.add_argument("--vars", type=str) + parser.add_argument("--tableID_ref", type=str) # run_type == "model_vs_model" only + parser.add_argument("--generate_sftlf", type=str) + parser.add_argument("--case_id", type=str) + parser.add_argument("--results_dir", type=str) + parser.add_argument("--debug", type=str) + + # For VariabilityModesParameters + parser.add_argument("--var_modes", type=str) # use either atm_mdoes or cpl_modes + + # Ignore the first arg + # (zi-pcmdi-variability-modes) + args: argparse.Namespace = parser.parse_args(sys.argv[1:]) + + if args.debug and args.debug.lower() == "true": + logger.setLevel("DEBUG") + logger.debug("Debug logging enabled") + + return vars(args) + + +def generate_varmode_cmds(modes, varOBS, reftyrs, reftyre, refname, refpath, case_id): + """Generates a list of command strings for variability modes processing.""" + + # EOF mode overrides for specific variability modes (default is 1) + eofn_map = {"NPO": 2, "NPGO": 2, "PSA1": 2, "PSA2": 3} + + commands = [] + + for var_mode in modes: + var_mode = var_mode.strip() + # Use specified EOF number if in map, otherwise default to 1 + eofn = eofn_map.get(var_mode, 1) + cmd = ( + f"variability_modes_driver.py -p parameterfile.py " + f"--variability_mode {var_mode} " + f"--eofn_mod {eofn} " + f"--eofn_obs {eofn} " + f"--varOBS {varOBS} " + f"--osyear {reftyrs} " + f"--oeyear {reftyre} " + f"--reference_data_name {refname} " + f"--reference_data_path {refpath} " + f"--case_id {case_id}" + ) + commands.append(cmd) + + return commands diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/__init__.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_merger.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_merger.py new file mode 100644 index 0000000..1e78281 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_merger.py @@ -0,0 +1,302 @@ +from collections.abc import MutableMapping +from copy import deepcopy +from typing import Dict, List, Optional + +import numpy as np +import pandas as pd + +from zppy_interfaces.multi_utils.logger import _setup_child_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import ( + get_highlight_models, + shift_row_to_bottom, +) + +logger = _setup_child_logger(__name__) + + +class ClimMetricsMerger: + def __init__(self, model_lib=None, cmip_lib=None, model_names=None): + self.model_lib = model_lib or {} + self.cmip_lib = cmip_lib or {} + self.model_names = model_names or [] + self.merged_lib = None + + def merge(self): + self._normalize_references() + self._filter_regions() + self._merge_and_standardize_units() + self._highlight_and_sort_models() + return self.merged_lib + + def _normalize_references(self): + if hasattr(self.model_lib, "references") and isinstance( + self.model_lib.references, dict + ): + self.model_lib.references = self._check_references( + self.model_lib.references + ) + if hasattr(self.cmip_lib, "references") and isinstance( + self.cmip_lib.references, dict + ): + self.cmip_lib.references = self._check_references(self.cmip_lib.references) + + def _check_references( + self, + data_dict: MutableMapping[str, Optional[List[str]]], + reference_alias: Optional[Dict[str, str]] = None, + ) -> MutableMapping[str, Optional[List[str]]]: + if reference_alias is None: + reference_alias = { + "ceres_ebaf_toa_v4.1": "ceres_ebaf_v4_1", + "ceres_ebaf_toa_v4.0": "ceres_ebaf_v4_0", + "ceres_ebaf_toa_v2.8": "ceres_ebaf_v2_8", + "ceres_ebaf_surface_v4.1": "ceres_ebaf_v4_1", + "ceres_ebaf_surface_v4.0": "ceres_ebaf_v4_0", + "ceres_ebaf_surface_v2.8": "ceres_ebaf_v2_8", + "CERES-EBAF-4-1": "ceres_ebaf_v4_1", + "CERES-EBAF-4-0": "ceres_ebaf_v4_0", + "CERES-EBAF-2-8": "ceres_ebaf_v2_8", + "GPCP_v2.3": "GPCP_v2_3", + "GPCP_v2.2": "GPCP_v2_2", + "GPCP_v3.2": "GPCP_v3_2", + "GPCP-2-3": "GPCP_v2_3", + "GPCP-2-2": "GPCP_v2_2", + "GPCP-3-2": "GPCP_v3_2", + "NOAA_20C": "NOAA-20C", + "ERA-INT": "ERA-Interim", + "ERA-5": "ERA5", + } + + for key, values in data_dict.items(): + if isinstance(values, list): + data_dict[key] = [reference_alias.get(val, val) for val in values] + elif values is not None: + data_dict[key] = reference_alias.get(values, values) + else: + raise ValueError("values is None") + + return data_dict + + def _filter_regions(self): + self.model_lib, self.cmip_lib = self._check_regions( + self.model_lib, self.cmip_lib + ) + + def _check_regions(self, data_lib, refr_lib): + shared_regions = [ + region for region in data_lib.regions if region in refr_lib.regions + ] + + for lib in [refr_lib, data_lib]: + for stat in lib.df_dict: + for season in lib.df_dict[stat]: + lib.df_dict[stat][season] = { + region: lib.df_dict[stat][season][region] + for region in shared_regions + if region in lib.df_dict[stat][season] + } + + data_lib.regions = shared_regions + refr_lib.regions = shared_regions + + return data_lib, refr_lib + + @staticmethod + def _prune_empty_dfs(lib): + for stat in lib.df_dict: + for season in lib.df_dict[stat]: + lib.df_dict[stat][season] = { + region: df + for region, df in lib.df_dict[stat][season].items() + if not df.empty and not df.isna().all().all() + } + return lib + + @staticmethod + def _safe_merge_libs(lib1, lib2): + """ + Merge two data libraries with nested dicts of DataFrames, + gracefully handling missing or inconsistent keys, while + avoiding FutureWarning due to all-NA/empty entries. + """ + merged = deepcopy(lib1) # Avoid modifying original + + for stat in lib2.df_dict: + if stat not in merged.df_dict: + merged.df_dict[stat] = {} + + for season in lib2.df_dict[stat]: + if season not in merged.df_dict[stat]: + merged.df_dict[stat][season] = {} + + for region, df2 in lib2.df_dict[stat][season].items(): + df1 = merged.df_dict[stat][season].get(region) + + # Collect and clean valid DataFrames + valid_dfs = [] + for df in (df1, df2): + if ( + isinstance(df, pd.DataFrame) + and not df.empty + and not df.isna().all().all() + ): + # Drop columns that are entirely NaN + df_clean = df.dropna(axis=1, how="all") + if not df_clean.empty and df_clean.shape[1] > 0: + valid_dfs.append(df_clean) + + if valid_dfs: + merged_df = pd.concat(valid_dfs, ignore_index=True, sort=False) + else: + merged_df = pd.DataFrame() + + merged.df_dict[stat][season][region] = merged_df + + return merged + + def _merge_and_standardize_units(self): + # Prune empty or fully-NaN DataFrames from the model library + cleaned_model_lib = self._prune_empty_dfs(self.model_lib) + if hasattr(cleaned_model_lib, "var_list"): + logger.debug(f"cleaned_model_lib.var_list: {cleaned_model_lib.var_list}") + + # Safe merge with fallback for missing stats/seasons/regions + self.merged_lib = self._safe_merge_libs(self.cmip_lib, cleaned_model_lib) + if hasattr(self.merged_lib, "var_list"): + logger.debug(f"merged_lib.var_list: {self.merged_lib.var_list}") + if hasattr(cleaned_model_lib, "var_list") and hasattr( + self.cmip_lib, "var_list" + ): + var_set_cleaned_model_lib = set(cleaned_model_lib.var_list) + var_set_cmip_lib = set(self.cmip_lib.var_list) + logger.debug( + f"Var list sizes - cleaned_model_lib: {len(var_set_cleaned_model_lib)}, cmip_lib: {len(var_set_cmip_lib)}" + ) + logger.debug( + f"Var list differences - in cleaned_model_lib not in cmip_lib: {var_set_cleaned_model_lib - var_set_cmip_lib}" + ) + logger.debug( + f"Var list differences - in cmip_lib not in cleaned_model_lib: {var_set_cmip_lib - var_set_cleaned_model_lib}" + ) + + # Standardize units after merging + self.merged_lib = self._check_units(self.merged_lib) + if hasattr(self.merged_lib, "var_list"): + logger.debug( + f"Post-unit-check merged_lib.var_list: {self.merged_lib.var_list}" + ) + + def _check_units(self, data_lib, verbose=False): + units_all = { + "prw": "[kg m$^{-2}$]", + "pr": "[mm d$^{-1}$]", + "prsn": "[mm d$^{-1}$]", + "prc": "[mm d$^{-1}$]", + "hfls": "[W m$^{-2}$]", + "hfss": "[W m$^{-2}$]", + "clivi": "[kg $m^{-2}$]", + "clwvi": "[kg $m^{-2}$]", + "psl": "[Pa]", + "rlds": "[W m$^{-2}$]", + "rldscs": "[W $m^{-2}$]", + "evspsbl": "[kg m$^{-2} s^{-1}$]", + "rtmt": "[W m$^{-2}$]", + "rsdt": "[W m$^{-2}$]", + "rlus": "[W m$^{-2}$]", + "rluscs": "[W m$^{-2}$]", + "rlut": "[W m$^{-2}$]", + "rlutcs": "[W m$^{-2}$]", + "rsds": "[W m$^{-2}$]", + "rsdscs": "[W m$^{-2}$]", + "rstcre": "[W m$^{-2}$]", + "rltcre": "[W m$^{-2}$]", + "rsus": "[W m$^{-2}$]", + "rsuscs": "[W m$^{-2}$]", + "rsut": "[W m$^{-2}$]", + "rsutcs": "[W m$^{-2}$]", + "ts": "[K]", + "tas": "[K]", + "tauu": "[Pa]", + "tauv": "[Pa]", + "zg-500": "[m]", + "ta-200": "[K]", + "sfcWind": "[m s$^{-1}$]", + "ta-850": "[K]", + "ua-200": "[m s$^{-1}$]", + "ua-850": "[m s$^{-1}$]", + "va-200": "[m s$^{-1}$]", + "va-850": "[m s$^{-1}$]", + "uas": "[m s$^{-1}$]", + "vas": "[m s$^{-1}$]", + "tasmin": "[K]", + "tasmax": "[K]", + "clt": "[%]", + } + + # Identify common variables and handle aliases like 'rt' or 'rmt' + common_vars = [var for var in data_lib.var_list if var in units_all] + if "rtmt" not in common_vars and any( + var in data_lib.var_list for var in ["rt", "rmt"] + ): + common_vars.append("rtmt") + + # Collect units for these variables + common_unts = [units_all[var] for var in common_vars if var in units_all] + + # Filter and correct reference list + new_var_ref_dict = {} + for var, ref in data_lib.var_ref_dict.items(): + if var in common_vars: + new_var_ref_dict[var] = ref + elif var in ["rt", "rmt"]: + new_var_ref_dict["rtmt"] = ref + if verbose: + logger.info(f"Alias {var} mapped to 'rtmt' in references.") + + data_lib.var_ref_dict = self._check_references(new_var_ref_dict) + + # Clean DataFrames + for stat, seasons in data_lib.df_dict.items(): + for season, regions in seasons.items(): + for region, df in regions.items(): + df = df.copy() + # Handle aliases + if "rt" in df.columns: + df["rtmt"] = df["rt"] + elif "rmt" in df.columns: + df["rtmt"] = df["rmt"] + + # Drop irrelevant variables + drop_cols = [ + var for var in df.columns[3:] if var not in common_vars + ] + if drop_cols and verbose: + logger.info( + f"Dropping variables in {stat}/{season}/{region}: {drop_cols}" + ) + df = df.drop(columns=drop_cols) + data_lib.df_dict[stat][season][region] = df + + logger.debug(f"Setting data_lib.var_list={common_vars}") + data_lib.var_list = common_vars + data_lib.var_unit_list = common_unts + + return data_lib + + def _highlight_and_sort_models(self): + if self.merged_lib: + for stat, seasons in self.merged_lib.df_dict.items(): + for season, regions in seasons.items(): + for region, df in regions.items(): + df = pd.DataFrame(df) + highlight_models = get_highlight_models( + df.get("model", []), self.model_names + ) + for model in highlight_models: + for idx in df[df["model"] == model].index: + df = shift_row_to_bottom(df, idx) + self.merged_lib.df_dict[stat][season][region] = df.fillna( + np.nan + ) + else: + raise ValueError("merged_lib is None") diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_reader.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_reader.py new file mode 100644 index 0000000..144728b --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/clim_metrics_reader.py @@ -0,0 +1,215 @@ +import glob +import json +import os +import re + +import numpy as np +import pandas as pd +from pcmdi_metrics.graphics import Metrics + +from zppy_interfaces.multi_utils.logger import _setup_child_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.clim_metrics_merger import ( + ClimMetricsMerger, +) +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import find_latest_file_list + +logger = _setup_child_logger(__name__) + + +class ClimMetricsReader: + def __init__(self, parameter, unit_check=True): + """ + Initialize the climate metrics collector. + + Args: + parameter (dict): Contains path, model info, and identifiers. + unit_check (bool): Whether to apply unit consistency check. + """ + self.parameter = parameter + self.unit_check = unit_check + self.cmip_lib = None + self.all_lib = None + self.all_names = [] + + self.var_pattern = re.compile(r"^([A-Za-z0-9\-]+)\.") + self.time_pattern = re.compile(r"\.v(\d{8})\.json$") + + def _load_clim_metrics_from_files(self, file_paths): + """ + Loads and processes synthetic climate metric data from JSON files. + + Parameters: + file_paths (list): List of file paths to load. + + Returns: + Metrics: Processed Metrics object. + """ + logger.info("file_paths=") + for i, fp in enumerate(file_paths): + logger.info(f"{i}. {fp}") + """ + FAILURE: + + list(results_dict_var["RESULTS"][model_list[0]]["default"][run_list[0]].keys()) + IndexError: list index out of range + + cat /lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zppy_pr719_output/unique_id_21/v3.LR.amip_0101/pcmdi_diags/model_vs_obs/metrics_data/mean_climate/rlus.2.5x2.5.e3sm.amip.v3-LR_0101.v20250725.json + + "RESULTS": { + "v3-LR": { + "default": { + "source": "ceres_ebaf_v4_1" + } + } + }, + + SYNTHETIC PLOTS ERROR #1: "source" is supposed to be a dictionary itself, even though mean_climate job completed successfully! + """ + lib = Metrics(file_paths) + lib = check_badvals(lib) + if self.unit_check: + a = ClimMetricsMerger() + lib = a._check_units(lib) + return lib + + def _load_cmip_metrics(self): + cmip_id_parts = self.parameter["cmip_name"].split(".") + cmip_dir = os.path.join( + self.parameter["cmip_path"], + cmip_id_parts[0], + cmip_id_parts[1], + cmip_id_parts[2], + ) + + cmip_files = sorted( + glob.glob(os.path.join(cmip_dir, f"*.{cmip_id_parts[2]}.json")) + ) + if not cmip_files: + raise FileNotFoundError(f"No CMIP metrics found in: {cmip_dir}") + + logger.info(f"Loading CMIP metrics from {len(cmip_files)} files...") + self.cmip_lib = self._load_clim_metrics_from_files(cmip_files) + + def _process_test_model(self, test_name, model_name): + test_key = test_name.split(".")[1] + test_path = self.parameter["test_path"].replace("put_model_here", model_name) + + model_files = find_latest_file_list( + path=test_path, + file_pattern="*.v*.json", + var_pattern=self.var_pattern, + time_pattern=self.time_pattern, + ) + + if not model_files or not os.path.exists(model_files[0]): + raise FileNotFoundError( + f"No synthetic mean climate metrics found for model: {model_name}" + ) + + logger.info( + f"Reading metrics for model: {model_name} from {len(model_files)} file(s)..." + ) + + valid_model_files = [] + + for file_path in model_files: + try: + with open(file_path, "r") as f: + data = json.load(f) + + results = data.get("RESULTS", {}) + modified = False + + for model, model_data in results.items(): + if test_key in model_data: + model_data["default"] = model_data.pop(test_key) + modified = True + + if modified: + with open(file_path, "w") as f: + json.dump(data, f, indent=2) + logger.info(f"Updated file: {file_path}") + + valid_model_files.append(file_path) + + except (FileNotFoundError, json.JSONDecodeError) as e: + logger.info(f"Warning: Could not load {file_path}: {e}") + + # Load metrics from valid files + model_lib = self._load_clim_metrics_from_files(valid_model_files) + + # Standardize model name in metric DataFrames + for stat, seasons in model_lib.df_dict.items(): + for season, regions in seasons.items(): + for region, df in regions.items(): + df = pd.DataFrame(df) + if "model" in df.columns: + df["model"] = model_name + model_lib.df_dict[stat][season][region] = df + + return model_lib + + def collect(self): + self._load_cmip_metrics() + if self.cmip_lib and hasattr(self.cmip_lib, "var_list"): + logger.debug( + f"ClimMetricsReader.cmip_lib.vars_list: {self.cmip_lib.var_list}" + ) + + for i, (test_name, model_name) in enumerate( + zip(self.parameter["test_name"], self.parameter["model_name"]) + ): + logger.debug( + f"Processing model {i + 1}: test_name={test_name}, model_name={model_name}" + ) + model_lib = self._process_test_model(test_name, model_name) + if hasattr(model_lib, "var_list"): + logger.debug(f"model_lib.vars_list: {model_lib.var_list}") + self.all_lib = ( + model_lib.copy() + if self.all_lib is None + else self.all_lib.merge(model_lib) + ) + self.all_names.append(model_name) + + logger.info("Merging model metrics with CMIP reference metrics...") + if self.all_lib and hasattr(self.all_lib, "var_list"): + logger.debug( + f"ClimMetricsReader.all_lib.vars_list: {self.all_lib.var_list}" + ) + merger = ClimMetricsMerger( + model_lib=self.all_lib, cmip_lib=self.cmip_lib, model_names=self.all_names + ) + merged_metrics = merger.merge() # Returns a new merged metrics library + + return merged_metrics + + +def check_badvals(data_lib): + """ + Replaces known bad values in the data library with NaN. + + Parameters: + data_lib (Metrics): Metrics object containing diagnostic DataFrames. + + Returns: + Metrics: Updated metrics with bad values replaced by NaN. + """ + # Define known bad values (model → variable with suspect data) + bad_model_vars = {"E3SM-1-0": "ta-850", "E3SM-1-1-ECA": "ta-850", "CIESM": "pr"} + + for stat in data_lib.df_dict: + for season in data_lib.df_dict[stat]: + for region in data_lib.df_dict[stat][season]: + df = pd.DataFrame(data_lib.df_dict[stat][season][region]) + + for model, bad_var in bad_model_vars.items(): + if bad_var in df.columns: + # Find all rows matching this model + bad_idx = df[df["model"] == model].index + df.loc[bad_idx, bad_var] = np.nan + + # Save cleaned DataFrame back + data_lib.df_dict[stat][season][region] = df + + return data_lib diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/enso_metrics_reader.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/enso_metrics_reader.py new file mode 100644 index 0000000..b07d9f4 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/enso_metrics_reader.py @@ -0,0 +1,93 @@ +import glob +import json +import os +import re + +from zppy_interfaces.multi_utils.logger import _setup_child_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import find_latest_file_list + +logger = _setup_child_logger(__name__) + + +class EnsoMetricsReader: + def __init__(self, parameter, stat, metric_dict, mips, collections): + self.parameter = parameter + self.stat = stat + self.mips = mips + self.metric_dict = metric_dict + self.metrics_collections = collections + self.dict_json_path = {} + + self.var_pattern = re.compile(r"\.(\w+)\..*\.v(\d{8})\.json$") + self.time_pattern = re.compile(r"\.v(\d{8})\.json$") + + def run(self): + """Collect paths to ENSO metrics JSON files and return the mapping.""" + for mip in self.mips: + self.dict_json_path[mip] = {} + for metrics_collection in self.metrics_collections: + if "cmip" in mip: + self.dict_json_path[mip][metrics_collection] = ( + self._get_cmip_json_path(mip, metrics_collection) + ) + else: + self.dict_json_path[mip][metrics_collection] = ( + self._get_test_json_path(mip, metrics_collection) + ) + + if len(self.dict_json_path[mip]) < 1: + raise FileNotFoundError( + f"No Synthetic ENSO Metrics Data for {mip}, aborting..." + ) + + return self.dict_json_path + + def _get_cmip_json_path(self, mip, metrics_collection): + path = os.path.join( + self.parameter["cmip_path"], + self.parameter["cmip_name"].split(".")[0], + self.parameter["cmip_name"].split(".")[1], + self.parameter["cmip_name"].split(".")[2], + metrics_collection, + f"{mip.lower()}_{self.parameter['cmip_name'].split('.')[1]}_{metrics_collection}_*.json", + ) + matches = glob.glob(path) + if not matches: + raise FileNotFoundError( + f"CMIP metrics file not found for {mip} and {metrics_collection}" + ) + return matches[0] + + def _get_test_json_path(self, mip, metrics_collection): + for i, model_name in enumerate(self.parameter["model_name"]): + model_path = self.parameter["test_path"].replace( + "put_model_here", model_name + ) + model_files = find_latest_file_list( + path=f"{model_path}/{metrics_collection}", + file_pattern="*.json", + var_pattern=self.var_pattern, + time_pattern=self.time_pattern, + ) + logger.info(f"{model_path}/{metrics_collection}") + if not model_files or not os.path.exists(model_files[0]): + raise FileNotFoundError( + f"No Synthetic ENSO Metrics Data For {mip} {model_name}, Aborting." + ) + + for json_path in model_files: + with open(json_path) as ff: + data_json = json.load(ff) + + old_key = list(data_json["RESULTS"]["model"].keys())[0] + + data_json["RESULTS"]["model"][mip] = data_json["RESULTS"]["model"].pop( + old_key + ) + + with open(json_path, "w", encoding="utf8") as ff: + json.dump( + data_json, ff, indent=4, separators=(",", ": "), sort_keys=True + ) + + return json_path diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/movs_metrics_reader.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/movs_metrics_reader.py new file mode 100644 index 0000000..d71fe46 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/movs_metrics_reader.py @@ -0,0 +1,175 @@ +import glob +import json +import os +import re + +import numpy as np +import pandas as pd +from pcmdi_metrics.utils import sort_human + +from zppy_interfaces.multi_utils.logger import _setup_child_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import ( + find_latest_file_list, + get_highlight_models, + shift_row_to_bottom, +) + +logger = _setup_child_logger(__name__) + + +class MoVsMetricsReader: + def __init__(self, parameter): + self.parameter = parameter + self.cmip_group, self.cmip_model, self.cmip_version = self.parameter[ + "cmip_name" + ].split(".") + self.movs_mode = parameter["movs_mode"] + self.var_pattern = re.compile(r"var_mode_(\w+)\.EOF\d+\..*\.json$") + self.time_pattern = re.compile(r"\.v(\d{8})\.json$") + + def collect_metrics(self): + cmip_files = self._get_cmip_files() + if not cmip_files or not os.path.exists(cmip_files[0]): + raise FileNotFoundError( + "ERROR: No Synthetic MoVs Metrics Data For CMIP, Aborting." + ) + + logger.info("Found Synthetic MoVs Metrics Data For CMIP, Reading...") + cmip_lib = self._load_movs_files(cmip_files) + + merge_lib = {} + for stat, diag_vars in self.parameter["diag_vars"].items(): + merge_df, mode_season_list = self._movs_dict_to_df(cmip_lib, stat) + + for i, model_name in enumerate(self.parameter["model_name"]): + model_path = self.parameter["test_path"].replace( + "put_model_here", model_name + ) + model_files = find_latest_file_list( + path=f"{model_path}/*/*", + file_pattern="var_mode_*.json", + var_pattern=self.var_pattern, + time_pattern=self.time_pattern, + ) + if not model_files or not os.path.exists(model_files[0]): + raise FileNotFoundError( + f"No Synthetic MoVs Metrics Data For {model_name}, Aborting." + ) + + logger.info( + f"Found Synthetic MoVs Metrics for {model_name}, Reading..." + ) + model_lib = self._load_movs_files(model_files) + + # Normalize model name key to match targets + model_lib = { + mode: {model_name: next(iter(model_data.values()))} + for mode, model_data in model_lib.items() + } + + # Convert dictionary to DataFrame + model_df, _ = self._movs_dict_to_df(model_lib, stat) + + # Append to the merged DataFrame + merge_df = pd.concat([merge_df, model_df], ignore_index=True) + + # Highlight and reorder models if applicable + highlight_models = get_highlight_models( + merge_df.get("model", []), self.parameter["model_name"] + ) + for model in merge_df["model"].tolist(): + if model in highlight_models: + for idx in merge_df[merge_df["model"] == model].index: + merge_df = shift_row_to_bottom(merge_df, idx) + + merge_lib[stat] = merge_df + + return merge_lib, mode_season_list + + def _get_cmip_files(self): + current_dir: str = os.path.abspath(os.getcwd()) + pattern: str = os.path.join( + self.parameter["cmip_path"], + self.cmip_group, + self.cmip_model, + self.cmip_version, + "*/*/var_mode_*.json", + ) + """ + Checking in /lcrc/group/e3sm/ac.forsyth2/zppy_pr719_output/unique_id_21/v3.LR.amip_0101/post/scripts/tmp.818290.CJbY for cmip files matching /lcrc/group/e3sm/diagnostics/pcmdi_data/metrics_data/variability_modes/cmip6/historical/v20220825/*/*/var_mode_*.json + """ + logger.info(f"From {current_dir}, checking for cmip files matching {pattern}") + matching_files = glob.glob(pattern) + num_matching_files = len(matching_files) + logger.info(f"Found {num_matching_files} matching files") + for file_name in matching_files: + logger.debug(f" - {file_name}") + if not matching_files: + logger.warning( + f"No matching files found for pattern: {pattern}. " + "Ensure the path and pattern are correct." + ) + return matching_files + + def _load_movs_files(self, file_lists): + json_lib = {} + for mode in self.movs_mode: + eof = {"PSA1": "EOF2", "NPO": "EOF2", "NPGO": "EOF2", "PSA2": "EOF3"}.get( + mode, "EOF1" + ) + for json_file in file_lists: + if mode in json_file and eof in json_file: + try: + with open(json_file, "r") as fj: + data = json.load(fj) + json_lib[mode] = data.get("RESULTS", {}) + except (FileNotFoundError, json.JSONDecodeError) as e: + logger.info(f"Warning: Could not load {json_file}: {e}") + break + return json_lib + + def _movs_dict_to_df(self, movs_dict, stat): + models = sorted(movs_dict.get("NAM", {}).keys()) + df = pd.DataFrame({"model": models, "num_runs": np.nan}) + mode_season_list = [] + + for mode in self.movs_mode: + seasons = ( + ["monthly"] + if mode in ["PDO", "NPGO"] + else ["yearly"] if mode == "AMO" else ["DJF", "MAM", "JJA", "SON"] + ) + + for season in seasons: + col_name = f"{mode}_{season}" + df[col_name] = np.nan + mode_season_list.append(col_name) + + for idx, model in enumerate(models): + value = np.nan + num_runs = 0 + + if mode in movs_dict and model in movs_dict[mode]: + runs = sort_human(list(movs_dict[mode][model].keys())) + stat_values = [] + + for run in runs: + try: + run_stat = movs_dict[mode][model][run][ + "defaultReference" + ][mode][season]["cbf"][stat] + stat_values.append(run_stat) + except KeyError: + continue + + if stat_values: + value = np.mean(stat_values) + num_runs = len(stat_values) + + df.at[idx, col_name] = value + if np.isnan(df.at[idx, "num_runs"]): + df.at[idx, "num_runs"] = num_runs + elif num_runs > 0: + df.at[idx, "num_runs"] = max(df.at[idx, "num_runs"], num_runs) + + return df, mode_season_list diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/synthetic_metrics_plotter.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/synthetic_metrics_plotter.py new file mode 100644 index 0000000..fce5645 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/synthetic_metrics_plotter.py @@ -0,0 +1,994 @@ +import os +from collections import OrderedDict +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +from matplotlib import pyplot as plt +from pcmdi_metrics.enso.lib import enso_portrait_plot +from pcmdi_metrics.graphics import ( + normalize_by_median, + parallel_coordinate_plot, + portrait_plot, +) + +from zppy_interfaces.multi_utils.logger import _setup_child_logger +from zppy_interfaces.pcmdi_diags.synthetic_plots.clim_metrics_reader import ( + ClimMetricsReader, +) +from zppy_interfaces.pcmdi_diags.synthetic_plots.enso_metrics_reader import ( + EnsoMetricsReader, +) +from zppy_interfaces.pcmdi_diags.synthetic_plots.movs_metrics_reader import ( + MoVsMetricsReader, +) +from zppy_interfaces.pcmdi_diags.synthetic_plots.utils import get_highlight_models + +logger = _setup_child_logger(__name__) + + +class SyntheticMetricsPlotter: + def __init__( + self, + case_name: str, + test_name: str, + table_id: str, + figure_format: str, + metric_dict: Dict[str, Any], + save_data: bool, + base_test_input_path: str, + results_dir: Optional[str] = None, + # Mean-climate viewer + clim_viewer: bool = True, + clim_vars: Optional[Union[List[str], str]] = None, + clim_regions: Optional[Union[List[str], str]] = None, + cmip_clim_dir: Optional[str] = None, + cmip_clim_set: Optional[str] = None, + # Atmosphere modes (MOVA) + mova_viewer: bool = True, + mova_modes: Optional[Union[List[str], str]] = None, + # Coupled modes (MOVC) + movc_viewer: bool = True, + movc_modes: Optional[Union[List[str], str]] = None, + cmip_movs_dir: Optional[str] = None, + cmip_movs_set: Optional[str] = None, + # ENSO viewer + enso_viewer: bool = True, + cmip_enso_dir: Optional[str] = None, + cmip_enso_set: Optional[str] = None, + ): + # Core + self.case_name = case_name + self.test_name = test_name + self.table_id = table_id + self.figure_format = figure_format + self.metric_dict = metric_dict + self.save_data = bool(save_data) + self.base_test_input_path = base_test_input_path + self.results_dir = results_dir or "." + + # Mean climate + self.clim_viewer = bool(clim_viewer) + self.clim_vars = self._to_list(clim_vars) # [] => all available + self.clim_regions = self._to_list(clim_regions) # [] => all regions + self.cmip_clim_dir = cmip_clim_dir + self.cmip_clim_set = cmip_clim_set + + # MOVA + self.mova_viewer = bool(mova_viewer) + self.mova_modes = self._to_list(mova_modes) if self.mova_viewer else [] + self.movc_viewer = bool(movc_viewer) + self.movc_modes = self._to_list(movc_modes) if self.movc_viewer else [] + self.movs_viewer = self.mova_viewer or self.movc_viewer + + self.cmip_movs_dir = cmip_movs_dir + self.cmip_movs_set = cmip_movs_set + + # ENSO + self.enso_viewer = bool(enso_viewer) + self.cmip_enso_dir = cmip_enso_dir + self.cmip_enso_set = cmip_enso_set + + # Final bundle for downstream readers/builders + self.parameter: Dict[str, Any] = self._initialize_parameter() + + # ---------- helpers ---------- + @staticmethod + def _to_list(value: Optional[Union[List[str], str]]) -> List[str]: + """Accept None | list[str] | comma/space-separated str -> List[str].""" + if value is None: + return [] + if isinstance(value, list): + return [str(v).strip() for v in value if str(v).strip()] + s = str(value).strip() + if not s: + return [] + parts = s.split(",") if "," in s else s.split() + return [p.strip() for p in parts if p.strip()] + + def _initialize_parameter(self): + # Parse comma-separated lists + tests = [t.strip() for t in str(self.test_name).split(",") if t.strip()] + cases = [c.strip() for c in str(self.case_name).split(",") if c.strip()] + + if len(tests) != len(cases): + raise ValueError( + f"test_name count ({len(tests)}) != case_name count ({len(cases)}). " + "They must align positionally." + ) + + parsed_test_names = [] + parsed_model_names = [] + + for raw_test, raw_case in zip(tests, cases): + parts = raw_test.split(".") + if len(parts) != 4: + raise ValueError( + f"Invalid test format '{raw_test}'. Expected 'a.b.c.d'" + ) + # Re-map 'a.b.c.d' -> 'a.b.c_d' (your original behavior) + test_id = f"{parts[0]}.{parts[1]}.{parts[2]}_{parts[3]}" + parsed_test_names.append(test_id) + parsed_model_names.append(raw_case) + + # Keep the exact keys your pipeline expects + param = OrderedDict( + { + "save_data": self.save_data, + "out_dir": os.path.join(self.results_dir, "ERROR_metric"), + "test_name": parsed_test_names, + "model_name": parsed_model_names, + "tableID": [self.table_id], + } + ) + + return param + + def generate(self) -> None: + logger.info("Generating synthetic metrics plots ...") + tasks = [ + (self.clim_viewer, "mean_climate", self._handle_mean_climate), + (self.movs_viewer, "variability_modes", self._handle_variability_modes), + (self.enso_viewer, "enso_metric", self._handle_enso_metric), + ] + + at_least_one_success = False + failures = [] + + for enabled, metric, handler in tasks: + self.parameter["test_path"] = self.base_test_input_path.replace( + "%(group_type)", metric + ) + self.parameter["diag_vars"] = self.metric_dict[metric] + if not enabled: + continue + logger.info("Processing metric: %s", metric) + try: + handler(metric) + at_least_one_success = True + except Exception as e: + logger.error("Failed to handle metric=%s: %s", metric, e, exc_info=True) + failures.append(metric) + + if not at_least_one_success: + raise RuntimeError("No synthetic metrics plots could be generated.") + + if failures: + logger.warning("Completed with partial failures: %s", ", ".join(failures)) + + def _handle_mean_climate(self, metric: str) -> None: + logger.info("Handling mean climate…") + self.parameter.update( + {"cmip_path": self.cmip_clim_dir, "cmip_name": self.cmip_clim_set} + ) + + collector = ClimMetricsReader(self.parameter) + merge_lib = collector.collect() + + # Variables (preserve original behavior unless filters are provided) + var_list = list(merge_lib.var_list) + var_unit_list = list(merge_lib.var_unit_list) + if self.clim_vars is not None: + name_to_unit = dict(zip(merge_lib.var_list, merge_lib.var_unit_list)) + missing = [v for v in self.clim_vars if v not in name_to_unit] + if missing: + logger.warning( + f"[mean_climate] Requested variables not found and will be skipped: {missing}" + ) + var_list = [v for v in self.clim_vars if v in name_to_unit] + var_unit_list = [name_to_unit[v] for v in var_list] + + # Regions (preserve order) + regions = list(merge_lib.regions) + if self.clim_regions is not None: + missing_r = [r for r in self.clim_regions if r not in merge_lib.regions] + if missing_r: + logger.warning( + f"[mean_climate] Requested regions not found and will be skipped: {missing_r}" + ) + regions = [r for r in self.clim_regions if r in merge_lib.regions] + + # Use the same `metric` variable as before (assuming it's defined in scope) + for stat, vars_ in self.metric_dict[metric].items(): + logger.debug(f"[mean_climate] Running plot driver: stat={stat}") + # Keep the exact positional calling convention you had before + mean_climate_plot_driver( + metric, + stat, + regions, + self.parameter["model_name"], + vars_, + merge_lib.df_dict[stat], + var_list, + var_unit_list, + self.parameter["save_data"], + self.parameter["out_dir"], + self.figure_format, + ) + + def _handle_variability_modes(self, metric: str) -> None: + logger.info("Handling modes variability …") + + # Combine atmospheric and coupled modes (already lists) + modes_list = (self.mova_modes or []) + (self.movc_modes or []) + + if not modes_list: + logger.warning( + "[variability_modes] No modes specified; skipping variability mode plots." + ) + return + + # Update parameters for reader + self.parameter.update( + { + "cmip_path": self.cmip_movs_dir, + "cmip_name": self.cmip_movs_set, + "movs_mode": modes_list, + } + ) + + # Collect metrics + reader = MoVsMetricsReader(self.parameter) + merge_lib, mode_season_list = reader.collect_metrics() + + # Ensure metric exists in dictionary + if metric not in self.metric_dict: + logger.error( + f"[variability_modes] Metric '{metric}' not found in metric_dict keys={list(self.metric_dict.keys())}" + ) + return + + # Loop through stats and plot + for stat, vars_ in self.metric_dict[metric].items(): + if stat not in merge_lib: + logger.warning( + f"[variability_modes] stat='{stat}' not found in merge_lib; available={list(merge_lib.keys())}" + ) + continue + + logger.debug(f"[variability_modes] Running plot driver for stat={stat}") + variability_modes_plot_driver( + metric, + stat, + self.parameter["model_name"], + vars_, + merge_lib[stat], + mode_season_list, + self.parameter["save_data"], + self.parameter["out_dir"], + self.figure_format, + ) + + def _handle_enso_metric(self, metric: str) -> None: + logger.info("Handling ENSO metrics…") + + # Update paths + self.parameter.update( + { + "cmip_path": self.cmip_enso_dir, + "cmip_name": self.cmip_enso_set, + } + ) + + # --- Build enso_mips: [] + model_name(s) --- + cmip_name = self.parameter.get("cmip_name", "") + mip_tag = ( + cmip_name.split(".")[0] + if isinstance(cmip_name, str) and cmip_name + else None + ) + + model_name = self.parameter.get("model_name", []) + if isinstance(model_name, str): + model_name = [model_name] + elif isinstance(model_name, tuple): + model_name = list(model_name) + elif not isinstance(model_name, list): + logger.warning( + f"[enso] Unexpected model_name type: {type(model_name).__name__}; coercing to list if possible." + ) + model_name = list(model_name) if model_name is not None else [] + + enso_mips = ([mip_tag] if mip_tag else []) + model_name + if not enso_mips: + logger.warning( + "[enso] No MIP/model names resolved for ENSO; continuing with empty list." + ) + + # --- Collections (optional config) --- + enso_collections = self.metric_dict.get("collection", []) + if not isinstance(enso_collections, (list, tuple)): + logger.warning( + f"[enso] 'collection' should be list/tuple; got {type(enso_collections).__name__}. Using empty list." + ) + enso_collections = [] + + # --- Validate metric entry --- + if metric not in self.metric_dict or not isinstance( + self.metric_dict[metric], dict + ): + logger.error( + f"[enso] metric_dict['{metric}'] missing or not a dict. Available: {list(self.metric_dict.keys())}" + ) + return + + diag_vars_all = self.parameter.get("diag_vars", {}) + if not isinstance(diag_vars_all, dict): + logger.error( + f"[enso] parameter['diag_vars'] must be a dict; got {type(diag_vars_all).__name__}." + ) + return + + # --- Main loop over stats --- + for stat in self.metric_dict[metric].keys(): + metric_dict = diag_vars_all.get(stat, {}) + if not metric_dict: + logger.warning( + f"[enso] No variables configured for stat='{stat}'. Skipping." + ) + continue + + logger.debug( + f"[enso] stat='{stat}', enso_mips={enso_mips}, collections={enso_collections}" + ) + try: + reader = EnsoMetricsReader( + self.parameter, stat, metric_dict, enso_mips, enso_collections + ) + dict_json_path = reader.run() + except Exception as e: + logger.exception(f"[enso] Reader failed for stat='{stat}': {e}") + continue + + if not dict_json_path: + logger.warning( + f"[enso] Reader returned empty path for stat='{stat}'. Skipping plot." + ) + continue + + try: + enso_plot_driver( + metric, stat, dict_json_path, self.parameter, self.figure_format + ) + logger.debug(f"[enso] Plotted stat='{stat}' successfully.") + except Exception as e: + logger.exception(f"[enso] Plot driver failed for stat='{stat}': {e}") + + +def mean_climate_plot_driver( + metric, + stat, + regions, + model_name, + metric_dict, + df_dict, + var_list, + var_unit_list, + save_data, + out_path, + fig_format, +): + """Driver Function for the mean climate metrics plot""" + if len(model_name) > 1: + mout_name = model_name[0].split("_")[0] + else: + mout_name = model_name[0] + + for region in regions: + for mtype in metric_dict["type"]: + if region in metric_dict["region"]: + do_plot = True + else: + do_plot = False + if do_plot and mtype == "portrait": + logger.info( + "Processing Portrait Plots for {} {} {}....".format( + metric, region, stat + ) + ) + var_names = sorted(var_list.copy()) + # label information + var_units = [] + for i, var in enumerate(var_names): + index = var_list.index(var) + var_units.append(var_unit_list[index]) + data_nor = dict() + for season in metric_dict["season"]: + data_dict = df_dict[season][region].copy() + if stat == "cor_xy": + data_nor[season] = data_dict[var_names].to_numpy().T + else: + logger.debug( + f"var_names={var_names} derived from var_list={var_list}." + ) + logger.debug(f"Available columns: {data_dict.columns.tolist()}") + try: + data_nor[season] = normalize_by_median( + data_dict[var_names].to_numpy().T, axis=1 + ) + except KeyError as e: + logger.error(f"KeyError on var_names={var_names}") + raise e + if save_data: + outdir = os.path.join(out_path, metric, region) + outdic = data_dict.drop(columns=["model_run"]).copy() + outdic[var_names] = data_nor[season].T + archive_data( + region, + stat, + season, + data_dict, + mout_name, + var_names, + var_units, + outdir, + ) + run_list = data_dict["model"].to_list() + stat_name = metric_dict["name"] + outdir = os.path.join(out_path, metric) + portrait_metric_plot( + region, + stat, + metric, + data_nor, + stat_name, + model_name, + var_names, + run_list, + outdir, + fig_format, + ) + elif do_plot and mtype == "parcoord": + logger.info( + "Processing Parallel Coordinate Plots for {} {} {}....".format( + metric, region, stat + ) + ) + for season in metric_dict["season"]: + if season in df_dict.keys(): + # drop data if all is NaNs + data_dict, var_names, var_units = drop_vars( + df_dict[season][region].copy(), + var_list.copy(), + var_unit_list.copy(), + ) + if save_data: + outdir = os.path.join(out_path, metric, region) + outdic = data_dict.drop(columns=["model_run"]).copy() + archive_data( + region, + stat, + season, + outdic, + mout_name, + var_list, + var_unit_list, + outdir, + ) + run_list = data_dict["model"].to_list() + stat_name = metric_dict["name"] + outdir = os.path.join(out_path, metric) + parcoord_metric_plot( + region, + stat, + metric, + data_dict, + stat_name, + model_name, + var_names, + var_units, + run_list, + outdir, + fig_format, + ) + return + + +def variability_modes_plot_driver( + metric, + stat, + model_name, + metric_dict, + df_dict, + mode_season_list, + save_data, + out_path, + fig_format, +): + """Driver Function for the modes variability metrics plot""" + season = "mon" + if len(model_name) > 1: + mout_name = model_name[0].split("_")[0] + else: + mout_name = model_name[0] + + for mtype in metric_dict["type"]: + if mtype == "portrait": + logger.info("Processing Portrait Plots for {} {}....".format(metric, stat)) + if stat not in ["stdv_pc_ratio_to_obs"]: + data_nor = normalize_by_median( + df_dict[mode_season_list].to_numpy().T, axis=1 + ) + else: + data_nor = df_dict[mode_season_list].to_numpy().T + if save_data: + df_dict[mode_season_list] = data_nor.T + outdir = os.path.join(out_path, metric) + archive_data( + metric, + stat, + season, + df_dict, + mout_name, + mode_season_list, + None, + outdir, + ) + run_list = df_dict["model"].to_list() + stat_name = metric_dict["name"] + portrait_metric_plot( + metric, + stat, + season, + data_nor, + stat_name, + model_name, + mode_season_list, + run_list, + out_path, + fig_format, + ) + elif mtype == "parcoord": + logger.info( + "Processing Parallel Coordinate Plots for {} {}....".format( + metric, stat + ) + ) + # drop data if all is NaNs + data_dict, var_names, var_units = drop_vars( + df_dict.copy(), mode_season_list.copy(), None + ) + if save_data: + outdir = os.path.join(out_path, metric) + archive_data( + metric, + stat, + season, + data_dict, + mout_name, + mode_season_list, + None, + outdir, + ) + run_list = data_dict["model"].to_list() + stat_name = metric_dict["name"] + parcoord_metric_plot( + metric, + stat, + season, + data_dict, + stat_name, + model_name, + var_names, + var_units, + run_list, + out_path, + fig_format, + ) + + return + + +def enso_plot_driver(metric, stat, dict_json_path, parameter, fig_format): + """ + Driver function to plot ENSO metrics based on specified type (e.g., portrait). + """ + metric_dict = parameter["diag_vars"][stat] + metrics_collections = metric_dict["collection"] + mips = [parameter["cmip_name"].split(".")[0]] + parameter["model_name"] + + for mtype in metric_dict["type"]: + if mtype == "portrait": + logger.info(f"Processing Portrait Plots for {metric} {stat}...") + + list_project = mips + list_obs: List[object] = ( + [] + ) # fill in if observational references are needed + outdir = os.path.join(parameter["out_dir"], metric) + os.makedirs(outdir, exist_ok=True) + + outfile = f"{metric}_{stat}_portrait.{fig_format}" + figure_name = os.path.join(outdir, outfile) + + fig, ref_info_dict = enso_portrait_plot( + metrics_collections, + list_project, + list_obs, + dict_json_path, + figure_name=figure_name, + reduced_set=True, + ) + + return + + +def archive_data( + region, stat, season, data_dict, model_name, var_names, var_units, outdir +): + """ + Archive processed data into a CSV file with variable units in column headers if available. + + Parameters: + region (str): Region name. + stat (str): Statistic type (e.g., mean, std). + season (str): Season name. + data_dict (dict or DataFrame): Data to archive. + model_name (str): Model identifier. + var_names (list): List of variable names. + var_units (list): List of variable units (optional, same order as var_names). + outdir (str): Directory to save the CSV file. + """ + df = pd.DataFrame(data_dict) + + # Determine the index of the first variable column (assumes first 3 are metadata) + metadata_cols = df.columns[:3].tolist() + variable_cols = df.columns[3:] + + filtered_cols = [] + new_column_names = df.columns.tolist() + + for var in variable_cols: + if var in var_names: + filtered_cols.append(var) + if var_units: + idx = df.columns.get_loc(var) + unit_label = var_units[var_names.index(var)] + new_column_names[idx] = f"{var} ({unit_label})" + + # Subset dataframe and rename columns if units provided + df = df[metadata_cols + filtered_cols] + df.columns = new_column_names[: len(df.columns)] + + # Ensure output directory exists + os.makedirs(outdir, exist_ok=True) + + # Construct and save the output filename + outfile = f"{stat}_{region}_{season}_{model_name}.csv" + df.to_csv(os.path.join(outdir, outfile), index=False) + + return + + +def portrait_metric_plot( + region, + stat, + group, + data_dict, + stat_name, + model_name, + var_list, + model_list, + out_path, + fig_format, + base_fontsize=20, + base_figsize=(40, 18), + base_legend_lw=1.5, + box_as_square=True, + missing_color="white", + logo_rect=[0, 0, 0, 0], + logo_off=True, +): + # === Figure scaling setup === + fscale = len(var_list) / 30.0 + fscale = max(0.5, min(fscale, 1.5)) # clamp to avoid extremes + + # Apply scaled parameters + fontsize = base_fontsize + figsize = (base_figsize[0], base_figsize[1] * fscale) + legend_box_xy = (1.08, 1.20) + legend_box_size = 4 * fscale + legend_lw = base_legend_lw * fscale + shrink = 0.8 * fscale + legend_fontsize = fontsize * 0.8 + + if group == "mean_climate": + # data for final plot + data_all_nor = np.stack( + [data_dict["djf"], data_dict["mam"], data_dict["jja"], data_dict["son"]] + ) + legend_on = True + legend_labels = ["DJF", "MAM", "JJA", "SON"] + else: + data_all_nor = data_dict + legend_on = False + legend_labels = [] + + highlight_models = get_highlight_models(model_list, model_name) + lable_colors = [] + for model in model_list: + if model in model_name: + lable_colors.append("#FC5A50") + elif "e3sm" in model.lower(): + lable_colors.append("#5170d7") + else: + lable_colors.append("#000000") + + var_range: Tuple[float, float] + if stat in ["cor_xy"]: + var_range = (0, 1.0) + cmap_color = "viridis" + cmap_bounds = np.linspace(0, 1, 21) + elif stat in ["stdv_pc_ratio_to_obs"]: + var_range = (0.5, 1.5) + cmap_color = "jet" + cmap_bounds = [0.5, 0.7, 0.9, 1.1, 1.3, 1.5] + cmap_bounds = [r / 10 for r in range(5, 16, 1)] + else: + var_range = (-0.5, 0.5) + cmap_color = "RdYlBu_r" + cmap_bounds = np.linspace(-0.5, 0.5, 11) + + fig, ax, cbar = portrait_plot( + data_all_nor, + xaxis_labels=model_list, + yaxis_labels=var_list, + cbar_label=stat, + cbar_label_fontsize=fontsize * 0.95, + cbar_tick_fontsize=fontsize * 0.95, + box_as_square=box_as_square, + vrange=var_range, + figsize=figsize, + cmap=cmap_color, + cmap_bounds=cmap_bounds, + cbar_kw={"extend": "both", "shrink": shrink}, + missing_color=missing_color, + legend_on=legend_on, + legend_labels=legend_labels, + legend_box_xy=legend_box_xy, + legend_box_size=legend_box_size, + legend_lw=legend_lw, + legend_fontsize=legend_fontsize, + logo_rect=logo_rect, + logo_off=logo_off, + ) + + ax.axvline(x=len(model_list) - len(highlight_models), color="k", linewidth=3) + ax.set_xticklabels(model_list, rotation=45, va="bottom", ha="left") + ax.set_yticklabels(var_list, rotation=0, va="center", ha="right") + for xtick, color in zip(ax.get_xticklabels(), lable_colors): + xtick.set_color(color) + ax.yaxis.label.set_color(lable_colors[0]) + + # Add title + fig.suptitle( + f"{region} — {group} ({stat_name})", fontsize=fontsize * 1.1, fontweight="bold" + ) + fig.tight_layout(rect=[0, 0, 1, 0.95]) # leave top 5 % free for title + + # Save figure as an image file + outdir = os.path.join(out_path, region) + if not os.path.exists(outdir): + os.makedirs(outdir) + outfile = "{}_{}_portrait_{}.{}".format(stat, region, group, fig_format) + fig.savefig(os.path.join(outdir, outfile), facecolor="w", bbox_inches="tight") + plt.close(fig) + + return + + +def drop_vars(data_dict, var_names, var_units=None): + """ + Drop variables (columns) from data_dict where more than 90% of the values are NaN. + + Parameters: + data_dict (pd.DataFrame): Data containing variable columns. + var_names (list): List of variable names matching data_dict columns. + var_units (list, optional): List of units for variables. Must match var_names in order. + + Returns: + Tuple of (filtered_data_dict, updated_var_names, updated_var_units) + """ + protected_columns = {"model", "run", "model_run", "num_runs"} + columns_to_drop = [] + + for column in data_dict.columns: + if column in protected_columns: + continue + nan_ratio = data_dict[column].isna().mean() + if nan_ratio > 0.9: + columns_to_drop.append(column) + + # Drop columns from DataFrame + data_dict = data_dict.drop(columns=columns_to_drop) + + # Update var_names and var_units if applicable + updated_var_names = [v for v in var_names if v not in columns_to_drop] + updated_var_units = None + if var_units is not None: + # Keep units only for remaining variables + name_to_unit = dict(zip(var_names, var_units)) + updated_var_units = [ + name_to_unit[v] for v in updated_var_names if v in name_to_unit + ] + + return data_dict, updated_var_names, updated_var_units + + +def parcoord_metric_plot( + region, + stat, + group, + data_dict, + stat_name, + model_name, + var_names, + var_units, + model_list, + out_path, + fig_format, + base_fontsize=20, + base_figsize=(60, 20), + base_legend_lw=1.5, + color_map="tab20_r", + xcolors=None, + group1_name="CMIP", + mean1_name="CMIP (Mean)", + group2_name="E3SM", + mean2_name="E3SM (Mean)", + identify_all_models=True, + vertical_center="median", + vertical_center_line=True, + show_boxplot=False, + show_violin=True, + violin_colors=("lightgrey", "pink"), + logo_rect=[0, 0, 0, 0], + logo_off=True, +): + """Function for parallel coordinate plots""" + # === Figure scaling setup === + fscale = len(var_names) / 30.0 + fscale = max(0.6, min(fscale, 1.5)) # clamp to avoid extremes + + fontsize = base_fontsize + figsize = (base_figsize[0] * fscale, base_figsize[1] * fscale) + + legend_ncol = int(7 * figsize[0] / 40.0) + legend_posistion = (0.50, -0.14) + + # colors for highlight lines + if xcolors is None: + xcolors = [ + "#e41a1c", + "#ff7f00", + "#4daf4a", + "#f781bf", + "#a65628", + "#984ea3", + "#377eb8", + "#dede00", + ] + + # ensemble mean for E3SM group + highlight_model1 = get_highlight_models(data_dict.get("model", []), model_name) + irow_str = data_dict[data_dict["model"] == highlight_model1[0]].index[0] + irow_end = data_dict[data_dict["model"] == highlight_model1[-1]].index[0] + 1 + data_dict.loc[mean2_name] = data_dict[irow_str:irow_end].mean( + numeric_only=True, skipna=True + ) + data_dict.at[mean2_name, "model"] = mean2_name + + # ensemble mean for CMIP group + irow_sub = data_dict[data_dict["model"] == highlight_model1[0]].index[0] + data_dict.loc[mean1_name] = data_dict[:irow_sub].mean( + numeric_only=True, skipna=True + ) + data_dict.at[mean1_name, "model"] = mean1_name + data_dict.loc[mean2_name] = data_dict[irow_sub:].mean( + numeric_only=True, skipna=True + ) + data_dict.at[mean2_name, "model"] = mean2_name + + model_list = data_dict["model"].to_list() + highlight_model2 = highlight_model1 + [mean1_name, mean2_name] + + # colors for highlight lines + lncolors = [] + for i, model in enumerate(highlight_model2): + if model == mean1_name: + lncolors.append("#000000") + elif model == mean2_name: + lncolors.append("#5b5b5b") # ("#999999") + else: + lncolors.append(xcolors[i % len(xcolors)]) + + var_name1 = sorted(var_names.copy()) + # label information + var_labels = [] + for i, var in enumerate(var_name1): + index = var_names.index(var) + if var_units is not None: + var_labels.append(var_names[index] + "\n" + var_units[index]) + else: + var_labels.append(var_names[index]) + + # final plot data + data_var = data_dict[var_name1].to_numpy() + + xlabel = "Metric" + ylabel = "{} ({})".format(stat_name, stat.upper()) + + if "mean_climate" in [group, region]: + title = "Model Performance of Annual Climatology ({}, {})".format( + stat.upper(), region.upper() + ) + elif "variability_modes" in [group, region]: + title = "Model Performance of Modes Variability ({})".format(stat.upper()) + elif "enso" in [group, region]: + title = "Model Performance of ENSO ({})".format(stat.upper()) + + fig, ax = parallel_coordinate_plot( + data_var, + var_labels, + model_list, + model_names2=highlight_model1, + group1_name=group1_name, + group2_name=group2_name, + models_to_highlight=highlight_model2, + models_to_highlight_colors=lncolors, + models_to_highlight_labels=highlight_model2, + identify_all_models=identify_all_models, + vertical_center=vertical_center, + vertical_center_line=vertical_center_line, + title="", + figsize=figsize, + colormap=color_map, + show_boxplot=show_boxplot, + show_violin=show_violin, + violin_colors=violin_colors, + legend_ncol=legend_ncol, + legend_bbox_to_anchor=legend_posistion, + legend_fontsize=fontsize * 0.85, + xtick_labelsize=fontsize * 0.95, + ytick_labelsize=fontsize * 0.95, + logo_rect=logo_rect, + logo_off=logo_off, + ) + + ax.set_xlabel(xlabel, fontsize=fontsize * 1.05) + ax.set_ylabel(ylabel, fontsize=fontsize * 1.05) + # ax.set_title(title, fontsize=fontsize * 1.05) + + # Add title + fig.suptitle(f"{title}", fontsize=fontsize * 1.05, fontweight="bold") + fig.tight_layout(rect=[0, 0, 1, 0.95]) # leave top 5 % free for title + + # Save figure as an image file + outdir = os.path.join(out_path, region) + if not os.path.exists(outdir): + os.makedirs(outdir) + outfile = "{}_{}_parcoord_{}.{}".format(stat, region, group, fig_format) + fig.savefig(os.path.join(outdir, outfile), facecolor="w", bbox_inches="tight") + plt.close(fig) + + return diff --git a/zppy_interfaces/pcmdi_diags/synthetic_plots/utils.py b/zppy_interfaces/pcmdi_diags/synthetic_plots/utils.py new file mode 100644 index 0000000..a205a28 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/synthetic_plots/utils.py @@ -0,0 +1,115 @@ +import glob +import os +import re +from datetime import datetime +from typing import Dict, List, Tuple + +import pandas as pd + +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +logger = _setup_child_logger(__name__) + + +def find_latest_file_list( + path: str, + file_pattern: str, + var_pattern=r"\.(\w+)\.\d{8}\.nc$", + time_pattern=r"\.(\d{8})\.nc$", +) -> List[str]: + """ + Find the latest NetCDF file for each variable in the directory based on timestamps in filenames. + + Args: + path (str): Directory to search. + file_pattern (str): Regex to search file lists. + var_pattern (str): Regex to extract variable name. + time_pattern (str): Regex to extract date. + + Returns: + List[str]: List of file paths, one for each variable (latest by timestamp). + """ + latest_files: Dict[str, Tuple[datetime, str]] = {} + files = glob.glob(os.path.join(path, file_pattern)) + if not files: + """ + FAILURE + + No files found in /lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zppy_pr719_output/unique_id_21/v3.LR.amip_0101/pcmdi_diags/model_vs_obs/metrics_data/variability_modes/*/* that match pattern: var_mode_*.json + + ls /lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zppy_pr_719_output/unique_id_21/v3.LR.amip_0101/pcmdi_diags/model_vs_obs/metrics_data/variability_modes/ + AMO NAM NAO NPGO NPO PDO PNA PSA1 PSA2 SAM + + ls /lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zppy_pr719_output/unique_id_21/v3.LR.amip_0101/pcmdi_diags/model_vs_obs/metrics_data/variability_modes/AMO/HadISST2/ + AMO_ts_EOF1_monthly_obs_1869-2014.nc AMO_ts_EOF1_yearly_obs_1869-2014.nc + + SYNTHETIC PLOTS ERROR #2: No json files produced by variability modes, even though those jobs completed successfully! + """ + logger.error(f"No files found in {path} that match pattern: {file_pattern}") + for f in files: + fname = os.path.basename(f) + var_match = re.search(var_pattern, fname) + time_match = re.search(time_pattern, fname) + + if var_match and time_match: + logger.info(f"{fname} matched var and time patterns") + var = var_match.group(1) + try: + timestamp = datetime.strptime(time_match.group(1), "%Y%m%d") + except ValueError: + continue + + if var not in latest_files or timestamp > latest_files[var][0]: + latest_files[var] = (timestamp, f) + else: + logger.warning(f"{fname} failed to match both var and time patterns") + + return [file for _, file in latest_files.values()] + + +def get_highlight_models(all_models, model_name): + """ + Prioritize models containing 'e3sm' and then any additional specified models. + + Parameters: + data_dict (dict): Dictionary with a 'model' key containing a list of model names. + model_name (list): List of models to also highlight (after e3sm models). + + Returns: + list: Ordered list of unique models to highlight. + """ + highlight_model1 = [] + + # First, collect all models that contain "e3sm" (case-insensitive) + e3sm_models = [m for m in all_models if "e3sm" in m.lower()] + + # Then collect models in model_name that are not already in e3sm_models + additional_models = [ + m for m in all_models if m in model_name and m not in e3sm_models + ] + + # Combine both lists + highlight_model1 = e3sm_models + additional_models + + return highlight_model1 + + +def shift_row_to_bottom(df, index_to_shift): + """ + Moves the specified row to the bottom of the DataFrame and resets the index. + + Parameters: + df (pd.DataFrame): The input DataFrame. + index_to_shift (int): The index of the row to move to the bottom. + + Returns: + pd.DataFrame: A new DataFrame with the row moved to the bottom and index reset. + """ + if index_to_shift not in df.index: + raise IndexError(f"Index {index_to_shift} not found in DataFrame.") + + df_top = df.drop(index=index_to_shift) + df_bottom = df.loc[[index_to_shift]] + + new_df = pd.concat([df_top, df_bottom], ignore_index=True) + return new_df diff --git a/zppy_interfaces/pcmdi_diags/utils.py b/zppy_interfaces/pcmdi_diags/utils.py new file mode 100644 index 0000000..10c3290 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/utils.py @@ -0,0 +1,104 @@ +import time +from subprocess import PIPE, Popen +from typing import Dict, List, Tuple + +import psutil + +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +logger = _setup_child_logger(__name__) + +# Mapping from observational variable names to CMIP-standard +ALT_OBS_MAP: Dict[str, str] = { + "pr": "PRECT", + "sst": "ts", + "sfcWind": "si10", + "taux": "tauu", + "tauy": "tauv", + "rltcre": "toa_cre_lw_mon", + "rstcre": "toa_cre_sw_mon", + "rtmt": "toa_net_all_mon", +} + + +def count_child_processes(process=None): + """ + Count the number of child processes for a given process. + + Parameters: + - process (psutil.Process, optional): The process to check. If None, uses the current process. + + Returns: + - int: Number of child processes. + """ + if process is None: + process = psutil.Process() + + children = process.children() + return len(children) + + +def run_parallel_jobs(cmds: List[str], num_workers: int) -> List[Tuple[str, str, int]]: + """ + Execute shell commands in parallel batches. + + Parameters: + - cmds: List of command strings to run. + - num_workers: Maximum number of subprocesses to run concurrently. + + Returns: + - List of tuples: (stdout, stderr, return_code) for each command. + """ + results = [] + procs = [] + + for i, cmd in enumerate(cmds): + proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True, text=True) + procs.append((cmd, proc)) + + # Run the batch if full or if it's the last command + if len(procs) >= num_workers or i == len(cmds) - 1: + logger.info(f"Running {count_child_processes()} subprocesses...") + for cmd, proc in procs: + stdout, stderr = proc.communicate() + return_code = proc.returncode + + if return_code != 0: + logger.error( + f"ERROR: Process failed: '{cmd}'\nError: {stderr.strip()}" + ) + raise RuntimeError(f"Subprocess failed: {cmd}") + + results.append((stdout.strip(), stderr.strip(), return_code)) + + time.sleep(0.25) # Throttle before starting the next batch + procs = [] + + return results + + +def run_serial_jobs(cmds: List[str]) -> List[Tuple[str, str, int]]: + """ + Execute shell commands one at a time (serially). + + Parameters: + - cmds: List of command strings to run. + + Returns: + - List of tuples: (stdout, stderr, return_code) for each command. + """ + results = [] + + for i, cmd in enumerate(cmds): + logger.info(f"Running [{i + 1}/{len(cmds)}]: {cmd}") + proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True, text=True) + stdout, stderr = proc.communicate() + return_code = proc.returncode + + if return_code != 0: + logger.error(f"ERROR: Process failed: '{cmd}'\nError: {stderr.strip()}") + raise RuntimeError(f"Subprocess failed: {cmd}") + + results.append((stdout.strip(), stderr.strip(), return_code)) + + return results diff --git a/zppy_interfaces/pcmdi_diags/viewer.py b/zppy_interfaces/pcmdi_diags/viewer.py new file mode 100644 index 0000000..9d80c85 --- /dev/null +++ b/zppy_interfaces/pcmdi_diags/viewer.py @@ -0,0 +1,1513 @@ +import glob +import os +from datetime import datetime +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +from jinja2 import Environment, FileSystemLoader + + +def collect_config( + title: str = "E3SM-PMP Diagnostics", + subtitle: str = "Model vs Obs.", + version: str = "v3.8.2", + case_id: str = "v3.LR.amip", + run_dir: str = "model_vs_obs", + out_dir: str = "./pcmdi_diags/viewer", + diag_dir: str = "./pcmdi_diags", + obs_dir: str = "./observations/Atm/time-series", + pmp_dir: str = "./pcmdi_data", + clim_viewer: bool = True, + clim_period: str = "1985-2014", + clim_vars: Optional[List[str]] = None, + clim_regions: Optional[List[str]] = None, + clim_reference: str = "obs", + mova_viewer: bool = True, + mova_modes: Optional[List[str]] = None, + mova_seasons: Optional[List[str]] = None, + mova_period: str = "1985-2014", + mova_vars: Optional[List[str]] = None, + mova_reference: str = "obs", + movc_viewer: bool = True, + movc_modes: Optional[List[str]] = None, + movc_seasons: Optional[List[str]] = None, + movc_period: str = "1985-2014", + movc_vars: Optional[List[str]] = None, + movc_reference: str = "obs", + enso_viewer: bool = True, + enso_period: str = "1985-2014", + enso_vars: Optional[List[str]] = None, + enso_reference: str = "obs", +) -> Dict[str, object]: + """ + Collects and returns configuration information for the diagnostics viewer. + Uses default settings when specific arguments are not provided. + + Returns: + config (dict): A dictionary containing all configuration settings. + """ + + # Define default regions and variable lists + if clim_regions is None: + clim_regions = ["global", "ocean", "land", "NHEX", "SHEX", "TROPICS"] + + if clim_vars is None: + # This list appears to dictate which variables show up under "Mean Climate Map" + # That is, if plots for the variable don't exist, the links will be grayed out. + clim_vars = [ + "pr", + "prw", + "psl", + "rlds", + "rldscs", + "rltcre", + "rstcre", + "rsus", + "rsuscs", + "rlus", + "rlut", + "rlutcs", + "rsds", + "rsdscs", + "rsdt", + "rsut", + "rsutcs", + "rtmt", + "sfcWind", + "tas", + "tauu", + "tauv", + "ts", + "ta-200", + "ta-850", + "ua-200", + "ua-850", + "va-200", + "va-850", + "zg-500", + ] + + if mova_vars is None: + mova_vars = ["psl"] + + if mova_modes is None: + mova_modes = ["NAM", "PNA", "NPO", "NAO", "SAM", "PSA1", "PSA2"] + + if mova_seasons is None: + mova_seasons = ["DJF", "MAM", "JJA", "SON", "yearly", "monthly"] + + if movc_vars is None: + movc_vars = ["ts"] + + if movc_modes is None: + movc_modes = ["PDO", "NPGO", "AMO"] + + if movc_seasons is None: + movc_seasons = ["yearly", "monthly"] + + if enso_vars is None: + enso_vars = [ + "psl", + "pr", + "prsn", + "ts", + "tas", + "tauu", + "tauv", + "hfls", + "hfss", + "rlds", + "rsds", + "rlus", + "rlut", + "rsdt", + ] + + # Derive additional paths + template_dir = os.path.join(pmp_dir, "viewer") + fig_dir = os.path.join(diag_dir, run_dir) + + # Consolidated configuration dictionary + config: Dict[str, object] = { + "title": title, + "subtitle": subtitle, + "version": version, + "case_id": case_id, + "clim_viewer": clim_viewer, + "clim_period": clim_period, + "clim_vars": clim_vars, + "clim_reference": clim_reference, + "clim_regions": clim_regions, + "mova_viewer": mova_viewer, + "mova_period": mova_period, + "mova_vars": mova_vars, + "mova_modes": mova_modes, + "mova_seasons": mova_seasons, + "mova_reference": mova_reference, + "movc_viewer": movc_viewer, + "movc_period": movc_period, + "movc_vars": movc_vars, + "movc_modes": movc_modes, + "movc_seasons": movc_seasons, + "movc_reference": movc_reference, + "enso_viewer": enso_viewer, + "enso_period": enso_period, + "enso_vars": enso_vars, + "enso_reference": enso_reference, + "diag_dir": fig_dir, + "obs_dir": obs_dir, + "pmp_dir": pmp_dir, + "template_dir": template_dir, + "fig_dir": fig_dir, + "out_dir": out_dir, + } + + return config + + +def setup_jinja_env(template_dir): + """ + Set up the Jinja2 environment + """ + return Environment(loader=FileSystemLoader(template_dir)) + + +def create_section(title, rows): + """ + Prepare section data as a dictionary with title and rows. + Each row is a tuple of (label, description). + """ + return {"title": title, "rows": rows} + + +def add_section(sections: List[Dict[str, object]], title, rows): + sections.append( + { + "title": title, + "rows": [ + {"label": label, "description": description} + for label, description in rows + ], + } + ) + + +def generate_methodology_html(config): + """ + Generate the Methodology and Definitions HTML page for the E3SM-PMP Diagnostics Package. + Expects: + - config["template_dir"], config["out_dir"] + - booleans: "clim_viewer", "mova_viewer", "movc_viewer", "enso_viewer" + """ + + # --- Safe getters with defaults + def cfg(key, default=None): + return config.get(key, default) + + # Data for content sections + general_notes = [ + ( + "Source Code", + ( + "See GitHub. " + "The Program for Climate Model Diagnosis & Intercomparison (PCMDI) Metrics Package (PMP) emphasizes metrics " + "of large- to global-scale annual cycle and both tropical and extra-tropical modes of variability. " + "The package expects model data to be CF-compliant. Sample usage: " + "PMP Website." + ), + ), + ( + "Reference", + ( + "The observations used were collected and processed by the " + "E3SM Diagnostics. " + "More info: Documentation." + ), + ), + ( + "Workflow", + ( + "The diagnostics were generated using the " + "E3SM-zppy workflow." + ), + ), + ] + + # -------------------- PMP summary metrics -------------------- + pmp_metrics = [] + if cfg("clim_viewer", False): + pmp_metrics.append( + ( + "Mean Climate", + "Mean climate summary statistics. See PMP mean climate webpage.", + ) + ) + if cfg("mova_viewer", False) or cfg("movc_viewer", False): + pmp_metrics.append( + ( + "EMoV Metrics", + "Summary metric for extra-tropical modes of variability (EMoV). See PMP EMoV webpage.", + ) + ) + if cfg("enso_viewer", False): + pmp_metrics.append( + ( + "ENSO Metrics", + "Summary metric for El Niño–Southern Oscillation (ENSO). See PMP ENSO webpage.", + ) + ) + pmp_metrics.extend( + [ + ( + "Method", + "Summary metrics are featured with Portrait Plot and " + "Parallel Coordinate Plot.", + ), + ( + "References", + "See PMP Metrics webpage.", + ), + ] + ) + + # -------------------- ENSO metrics -------------------- + enso_metrics = [] + if cfg("enso_viewer", False): + enso_metrics.extend( + [ + ( + "Region", + "Diagnostics are based on monthly anomalies in the Niño 3.4 region (5°S–5°N, 170°–120°W).", + ), + ( + "ENSO Perf", + "Performance collection (15 metrics) covering: (a) background climatology (double ITCZ, equator too dry, cold tongue bias, shifted trade winds); " + "(b) basic ENSO characteristics (amplitude, skewness, seasonality, SSTA pattern, lifecycle, duration, diversity).", + ), + ( + "ENSO Proc", + "Processes collection (11 metrics) covering: (a) background climatology (cold tongue bias, shifted trade winds); " + "(b) basic characteristics (amplitude, skewness, seasonality, SSTA pattern); " + "(c) feedbacks (SSH–SST, SST–heat fluxes, SST–τx, τx–SSH); " + "(d) ocean-driven SST change.", + ), + ( + "ENSO Telec", + "Teleconnections collection (7 metrics) covering: (a) basic characteristics (amplitude, seasonality, SSTA pattern); " + "(b) ENSO-related anomalies (precipitation and surface temperature) outside the equatorial Pacific during events.", + ), + ( + "Method", + "Developed by the ENSO Metrics Working Group of the International CLIVAR Pacific Panel.", + ), + ( + "References", + "See the ENSO metrics documentation for algorithms and examples.", + ), + ] + ) + + # -------------------- EMoV (Atmos/Coupled modes) -------------------- + emov_metrics = [] + if cfg("mova_viewer", False): + emov_metrics.extend( + [ + ( + "NAM", + "Northern Annular Mode. 1st EOF mode of PSL over 20–90°N, 0–360°.", + ), + ( + "PNA", + "Pacific–North American pattern. 1st EOF mode of PSL over 20–85°N, 120°E–120°W.", + ), + ( + "NPO", + "North Pacific Oscillation. 2nd EOF mode of PSL over 20–85°N, 120°E–120°W.", + ), + ( + "NAO", + "North Atlantic Oscillation. 1st EOF mode of PSL over 20–80°N, 40°E–90°W.", + ), + ( + "SAM", + "Southern Annular Mode. 1st EOF mode of PSL over 20–90°S, 0–360°.", + ), + ( + "PSA1", + "Pacific–South American pattern 1. 2nd EOF mode of PSL over 20–90°S, 0–360°.", + ), + ( + "PSA2", + "Pacific–South American pattern 2. 3rd EOF mode of PSL over 20–90°S, 0–360°.", + ), + ] + ) + if cfg("movc_viewer", False): + emov_metrics.extend( + [ + ( + "PDO", + "Pacific Decadal Oscillation. 1st EOF mode of SST over the North Pacific (poleward of 20°N).", + ), + ( + "NPGO", + "North Pacific Gyre Oscillation. 2nd EOF mode of SST over the Northeast Pacific.", + ), + ( + "AMO", + "Atlantic Multidecadal Oscillation, typically derived from detrended North Atlantic SST averages (method varies by index).", + ), + ] + ) + emov_metrics.extend( + [ + ( + "Method", + "The Common Basis Function (CBF) approach is employed in addition to traditional EOFs, projecting model anomalies onto observed modes.", + ), + ( + "References", + "See the EMoV metrics documentation.", + ), + ] + ) + + # -------------------- Mean climate metrics -------------------- + clim_metrics = [ + ( + "Mean Bias", + "Climatological annual/seasonal mean differences between model and observations.", + ), + ( + "RMSE", + "Root-Mean-Square Error (L2 norm) against observations/reanalyses for seasonal and mean-state climatologies.", + ), + ("Centered RMSE", "RMSE of anomalies after removing the mean bias."), + ( + "Region", + "Metrics computed over global, hemispheric, tropical, extra-tropical, and other selected domains.", + ), + ( + "References", + "See PMP mean climate.", + ), + ] + + # Build sections (skip empty bodies gracefully if your create_section handles it) + sections = [ + create_section("General Notes", general_notes), + create_section("Summary Metrics", pmp_metrics), + ] + + # ENSO + if cfg("enso_viewer", False): + sections.append( + create_section("El Niño–Southern Oscillation (ENSO)", enso_metrics) + ) + + # MOVs (MOVA/MOVC) + movs_on = cfg("movs_viewer", None) + if movs_on is None: + movs_on = cfg("mova_viewer", False) or cfg("movc_viewer", False) + + if movs_on: + sections.append( + create_section("Extra-Tropical Modes of Variability (EMoV)", emov_metrics) + ) + + # Mean climate + if cfg("clim_viewer", False): + sections.append(create_section("Mean Climate", clim_metrics)) + + # Setup Jinja2 environment and load template + env = setup_jinja_env(cfg("template_dir")) + template = env.get_template("methodology_template.html") + + # Render and write + rendered_html = template.render(sections=sections) + out_path = os.path.join(cfg("out_dir"), "methodology.html") + Path(out_path).write_text(rendered_html) + print(f"HTML file written to: {cfg('out_dir')}") + + return out_path + + +def generate_data_html(config): + """ + Generate diagnostic output HTML pages for the E3SM-PMP Diagnostics Package. + """ + + # Safe getters with defaults + def cfg(key, default=None): + return config.get(key, default) + + # Join lists safely + def join_list(key): + vals = cfg(key, []) or [] + return ", ".join(vals) + + clim_vars = join_list("clim_vars") + mova_vars = join_list("mova_vars") + movc_vars = join_list("movc_vars") + enso_vars = join_list("enso_vars") + + clim_reference = cfg("clim_reference", "") + mova_reference = cfg("mova_reference", "") + movc_reference = cfg("movc_reference", "") + enso_reference = cfg("enso_reference", "") + + sections: List[Dict[str, object]] = [] + + # ---------------- General Notes ---------------- + general_notes: List[Tuple[str, str]] = [ + ( + "Source Code", + ( + "Diagnosis & Intercomparison (PCMDI) Metrics Package (PMP) " + f"Version {cfg('version', '')}." + ), + ), + ( + "Reference", + "The observations used were collected and processed by the " + "E3SM Diagnostics. " + "More info: Documentation.", + ), + ("Experiment", f"{cfg('case_id', '')}"), + ("Output Path", f"{cfg('diag_dir', '')}/metrics_data"), + ("Reference Path", f"{cfg('obs_dir', '')}"), + ("PMP Path", f"{cfg('pmp_dir', '')}"), + ] + add_section(sections, "General Notes", general_notes) + + # ---------------- Mean Climate ---------------- + if cfg("clim_viewer", False): + clim_metrics = [ + ( + "Metrics", + "Mean Bias, RMSE, Centered RMSE, etc.; see full list in the " + "document.", + ), + ( + "Region", + "Global, Northern Hemisphere, Southern Hemisphere, and Tropics. " + "See region/regions_specs.json under PMP Path.", + ), + ( + "Variables", + f"CMIP conventions: {clim_vars}.", + ), + ( + "References", + f"All use {clim_reference}, defined in reference/reference_alias.json under PMP Path. " + "Source data linked from Reference Path.", + ), + ( + "Model Diagnostics", + "JSON files in mean_climate under Output Path.", + ), + ( + "CMIP Diagnostics", + "Pre-generated datasets: " + "CMIP mean climate.", + ), + ] + add_section( + sections, + f"Mean Climate Metrics Data ({cfg('clim_period', '')})", + clim_metrics, + ) + + # ---------------- EMoV (Atmos/Coupled modes) ---------------- + metric_string = variable_string = reference_string = "" + mova_view = cfg("mova_viewer", False) + movc_view = cfg("movc_viewer", False) + + if mova_view and movc_view: + metric_string = ( + "The coupled modes considered include the Pacific Decadal Oscillation (PDO), " + "North Pacific Gyre Oscillation (NPGO), and Atlantic Multidecadal Oscillation (AMO), " + "while the atmospheric modes include the Northern Annular Mode (NAM), Pacific–North American pattern (PNA), " + "North Pacific Oscillation (NPO), North Atlantic Oscillation (NAO), Southern Annular Mode (SAM), " + "and the Pacific–South American patterns (PSA1 and PSA2). " + "Metrics were derived from empirical orthogonal function (EOF) analysis, using sea-level pressure (PSL) for atmospheric modes " + "and sea surface temperature (TS) for coupled modes." + ) + variable_string = f"{mova_vars} (ATM) and {movc_vars} (CPL)" + reference_string = f"{mova_reference} (ATM) and {movc_reference} (CPL)" + elif mova_view: + metric_string = ( + "Atmospheric modes include NAM, PNA, NPO, NAO, SAM, and PSA1/PSA2. " + "Metrics were derived from EOF analysis based on sea-level pressure (PSL)." + ) + variable_string = f"{mova_vars} (ATM)" + reference_string = f"{mova_reference} (ATM)" + elif movc_view: + metric_string = ( + "The coupled modes considered include PDO, NPGO, and AMO. " + "Metrics were derived from EOF analysis based on sea surface temperature (TS)." + ) + variable_string = f"{movc_vars} (CPL)" + reference_string = f"{movc_reference} (CPL)" + + if mova_view or movc_view: + emov_metrics = [ + ( + "Metrics", + f"{metric_string} See full descriptions in the " + "document.", + ), + ( + "Region", + "Regions defined in region/regions_specs.json under PMP Path.", + ), + ( + "Variables", + f"CMIP conventions: {variable_string}.", + ), + ( + "References", + f"Following Variables: {reference_string}, defined in " + "reference/reference_alias.json under PMP Path. " + "Source data linked from Reference Path.", + ), + ( + "Model Diagnostics", + "JSON files in variability_modes under Output Path.", + ), + ( + "CMIP Diagnostics", + "Pre-generated datasets: " + "CMIP modes variability.", + ), + ] + add_section( + sections, + f"EMoV Metrics Data (coupled modes: {cfg('movc_period', '')}; atmospheric modes: {cfg('mova_period', '')})", + emov_metrics, + ) + + # ---------------- ENSO ---------------- + if cfg("enso_viewer", False): + enso_metrics = [ + ( + "Metrics", + "Three groups (ENSO_perf, ENSO_proc, ENSO_tel); see full list in the " + "document.", + ), + ("Region", "Niño 3.4 region (5°S–5°N, 170°–120°W)."), + ("Variables", f"CMIP conventions: {enso_vars}."), + ( + "References", + f"Following Variables: {enso_reference}, defined in " + "reference/reference_alias.json under PMP Path. " + "Source data linked from Reference Path.", + ), + ( + "Model Diagnostics", + "JSON files in enso_metric/ENSO_perf, " + "enso_metric/ENSO_proc, and enso_metric/ENSO_tel under Output Path.", + ), + ( + "CMIP Diagnostics", + "Pre-generated datasets: " + "CMIP ENSO metrics.", + ), + ] + add_section( + sections, f"ENSO Metrics Data ({cfg('enso_period', '')})", enso_metrics + ) + + # ---------------- Render ---------------- + env = setup_jinja_env(cfg("template_dir")) + template = env.get_template("data_template.html") + + output_html = template.render( + title="E3SM-PMP Diagnostics Package", sections=sections + ) + + out_path = os.path.join(cfg("out_dir"), "diag_data.html") + Path(out_path).write_text(output_html) + print(f"HTML file written to: {cfg('out_dir')}") + + return out_path + + +def to_relative_path(absolute_path, base_path=None): + """ + Converts an absolute path to a relative path. + + Parameters: + absolute_path (str): The absolute file path to convert. + base_path (str): The base directory to make the path relative to. + If None, uses the current working directory. + + Returns: + str: Relative path. + """ + if base_path is None: + base_path = os.getcwd() + return os.path.join("..", os.path.relpath(absolute_path, start=base_path)) + + +def safe_join(base, filename): + if base.endswith("/"): + return f"{base}{filename}" + else: + return f"{base}/{filename}" + + +def create_image_link( + fig_dir, diag_dir, subdirs, filename_pattern, label, fallback_filename=None +): + sub_path = Path(*subdirs) + search_path = Path(fig_dir) / sub_path / filename_pattern + matches = glob.glob(str(search_path)) + + if matches: + file_name = Path(matches[0]).name + else: + file_name = fallback_filename or filename_pattern + + full_path = Path(fig_dir) / sub_path / file_name + href = Path(to_relative_path(fig_dir, diag_dir)) / sub_path / file_name + href_str = str(href).replace("\\", "/") + if full_path.is_file(): + return ( + f'' + f'{label}' + ) + else: + return f'{label}' + + +class SummaryTableBuilder: + def __init__(self, diag_dir, fig_dir): + self.diag_dir = diag_dir + self.fig_dir = fig_dir + self.regions = ["global", "ocean", "land", "NHEX", "SHEX", "TROPICS"] + self.metrics = [ + ("Mean Bias", "mae_xy", "Portrait"), + ("Pattern Corr.", "cor_xy", "Portrait"), + ("RMSE", "rms_xy", "Portrait", "rms_xyt", "ParCoord"), + ] + + def build_summary_table(self, regions=None, metrics=None): + clim_path = safe_join(str(self.fig_dir), "ERROR_metric/mean_climate") + metric_table = [] + + if regions is None: + regions = self.regions + + if metrics is None: + metrics = self.metrics + + for i, region in enumerate(regions): + row = [] + if i == 0: + row.append( + {"content": "Mean Climate", "rowspan": len(self.regions)} + ) + row.append({"content": region.upper()}) + + for metric in metrics: + if len(metric) == 3: + name, prefix, mode = metric + filename_pattern = f"{prefix}_{region}_portrait_mean_climate.png" + link = create_image_link( + fig_dir=clim_path, + diag_dir=self.diag_dir, + subdirs=[region], + filename_pattern=filename_pattern, + label=mode, + ) + row.append({"colspan": 4, "content": f"{name}
{link}"}) + + elif len(metric) == 5: + name, prefix1, mode1, prefix2, mode2 = metric + + link1 = create_image_link( + fig_dir=clim_path, + diag_dir=self.diag_dir, + subdirs=[region], + filename_pattern=f"{prefix1}_{region}_portrait_mean_climate.png", + label=mode1, + ) + link2 = create_image_link( + fig_dir=clim_path, + diag_dir=self.diag_dir, + subdirs=[region], + filename_pattern=f"{prefix2}_{region}_parcoord_mean_climate.png", + label=mode2, + ) + + row.append({"colspan": 4, "content": f"{name}
{link1} {link2}"}) + + metric_table.append(row) + + return metric_table + + def build_enso_row(self): + row: List[Dict[str, object]] = [] + row.append({"content": "ENSO"}) + row.append({"content": "TROPICS"}) + enso_path = safe_join(str(self.fig_dir), "ERROR_metric/enso_metric") + + link = create_image_link( + fig_dir=enso_path, + diag_dir=self.diag_dir, + subdirs=[], + filename_pattern="enso_metric_skill_portrait.png", + label="Portrait", + ) + + row.append({"colspan": 12, "content": f"Performance Skill
{link}"}) + return row + + def build_emov_row(self): + row: List[Dict[str, object]] = [] + row.append({"content": "EMoVs"}) + row.append({"content": "Extra-TROPICS"}) + emov_path = safe_join(str(self.fig_dir), "ERROR_metric/variability_modes") + + modes_metrics = [ + ( + "PC_Std_Dev", + "stdv_pc_ratio_to_obs_variability_modes_portrait_mon.png", + "Portrait", + "stdv_pc_ratio_to_obs_variability_modes_parcoord_mon.png", + "ParCoord", + ), + ( + "Centered RMSE", + "rmsc_variability_modes_portrait_mon.png", + "Portrait", + "rmsc_variability_modes_parcoord_mon.png", + "ParCoord", + ), + ( + "RMSE", + "rms_variability_modes_portrait_mon.png", + "Portrait", + "rms_variability_modes_parcoord_mon.png", + "ParCoord", + ), + ] + + for name, f1, mode1, f2, mode2 in modes_metrics: + link1 = create_image_link(emov_path, self.diag_dir, [], f1, mode1) + link2 = create_image_link(emov_path, self.diag_dir, [], f2, mode2) + row.append({"colspan": 4, "content": f"{name}
{link1} {link2}"}) + return row + + +def generate_summary_table( + diag_dir: str, + fig_dir: str, + clim_show: Union[bool, str] = True, + clim_regions: Optional[Union[List[str], str]] = None, + clim_metrics: Optional[List[Tuple[str, ...]]] = None, + mova_show: Union[bool, str] = True, + movc_show: Union[bool, str] = True, + enso_show: Union[bool, str] = True, +) -> List[list]: + """ + Build the summary metrics table (Mean Climate, ENSO, EMoV). + Returns a list of rows (each row is a list of cell dicts). + """ + # Coerce possible string flags + for name in ("clim_show", "mova_show", "movc_show", "enso_show"): + val = locals()[name] + if not isinstance(val, bool): + locals()[name] = str(val).strip().lower() in { + "1", + "true", + "t", + "yes", + "y", + "on", + } + clim_show, mova_show, movc_show, enso_show = ( + clim_show, + mova_show, + movc_show, + enso_show, + ) + + builder = SummaryTableBuilder(diag_dir, fig_dir) + table: List[list] = [] + + if clim_show: + table.extend( + builder.build_summary_table(regions=clim_regions, metrics=clim_metrics) + ) + + if enso_show: + table.append(builder.build_enso_row()) + + if mova_show or movc_show: + table.append(builder.build_emov_row()) + + return table + + +class CMVARGroupBuilder: + def __init__(self): + self.regions = { + "Global": {"All": "01", "El/La": "02"}, + "Africa": {"All": "03", "El/La": "08"}, + "CONUS": {"All": "04", "El/La": "09"}, + "SA": {"All": "05", "El/La": "10"}, + "SCS": {"All": "06", "El/La": "11"}, + "AUS": {"All": "07", "El/La": "12"}, + } + + def create_metric_group(self, *metrics): + return { + name: f"divedown{str(i + 1).zfill(2)}" for i, name in enumerate(metrics) + } + + def create_region_maps(self): + region_metrics = {} + for reg, types in self.regions.items(): + for typ, val in types.items(): + region_metrics[f"{reg}({typ})"] = f"divedown{str(val).zfill(2)}" + return region_metrics + + def construct(self): + region_metrics = self.create_region_maps() + + return { + "Perf": { + "ENSO_perf": { + "BiasPrLat": self.create_metric_group("Skill(Lat)", "Pattern(SRF)"), + "BiasPrLon": self.create_metric_group("Skill(Lon)", "Pattern(SRF)"), + "BiasSstLon": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)" + ), + "BiasTauxLon": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoAmpl": self.create_metric_group( + "Skill(All)", "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoDuration": self.create_metric_group( + "Skill(All)", "Skill(AC)", "Skill(El/La)" + ), + "EnsoSstDiversity": self.create_metric_group( + "Skill(All)", "Skill(El/La)" + ), + "EnsoSstSkew": self.create_metric_group( + "Skill(All)", "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoSstLonRmse": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)", "Skill(El/La)", "Pattern(El/La)" + ), + "EnsoSstTsRmse": self.create_metric_group( + "Skill(All)", "Pattern(Hov)", "Sill(El/La)", "Pattern(El/La)" + ), + "EnsoSeasonality": self.create_metric_group( + "Skill(All)", + "Skill(AC)", + "Pattern(Hov)", + "Skill(El/La)", + "Pattern(El/La)", + ), + "SeasonalPrLat": self.create_metric_group( + "Skill(All)", "Pattern(SRF)", "Pattern(Hov)" + ), + "SeasonalPrLon": self.create_metric_group( + "Skill(All)", "Pattern(SRF)", "Pattern(Hov)" + ), + "SeasonalSstLon": self.create_metric_group( + "Skill(All)", "Pattern(SRF)", "Pattern(Hov)" + ), + "SeasonalTauxLon": self.create_metric_group( + "Skill(All)", "Pattern(SRF)", "Pattern(Hov)" + ), + } + }, + "Proc": { + "ENSO_proc": { + "BiasSstLon": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)" + ), + "BiasTauxLon": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoAmpl": self.create_metric_group( + "Skill(All)", "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoFbSstTaux": self.create_metric_group( + "Skill(CPL)", "Skill(NLIN)", "SKill(FDBK)", "Pattern(Hov)" + ), + "EnsoSeasonality": self.create_metric_group( + "Skill(All)", + "Skill(AC)", + "Pattern(Hov)", + "Skill(Lon)", + "Pattern(SRF)", + ), + "EnsoSstLonRmse": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)", "Skill(El/La)", "Pattern(El/La)" + ), + "EnsoSstSkew": self.create_metric_group( + "Skill(All)", "Skill(Lon)", "Pattern(SRF)" + ), + } + }, + "Telec": { + "ENSO_tel": { + "EnsoAmpl": self.create_metric_group( + "Skill(All)", "Skill(Lon)", "Pattern(SRF)" + ), + "EnsoSeasonality": self.create_metric_group( + "Skill(All)", + "Skill(AC)", + "Pattern(Hov)", + "Skill(El/La)", + "Pattern(El/La)", + ), + "EnsoSstLonRmse": self.create_metric_group( + "Skill(Lon)", "Pattern(SRF)", "Skill(El/La)", "Pattern(El/La)" + ), + "EnsoPrMapDjf": region_metrics, + "EnsoPrMapJja": region_metrics, + "EnsoSstMapDjf": region_metrics, + "EnsoSstMapJja": region_metrics, + } + }, + } + + @staticmethod + def reshape_1d_to_2d(lst, num_cols, fill_value=""): + num_rows = -(-len(lst) // num_cols) + padded = lst + [fill_value] * (num_rows * num_cols - len(lst)) + return [padded[i * num_cols : (i + 1) * num_cols] for i in range(num_rows)], ( + num_rows, + num_cols, + ) + + @staticmethod + def build_cmvar_cell(fig_dir, diag_dir, group, variable, keys_dict): + content = f"{variable}
" + keys = list(keys_dict.items()) + for i, (label, code) in enumerate(keys): + link = create_image_link( + fig_dir=fig_dir, + diag_dir=diag_dir, + subdirs=[group], + filename_pattern=f"{group}_{variable}*_{code}.png", + label=label, + fallback_filename=f"{group}_{variable}_{code}.png", + ) + if len(keys) > 2 and (i - 1) % 2 == 0: + content += f" {link}
" + else: + content += f" {link}" + return content + + def generate_mcpl_row(self, mode, diag_dir, fig_dir): + groups = { + "MOV_eoftest": { + "EOF Spectr": { + "EG Spec(Yearly)": {"EG_Spec": "yearly"}, + "EG Spec(Monthly)": {"EG_Spec": "monthly"}, + } + }, + "MOV_compose": { + "EOF Compos": { + "CBF(Yearly)": {"cbf": "yearly"}, + "CBF(Monthly)": {"cbf": "monthly"}, + "EOF(Yearly)": {"cbf": "yearly"}, + "EOF(Monthly)": {"cbf": "monthly"}, + } + }, + "MOV_pattern": { + "EOF Pattern": { + "EOF1(Yearly)": {"eof1": "yearly"}, + "EOF1(Monthly)": {"eof1": "monthly"}, + "EOF2(Yearly)": {"eof2": "yearly"}, + "EOF2(Monthly)": {"eof2": "monthly"}, + "EOF3(Yearly)": {"eof3": "yearly"}, + "EOF3(Monthly)": {"eof3": "monthly"}, + } + }, + "MOV_telecon": { + "EOF Telec": { + "CBF(Yearly)": {"cbf": "yearly"}, + "CBF(Monthly)": {"cbf": "monthly"}, + "EOF1(Yearly)": {"eof1": "yearly"}, + "EOF1(Monthly)": {"eof1": "monthly"}, + "EOF2(Yearly)": {"eof2": "yearly"}, + "EOF2(Monthly)": {"eof2": "monthly"}, + "EOF3(Yearly)": {"eof3": "yearly"}, + "EOF3(Monthly)": {"eof3": "monthly"}, + } + }, + } + + row = [{"content": f"{mode} (SST)", "rowspan": 1}] + for src_grp, cat_dict in groups.items(): + category_label = next(iter(cat_dict)) + display_items = cat_dict[category_label] + content = f"{category_label}
" + + for idx, (label, fileinfo) in enumerate(display_items.items()): + key_type, season = next(iter(fileinfo.items())) + pattern = f"{src_grp}_{mode}_*{key_type}.png" + fallback = f"{src_grp}_{mode}_{key_type}.png" + link = create_image_link( + fig_dir=fig_dir, + diag_dir=diag_dir, + subdirs=[src_grp, season], + filename_pattern=pattern, + fallback_filename=fallback, + label=label, + ) + + if len(display_items) > 2 and (idx % 2 == 1): + content += f"{link}
" + else: + content += f"{link} " + + row.append({"colspan": 4, "content": content.strip()}) + return row + + +def generate_cmvar_table( + diag_dir: str, + fig_dir: str, + enso_show: Union[bool, str] = True, + movc_show: Union[bool, str] = True, + movc_modes: Optional[Union[List[str], str]] = None, +) -> List[list]: + """ + Build the Coupled Modes Variability (CMVAR) table (ENSO + coupled ocean modes). + + Returns: + list: HTML-ready rows (list of lists of cell dicts). + """ + # Coerce flags if they might come from text config + if not isinstance(enso_show, bool): + enso_show = str(enso_show).strip().lower() in { + "1", + "true", + "t", + "yes", + "y", + "on", + } + if not isinstance(movc_show, bool): + movc_show = str(movc_show).strip().lower() in { + "1", + "true", + "t", + "yes", + "y", + "on", + } + + # Accept comma-separated modes + if movc_modes is None: + movc_modes = ["PDO", "NPGO", "AMO"] + elif isinstance(movc_modes, str): + movc_modes = [s.strip() for s in movc_modes.split(",") if s.strip()] + + cmvar_table: List[list] = [] + + builder = CMVARGroupBuilder() + + # --- ENSO block --- + if enso_show: + # Prefer existing dir; fallback keeps old behavior + enso_path_candidates = ( + safe_join(str(fig_dir), "enso_metric"), + safe_join(str(fig_dir), "ENSO_metric"), + ) + enso_path = next( + (p for p in enso_path_candidates if os.path.isdir(p)), + enso_path_candidates[-1], + ) + + cmvar_groups = builder.construct() or {} + for mode_label, group_data in cmvar_groups.items(): + if not isinstance(group_data, dict) or not group_data: + continue + group_name = next(iter(group_data)) + var_dict = group_data[group_name] or {} + + variables, (nrows, _ncols) = builder.reshape_1d_to_2d( + list(var_dict.keys()), 4 + ) + + for i, row_vars in enumerate(variables): + row = [] + if i == 0: + row.append( + {"content": f"ENSO {mode_label}", "rowspan": nrows} + ) + + for var in row_vars: + if not var: + row.append({"colspan": 4, "content": "--"}) + continue + content = builder.build_cmvar_cell( + fig_dir=enso_path, + diag_dir=diag_dir, + group=group_name, + variable=var, + keys_dict=var_dict.get(var, {}), + ) + row.append({"colspan": 4, "content": content}) + + cmvar_table.append(row) + + # --- Coupled modes (PDO/NPGO/AMO) block --- + if movc_show: + emov_path = safe_join(str(fig_dir), "MOV_metric") + for mode in movc_modes: + mcpl_row = builder.generate_mcpl_row(mode, diag_dir, emov_path) + if mcpl_row: + cmvar_table.append(mcpl_row) + + return cmvar_table + + +class EMOVGroupBuilder: + def __init__( + self, + diag_dir, + fig_dir, + modes_names=None, + modes_seasons=None, + ): + self.diag_dir = diag_dir + self.fig_dir = fig_dir + self.modes = self.map_modes(modes_names) + self.seasons = self.map_seasons(modes_seasons) + + # (UI label, filename template) + self.rowspecs = [ + ("Composite (CBF)", "MOV_compose_{}_{}_cbf.png"), + ("Composite (EOF)", "MOV_compose_{}_{}_{}.png"), # needs 3 args + ("North Test", "MOV_eoftest_{}_{}_EG_Spec.png"), + ("EOF1 Pattern", "MOV_pattern_{}_{}_eof1.png"), + ("EOF1 Telecon.", "MOV_telecon_{}_{}_eof1.png"), + ("EOF2 Pattern", "MOV_pattern_{}_{}_eof2.png"), + ("EOF2 Telecon.", "MOV_telecon_{}_{}_eof2.png"), + ("EOF3 Pattern", "MOV_pattern_{}_{}_eof3.png"), + ("EOF3 Telecon.", "MOV_telecon_{}_{}_eof3.png"), + ] + + def map_modes(self, names): + default_modes = { + "NAM": "EOF1", + "PNA": "EOF1", + "NAO": "EOF1", + "SAM": "EOF1", + "NPO": "EOF2", + "PSA1": "EOF2", + "PSA2": "EOF3", + } + if names is None: + return default_modes + return {k: default_modes.get(k, "EOF1") for k in names} + + def map_seasons(self, names): + default_seasons = ["DJF", "MAM", "JJA", "SON", "yearly", "monthly"] + if names is None: + return default_seasons + + seasons = [] + for sea in names: + s = str(sea).strip().lower() + if s in {"ann", "year", "yearly"}: + seasons.append("yearly") + elif s in {"mon", "month", "monthly"}: + seasons.append("monthly") + else: + # preserve canonical case for DJF/MAM/JJA/SON if provided + seasons.append(sea) + return seasons + + def _format_filename(self, pattern, mode, season, eof_tag): + n = pattern.count("{}") + if n == 3: + return pattern.format(mode, season, eof_tag) + elif n == 2: + return pattern.format(mode, season) + else: + # Unexpected template; return as-is to avoid crashing + return pattern + + def build(self): + emov_path = safe_join(str(self.fig_dir), "MOV_metric") + table = [] + + for mode, eof in self.modes.items(): + eof_tag = str(eof).lower() # e.g., "eof1" + for i, (label, pattern) in enumerate(self.rowspecs): + row = [] + if i == 0: + row.append( + { + "content": f"{mode} (PSL)", + "rowspan": len(self.rowspecs), + } + ) + + row.append( + { + "colspan": 4, + "content": f'{label}', + } + ) + + for season in self.seasons: + filename = self._format_filename(pattern, mode, season, eof_tag) + + # subdir convention: e.g., "MOV_compose/" + prefix = "_".join(filename.split("_")[:2]) # "MOV_compose" + subdir_parts = [prefix, season] + + link = create_image_link( + fig_dir=emov_path, + diag_dir=self.diag_dir, + subdirs=subdir_parts, + filename_pattern=filename, + fallback_filename=filename, + label=str(season).upper(), + ) + + row.append({"colspan": 4, "content": link}) + + table.append(row) + + return table + + +def generate_emovs_table( + diag_dir: str, + fig_dir: str, + show: Union[bool, str] = True, + modes: Optional[Union[List[str], str]] = None, +) -> List[list]: + """ + Build the Extratropical Modes of Variability (EMoV) table. + + Args: + diag_dir: Path to diagnostics directory (for relative links). + fig_dir: Path to figures directory. + show: Whether to build the table. If a string (e.g., "false"), it will be coerced. + modes: List of modes or comma-separated string (e.g., "PDO,NPGO,AMO"). + + Returns: + list: HTML-ready rows (list of lists of cell dicts). + """ + # Coerce show if it might come in as a string + if not isinstance(show, bool): + show = str(show).strip().lower() in {"1", "true", "t", "yes", "y", "on"} + if not show: + return [] + + # Accept comma-separated string for modes + if modes is None: + modes = ["PDO", "NPGO", "AMO"] + elif isinstance(modes, str): + modes = [s.strip() for s in modes.split(",") if s.strip()] + + builder = EMOVGroupBuilder( + diag_dir=diag_dir, + fig_dir=fig_dir, + modes_names=modes, + ) + return builder.build() + + +class MeanClimateTableBuilder: + def __init__(self, diag_dir, fig_dir, variables=None, regions=None, seasons=None): + self.diag_dir = Path(diag_dir) + self.fig_dir = Path(fig_dir) + self.variables = self.map_vars(variables) + self.regions = self.map_regions(regions) + self.seasons = self.map_seasons(seasons) + + def map_vars(self, names): + default_vars = ["pr", "psl", "tas", "ts", "rlds", "rlut"] + if names is None: + variables = default_vars + else: + variables = names + return variables + + def map_regions(self, names): + default_regions = ["global", "land", "ocean", "TROPICS", "NHEX", "SHEX"] + if names is None: + seasons = default_regions + else: + seasons = names + return seasons + + def map_seasons(self, names): + default_seasons = ["DJF", "MAM", "JJA", "SON", "AC"] # AC = Annual Cycle + if names is None: + return default_seasons + seasons = [] + for sea in names: + s = str(sea).strip().lower() + if s in {"ann", "year", "yearly"}: + seasons.append("AC") + else: + # preserve canonical case for DJF/MAM/JJA/SON if provided + seasons.append(sea) + return seasons + + def build_table(self): + """ + Constructs a list of table row dictionaries for use in an HTML diagnostic viewer. + """ + clim_path = safe_join(str(self.fig_dir), "CLIM_patttern") + table = [] + + for var in self.variables: + for i, region in enumerate(self.regions): + row = [] + if i == 0: + row.append( + {"content": f"{var}", "rowspan": len(self.regions)} + ) + + row.append( + { + "colspan": 4, + "content": f'{region}', + } + ) + + for season in self.seasons: + label = "Yearly" if season == "AC" else season + filename = f"{var.strip()}_{region}_{season}.png" + subdir = safe_join(str(region), str(season)) + link = create_image_link( + fig_dir=clim_path, + diag_dir=self.diag_dir, + subdirs=[subdir], + filename_pattern=filename, + label=label, + fallback_filename=filename, + ) + row.append({"colspan": 4, "content": link}) + + table.append(row) + + return table + + +def generate_clim_table( + diag_dir: str, + fig_dir: str, + show: Union[bool, str] = True, + variables: Optional[Union[List[str], str]] = None, + regions: Optional[Union[List[str], str]] = None, +) -> List[list]: + """ + Build the climatology diagnostics table. + + Args: + diag_dir: Path to the diagnostics directory (used to compute relative links). + fig_dir: Path to the figures directory. + show: Whether to build the table. If a string (e.g., "true"/"false"), it will be coerced. + variables: List of variable names (or a comma-separated string). If None, builder defaults are used. + regions: List of region names (or a comma-separated string). If None, builder defaults are used. + + Returns: + A list of table rows (each row is a list of cell dicts) ready for HTML rendering. + """ + # Coerce 'show' if it might come from text config + # In case some unexpected guess values are passed + if not isinstance(show, bool): + show = str(show).strip().lower() in {"1", "true", "t", "yes", "y", "on"} + + # In case comma-separated strings is passed + if isinstance(variables, str): + variables = [s.strip() for s in variables.split(",") if s.strip()] + if isinstance(regions, str): + regions = [s.strip() for s in regions.split(",") if s.strip()] + + if not show: + return [] + + builder = MeanClimateTableBuilder( + diag_dir, + fig_dir, + variables=variables, + regions=regions, + ) + return builder.build_table() + + +def generate_viewer_html(config): + """ + Generate overview HTML page for the E3SM-PMP Diagnostics + """ + env = setup_jinja_env(config["template_dir"]) + template = env.get_template("index_template.html") + + # Add 'summary metrics' row + metric_table = generate_summary_table( + config["diag_dir"], + config["fig_dir"], + clim_show=config["clim_viewer"], + clim_regions=config["clim_regions"], + mova_show=config["mova_viewer"], + movc_show=config["movc_viewer"], + enso_show=config["enso_viewer"], + ) + + # Add 'Coupled Modes Variability' table + cmvars_table = generate_cmvar_table( + config["diag_dir"], + config["fig_dir"], + enso_show=config["enso_viewer"], + movc_show=config["movc_viewer"], + movc_modes=config["movc_modes"], + ) + + # Add 'Extratropical Modes Variability' tabel + emovs_table = generate_emovs_table( + config["diag_dir"], + config["fig_dir"], + show=config["mova_viewer"], + modes=config["mova_modes"], + ) + + # Add 'Mean Climate Map' tabel + clim_table = generate_clim_table( + config["diag_dir"], + config["fig_dir"], + show=config["clim_viewer"], + variables=config["clim_vars"], + regions=config["clim_regions"], + ) + + # Render final HTML + output_html = template.render( + title=config["title"], + subtitle=config["subtitle"], + version=config["version"], + clim_period=config["clim_period"], + mova_period=config["mova_period"], + movc_period=config["movc_period"], + enso_period=config["enso_period"], + created=datetime.now().strftime("%Y-%m-%d"), + metric_table=metric_table, + cmvars_table=cmvars_table, + emovs_table=emovs_table, + clim_table=clim_table, + ) + + # Write the generated HTML to the specified file + Path(os.path.join(config["out_dir"], "index.html")).write_text(output_html) + print(f"HTML file written to: {config['out_dir']}") + return