From ee1eb432b01b6495c00913b370e5abdb06627579 Mon Sep 17 00:00:00 2001 From: Ryan Forsyth Date: Mon, 2 Jun 2025 18:59:03 -0500 Subject: [PATCH 01/17] Split coupled_global --- conda/dev.yml | 5 + .../cases_global_time_series.py | 184 ------ .../test_global_time_series.py | 76 ++- .../global_time_series/__main__.py | 80 +-- .../global_time_series/classic/__init__.py | 0 .../classic/coupled_global.py | 166 +++++ .../classic/coupled_global_plotting.py | 368 +++++++++++ .../global_time_series/classic/driver.py | 59 ++ .../{ => classic}/ocean_month.py | 9 +- .../global_time_series/coupled_global.py | 366 ----------- .../coupled_global_dataset_wrapper.py | 192 ------ .../coupled_global_plotting.py | 609 +++--------------- .../coupled_global_utils.py | 297 +++++++++ zppy_interfaces/global_time_series/utils.py | 85 ++- .../global_time_series/viewer/__init__.py | 0 .../viewer/coupled_global.py | 234 +++++++ .../viewer/coupled_global_plotting.py | 92 +++ .../{ => viewer}/coupled_global_viewer.py | 0 .../global_time_series/viewer/driver.py | 12 + 19 files changed, 1466 insertions(+), 1368 deletions(-) delete mode 100755 tests/integration/global_time_series/cases_global_time_series.py create mode 100644 zppy_interfaces/global_time_series/classic/__init__.py create mode 100644 zppy_interfaces/global_time_series/classic/coupled_global.py create mode 100644 zppy_interfaces/global_time_series/classic/coupled_global_plotting.py create mode 100644 zppy_interfaces/global_time_series/classic/driver.py rename zppy_interfaces/global_time_series/{ => classic}/ocean_month.py (95%) delete mode 100644 zppy_interfaces/global_time_series/coupled_global.py delete mode 100644 zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py create mode 100644 zppy_interfaces/global_time_series/viewer/__init__.py create mode 100644 zppy_interfaces/global_time_series/viewer/coupled_global.py create mode 100644 zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py rename zppy_interfaces/global_time_series/{ => viewer}/coupled_global_viewer.py (100%) create mode 100644 zppy_interfaces/global_time_series/viewer/driver.py diff --git a/conda/dev.yml b/conda/dev.yml index 088c786..b0d25fc 100644 --- a/conda/dev.yml +++ b/conda/dev.yml @@ -1,4 +1,9 @@ # Conda development environment for testing local source code changes to zppy-interfaces before merging them to production (main branch). +# To set up: +# conda clean --all --y +# conda env create -f conda/dev.yml -n env_name +# pre-commit run --all-files +# pip install . name: zppy-interfaces-dev channels: - conda-forge diff --git a/tests/integration/global_time_series/cases_global_time_series.py b/tests/integration/global_time_series/cases_global_time_series.py deleted file mode 100755 index 19bb93f..0000000 --- a/tests/integration/global_time_series/cases_global_time_series.py +++ /dev/null @@ -1,184 +0,0 @@ -import os -import shutil - -from zppy_interfaces.global_time_series.__main__ import main -from zppy_interfaces.global_time_series.utils import Parameters - -CASE_DIR = "/lcrc/group/e3sm/ac.forsyth2/zi-test-input-data" # This is 44G. -WEB_DIR = "/lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zi-test-webdir/" -RESULTS_DIR_PREFIX = "global_time_series_1985-1995_results" - -plots_lnd_metric_average = "FSH,RH2M,LAISHA,LAISUN,QINTR,QOVER,QRUNOFF,QSOIL,QVEGE,QVEGT,SOILWATER_10CM,TSA,H2OSNO," -plots_lnd_metric_total = ( - "TOTLITC,CWDC,SOIL1C,SOIL2C,SOIL3C,SOIL4C,WOOD_HARVESTC,TOTVEGC,NBP,GPP,AR,HR" -) -plots_lnd_all = plots_lnd_metric_average + plots_lnd_metric_total - -parameters_viewers: Parameters = Parameters( - { - "use_ocn": "False", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/ocn/hist", - "moc_file": "None", - "case_dir": CASE_DIR, - "experiment_name": "v3.LR.historical_0051", - "figstr": "v3.LR.historical_0051", - "color": "Blue", - "ts_num_years": "5", - "plots_original": "None", - "plots_atm": "TREFHT", - "plots_ice": "None", - "plots_lnd": plots_lnd_all, - "plots_ocn": "None", - "nrows": "1", - "ncols": "1", - "results_dir": f"{RESULTS_DIR_PREFIX}_viewers", - "regions": "glb,n,s", - "make_viewer": "True", - "start_yr": "1985", - "end_yr": "1995", - } -) - -parameters_custom: Parameters = Parameters( - { - "use_ocn": "False", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/ocn/hist", - "moc_file": "None", - "case_dir": CASE_DIR, - "experiment_name": "v3.LR.historical_0051", - "figstr": "v3.LR.historical_0051", - "color": "Blue", - "ts_num_years": "5", - "plots_original": "None", - "plots_atm": "TREFHT", - "plots_ice": "None", - "plots_lnd": plots_lnd_all, - "plots_ocn": "None", - "nrows": "4", - "ncols": "2", - "results_dir": f"{RESULTS_DIR_PREFIX}_custom", - "regions": "glb,n,s", - "make_viewer": "False", - "start_yr": "1985", - "end_yr": "1995", - } -) - - -parameters_original_8_no_ocn: Parameters = Parameters( - { - "use_ocn": "False", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/ocn/hist", - "moc_file": "None", - "case_dir": CASE_DIR, - "experiment_name": "v3.LR.historical_0051", - "figstr": "v3.LR.historical_0051", - "color": "Blue", - "ts_num_years": "5", - "plots_original": "net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,net_atm_water_imbalance", - "plots_atm": "None", - "plots_ice": "None", - "plots_lnd": "None", - "plots_ocn": "None", - "nrows": "4", - "ncols": "2", - "results_dir": f"{RESULTS_DIR_PREFIX}_original_8_no_ocn", - "regions": "glb,n,s", - "make_viewer": "False", - "start_yr": "1985", - "end_yr": "1995", - } -) - -parameters_original_8: Parameters = Parameters( - { - "use_ocn": "True", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/ocn/hist", - "moc_file": "mocTimeSeries_1985-1995.nc", - "case_dir": CASE_DIR, - "experiment_name": "v3.LR.historical_0051", - "figstr": "v3.LR.historical_0051", - "color": "Blue", - "ts_num_years": "5", - "plots_original": "net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", - "plots_atm": "None", - "plots_ice": "None", - "plots_lnd": "None", - "plots_ocn": "None", - "nrows": "4", - "ncols": "2", - "results_dir": f"{RESULTS_DIR_PREFIX}_original_8", - "regions": "glb,n,s", - "make_viewer": "False", - "start_yr": "1985", - "end_yr": "1995", - } -) - -parameters_comprehensive_v3: Parameters = Parameters( - { - "use_ocn": "True", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/ocn/hist", - "moc_file": "mocTimeSeries_1985-1995.nc", - "case_dir": CASE_DIR, - "experiment_name": "v3.LR.historical_0051", - "figstr": "v3.LR.historical_0051", - "color": "Blue", - "ts_num_years": "5", - "plots_original": "net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", - "plots_atm": "None", - "plots_ice": "None", - "plots_lnd": plots_lnd_all, - "plots_ocn": "None", - "nrows": "4", - "ncols": "2", - "results_dir": f"{RESULTS_DIR_PREFIX}_comprehensive_v3", - "regions": "glb,n,s", - "make_viewer": "False", - "start_yr": "1985", - "end_yr": "1995", - } -) - - -def generate_results(parameters: Parameters): - print(f"Generating results for {parameters.results_dir}") - - # CASE_DIR is large, so we don't want to copy it every time. - # We also want to be able to reset it to the original state after running the test. - # The only modification `global_time_series` makes to CASE_DIR is add an `ocn` subdirectory. - # (It also adds a parameters.results_dir subdirectory to the current directory, which we then copy to WEB_DIR). - # So, we just have to remove those - for subdir in [ - f"{CASE_DIR}/ocn", - parameters.results_dir, - f"{WEB_DIR}/{parameters.results_dir}", - ]: - if os.path.exists(subdir): - print(f"Removing {subdir}") - shutil.rmtree(subdir) - - print("Running main") - main(parameters) - - print(f"Copying {parameters.results_dir} to {WEB_DIR}/{parameters.results_dir}") - shutil.copytree(parameters.results_dir, f"{WEB_DIR}/{parameters.results_dir}") - - -def run_all_cases(): - generate_results(parameters_viewers) - generate_results(parameters_custom) - generate_results(parameters_original_8_no_ocn) - generate_results(parameters_original_8) - generate_results(parameters_comprehensive_v3) - - -if __name__ == "__main__": - # TODO: Create actual pytest cases, including image comparison checks - # (See https://github.com/E3SM-Project/zppy-interfaces/issues/5) - run_all_cases() diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index a63ca46..4eae963 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -2,24 +2,32 @@ import pytest -from zppy_interfaces.global_time_series.coupled_global import ( - construct_generic_variables, - get_data_dir, - get_exps, - get_vars_original, - land_csv_row_to_var, +from zppy_interfaces.global_time_series.classic.coupled_global import ( + get_exps as classic_get_exps, ) +from zppy_interfaces.global_time_series.classic.coupled_global import get_vars_original from zppy_interfaces.global_time_series.coupled_global_plotting import get_ylim -from zppy_interfaces.global_time_series.coupled_global_utils import Metric, Variable -from zppy_interfaces.global_time_series.coupled_global_viewer import ( - VariableGroup, - get_variable_groups, +from zppy_interfaces.global_time_series.coupled_global_utils import ( + Metric, + Variable, + get_data_dir, ) from zppy_interfaces.global_time_series.utils import ( Parameters, get_region, param_get_list, ) +from zppy_interfaces.global_time_series.viewer.coupled_global import ( + construct_generic_variables, +) +from zppy_interfaces.global_time_series.viewer.coupled_global import ( + get_exps as viewer_get_exps, +) +from zppy_interfaces.global_time_series.viewer.coupled_global import land_csv_row_to_var +from zppy_interfaces.global_time_series.viewer.coupled_global_viewer import ( + VariableGroup, + get_variable_groups, +) # Run tests with `pytest tests/unit/global_time_series/test_*.py` @@ -67,41 +75,48 @@ def test_get_region(): def test_Parameters_and_related_functions(): # Consider the following parameters given by a user. args: Dict[str, str] = { - "use_ocn": "True", - "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "input_subdir": "archive/atm/hist", - "moc_file": "mocTimeSeries_1985-1995.nc", + "make_viewer": "True", "case_dir": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051", "experiment_name": "v3.LR.historical_0051", "figstr": "v3.LR.historical_0051", "color": "Blue", "ts_num_years": "5", + "results_dir": "results", + "regions": "glb,n,s", + "start_yr": "1985", + "end_yr": "1989", + "subsection": "subsection_name", + "use_ocn": "True", + "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", + "input_subdir": "archive/atm/hist", + "moc_file": "mocTimeSeries_1985-1995.nc", "plots_original": "None", + "nrows": "1", + "ncols": "1", "plots_atm": "TREFHT", "plots_ice": "None", "plots_lnd": "FSH,RH2M,LAISHA,LAISUN,QINTR,QOVER,QRUNOFF,QSOIL,QVEGE,QVEGT,SOILWATER_10CM,TSA,H2OSNO,TOTLITC,CWDC,SOIL1C,SOIL2C,SOIL3C,SOIL4C,WOOD_HARVESTC,TOTVEGC,NBP,GPP,AR,HR", "plots_ocn": "None", - "nrows": "1", - "ncols": "1", - "results_dir": "results", - "regions": "glb,n,s", - "make_viewer": "True", - "start_yr": "1985", - "end_yr": "1989", } # Then: parameters: Parameters = Parameters(args) + assert parameters.make_viewer assert ( parameters.case_dir == "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051" ) assert parameters.experiment_name == "v3.LR.historical_0051" assert parameters.figstr == "v3.LR.historical_0051" - assert parameters.year1 == 1985 - assert parameters.year2 == 1989 assert parameters.color == "Blue" assert parameters.ts_num_years_str == "5" + assert parameters.results_dir == "results" + assert parameters.regions == ["glb", "n", "s"] + assert parameters.year1 == 1985 + assert parameters.year2 == 1989 + assert parameters.subtask_name == "global_time_series_subsection_name" assert parameters.plots_original == [] + assert parameters.nrows == 1 + assert parameters.ncols == 1 assert parameters.plots_atm == ["TREFHT"] assert parameters.plots_ice == [] assert parameters.plots_lnd == [ @@ -132,9 +147,6 @@ def test_Parameters_and_related_functions(): "HR", ] assert parameters.plots_ocn == [] - assert parameters.nrows == 1 - assert parameters.ncols == 1 - assert parameters.regions == ["glb", "n", "s"] # test_get_data_dir assert ( @@ -159,15 +171,13 @@ def test_Parameters_and_related_functions(): assert get_data_dir(parameters, "ocn", False) == "" # test_get_exps - exps: List[Dict[str, Any]] = get_exps(parameters) + exps: List[Dict[str, Any]] = viewer_get_exps(parameters) assert len(exps) == 1 expected = { "atmos": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/atm/glb/ts/monthly/5yr/", "ice": "", "land": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/lnd/glb/ts/monthly/5yr/", "ocean": "", - "moc": "", - "vol": "", "name": "v3.LR.historical_0051", "yoffset": 0.0, "yr": ([1985, 1989],), @@ -180,15 +190,13 @@ def test_Parameters_and_related_functions(): ) parameters.plots_atm = [] parameters.plots_lnd = [] - exps = get_exps(parameters) + exps = classic_get_exps(parameters) assert len(exps) == 1 expected = { "atmos": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/atm/glb/ts/monthly/5yr/", - "ice": "", - "land": "", "ocean": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/ocn/glb/ts/monthly/5yr/", - "moc": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/ocn/glb/ts/monthly/5yr/", - "vol": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/ocn/glb/ts/monthly/5yr/", + "moc": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/global_time_series_subsection_name/ocn/glb/ts/monthly/5yr/", + "vol": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/global_time_series_subsection_name/ocn/glb/ts/monthly/5yr/", "name": "v3.LR.historical_0051", "yoffset": 0.0, "yr": ([1985, 1989],), diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 0cad995..314c4e4 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -1,10 +1,8 @@ import argparse -import os -import shutil import sys -from zppy_interfaces.global_time_series.coupled_global import coupled_global -from zppy_interfaces.global_time_series.ocean_month import ocean_month +import zppy_interfaces.global_time_series.classic.driver as classic_driver +import zppy_interfaces.global_time_series.viewer.driver as viewer_driver from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger @@ -17,38 +15,11 @@ def main(parameters=None): if not parameters: parameters = _get_args() - - if parameters.use_ocn: - logger.info("Create ocean time series") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Creates the directory post/ocn - os.makedirs( - f"{parameters.case_dir}/post/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr", - exist_ok=True, - ) - input: str = f"{parameters.input}/{parameters.input_subdir}" - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Modifies post/ocn (which we just created in the first place) - ocean_month( - input, - parameters.case_dir, - parameters.year1, - parameters.year2, - int(parameters.ts_num_years_str), - ) - - logger.info("Copy moc file") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Copies files to post/ocn (which we just created in the first place) - shutil.copy( - f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}", - f"{parameters.case_dir}/post/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/", - ) - - logger.info("Update time series figures") - # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) - # Creates the directory parameters.results_dir - coupled_global(parameters) + # Determine if we want the Classic PDF or the Viewer + if parameters.make_viewer: + viewer_driver.run(parameters) + else: + classic_driver.run(parameters) def _get_args() -> Parameters: @@ -58,32 +29,39 @@ def _get_args() -> Parameters: description="Generate Global Time Series plots", ) - # For ocean_month - parser.add_argument("--use_ocn", type=str, help="Use ocean") - parser.add_argument("--input", type=str, help="Input directory") - parser.add_argument("--input_subdir", type=str, help="Input subdirectory") - parser.add_argument("--moc_file", type=str, help="MOC file") - + # Used by both Classic PDF and Viewer + # For determining which output type to produce + parser.add_argument("--make_viewer", type=str, help="Make viewer") # For coupled_global parser.add_argument("--case_dir", type=str, help="Case directory") parser.add_argument("--experiment_name", type=str, help="Experiment name") parser.add_argument("--figstr", type=str, help="Figure string") parser.add_argument("--color", type=str, help="Color") parser.add_argument("--ts_num_years", type=str, help="Time series number of years") + parser.add_argument("--results_dir", type=str, help="Results directory") + parser.add_argument("--regions", type=str, help="Regions") + # For both ocean_month and coupled_global + parser.add_argument("--start_yr", type=str, help="Start year") + parser.add_argument("--end_yr", type=str, help="End year") + + # Classic PDF only + # For ocean_month + parser.add_argument("--subsection", type=str, help="Subtask name") + parser.add_argument("--use_ocn", type=str, help="Use ocean") + parser.add_argument("--input", type=str, help="Input directory") + parser.add_argument("--input_subdir", type=str, help="Input subdirectory") + parser.add_argument("--moc_file", type=str, help="MOC file") + # For coupled_global parser.add_argument("--plots_original", type=str, help="Plots original") + parser.add_argument("--nrows", type=str, help="Number of rows in pdf") + parser.add_argument("--ncols", type=str, help="Number of columns in pdf") + + # Viewer only + # For coupled_global parser.add_argument("--plots_atm", type=str, help="Plots atmosphere") parser.add_argument("--plots_ice", type=str, help="Plots ice") parser.add_argument("--plots_lnd", type=str, help="Plots land") parser.add_argument("--plots_ocn", type=str, help="Plots ocean") - parser.add_argument("--nrows", type=str, help="Number of rows in pdf") - parser.add_argument("--ncols", type=str, help="Number of columns in pdf") - parser.add_argument("--results_dir", type=str, help="Results directory") - parser.add_argument("--regions", type=str, help="Regions") - parser.add_argument("--make_viewer", type=str, help="Make viewer") - - # For both - parser.add_argument("--start_yr", type=str, help="Start year") - parser.add_argument("--end_yr", type=str, help="End year") # Ignore the first arg # (zi-global-time-series) diff --git a/zppy_interfaces/global_time_series/classic/__init__.py b/zppy_interfaces/global_time_series/classic/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/zppy_interfaces/global_time_series/classic/coupled_global.py b/zppy_interfaces/global_time_series/classic/coupled_global.py new file mode 100644 index 0000000..5f28ef3 --- /dev/null +++ b/zppy_interfaces/global_time_series/classic/coupled_global.py @@ -0,0 +1,166 @@ +from typing import Any, Dict, List + +from zppy_interfaces.global_time_series.classic.coupled_global_plotting import ( + make_plot_pdfs, +) +from zppy_interfaces.global_time_series.coupled_global_utils import ( + DatasetWrapper, + Variable, + get_data_dir, + set_var, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + + +# Classes ##################################################################### +class RequestedVariables(object): + def __init__(self, parameters: Parameters): + self.vars_original: List[Variable] = get_vars_original( + parameters.plots_original + ) + + +def get_vars_original(plots_original: List[str]) -> List[Variable]: + # NOTE: These are ALL atmosphere variables + vars_original: List[Variable] = [] + if ("net_toa_flux_restom" in plots_original) or ( + "net_atm_energy_imbalance" in plots_original + ): + vars_original.append(Variable("RESTOM")) + if "net_atm_energy_imbalance" in plots_original: + vars_original.append(Variable("RESSURF")) + if "global_surface_air_temperature" in plots_original: + vars_original.append(Variable("TREFHT")) + if "toa_radiation" in plots_original: + vars_original.append(Variable("FSNTOA")) + vars_original.append(Variable("FLUT")) + if "net_atm_water_imbalance" in plots_original: + vars_original.append(Variable("PRECC")) + vars_original.append(Variable("PRECL")) + vars_original.append(Variable("QFLX")) + return vars_original + + +# Main functionality ########################################################## + + +def run_coupled_global(parameters: Parameters) -> None: + requested_variables = RequestedVariables(parameters) + run(parameters, requested_variables) + + +def run(parameters: Parameters, requested_variables: RequestedVariables): + # Experiments + exps: List[Dict[str, Any]] = process_data(parameters, requested_variables) + + xlim: List[float] = [float(parameters.year1), float(parameters.year2)] + + # Note: we use `parameters.plots_original` rather than `requested_variables.vars_original` + # because the "original" plots are expecting plot names that are not variable names. + # The model components however are expecting plot names to be variable names. + for rgn in parameters.regions: + valid_plots: List[str] = [] + invalid_plots: List[str] = [] + make_plot_pdfs( + parameters, + rgn, + "original", + xlim, + exps, + parameters.plots_original, + valid_plots, + invalid_plots, + ) + logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") + if invalid_plots: + logger.error( + f"These {rgn} region plots could not be generated successfully: {invalid_plots}" + ) + + +def process_data( + parameters: Parameters, requested_variables: RequestedVariables +) -> List[Dict[str, Any]]: + exps: List[Dict[str, Any]] = get_exps(parameters) + valid_vars: List[str] = [] + invalid_vars: List[str] = [] + exp: Dict[str, Any] + + logger.info("Processing data for Classic PDF.") + for exp in exps: + exp["annual"] = {} + + logger.info("Setting requested variables") + requested_variables.vars_original = set_var( + exp, + "atmos", + requested_variables.vars_original, + valid_vars, + invalid_vars, + ) + # Optionally read ohc + logger.info("Reading ohc") + if exp["moc"] != "": + ohc_variable = Variable("ohc") + dataset_wrapper = DatasetWrapper(exp["moc"]) + exp["annual"]["ohc"], _ = dataset_wrapper.globalAnnual(ohc_variable) + # anomalies with respect to first year + exp["annual"]["ohc"][:] = exp["annual"]["ohc"][:] - exp["annual"]["ohc"][0] + + logger.info("Reading vol") + if exp["vol"] != "": + vol_variable = Variable("volume") + dataset_wrapper = DatasetWrapper(exp["vol"]) + exp["annual"]["volume"], _ = dataset_wrapper.globalAnnual(vol_variable) + # annomalies with respect to first year + exp["annual"]["volume"][:] = ( + exp["annual"]["volume"][:] - exp["annual"]["volume"][0] + ) + + logger.info( + f"globalAnnual was computed successfully for these variables: {valid_vars}" + ) + if invalid_vars: + logger.error( + f"globalAnnual could not be computed for these variables: {invalid_vars}" + ) + return exps + + +def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: + # Experiments + atm_set_intersection: set = set( + [ + "net_toa_flux_restom", + "global_surface_air_temperature", + "toa_radiation", + "net_atm_energy_imbalance", + "net_atm_water_imbalance", + ] + ) & set(parameters.plots_original) + # Use set intersection: check if any of these 3 plots were requested + ocn_set_intersection: set = set( + ["change_ohc", "max_moc", "change_sea_level"] + ) & set(parameters.plots_original) + ocean_dir = get_data_dir(parameters, "ocn", ocn_set_intersection != set()) + ocean_month_dir = get_data_dir( + parameters, + f"{parameters.subtask_name}/ocn", + ocn_set_intersection != set(), + ) + exps: List[Dict[str, Any]] = [ + { + "atmos": get_data_dir(parameters, "atm", atm_set_intersection != set()), + "ocean": ocean_dir, + "moc": ocean_month_dir, + "vol": ocean_month_dir, + "name": parameters.experiment_name, + "yoffset": 0.0, + "yr": ([parameters.year1, parameters.year2],), + "color": f"{parameters.color}", + } + ] + return exps diff --git a/zppy_interfaces/global_time_series/classic/coupled_global_plotting.py b/zppy_interfaces/global_time_series/classic/coupled_global_plotting.py new file mode 100644 index 0000000..43a6cda --- /dev/null +++ b/zppy_interfaces/global_time_series/classic/coupled_global_plotting.py @@ -0,0 +1,368 @@ +import math +import os +import traceback + +import matplotlib as mpl +import matplotlib.backends.backend_pdf +import matplotlib.pyplot as plt +import numpy as np + +from zppy_interfaces.global_time_series.coupled_global_plotting import plot +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +mpl.use("Agg") + +logger = _setup_custom_logger(__name__) + + +# FIXME: C901 'make_plot_pdfs' is too complex (20) +def make_plot_pdfs( # noqa: C901 + parameters: Parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, +): + logger.info(f"Global Time Series Classic PDF: make_plot_pdfs for rgn={rgn}") + num_plots = len(plot_list) + if num_plots == 0: + return + + plots_per_page = parameters.nrows * parameters.ncols + num_pages = math.ceil(num_plots / plots_per_page) + + counter = 0 + os.makedirs(parameters.results_dir, exist_ok=True) + # https://stackoverflow.com/questions/58738992/save-multiple-figures-with-subplots-into-a-pdf-with-multiple-pages + pdf = matplotlib.backends.backend_pdf.PdfPages( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.pdf" + ) + for page in range(num_pages): + fig = plt.figure(1, figsize=[13.5, 16.5]) + logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") + fig.suptitle(f"{parameters.figstr}_{rgn}_{component}") + for j in range(plots_per_page): + logger.info( + f"Plotting plot {j} on page {page}. This is plot {counter} in total." + ) + # The final page doesn't need to be filled out with plots. + if counter >= num_plots: + break + ax = plt.subplot( + parameters.nrows, + parameters.ncols, + j + 1, + ) + plot_name = plot_list[counter] + try: + plot_function = PLOT_DICT[plot_name] + except KeyError: + raise KeyError(f"Invalid plot name: {plot_name}") + try: + plot_function(ax, xlim, exps, rgn) + valid_plots.append(plot_name) + except Exception: + traceback.print_exc() + required_vars = [] + if plot_name == "net_toa_flux_restom": + required_vars = ["RESTOM"] + elif plot_name == "net_atm_energy_imbalance": + required_vars = ["RESTOM", "RESSURF"] + elif plot_name == "global_surface_air_temperature": + required_vars = ["TREFHT"] + elif plot_name == "toa_radiation": + required_vars = ["FSNTOA", "FLUT"] + elif plot_name == "net_atm_water_imbalance": + required_vars = ["PRECC", "PRECL", "QFLX"] + logger.error( + f"Failed plot_function for {plot_name}. Check that {required_vars} are available." + ) + invalid_plots.append(plot_name) + counter += 1 + + fig.tight_layout() + pdf.savefig(1) + if num_pages > 1: + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{page}.png", + dpi=150, + ) + else: + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.png", + dpi=150, + ) + plt.close(fig) + pdf.close() + + +# 1 +def plot_net_toa_flux_restom(ax, xlim, exps, rgn): + logger.info("Plot 1: plot_net_toa_flux_restom") + param_dict = { + "2nd_var": False, + "axhline_y": 0, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": True, + "default_ylim": [-1.5, 1.5], + "do_add_line": True, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": True, + "set_legend": True, + "shorten_year": False, + "title": "Net TOA flux (restom)", + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"]["RESTOM"][rgn][0]), + "verbose": False, + "vol": False, + "ylabel": "W m-2", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 2 +def plot_global_surface_air_temperature(ax, xlim, exps, rgn): + logger.info("Plot 2: plot_global_surface_air_temperature") + if rgn == "glb": + region_title = "Global" + elif rgn == "n": + region_title = "Northern Hemisphere" + elif rgn == "s": + region_title = "Southern Hemisphere" + else: + raise RuntimeError(f"Invalid rgn={rgn}") + param_dict = { + "2nd_var": False, + "axhline_y": None, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": True, + "default_ylim": [13, 15.5], + "do_add_line": True, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": False, + "title": f"{region_title} surface air temperature", + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"]["TREFHT"][rgn][0]) - 273.15, + "verbose": False, + "vol": False, + "ylabel": "degC", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 3 +def plot_toa_radiation(ax, xlim, exps, rgn): + logger.info("Plot 3: plot_toa_radiation") + param_dict = { + "2nd_var": True, + "axhline_y": None, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": False, + "default_ylim": [235, 245], + "do_add_line": False, + "do_add_trend": False, + "format": None, + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": False, + "shorten_year": False, + "title": "TOA radiation: SW (solid), LW (dashed)", + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"]["FSNTOA"][rgn][0]), + "verbose": None, + "vol": None, + "ylabel": "W m-2", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 4 +def plot_net_atm_energy_imbalance(ax, xlim, exps, rgn): + logger.info("Plot 4: plot_net_atm_energy_imbalance") + param_dict = { + "2nd_var": False, + "axhline_y": None, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": True, + "default_ylim": [-0.3, 0.3], + "do_add_line": True, + "do_add_trend": False, + "format": "%4.2f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": False, + "title": "Net atm energy imbalance (restom-ressurf)", + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"]["RESTOM"][rgn][0]) + - np.array(exp["annual"]["RESSURF"][rgn][0]), + "verbose": False, + "vol": False, + "ylabel": "W m-2", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 5 +def plot_change_ohc(ax, xlim, exps, rgn): + logger.info("Plot 5: plot_change_ohc") + param_dict = { + "2nd_var": False, + "axhline_y": 0, + "check_exp_ocean": True, + "check_exp_vol": False, + "check_exp_year": False, + "default_ylim": [-0.3e24, 0.9e24], + "do_add_line": False, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": True, + "lw": 1.5, + "ohc": True, + "set_axhline": True, + "set_legend": True, + "shorten_year": True, + "title": "Change in ocean heat content", + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"]["ohc"]), + "verbose": False, + "vol": False, + "ylabel": "J", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 6 +def plot_max_moc(ax, xlim, exps, rgn): + logger.info("Plot 6: plot_max_moc") + param_dict = { + "2nd_var": False, + "axhline_y": 10, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": False, + "default_ylim": [4, 22], + "do_add_line": False, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": True, + "lw": 1.5, + "ohc": False, + "set_axhline": True, + "set_legend": True, + "shorten_year": False, + "title": "Max MOC Atlantic streamfunction at 26.5N", + "use_getmoc": True, + "var": None, + "verbose": True, + "vol": None, + "ylabel": "Sv", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 7 +def plot_change_sea_level(ax, xlim, exps, rgn): + logger.info("Plot 7: plot_change_sea_level") + param_dict = { + "2nd_var": False, + "axhline_y": None, + "check_exp_ocean": False, + "check_exp_vol": True, + "check_exp_year": True, + "default_ylim": [4, 22], + "do_add_line": False, + "do_add_trend": True, + "format": "%5.3f", + "glb_only": True, + "lw": 1.5, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": True, + "title": "Change in sea level", + "use_getmoc": False, + "var": lambda exp: ( + 1e3 + * np.array(exp["annual"]["volume"]) + / (4.0 * math.pi * (6371229.0) ** 2 * 0.7) + ), + "verbose": True, + "vol": True, + "ylabel": "mm", + } + plot(ax, xlim, exps, param_dict, rgn) + + +# 8 +def plot_net_atm_water_imbalance(ax, xlim, exps, rgn): + logger.info("Plot 8: plot_net_atm_water_imbalance") + param_dict = { + "2nd_var": False, + "axhline_y": None, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": False, + "default_ylim": [-1, 1], + "do_add_line": True, + "do_add_trend": False, + "format": "%5.4f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": False, + "title": "Net atm water imbalance (evap-prec)", + "use_getmoc": False, + "var": lambda exp: ( + 365 + * 86400 + * ( + np.array(exp["annual"]["QFLX"][rgn][0]) + - 1e3 + * ( + np.array(exp["annual"]["PRECC"][rgn][0]) + + np.array(exp["annual"]["PRECL"][rgn][0]) + ) + ) + ), + "verbose": False, + "vol": False, + "ylabel": "mm yr-1", + } + plot(ax, xlim, exps, param_dict, rgn) + + +PLOT_DICT = { + "net_toa_flux_restom": plot_net_toa_flux_restom, + "global_surface_air_temperature": plot_global_surface_air_temperature, + "toa_radiation": plot_toa_radiation, + "net_atm_energy_imbalance": plot_net_atm_energy_imbalance, + "change_ohc": plot_change_ohc, # only glb + "max_moc": plot_max_moc, # only glb + "change_sea_level": plot_change_sea_level, # only glb + "net_atm_water_imbalance": plot_net_atm_water_imbalance, +} diff --git a/zppy_interfaces/global_time_series/classic/driver.py b/zppy_interfaces/global_time_series/classic/driver.py new file mode 100644 index 0000000..83a0f03 --- /dev/null +++ b/zppy_interfaces/global_time_series/classic/driver.py @@ -0,0 +1,59 @@ +import os +import shutil + +from zppy_interfaces.global_time_series.classic.coupled_global import run_coupled_global +from zppy_interfaces.global_time_series.classic.ocean_month import ocean_month +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + + +def run(parameters: Parameters): + # From zppy's default.ini: + # Remove the 3 ocean plots (change_ohc,max_moc,change_sea_level) if you don't have ocean data. + # plots_original = string(default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance") + if parameters.use_ocn: + if set(["change_ohc", "max_moc", "change_sea_level"]) & set( + parameters.plots_original + ): + logger.info("Create ocean time series") + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Creates the directory post/ocn + os.makedirs( + f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr", + exist_ok=True, + ) + input_dir: str = f"{parameters.input}/{parameters.input_subdir}" + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Modifies post/ocn (which we just created in the first place) + ocean_month( + input_dir, + parameters.subtask_name, + parameters.case_dir, + parameters.year1, + parameters.year2, + int(parameters.ts_num_years_str), + ) + + src: str = ( + f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}" + ) + dst: str = ( + f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/" + ) + logger.info(f"Copy moc file from {src} to {dst}") + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Copies files to post//ocn (which we just created in the first place) + shutil.copy( + src, + dst, + ) + else: + logger.info( + "use_ocn is set unecessarily. ocn plots have not been requested" + ) + logger.info("Update time series figures") + # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) + # Creates the directory parameters.results_dir + run_coupled_global(parameters) diff --git a/zppy_interfaces/global_time_series/ocean_month.py b/zppy_interfaces/global_time_series/classic/ocean_month.py similarity index 95% rename from zppy_interfaces/global_time_series/ocean_month.py rename to zppy_interfaces/global_time_series/classic/ocean_month.py index 4295203..54279d0 100644 --- a/zppy_interfaces/global_time_series/ocean_month.py +++ b/zppy_interfaces/global_time_series/classic/ocean_month.py @@ -13,9 +13,14 @@ def ocean_month( - path_in: str, case_dir: str, start_yr: int, end_yr: int, ts_num_years: int + path_in: str, + subtask_name: str, + case_dir: str, + start_yr: int, + end_yr: int, + ts_num_years: int, ): - path_out = f"{case_dir}/post/ocn/glb/ts/monthly/{ts_num_years}yr" + path_out = f"{case_dir}/post/{subtask_name}/ocn/glb/ts/monthly/{ts_num_years}yr" # Ocean constants # specific heat [J/(kg*degC)] diff --git a/zppy_interfaces/global_time_series/coupled_global.py b/zppy_interfaces/global_time_series/coupled_global.py deleted file mode 100644 index e7b6724..0000000 --- a/zppy_interfaces/global_time_series/coupled_global.py +++ /dev/null @@ -1,366 +0,0 @@ -# Script to plot some global atmosphere and ocean time series -import csv -import importlib.resources as imp_res -from typing import Any, Dict, List, Tuple - -import cftime -import numpy as np -import xarray - -from zppy_interfaces.global_time_series.coupled_global_dataset_wrapper import ( - DatasetWrapper, -) -from zppy_interfaces.global_time_series.coupled_global_plotting import make_plot_pdfs -from zppy_interfaces.global_time_series.coupled_global_utils import Metric, Variable -from zppy_interfaces.global_time_series.coupled_global_viewer import ( - create_viewer, - create_viewer_index, -) -from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_child_logger - -logger = _setup_child_logger(__name__) - - -# Useful helper functions and classes ######################################### - - -def get_vars_original(plots_original: List[str]) -> List[Variable]: - # NOTE: These are ALL atmosphere variables - vars_original: List[Variable] = [] - if ("net_toa_flux_restom" in plots_original) or ( - "net_atm_energy_imbalance" in plots_original - ): - vars_original.append(Variable("RESTOM")) - if "net_atm_energy_imbalance" in plots_original: - vars_original.append(Variable("RESSURF")) - if "global_surface_air_temperature" in plots_original: - vars_original.append(Variable("TREFHT")) - if "toa_radiation" in plots_original: - vars_original.append(Variable("FSNTOA")) - vars_original.append(Variable("FLUT")) - if "net_atm_water_imbalance" in plots_original: - vars_original.append(Variable("PRECC")) - vars_original.append(Variable("PRECL")) - vars_original.append(Variable("QFLX")) - return vars_original - - -def land_csv_row_to_var(csv_row: List[str]) -> Variable: - # “A” or “T” for global average over land area or global total, respectively - metric: Metric - if csv_row[1] == "A": - metric = Metric.AVERAGE - elif csv_row[1] == "T": - metric = Metric.TOTAL - else: - raise ValueError(f"Invalid metric={csv_row[1]}") - return Variable( - variable_name=csv_row[0], - metric=metric, - scale_factor=float(csv_row[2]), - original_units=csv_row[3], - final_units=csv_row[4], - group=csv_row[5], - long_name=csv_row[6], - ) - - -def construct_land_variables(requested_vars: List[str]) -> List[Variable]: - var_list: List[Variable] = [] - header = True - csv_filename = str( - imp_res.files("zppy_interfaces.global_time_series") / "zppy_land_fields.csv" - ) - with open(csv_filename, newline="") as csv_file: - logger.debug("Reading zppy_land_fields.csv") - var_reader = csv.reader(csv_file) - for row in var_reader: - # logger.debug(f"row={row}") - # Skip the header row - if header: - header = False - else: - # If set to "all" then we want all variables. - # Design note: we can't simply run all variables if requested_vars is empty because - # that would actually mean the user doesn't want to make *any* land plots. - if (requested_vars == ["all"]) or (row[0] in requested_vars): - row_elements_strip_whitespace: List[str] = list( - map(lambda x: x.strip(), row) - ) - var_list.append(land_csv_row_to_var(row_elements_strip_whitespace)) - return var_list - - -def construct_generic_variables(requested_vars: List[str]) -> List[Variable]: - var_list: List[Variable] = [] - for var_name in requested_vars: - var_list.append(Variable(var_name)) - return var_list - - -class RequestedVariables(object): - def __init__(self, parameters: Parameters): - self.vars_original: List[Variable] = get_vars_original( - parameters.plots_original - ) - self.vars_land: List[Variable] = construct_land_variables(parameters.plots_lnd) - - # Use generic constructor - self.vars_atm: List[Variable] = construct_generic_variables( - parameters.plots_atm - ) - self.vars_ice: List[Variable] = construct_generic_variables( - parameters.plots_ice - ) - self.vars_ocn: List[Variable] = construct_generic_variables( - parameters.plots_ocn - ) - - -# Setup ####################################################################### -def get_data_dir(parameters: Parameters, component: str, conditional: bool) -> str: - return ( - f"{parameters.case_dir}/post/{component}/glb/ts/monthly/{parameters.ts_num_years_str}yr/" - if conditional - else "" - ) - - -def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: - # Experiments - use_atmos: bool = (parameters.plots_atm != []) or (parameters.plots_original != []) - # Use set intersection: check if any of these 3 plots were requested - set_intersection: set = set(["change_ohc", "max_moc", "change_sea_level"]) & set( - parameters.plots_original - ) - has_original_ocn_plots: bool = set_intersection != set() - use_ocn: bool = (parameters.plots_ocn != []) or has_original_ocn_plots - ocean_dir = get_data_dir(parameters, "ocn", use_ocn) - exps: List[Dict[str, Any]] = [ - { - "atmos": get_data_dir(parameters, "atm", use_atmos), - "ice": get_data_dir(parameters, "ice", parameters.plots_ice != []), - "land": get_data_dir(parameters, "lnd", parameters.plots_lnd != []), - "ocean": ocean_dir, - "moc": ocean_dir, - "vol": ocean_dir, - "name": parameters.experiment_name, - "yoffset": 0.0, - "yr": ([parameters.year1, parameters.year2],), - "color": f"{parameters.color}", - } - ] - return exps - - -def set_var( - exp: Dict[str, Any], - exp_key: str, - var_list: List[Variable], - valid_vars: List[str], - invalid_vars: List[str], - rgn: str, -) -> List[Variable]: - new_var_list: List[Variable] = [] - if exp[exp_key] != "": - try: - dataset_wrapper: DatasetWrapper = DatasetWrapper(exp[exp_key]) - except Exception as e: - logger.critical(e) - logger.critical( - f"DatasetWrapper object could not be created for {exp_key}={exp[exp_key]}" - ) - raise e - for var in var_list: - var_str: str = var.variable_name - try: - data_array: xarray.core.dataarray.DataArray - units: str - data_array, units = dataset_wrapper.globalAnnual(var) - valid_vars.append(str(var_str)) # Append the name - new_var_list.append(var) # Append the variable itself - except Exception as e: - logger.error(e) - logger.error(f"globalAnnual failed for {var_str}") - invalid_vars.append(str(var_str)) - continue - if data_array.sizes["rgn"] > 1: - # number of years x 3 regions = data_array.shape - # 3 regions = global, northern hemisphere, southern hemisphere - # We get here if we used the updated `ts` task - # (using `rgn_avg` rather than `glb_avg`). - if rgn == "glb": - n = 0 - elif rgn == "n": - n = 1 - elif rgn == "s": - n = 2 - else: - raise RuntimeError(f"Invalid rgn={rgn}") - data_array = data_array.isel(rgn=n) # Just use nth region - elif rgn != "glb": - # data_array only has one dimension -- glb. - # Therefore it is not possible to get n or s plots. - raise RuntimeError( - f"var={var_str} only has global data. Cannot process rgn={rgn}" - ) - exp["annual"][var_str] = (data_array, units) - if "year" not in exp["annual"]: - years: np.ndarray[cftime.DatetimeNoLeap] = data_array.coords[ - "time" - ].values - exp["annual"]["year"] = [x.year for x in years] - del dataset_wrapper - return new_var_list - - -def process_data( - parameters: Parameters, requested_variables: RequestedVariables, rgn: str -) -> List[Dict[str, Any]]: - exps: List[Dict[str, Any]] = get_exps(parameters) - valid_vars: List[str] = [] - invalid_vars: List[str] = [] - exp: Dict[str, Any] - for exp in exps: - exp["annual"] = {} - - requested_variables.vars_original = set_var( - exp, - "atmos", - requested_variables.vars_original, - valid_vars, - invalid_vars, - rgn, - ) - requested_variables.vars_atm = set_var( - exp, "atmos", requested_variables.vars_atm, valid_vars, invalid_vars, rgn - ) - requested_variables.vars_ice = set_var( - exp, "ice", requested_variables.vars_ice, valid_vars, invalid_vars, rgn - ) - requested_variables.vars_land = set_var( - exp, - "land", - requested_variables.vars_land, - valid_vars, - invalid_vars, - rgn, - ) - requested_variables.vars_ocn = set_var( - exp, "ocean", requested_variables.vars_ocn, valid_vars, invalid_vars, rgn - ) - - # Optionally read ohc - if exp["ocean"] != "": - dataset_wrapper = DatasetWrapper(exp["ocean"]) - exp["annual"]["ohc"], _ = dataset_wrapper.globalAnnual(Variable("ohc")) - # anomalies with respect to first year - exp["annual"]["ohc"][:] = exp["annual"]["ohc"][:] - exp["annual"]["ohc"][0] - - if exp["vol"] != "": - dataset_wrapper = DatasetWrapper(exp["vol"]) - exp["annual"]["volume"], _ = dataset_wrapper.globalAnnual( - Variable("volume") - ) - # annomalies with respect to first year - exp["annual"]["volume"][:] = ( - exp["annual"]["volume"][:] - exp["annual"]["volume"][0] - ) - - logger.info( - f"{rgn} region globalAnnual was computed successfully for these variables: {valid_vars}" - ) - logger.error( - f"{rgn} region globalAnnual could not be computed for these variables: {invalid_vars}" - ) - return exps - - -# Run coupled_global ########################################################## -def run(parameters: Parameters, requested_variables: RequestedVariables, rgn: str): - # Experiments - exps: List[Dict[str, Any]] = process_data(parameters, requested_variables, rgn) - - xlim: List[float] = [float(parameters.year1), float(parameters.year2)] - - valid_plots: List[str] = [] - invalid_plots: List[str] = [] - - # Use list of tuples rather than a dict, to keep order - # Note: we use `parameters.plots_original` rather than `requested_variables.vars_original` - # because the "original" plots are expecting plot names that are not variable names. - # The model components however are expecting plot names to be variable names. - mapping: List[Tuple[str, List[str]]] = [ - ("original", parameters.plots_original), - ("atm", list(map(lambda v: v.variable_name, requested_variables.vars_atm))), - ("ice", list(map(lambda v: v.variable_name, requested_variables.vars_ice))), - ("lnd", list(map(lambda v: v.variable_name, requested_variables.vars_land))), - ("ocn", list(map(lambda v: v.variable_name, requested_variables.vars_ocn))), - ] - for component, plot_list in mapping: - make_plot_pdfs( - parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, - ) - logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") - logger.error( - f"These {rgn} region plots could not be generated successfully: {invalid_plots}" - ) - - -def get_vars(requested_variables: RequestedVariables, component: str) -> List[Variable]: - vars: List[Variable] - if component == "original": - vars = requested_variables.vars_original - elif component == "atm": - vars = requested_variables.vars_atm - elif component == "ice": - vars = requested_variables.vars_ice - elif component == "lnd": - vars = requested_variables.vars_land - elif component == "ocn": - vars = requested_variables.vars_ocn - else: - raise ValueError(f"Invalid component={component}") - return vars - - -def coupled_global(parameters: Parameters) -> None: - requested_variables = RequestedVariables(parameters) - for rgn in parameters.regions: - run(parameters, requested_variables, rgn) - if parameters.make_viewer: - # In this case, we don't want the summary PDF. - # Rather, we want to construct a viewer similar to E3SM Diags. - title_and_url_list: List[Tuple[str, str]] = [] - for component in [ - "atm", - "ice", - "lnd", - "ocn", - ]: # Don't create viewer for original component - vars = get_vars(requested_variables, component) - if vars: - url = create_viewer(parameters, vars, component) - logger.info(f"Viewer URL for {component}: {url}") - title_and_url_list.append((component, url)) - # Special case for original plots: always use user-provided dimensions. - vars = get_vars(requested_variables, "original") - if vars: - logger.info("Using user provided dimensions for original plots PDF") - title_and_url_list.append( - ( - "original", - f"{parameters.figstr}_glb_original.pdf", - ) - ) - - index_url: str = create_viewer_index(parameters.results_dir, title_and_url_list) - logger.info(f"Viewer index URL: {index_url}") diff --git a/zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py b/zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py deleted file mode 100644 index e9b3cce..0000000 --- a/zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py +++ /dev/null @@ -1,192 +0,0 @@ -from typing import Tuple - -import xarray -import xcdat - -from zppy_interfaces.global_time_series.coupled_global_utils import Metric, Variable -from zppy_interfaces.multi_utils.logger import _setup_child_logger - -logger = _setup_child_logger(__name__) - - -class DatasetWrapper(object): - def __init__(self, directory): - - self.directory: str = directory - - # `directory` will be of the form `{case_dir}/post//glb/ts/monthly/{ts_num_years_str}yr/` - self.dataset: xarray.core.dataset.Dataset = xcdat.open_mfdataset( - f"{directory}*.nc", center_times=True - ) - - self.area_tuple = None - - def set_area_tuple(self): - keys = list(self.dataset.keys()) - if "valid_area_per_gridcell" in keys: - logger.debug("Setting area_tuple, using valid_area_per_gridcell") - land_area_per_gridcell = self.dataset["valid_area_per_gridcell"] - # land_area_per_gridcell.shape = (360, 720) - logger.debug(f"land_area_per_gridcell.shape={land_area_per_gridcell.shape}") - total_land_area = land_area_per_gridcell.sum() # Sum over all dimensions - # Account for hemispheric plots: - north_land_area = land_area_per_gridcell.where( - land_area_per_gridcell.lat >= 0 - ).sum() - south_land_area = land_area_per_gridcell.where( - land_area_per_gridcell.lat < 0 - ).sum() - else: - logger.debug("Setting area_tuple, using area and landfrac") - area: xarray.core.dataarray.DataArray = self.dataset["area"] - landfrac: xarray.core.dataarray.DataArray = self.dataset["landfrac"] - - # area.shape = (180, 360) - logger.debug(f"area.shape={area.shape}") - # landfrac.shape = (180, 360) - logger.debug(f"landfrac.shape={landfrac.shape}") - - total_land_area = (area * landfrac).sum() # Sum over all dimensions - - # Account for hemispheric plots: - north_area = area.where(area.lat >= 0) - north_landfrac = landfrac.where(landfrac.lat >= 0) - north_land_area = (north_area * north_landfrac).sum() - - south_area = area.where(area.lat < 0) - south_landfrac = landfrac.where(landfrac.lat < 0) - south_land_area = (south_area * south_landfrac).sum() - - logger.debug(f"total_land_area.shape={total_land_area.shape}") - logger.debug(f"north_land_area.shape={north_land_area.shape}") - logger.debug(f"south_land_area.shape={south_land_area.shape}") - - # logger.debug(f"total_land_area={total_land_area.item()}") - # logger.debug(f"north_land_area={north_land_area.item()}") - # logger.debug(f"south_land_area={south_land_area.item()}") - - self.area_tuple = (total_land_area, north_land_area, south_land_area) - # logger.debug(f"For Metric.TOTAL, data_array's glb,n,s will be scaled respectively by {self.area_tuple}") - - def __del__(self): - - self.dataset.close() - - def globalAnnualHelper( - self, - var: str, - metric: Metric, - scale_factor: float, - original_units: str, - final_units: str, - ) -> Tuple[xarray.core.dataarray.DataArray, str]: - - data_array: xarray.core.dataarray.DataArray - units: str = "" - - # Constants, from AMWG diagnostics - Lv = 2.501e6 - Lf = 3.337e5 - - # Is this a derived variable? - if var == "RESTOM": - FSNT, _ = self.globalAnnualHelper( - "FSNT", metric, scale_factor, original_units, final_units - ) - FLNT, _ = self.globalAnnualHelper( - "FLNT", metric, scale_factor, original_units, final_units - ) - data_array = FSNT - FLNT - elif var == "RESTOA": - logger.warning("NOT READY") - FSNTOA, _ = self.globalAnnualHelper( - "FSNTOA", metric, scale_factor, original_units, final_units - ) - FLUT, _ = self.globalAnnualHelper( - "FLUT", metric, scale_factor, original_units, final_units - ) - data_array = FSNTOA - FLUT - elif var == "LHFLX": - QFLX, _ = self.globalAnnualHelper( - "QFLX", metric, scale_factor, original_units, final_units - ) - PRECC, _ = self.globalAnnualHelper( - "PRECC", metric, scale_factor, original_units, final_units - ) - PRECL, _ = self.globalAnnualHelper( - "PRECL", metric, scale_factor, original_units, final_units - ) - PRECSC, _ = self.globalAnnualHelper( - "PRECSC", metric, scale_factor, original_units, final_units - ) - PRECSL, _ = self.globalAnnualHelper( - "PRECSL", metric, scale_factor, original_units, final_units - ) - data_array = (Lv + Lf) * QFLX - Lf * 1.0e3 * ( - PRECC + PRECL - PRECSC - PRECSL - ) - elif var == "RESSURF": - FSNS, _ = self.globalAnnualHelper( - "FSNS", metric, scale_factor, original_units, final_units - ) - FLNS, _ = self.globalAnnualHelper( - "FLNS", metric, scale_factor, original_units, final_units - ) - SHFLX, _ = self.globalAnnualHelper( - "SHFLX", metric, scale_factor, original_units, final_units - ) - LHFLX, _ = self.globalAnnualHelper( - "LHFLX", metric, scale_factor, original_units, final_units - ) - data_array = FSNS - FLNS - SHFLX - LHFLX - elif var == "PREC": - PRECC, _ = self.globalAnnualHelper( - "PRECC", metric, scale_factor, original_units, final_units - ) - PRECL, _ = self.globalAnnualHelper( - "PRECL", metric, scale_factor, original_units, final_units - ) - data_array = 1.0e3 * (PRECC + PRECL) - else: - # Non-derived variables - annual_average_dataset_for_var: xarray.core.dataset.Dataset = ( - self.dataset.temporal.group_average(var, "year") - ) - data_array = annual_average_dataset_for_var.data_vars[var] - if metric == Metric.TOTAL: - if not self.area_tuple: - self.set_area_tuple() - # Appease the type checker (avoid `Value of type "Optional[Any]" is not indexable`) - if not self.area_tuple: - raise ValueError("area_tuple still not set") - # data_array.shape = (number of years, number of regions) - # We want to keep those dimensions, but with these values: - # (glb*total_land_area, n*north_land_area, s*south_land_area) - try: - data_array[:, 0] *= self.area_tuple[0] - data_array[:, 1] *= self.area_tuple[1] - data_array[:, 2] *= self.area_tuple[2] - except Exception as e: - logger.error(f"Error while scaling data_array: {e}") - raise e - units = data_array.units - # `units` will be "1" if it's a dimensionless quantity - if (units != "1") and (original_units != "") and original_units != units: - raise ValueError( - f"Units don't match up: Have {units} but expected {original_units}. This renders the supplied scale_factor ({scale_factor}) unusable." - ) - if (scale_factor != 1) and (final_units != ""): - data_array *= scale_factor - units = final_units - return data_array, units - - def globalAnnual( - self, var: Variable - ) -> Tuple[xarray.core.dataarray.DataArray, str]: - return self.globalAnnualHelper( - var.variable_name, - var.metric, - var.scale_factor, - var.original_units, - var.final_units, - ) diff --git a/zppy_interfaces/global_time_series/coupled_global_plotting.py b/zppy_interfaces/global_time_series/coupled_global_plotting.py index 76516d4..dec4810 100644 --- a/zppy_interfaces/global_time_series/coupled_global_plotting.py +++ b/zppy_interfaces/global_time_series/coupled_global_plotting.py @@ -1,15 +1,10 @@ import glob import math -import os -import traceback import matplotlib as mpl -import matplotlib.backends.backend_pdf -import matplotlib.pyplot as plt import numpy as np from netCDF4 import Dataset -from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger mpl.use("Agg") @@ -17,6 +12,97 @@ logger = _setup_child_logger(__name__) +# FIXME: C901 'plot' is too complex (19) +def plot(ax, xlim, exps, param_dict, rgn): # noqa: C901 + if param_dict["glb_only"] and (rgn != "glb"): + return + ax.set_xlim(xlim) + extreme_values = [] + for exp in exps: + # Relevant to "Plot 5: plot_change_ohc" + if param_dict["check_exp_ocean"] and (exp["ocean"] == ""): + continue + # Relevant to "Plot 7: plot_change_sea_level" + # This must be checked before plot 6, + # otherwise, `param_dict["var"]` will be run, + # but `exp["annual"]["volume"]` won't exist. + if param_dict["check_exp_vol"] and (exp["vol"] == ""): + continue + # Relevant to "Plot 6: plot_max_moc" + if param_dict["use_getmoc"]: + if exp["moc"]: + [year, var] = getmoc(exp["moc"]) + else: + continue + else: + year = np.array(exp["annual"]["year"]) + exp["yoffset"] + var = param_dict["var"](exp) + extreme_values.append(np.amax(var)) + extreme_values.append(np.amin(var)) + if param_dict["shorten_year"]: + year = year[: len(var)] + try: + ax.plot( + year, + var, + lw=param_dict["lw"], + marker=None, + c=exp["color"], + label=exp["name"], + ) + except Exception: + raise RuntimeError(f"{param_dict['title']} could not be plotted.") + if param_dict["2nd_var"]: + # Specifically for plot_toa_radiation + # TODO: if more plots require a 2nd variable, we can change `var` to be a list, + # but that will be a more significant refactoring. + var = np.array(exp["annual"]["FLUT"][rgn][0]) + ax.plot(year, var, lw=1.0, marker=None, ls=":", c=exp["color"]) + continue + if param_dict["check_exp_year"] and exp["yr"] is None: + continue + elif param_dict["do_add_line"] or param_dict["do_add_trend"]: + for yrs in exp["yr"]: + if param_dict["do_add_line"]: + add_line( + year, + var, + yrs[0], + yrs[1], + format=param_dict["format"], + ax=ax, + lw=2 * param_dict["lw"], + color=exp["color"], + ) + if param_dict["do_add_trend"]: + add_trend( + year, + var, + yrs[0], + yrs[1], + format=param_dict["format"], + ax=ax, + lw=2 * param_dict["lw"], + color=exp["color"], + ohc=param_dict["ohc"], + verbose=param_dict["verbose"], + vol=param_dict["vol"], + ) + ylim = get_ylim(param_dict["default_ylim"], extreme_values) + ax.set_ylim(ylim) + if param_dict["set_axhline"]: + ax.axhline(y=param_dict["axhline_y"], lw=1, c="0.5") + ax.set_title(param_dict["title"]) + ax.set_xlabel("Year") + units = param_dict["ylabel"] + c = callable(units) + if c: + units = units(exps[0]) + ax.set_ylabel(units) + if param_dict["set_legend"]: + ax.legend(loc="best") + + # ---additional function to get moc time series def getmoc(dir_in): files = sorted(glob.glob(dir_in + "mocTimeSeries*.nc")) @@ -45,7 +131,6 @@ def getmoc(dir_in): return time, var -# ----------------------------------------------------------------------------- # Function to add horizontal line showing average value over a specified period def add_line(year, var, year1, year2, ax, format="%4.2f", lw=1, color="b"): @@ -59,7 +144,6 @@ def add_line(year, var, year1, year2, ax, format="%4.2f", lw=1, color="b"): return -# ----------------------------------------------------------------------------- # Function to add line showing linear trend over a specified period def add_trend( year, @@ -107,7 +191,6 @@ def add_trend( return -# ----------------------------------------------------------------------------- # Function to get ylim def get_ylim(standard_range, extreme_values): if len(extreme_values) > 0: @@ -146,513 +229,3 @@ def get_ylim(standard_range, extreme_values): else: raise ValueError("Not enough range information supplied") return [ylim_min, ylim_max] - - -# ----------------------------------------------------------------------------- -# Plotting functions - - -# 1 -def plot_net_toa_flux_restom(ax, xlim, exps, rgn): - logger.info("Plot 1: plot_net_toa_flux_restom") - param_dict = { - "2nd_var": False, - "axhline_y": 0, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": True, - "default_ylim": [-1.5, 1.5], - "do_add_line": True, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": True, - "set_legend": True, - "shorten_year": False, - "title": "Net TOA flux (restom)", - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"]["RESTOM"][0]), - "verbose": False, - "vol": False, - "ylabel": "W m-2", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 2 -def plot_global_surface_air_temperature(ax, xlim, exps, rgn): - logger.info("Plot 2: plot_global_surface_air_temperature") - if rgn == "glb": - region_title = "Global" - elif rgn == "n": - region_title = "Northern Hemisphere" - elif rgn == "s": - region_title = "Southern Hemisphere" - else: - raise RuntimeError(f"Invalid rgn={rgn}") - param_dict = { - "2nd_var": False, - "axhline_y": None, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": True, - "default_ylim": [13, 15.5], - "do_add_line": True, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": False, - "title": f"{region_title} surface air temperature", - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"]["TREFHT"][0]) - 273.15, - "verbose": False, - "vol": False, - "ylabel": "degC", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 3 -def plot_toa_radiation(ax, xlim, exps, rgn): - logger.info("Plot 3: plot_toa_radiation") - param_dict = { - "2nd_var": True, - "axhline_y": None, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": False, - "default_ylim": [235, 245], - "do_add_line": False, - "do_add_trend": False, - "format": None, - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": False, - "shorten_year": False, - "title": "TOA radiation: SW (solid), LW (dashed)", - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"]["FSNTOA"][0]), - "verbose": None, - "vol": None, - "ylabel": "W m-2", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 4 -def plot_net_atm_energy_imbalance(ax, xlim, exps, rgn): - logger.info("Plot 4: plot_net_atm_energy_imbalance") - param_dict = { - "2nd_var": False, - "axhline_y": None, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": True, - "default_ylim": [-0.3, 0.3], - "do_add_line": True, - "do_add_trend": False, - "format": "%4.2f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": False, - "title": "Net atm energy imbalance (restom-ressurf)", - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"]["RESTOM"][0]) - - np.array(exp["annual"]["RESSURF"][0]), - "verbose": False, - "vol": False, - "ylabel": "W m-2", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 5 -def plot_change_ohc(ax, xlim, exps, rgn): - logger.info("Plot 5: plot_change_ohc") - param_dict = { - "2nd_var": False, - "axhline_y": 0, - "check_exp_ocean": True, - "check_exp_vol": False, - "check_exp_year": False, - "default_ylim": [-0.3e24, 0.9e24], - "do_add_line": False, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": True, - "lw": 1.5, - "ohc": True, - "set_axhline": True, - "set_legend": True, - "shorten_year": True, - "title": "Change in ocean heat content", - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"]["ohc"]), - "verbose": False, - "vol": False, - "ylabel": "J", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 6 -def plot_max_moc(ax, xlim, exps, rgn): - logger.info("Plot 6: plot_max_moc") - param_dict = { - "2nd_var": False, - "axhline_y": 10, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": False, - "default_ylim": [4, 22], - "do_add_line": False, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": True, - "lw": 1.5, - "ohc": False, - "set_axhline": True, - "set_legend": True, - "shorten_year": False, - "title": "Max MOC Atlantic streamfunction at 26.5N", - "use_getmoc": True, - "var": None, - "verbose": True, - "vol": None, - "ylabel": "Sv", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 7 -def plot_change_sea_level(ax, xlim, exps, rgn): - logger.info("Plot 7: plot_change_sea_level") - param_dict = { - "2nd_var": False, - "axhline_y": None, - "check_exp_ocean": False, - "check_exp_vol": True, - "check_exp_year": True, - "default_ylim": [4, 22], - "do_add_line": False, - "do_add_trend": True, - "format": "%5.3f", - "glb_only": True, - "lw": 1.5, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": True, - "title": "Change in sea level", - "use_getmoc": False, - "var": lambda exp: ( - 1e3 - * np.array(exp["annual"]["volume"]) - / (4.0 * math.pi * (6371229.0) ** 2 * 0.7) - ), - "verbose": True, - "vol": True, - "ylabel": "mm", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# 8 -def plot_net_atm_water_imbalance(ax, xlim, exps, rgn): - logger.info("Plot 8: plot_net_atm_water_imbalance") - param_dict = { - "2nd_var": False, - "axhline_y": None, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": False, - "default_ylim": [-1, 1], - "do_add_line": True, - "do_add_trend": False, - "format": "%5.4f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": False, - "title": "Net atm water imbalance (evap-prec)", - "use_getmoc": False, - "var": lambda exp: ( - 365 - * 86400 - * ( - np.array(exp["annual"]["QFLX"][0]) - - 1e3 - * ( - np.array(exp["annual"]["PRECC"][0]) - + np.array(exp["annual"]["PRECL"][0]) - ) - ) - ), - "verbose": False, - "vol": False, - "ylabel": "mm yr-1", - } - plot(ax, xlim, exps, param_dict, rgn) - - -# Generic plot function -def plot_generic(ax, xlim, exps, var_name, rgn): - logger.info(f"plot_generic for {var_name}, rgn={rgn}") - param_dict = { - "2nd_var": False, - "axhline_y": 0, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": True, - "default_ylim": [], - "do_add_line": True, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": False, - "title": var_name, - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"][var_name][0]), - "verbose": False, - "vol": False, - "ylabel": lambda exp: np.array(exp["annual"][var_name][1]), - } - plot(ax, xlim, exps, param_dict, rgn) - - -# FIXME: C901 'plot' is too complex (19) -def plot(ax, xlim, exps, param_dict, rgn): # noqa: C901 - if param_dict["glb_only"] and (rgn != "glb"): - return - ax.set_xlim(xlim) - extreme_values = [] - for exp in exps: - # Relevant to "Plot 5: plot_change_ohc" - if param_dict["check_exp_ocean"] and (exp["ocean"] == ""): - continue - # Relevant to "Plot 7: plot_change_sea_level" - # This must be checked before plot 6, - # otherwise, `param_dict["var"]` will be run, - # but `exp["annual"]["volume"]` won't exist. - if param_dict["check_exp_vol"] and (exp["vol"] == ""): - continue - # Relevant to "Plot 6: plot_max_moc" - if param_dict["use_getmoc"]: - if exp["moc"]: - [year, var] = getmoc(exp["moc"]) - else: - continue - else: - year = np.array(exp["annual"]["year"]) + exp["yoffset"] - var = param_dict["var"](exp) - extreme_values.append(np.amax(var)) - extreme_values.append(np.amin(var)) - if param_dict["shorten_year"]: - year = year[: len(var)] - try: - ax.plot( - year, - var, - lw=param_dict["lw"], - marker=None, - c=exp["color"], - label=exp["name"], - ) - except Exception: - raise RuntimeError(f"{param_dict['title']} could not be plotted.") - if param_dict["2nd_var"]: - # Specifically for plot_toa_radiation - # TODO: if more plots require a 2nd variable, we can change `var` to be a list, - # but that will be a more significant refactoring. - var = np.array(exp["annual"]["FLUT"][0]) - ax.plot(year, var, lw=1.0, marker=None, ls=":", c=exp["color"]) - continue - if param_dict["check_exp_year"] and exp["yr"] is None: - continue - elif param_dict["do_add_line"] or param_dict["do_add_trend"]: - for yrs in exp["yr"]: - if param_dict["do_add_line"]: - add_line( - year, - var, - yrs[0], - yrs[1], - format=param_dict["format"], - ax=ax, - lw=2 * param_dict["lw"], - color=exp["color"], - ) - if param_dict["do_add_trend"]: - add_trend( - year, - var, - yrs[0], - yrs[1], - format=param_dict["format"], - ax=ax, - lw=2 * param_dict["lw"], - color=exp["color"], - ohc=param_dict["ohc"], - verbose=param_dict["verbose"], - vol=param_dict["vol"], - ) - ylim = get_ylim(param_dict["default_ylim"], extreme_values) - ax.set_ylim(ylim) - if param_dict["set_axhline"]: - ax.axhline(y=param_dict["axhline_y"], lw=1, c="0.5") - ax.set_title(param_dict["title"]) - ax.set_xlabel("Year") - units = param_dict["ylabel"] - c = callable(units) - if c: - units = units(exps[0]) - ax.set_ylabel(units) - if param_dict["set_legend"]: - ax.legend(loc="best") - - -PLOT_DICT = { - "net_toa_flux_restom": plot_net_toa_flux_restom, - "global_surface_air_temperature": plot_global_surface_air_temperature, - "toa_radiation": plot_toa_radiation, - "net_atm_energy_imbalance": plot_net_atm_energy_imbalance, - "change_ohc": plot_change_ohc, # only glb - "max_moc": plot_max_moc, # only glb - "change_sea_level": plot_change_sea_level, # only glb - "net_atm_water_imbalance": plot_net_atm_water_imbalance, -} - - -# FIXME: C901 'make_plot_pdfs' is too complex (20) -def make_plot_pdfs( # noqa: C901 - parameters: Parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, -): - logger.info(f"make_plot_pdfs for rgn={rgn}, component={component}") - num_plots = len(plot_list) - if num_plots == 0: - return - - # If make_viewer, then we want to do 1 plot per page. - # However, the original plots are excluded from this restriction. - # Note: if the user provides nrows=ncols=1, there will still be a single plot per page - keep_user_dims = (not parameters.make_viewer) or (component == "original") - if keep_user_dims: - nrows = parameters.nrows - ncols = parameters.ncols - else: - nrows = 1 - ncols = 1 - - plots_per_page = nrows * ncols - num_pages = math.ceil(num_plots / plots_per_page) - - counter = 0 - os.makedirs(parameters.results_dir, exist_ok=True) - # https://stackoverflow.com/questions/58738992/save-multiple-figures-with-subplots-into-a-pdf-with-multiple-pages - pdf = matplotlib.backends.backend_pdf.PdfPages( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.pdf" - ) - for page in range(num_pages): - if plots_per_page == 1: - logger.info("Using reduced figsize") - fig = plt.figure(1, figsize=[13.5 / 2, 16.5 / 4]) - else: - logger.info("Using standard figsize") - fig = plt.figure(1, figsize=[13.5, 16.5]) - logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") - fig.suptitle(f"{parameters.figstr}_{rgn}_{component}") - for j in range(plots_per_page): - logger.info( - f"Plotting plot {j} on page {page}. This is plot {counter} in total." - ) - # The final page doesn't need to be filled out with plots. - if counter >= num_plots: - break - ax = plt.subplot( - nrows, - ncols, - j + 1, - ) - plot_name = plot_list[counter] - if component == "original": - try: - plot_function = PLOT_DICT[plot_name] - except KeyError: - raise KeyError(f"Invalid plot name: {plot_name}") - try: - plot_function(ax, xlim, exps, rgn) - valid_plots.append(plot_name) - except Exception: - traceback.print_exc() - required_vars = [] - if plot_name == "net_toa_flux_restom": - required_vars = ["RESTOM"] - elif plot_name == "net_atm_energy_imbalance": - required_vars = ["RESTOM", "RESSURF"] - elif plot_name == "global_surface_air_temperature": - required_vars = ["TREFHT"] - elif plot_name == "toa_radiation": - required_vars = ["FSNTOA", "FLUT"] - elif plot_name == "net_atm_water_imbalance": - required_vars = ["PRECC", "PRECL", "QFLX"] - logger.error( - f"Failed plot_function for {plot_name}. Check that {required_vars} are available." - ) - invalid_plots.append(plot_name) - counter += 1 - else: - try: - plot_generic(ax, xlim, exps, plot_name, rgn) - valid_plots.append(plot_name) - except Exception: - traceback.print_exc() - logger.error( - f"plot_generic failed. Invalid plot={plot_name}, rgn={rgn}" - ) - invalid_plots.append(plot_name) - counter += 1 - - fig.tight_layout() - pdf.savefig(1) - # Always save individual PNGs for viewer mode - if plots_per_page == 1: - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{plot_name}.png", - dpi=150, - ) - elif num_pages > 1: - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{page}.png", - dpi=150, - ) - else: - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.png", - dpi=150, - ) - plt.close(fig) - pdf.close() diff --git a/zppy_interfaces/global_time_series/coupled_global_utils.py b/zppy_interfaces/global_time_series/coupled_global_utils.py index e32d496..c8587ff 100644 --- a/zppy_interfaces/global_time_series/coupled_global_utils.py +++ b/zppy_interfaces/global_time_series/coupled_global_utils.py @@ -1,4 +1,16 @@ +import os.path from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +import cftime +import numpy as np +import xarray +import xcdat + +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) class Metric(Enum): @@ -35,3 +47,288 @@ def __init__( self.group: str = group # Descriptive text to add to the plot page to help users identify the variable self.long_name: str = long_name + + +class DatasetWrapper(object): + def __init__( + self, + directory: str, + var_list: Optional[List[Variable]] = None, + parameters: Optional[Parameters] = None, + ): + + self.directory: str = directory + self.var_list: Optional[List[Variable]] = var_list + + # `directory` will be of the form `{case_dir}/post//glb/ts/monthly/{ts_num_years_str}yr/` + self.dataset: xarray.core.dataset.Dataset + if var_list and parameters: + file_path_list: List[str] = [] + for var in var_list: + num_years: int = int(parameters.ts_num_years_str) + y1: int = parameters.year1 + y2: int = parameters.year1 + num_years - 1 + while y2 <= parameters.year2: + # `var.variable_name` will be of the form `FSNS`, `FLNS`, etc. + file_path: str = ( + f"{directory}{var.variable_name}_{y1:04d}01_{y2:04d}12.nc" + ) + if os.path.exists(file_path): + file_path_list.append(file_path) + else: + logger.info(f"{file_path} does not exist.") + y1 += num_years + y2 += num_years + self.dataset = xcdat.open_mfdataset(file_path_list, center_times=True) + else: + self.dataset = xcdat.open_mfdataset(f"{directory}*.nc", center_times=True) + self.area_tuple: Optional[Tuple[Any, Any, Any]] = None + + def set_area_tuple(self): + keys = list(self.dataset.keys()) + if "valid_area_per_gridcell" in keys: + logger.debug("Setting area_tuple, using valid_area_per_gridcell") + land_area_per_gridcell = self.dataset["valid_area_per_gridcell"] + # land_area_per_gridcell.shape = (360, 720) + logger.debug(f"land_area_per_gridcell.shape={land_area_per_gridcell.shape}") + total_land_area = land_area_per_gridcell.sum() # Sum over all dimensions + # Account for hemispheric plots: + north_land_area = land_area_per_gridcell.where( + land_area_per_gridcell.lat >= 0 + ).sum() + south_land_area = land_area_per_gridcell.where( + land_area_per_gridcell.lat < 0 + ).sum() + else: + logger.debug("Setting area_tuple, using area and landfrac") + area: xarray.core.dataarray.DataArray = self.dataset["area"] + landfrac: xarray.core.dataarray.DataArray = self.dataset["landfrac"] + + # area.shape = (180, 360) + logger.debug(f"area.shape={area.shape}") + # landfrac.shape = (180, 360) + logger.debug(f"landfrac.shape={landfrac.shape}") + + total_land_area = (area * landfrac).sum() # Sum over all dimensions + + # Account for hemispheric plots: + north_area = area.where(area.lat >= 0) + north_landfrac = landfrac.where(landfrac.lat >= 0) + north_land_area = (north_area * north_landfrac).sum() + + south_area = area.where(area.lat < 0) + south_landfrac = landfrac.where(landfrac.lat < 0) + south_land_area = (south_area * south_landfrac).sum() + + logger.debug(f"total_land_area.shape={total_land_area.shape}") + logger.debug(f"north_land_area.shape={north_land_area.shape}") + logger.debug(f"south_land_area.shape={south_land_area.shape}") + + # logger.debug(f"total_land_area={total_land_area.item()}") + # logger.debug(f"north_land_area={north_land_area.item()}") + # logger.debug(f"south_land_area={south_land_area.item()}") + + self.area_tuple = (total_land_area, north_land_area, south_land_area) + # logger.debug(f"For Metric.TOTAL, data_array's glb,n,s will be scaled respectively by {self.area_tuple}") + + def __del__(self): + + try: + self.dataset.close() + except AttributeError: + raise AttributeError( + "DatasetWrapper.dataset was not set. This could be because of a failure producing the dataset." + ) + + def globalAnnualHelper( + self, + var: str, + metric: Metric, + scale_factor: float, + original_units: str, + final_units: str, + ) -> Tuple[xarray.core.dataarray.DataArray, str]: + + data_array: xarray.core.dataarray.DataArray + units: str = "" + + # Constants, from AMWG diagnostics + Lv = 2.501e6 + Lf = 3.337e5 + + if (not self.var_list) and ( + var in ["RESTOM", "RESTOA", "LHFLX", "RESSURF", "PREC"] + ): + # We've loaded ALL variables. + # That means we can attempt derivations from other variables + # not explicitally requested. + if var == "RESTOM": + FSNT, _ = self.globalAnnualHelper( + "FSNT", metric, scale_factor, original_units, final_units + ) + FLNT, _ = self.globalAnnualHelper( + "FLNT", metric, scale_factor, original_units, final_units + ) + data_array = FSNT - FLNT + elif var == "RESTOA": + logger.warning("NOT READY") + FSNTOA, _ = self.globalAnnualHelper( + "FSNTOA", metric, scale_factor, original_units, final_units + ) + FLUT, _ = self.globalAnnualHelper( + "FLUT", metric, scale_factor, original_units, final_units + ) + data_array = FSNTOA - FLUT + elif var == "LHFLX": + QFLX, _ = self.globalAnnualHelper( + "QFLX", metric, scale_factor, original_units, final_units + ) + PRECC, _ = self.globalAnnualHelper( + "PRECC", metric, scale_factor, original_units, final_units + ) + PRECL, _ = self.globalAnnualHelper( + "PRECL", metric, scale_factor, original_units, final_units + ) + PRECSC, _ = self.globalAnnualHelper( + "PRECSC", metric, scale_factor, original_units, final_units + ) + PRECSL, _ = self.globalAnnualHelper( + "PRECSL", metric, scale_factor, original_units, final_units + ) + data_array = (Lv + Lf) * QFLX - Lf * 1.0e3 * ( + PRECC + PRECL - PRECSC - PRECSL + ) + elif var == "RESSURF": + FSNS, _ = self.globalAnnualHelper( + "FSNS", metric, scale_factor, original_units, final_units + ) + FLNS, _ = self.globalAnnualHelper( + "FLNS", metric, scale_factor, original_units, final_units + ) + SHFLX, _ = self.globalAnnualHelper( + "SHFLX", metric, scale_factor, original_units, final_units + ) + LHFLX, _ = self.globalAnnualHelper( + "LHFLX", metric, scale_factor, original_units, final_units + ) + data_array = FSNS - FLNS - SHFLX - LHFLX + elif var == "PREC": + PRECC, _ = self.globalAnnualHelper( + "PRECC", metric, scale_factor, original_units, final_units + ) + PRECL, _ = self.globalAnnualHelper( + "PRECL", metric, scale_factor, original_units, final_units + ) + data_array = 1.0e3 * (PRECC + PRECL) + else: + raise ValueError(f"Invalid var={var}") + else: + # Non-derived variables + annual_average_dataset_for_var: xarray.core.dataset.Dataset = ( + self.dataset.temporal.group_average(var, "year") + ) + data_array = annual_average_dataset_for_var.data_vars[var] + if metric == Metric.TOTAL: + if not self.area_tuple: + self.set_area_tuple() + # Appease the type checker (avoid `Value of type "Optional[Any]" is not indexable`) + if not self.area_tuple: + raise ValueError("area_tuple still not set") + # data_array.shape = (number of years, number of regions) + # We want to keep those dimensions, but with these values: + # (glb*total_land_area, n*north_land_area, s*south_land_area) + try: + data_array[:, 0] *= self.area_tuple[0] + data_array[:, 1] *= self.area_tuple[1] + data_array[:, 2] *= self.area_tuple[2] + except Exception as e: + logger.error(f"Error while scaling data_array: {e}") + raise e + units = data_array.units + # `units` will be "1" if it's a dimensionless quantity + if (units != "1") and (original_units != "") and original_units != units: + raise ValueError( + f"Units don't match up: Have {units} but expected {original_units}. This renders the supplied scale_factor ({scale_factor}) unusable." + ) + if (scale_factor != 1) and (final_units != ""): + data_array *= scale_factor + units = final_units + return data_array, units + + def globalAnnual( + self, var: Variable + ) -> Tuple[xarray.core.dataarray.DataArray, str]: + return self.globalAnnualHelper( + var.variable_name, + var.metric, + var.scale_factor, + var.original_units, + var.final_units, + ) + + +# Helper functions ############################################################ +def get_data_dir(parameters: Parameters, component: str, conditional: bool) -> str: + return ( + f"{parameters.case_dir}/post/{component}/glb/ts/monthly/{parameters.ts_num_years_str}yr/" + if conditional + else "" + ) + + +def set_var( + exp: Dict[str, Any], + exp_key: str, + var_list: List[Variable], + valid_vars: List[str], + invalid_vars: List[str], + parameters: Optional[Parameters] = None, +) -> List[Variable]: + new_var_list: List[Variable] = [] + if parameters and (var_list == []): + # If we want to load specific variables, + # but none are specified, + # then we can just immediately return. + return new_var_list + if exp[exp_key] != "": + try: + dataset_wrapper: DatasetWrapper + if parameters: + # If this is passed in, then we want to load specific vars. + dataset_wrapper = DatasetWrapper(exp[exp_key], var_list, parameters) + else: + dataset_wrapper = DatasetWrapper(exp[exp_key]) + except Exception as e: + logger.critical(e) + logger.critical( + f"DatasetWrapper object could not be created for {exp_key}={exp[exp_key]}" + ) + raise e + for var in var_list: + var_str: str = var.variable_name + try: + data_array: xarray.core.dataarray.DataArray + units: str + data_array, units = dataset_wrapper.globalAnnual(var) + valid_vars.append(str(var_str)) # Append the name + new_var_list.append(var) # Append the variable itself + except Exception as e: + logger.error(e) + logger.error(f"globalAnnual failed for {var_str}") + invalid_vars.append(str(var_str)) + continue + exp["annual"][var_str] = {"glb": (data_array.isel(rgn=0), units)} + if data_array.sizes["rgn"] > 1: + # data_array.shape => number of years x 3 regions + # 3 regions = global, northern hemisphere, southern hemisphere + # We get here if we used the updated `ts` task + # (using `rgn_avg` rather than `glb_avg`). + exp["annual"][var_str]["n"] = (data_array.isel(rgn=1), units) + exp["annual"][var_str]["s"] = (data_array.isel(rgn=2), units) + if "year" not in exp["annual"]: + years: np.ndarray[cftime.DatetimeNoLeap] = data_array.coords[ + "time" + ].values + exp["annual"]["year"] = [x.year for x in years] + del dataset_wrapper + return new_var_list diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index 105f4fe..b792ff5 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -1,11 +1,35 @@ from typing import Dict, List +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + # Parameters ################################################################## class Parameters(object): def __init__(self, args: Dict[str, str]): + # Used by both Classic PDF and Viewer + # For determining which output type to produce + self.make_viewer: bool = _str2bool(args["make_viewer"]) + # For coupled_global + self.case_dir: str = args["case_dir"] + self.experiment_name: str = args["experiment_name"] + self.figstr: str = args["figstr"] + self.color: str = args["color"] + self.ts_num_years_str: str = args["ts_num_years"] + self.results_dir: str = args["results_dir"] + # These regions are used often as strings, + # so making an Enum Region={GLOBAL, NORTH, SOUTH} would be limiting. + self.regions: List[str] = list( + map(lambda rgn: get_region(rgn), args["regions"].split(",")) + ) + # For both ocean_month and coupled_global + self.year1: int = int(args["start_yr"]) + self.year2: int = int(args["end_yr"]) + # Used by Classic PDF only # For ocean_month + self.subtask_name: str = f"global_time_series_{args['subsection']}" self.use_ocn: bool = _str2bool(args["use_ocn"]) self.input: str = args["input"] self.input_subdir: str = args["input_subdir"] @@ -14,31 +38,53 @@ def __init__(self, args: Dict[str, str]): self.moc_file = "" else: self.moc_file = args["moc_file"] - # For coupled_global - self.case_dir: str = args["case_dir"] - self.experiment_name: str = args["experiment_name"] - self.figstr: str = args["figstr"] - self.color: str = args["color"] - self.ts_num_years_str: str = args["ts_num_years"] self.plots_original: List[str] = param_get_list(args["plots_original"]) + self.nrows: int = int(args["nrows"]) + self.ncols: int = int(args["ncols"]) + + # Used by Viewer only + # For coupled_global self.plots_atm: List[str] = param_get_list(args["plots_atm"]) self.plots_ice: List[str] = param_get_list(args["plots_ice"]) self.plots_lnd: List[str] = param_get_list(args["plots_lnd"]) self.plots_ocn: List[str] = param_get_list(args["plots_ocn"]) - self.nrows: int = int(args["nrows"]) - self.ncols: int = int(args["ncols"]) - self.results_dir: str = args["results_dir"] - # These regions are used often as strings, - # so making an Enum Region={GLOBAL, NORTH, SOUTH} would be limiting. - self.regions: List[str] = list( - map(lambda rgn: get_region(rgn), args["regions"].split(",")) - ) - self.make_viewer: bool = _str2bool(args["make_viewer"]) - # For both - self.year1: int = int(args["start_yr"]) - self.year2: int = int(args["end_yr"]) + # Input validation + if self.make_viewer: + if self.plots_original: + logger.warning( + f"plots_original={self.plots_original} will not be plotted in Viewer mode." + ) + else: + if self.use_ocn and (not self.moc_file): + raise ValueError( + "moc_file must be set for ocean plots in Classic PDF mode." + ) + if self.nrows != 4: + logger.warning( + f"nrows={self.nrows} is DEPRECATED. It will be overridden as 4." + ) + if self.ncols != 2: + logger.warning( + f"ncols={self.ncols} is DEPRECATED. It will be overridden as 2." + ) + if self.plots_atm: + logger.warning( + f"plots_atm={self.plots_atm} will not be plotted in Classic PDF mode." + ) + if self.plots_ice: + logger.warning( + f"plots_ice={self.plots_ice} will not be plotted in Classic PDF mode." + ) + if self.plots_lnd: + logger.warning( + f"plots_lnd={self.plots_lnd} will not be plotted in Classic PDF mode." + ) + if self.plots_ocn: + logger.warning( + f"plots_ocn={self.plots_ocn} will not be plotted in Classic PDF mode." + ) def _str2bool(s: str) -> bool: @@ -62,6 +108,3 @@ def get_region(rgn: str) -> str: else: raise ValueError(f"Invalid rgn={rgn}") return rgn - - -############################################################################### diff --git a/zppy_interfaces/global_time_series/viewer/__init__.py b/zppy_interfaces/global_time_series/viewer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/zppy_interfaces/global_time_series/viewer/coupled_global.py b/zppy_interfaces/global_time_series/viewer/coupled_global.py new file mode 100644 index 0000000..fd54305 --- /dev/null +++ b/zppy_interfaces/global_time_series/viewer/coupled_global.py @@ -0,0 +1,234 @@ +import csv +import importlib.resources as imp_res +from typing import Any, Dict, List, Tuple + +from zppy_interfaces.global_time_series.coupled_global_utils import ( + Metric, + Variable, + get_data_dir, + set_var, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.global_time_series.viewer.coupled_global_plotting import ( + make_plot_pdfs, +) +from zppy_interfaces.global_time_series.viewer.coupled_global_viewer import ( + create_viewer, + create_viewer_index, +) +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + + +# Classes ##################################################################### +class RequestedVariables(object): + def __init__(self, parameters: Parameters): + # Land variables are constructed differently + self.vars_land: List[Variable] = construct_land_variables(parameters.plots_lnd) + + # Other variables use the generic constructor + self.vars_atm: List[Variable] = construct_generic_variables( + parameters.plots_atm + ) + self.vars_ice: List[Variable] = construct_generic_variables( + parameters.plots_ice + ) + self.vars_ocn: List[Variable] = construct_generic_variables( + parameters.plots_ocn + ) + + +def construct_land_variables(requested_vars: List[str]) -> List[Variable]: + var_list: List[Variable] = [] + header = True + csv_filename = str( + imp_res.files("zppy_interfaces.global_time_series") / "zppy_land_fields.csv" + ) + with open(csv_filename, newline="") as csv_file: + logger.debug("Reading zppy_land_fields.csv") + var_reader = csv.reader(csv_file) + for row in var_reader: + # logger.debug(f"row={row}") + # Skip the header row + if header: + header = False + else: + # If set to "all" then we want all variables. + # Design note: we can't simply run all variables if requested_vars is empty because + # that would actually mean the user doesn't want to make *any* land plots. + if (requested_vars == ["all"]) or (row[0] in requested_vars): + row_elements_strip_whitespace: List[str] = list( + map(lambda x: x.strip(), row) + ) + var_list.append(land_csv_row_to_var(row_elements_strip_whitespace)) + return var_list + + +def land_csv_row_to_var(csv_row: List[str]) -> Variable: + # “A” or “T” for global average over land area or global total, respectively + metric: Metric + if csv_row[1] == "A": + metric = Metric.AVERAGE + elif csv_row[1] == "T": + metric = Metric.TOTAL + else: + raise ValueError(f"Invalid metric={csv_row[1]}") + return Variable( + variable_name=csv_row[0], + metric=metric, + scale_factor=float(csv_row[2]), + original_units=csv_row[3], + final_units=csv_row[4], + group=csv_row[5], + long_name=csv_row[6], + ) + + +def construct_generic_variables(requested_vars: List[str]) -> List[Variable]: + var_list: List[Variable] = [] + for var_name in requested_vars: + var_list.append(Variable(var_name)) + return var_list + + +# Main functionality ########################################################## + + +def run_coupled_global(parameters: Parameters) -> None: + requested_variables = RequestedVariables(parameters) + run(parameters, requested_variables) + title_and_url_list: List[Tuple[str, str]] = [] + for component in [ + "atm", + "ice", + "lnd", + "ocn", + ]: # Don't create viewer for original component + vars_list: List[Variable] = get_vars(requested_variables, component) + if vars_list: + url = create_viewer(parameters, vars_list, component) + logger.info(f"Viewer URL for {component}: {url}") + title_and_url_list.append((component, url)) + + index_url: str = create_viewer_index(parameters.results_dir, title_and_url_list) + logger.info(f"Viewer index URL: {index_url}") + + +def run(parameters: Parameters, requested_variables: RequestedVariables): + # Experiments + exps: List[Dict[str, Any]] = process_data(parameters, requested_variables) + + xlim: List[float] = [float(parameters.year1), float(parameters.year2)] + + # Use list of tuples rather than a dict, to keep order + mapping: List[Tuple[str, List[str]]] = [ + ("atm", list(map(lambda v: v.variable_name, requested_variables.vars_atm))), + ("ice", list(map(lambda v: v.variable_name, requested_variables.vars_ice))), + ("lnd", list(map(lambda v: v.variable_name, requested_variables.vars_land))), + ("ocn", list(map(lambda v: v.variable_name, requested_variables.vars_ocn))), + ] + for rgn in parameters.regions: + valid_plots: List[str] = [] + invalid_plots: List[str] = [] + for component, plot_list in mapping: + make_plot_pdfs( + parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, + ) + logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") + if invalid_plots: + logger.error( + f"These {rgn} region plots could not be generated successfully: {invalid_plots}" + ) + + +def get_vars(requested_variables: RequestedVariables, component: str) -> List[Variable]: + vars: List[Variable] + if component == "atm": + vars = requested_variables.vars_atm + elif component == "ice": + vars = requested_variables.vars_ice + elif component == "lnd": + vars = requested_variables.vars_land + elif component == "ocn": + vars = requested_variables.vars_ocn + else: + raise ValueError(f"Invalid component={component}") + return vars + + +def process_data( + parameters: Parameters, requested_variables: RequestedVariables +) -> List[Dict[str, Any]]: + exps: List[Dict[str, Any]] = get_exps(parameters) + valid_vars: List[str] = [] + invalid_vars: List[str] = [] + exp: Dict[str, Any] + for exp in exps: + exp["annual"] = {} + + requested_variables.vars_atm = set_var( + exp, + "atmos", + requested_variables.vars_atm, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ice = set_var( + exp, + "ice", + requested_variables.vars_ice, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_land = set_var( + exp, + "land", + requested_variables.vars_land, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ocn = set_var( + exp, + "ocean", + requested_variables.vars_ocn, + valid_vars, + invalid_vars, + parameters, + ) + + logger.info( + f"globalAnnual was computed successfully for these variables: {valid_vars}" + ) + if invalid_vars: + logger.error( + f"globalAnnual could not be computed for these variables: {invalid_vars}" + ) + return exps + + +def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: + # Experiments + exps: List[Dict[str, Any]] = [ + { + "atmos": get_data_dir(parameters, "atm", parameters.plots_atm != []), + "ice": get_data_dir(parameters, "ice", parameters.plots_ice != []), + "land": get_data_dir(parameters, "lnd", parameters.plots_lnd != []), + "ocean": get_data_dir(parameters, "ocn", parameters.plots_ocn != []), + "name": parameters.experiment_name, + "yoffset": 0.0, + "yr": ([parameters.year1, parameters.year2],), + "color": f"{parameters.color}", + } + ] + return exps diff --git a/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py b/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py new file mode 100644 index 0000000..e7fb03d --- /dev/null +++ b/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py @@ -0,0 +1,92 @@ +import os +import traceback + +import matplotlib as mpl +import matplotlib.backends.backend_pdf +import matplotlib.pyplot as plt +import numpy as np + +from zppy_interfaces.global_time_series.coupled_global_plotting import plot +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +mpl.use("Agg") + +logger = _setup_custom_logger(__name__) + + +# FIXME: C901 'make_plot_pdfs' is too complex (20) +def make_plot_pdfs( # noqa: C901 + parameters: Parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, +): + logger.info( + f"Global Time Series Viewer: make_plot_pdfs for rgn={rgn}, component={component}" + ) + num_plots = len(plot_list) + if num_plots == 0: + return + nrows = 1 + ncols = 1 + + os.makedirs(parameters.results_dir, exist_ok=True) + for i in range(num_plots): + fig = plt.figure(1, figsize=[13.5 / 2, 16.5 / 4]) + logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") + plot_name = plot_list[i] + fig.suptitle(plot_name) + ax = plt.subplot( + nrows, + ncols, + 1, + ) + try: + plot_generic(ax, xlim, exps, plot_name, rgn) + valid_plots.append(plot_name) + except Exception: + traceback.print_exc() + logger.error(f"plot_generic failed. Invalid plot={plot_name}, rgn={rgn}") + invalid_plots.append(plot_name) + + fig.tight_layout() + # Save individual PNGs + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{plot_name}.png", + dpi=150, + ) + plt.close(fig) + + +# Generic plot function +def plot_generic(ax, xlim, exps, var_name, rgn): + logger.info(f"plot_generic for {var_name}, rgn={rgn}") + param_dict = { + "2nd_var": False, + "axhline_y": 0, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": True, + "default_ylim": [], + "do_add_line": True, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": False, + "title": var_name, + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"][var_name][rgn][0]), + "verbose": False, + "vol": False, + "ylabel": lambda exp: np.array(exp["annual"][var_name][rgn][1]), + } + plot(ax, xlim, exps, param_dict, rgn) diff --git a/zppy_interfaces/global_time_series/coupled_global_viewer.py b/zppy_interfaces/global_time_series/viewer/coupled_global_viewer.py similarity index 100% rename from zppy_interfaces/global_time_series/coupled_global_viewer.py rename to zppy_interfaces/global_time_series/viewer/coupled_global_viewer.py diff --git a/zppy_interfaces/global_time_series/viewer/driver.py b/zppy_interfaces/global_time_series/viewer/driver.py new file mode 100644 index 0000000..e603a8a --- /dev/null +++ b/zppy_interfaces/global_time_series/viewer/driver.py @@ -0,0 +1,12 @@ +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.global_time_series.viewer.coupled_global import run_coupled_global +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + + +def run(parameters: Parameters): + logger.info("Update time series figures") + # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) + # Creates the directory parameters.results_dir + run_coupled_global(parameters) From 188031a4ffbc7a1f3e4bddb066f70977d26993e8 Mon Sep 17 00:00:00 2001 From: Ryan Forsyth Date: Wed, 18 Jun 2025 19:49:42 -0500 Subject: [PATCH 02/17] Redesign for backwards compatibility --- .../test_global_time_series.py | 38 ++- .../global_time_series/__main__.py | 75 ++++-- .../global_time_series/classic/driver.py | 59 ----- .../{classic => coupled_global}/__init__.py | 0 .../coupled_global/driver.py | 99 ++++++++ .../mix_pdf_original.py} | 108 ++------ .../mix_viewer_component.py} | 197 ++++++++------- .../coupled_global/mix_viewer_original.py | 40 +++ .../coupled_global/mode_pdf.py | 112 +++++++++ .../coupled_global/mode_viewer.py | 116 +++++++++ .../coupled_global/plots_component.py | 125 ++++++++++ .../plots_original.py} | 81 +----- .../plotting.py} | 8 +- .../utils.py} | 103 ++++++++ .../ocean_month.py => create_ocean_ts.py} | 38 +++ zppy_interfaces/global_time_series/utils.py | 46 +--- .../global_time_series/viewer/__init__.py | 0 .../viewer/coupled_global.py | 234 ------------------ .../viewer/coupled_global_plotting.py | 92 ------- .../global_time_series/viewer/driver.py | 12 - 20 files changed, 870 insertions(+), 713 deletions(-) delete mode 100644 zppy_interfaces/global_time_series/classic/driver.py rename zppy_interfaces/global_time_series/{classic => coupled_global}/__init__.py (100%) create mode 100644 zppy_interfaces/global_time_series/coupled_global/driver.py rename zppy_interfaces/global_time_series/{classic/coupled_global_plotting.py => coupled_global/mix_pdf_original.py} (70%) rename zppy_interfaces/global_time_series/{viewer/coupled_global_viewer.py => coupled_global/mix_viewer_component.py} (50%) create mode 100644 zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py create mode 100644 zppy_interfaces/global_time_series/coupled_global/mode_pdf.py create mode 100644 zppy_interfaces/global_time_series/coupled_global/mode_viewer.py create mode 100644 zppy_interfaces/global_time_series/coupled_global/plots_component.py rename zppy_interfaces/global_time_series/{classic/coupled_global.py => coupled_global/plots_original.py} (53%) rename zppy_interfaces/global_time_series/{coupled_global_plotting.py => coupled_global/plotting.py} (95%) rename zppy_interfaces/global_time_series/{coupled_global_utils.py => coupled_global/utils.py} (79%) rename zppy_interfaces/global_time_series/{classic/ocean_month.py => create_ocean_ts.py} (77%) delete mode 100644 zppy_interfaces/global_time_series/viewer/__init__.py delete mode 100644 zppy_interfaces/global_time_series/viewer/coupled_global.py delete mode 100644 zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py delete mode 100644 zppy_interfaces/global_time_series/viewer/driver.py diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index 4eae963..45155d8 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -2,32 +2,30 @@ import pytest -from zppy_interfaces.global_time_series.classic.coupled_global import ( - get_exps as classic_get_exps, +from zppy_interfaces.global_time_series.coupled_global.mix_viewer_component import ( + VariableGroup, + _get_variable_groups, +) +from zppy_interfaces.global_time_series.coupled_global.plots_component import ( + _get_exps as component_get_exps, +) +from zppy_interfaces.global_time_series.coupled_global.plots_original import ( + _get_exps as original_get_exps, ) -from zppy_interfaces.global_time_series.classic.coupled_global import get_vars_original -from zppy_interfaces.global_time_series.coupled_global_plotting import get_ylim -from zppy_interfaces.global_time_series.coupled_global_utils import ( +from zppy_interfaces.global_time_series.coupled_global.plotting import get_ylim +from zppy_interfaces.global_time_series.coupled_global.utils import ( Metric, Variable, + _land_csv_row_to_var, + construct_generic_variables, get_data_dir, + get_vars_original, ) from zppy_interfaces.global_time_series.utils import ( Parameters, get_region, param_get_list, ) -from zppy_interfaces.global_time_series.viewer.coupled_global import ( - construct_generic_variables, -) -from zppy_interfaces.global_time_series.viewer.coupled_global import ( - get_exps as viewer_get_exps, -) -from zppy_interfaces.global_time_series.viewer.coupled_global import land_csv_row_to_var -from zppy_interfaces.global_time_series.viewer.coupled_global_viewer import ( - VariableGroup, - get_variable_groups, -) # Run tests with `pytest tests/unit/global_time_series/test_*.py` @@ -171,7 +169,7 @@ def test_Parameters_and_related_functions(): assert get_data_dir(parameters, "ocn", False) == "" # test_get_exps - exps: List[Dict[str, Any]] = viewer_get_exps(parameters) + exps: List[Dict[str, Any]] = component_get_exps(parameters) assert len(exps) == 1 expected = { "atmos": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/atm/glb/ts/monthly/5yr/", @@ -190,7 +188,7 @@ def test_Parameters_and_related_functions(): ) parameters.plots_atm = [] parameters.plots_lnd = [] - exps = classic_get_exps(parameters) + exps = original_get_exps(parameters) assert len(exps) == 1 expected = { "atmos": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/atm/glb/ts/monthly/5yr/", @@ -256,7 +254,7 @@ def test_land_csv_row_to_var(): csv_row = "BCDEP,A,1.00000E+00,kg/m^2/s,kg/m^2/s,Aerosol Flux,total black carbon deposition (dry+wet) from atmosphere".split( "," ) - v: Variable = land_csv_row_to_var(csv_row) + v: Variable = _land_csv_row_to_var(csv_row) assert v.variable_name == "BCDEP" assert v.metric == Metric.AVERAGE assert v.scale_factor == 1.0 @@ -288,7 +286,7 @@ def test_get_variable_groups(): def get_group_names(groups: List[VariableGroup]) -> List[str]: return list(map(lambda g: g.group_name, groups)) - assert get_group_names(get_variable_groups([a, b, x, y])) == ["GroupA", "GroupX"] + assert get_group_names(_get_variable_groups([a, b, x, y])) == ["GroupA", "GroupX"] def test_get_ylim(): diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 314c4e4..316c6ec 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -1,8 +1,8 @@ import argparse import sys -import zppy_interfaces.global_time_series.classic.driver as classic_driver -import zppy_interfaces.global_time_series.viewer.driver as viewer_driver +from zppy_interfaces.global_time_series.coupled_global.driver import run_coupled_global +from zppy_interfaces.global_time_series.create_ocean_ts import create_ocean_ts from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger, _setup_root_logger @@ -15,11 +15,47 @@ def main(parameters=None): if not parameters: parameters = _get_args() - # Determine if we want the Classic PDF or the Viewer - if parameters.make_viewer: - viewer_driver.run(parameters) - else: - classic_driver.run(parameters) + """ + Determine if we want the Classic PDF or the Viewer + There are several cases to consider. In markdown table format: + + | case | make_viewer = | plots_original non-empty? | `plots_` non-empty? | results page shows | plots_original | `plots_` | + | --- | --- | --- | --- | --- | --- | --- | + | 1 | T | T | T | viewer list HTML | `_original` PDF & PNG | each component gets a Viewer, with rows=vars, cols=rgns | + | 2 | T | T | F | viewer list HTML | `_original` PDF & PNG | no Viewers | + | 3 | T | F | T | viewer list HTML | no original PDF/PNG links | each component gets a Viewer, with rows=vars, cols=rgns | + | 4 | T | F | F | viewer list HTML | no original PDF/PNG links | no Viewers | + | 5 | F | T | T | no-frills file list | `_original` PDF & PNG | each component gets a cumulative PDF | + | 6 | F | T | F | no-frills file list | `_original` PDF & PNG | no component PDFs | + | 7 | F | F | T | no-frills file list | no classic plots | each component gets a cumulative PDF | + | 8 | F | F | F | no-frills file list | no classic plots | no component PDFs | + + examples/post.v3.LR.historical_zppy_v3.cfg has: make_viewer = True, plots_original = 8 plots, plots_lnd = variable list + That is: | T | T | T |, or case 1 in the table above. + + By default: make_viewer = False, plots_original = 8 plots, plots_ = "" + That is: | F | T | F |, or case 6 in the table above. + + We can simplify the above table to: + + | make_viewer= | plots_original | `plots_` | + | --- | --- | --- | + | True | `_original` PDF & PNG | each component gets a Viewer, with rows=vars, cols=rgns | + | False | `_original` PDF & PNG | each component gets a cumulative PDF | + """ + if parameters.use_ocn: + # From zppy's default.ini: + # Remove the 3 ocean plots (change_ohc,max_moc,change_sea_level) if you don't have ocean data. + # plots_original = string(default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance") + if set(["change_ohc", "max_moc", "change_sea_level"]) & set( + parameters.plots_original + ): + create_ocean_ts(parameters) + logger.info("Update time series figures") + # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) + # Creates the directory parameters.results_dir + run_coupled_global(parameters) + # TODO: Add tests for all of the above cases on the zppy side def _get_args() -> Parameters: @@ -29,10 +65,10 @@ def _get_args() -> Parameters: description="Generate Global Time Series plots", ) - # Used by both Classic PDF and Viewer - # For determining which output type to produce + # Used in all cases + # > For determining which output type to produce parser.add_argument("--make_viewer", type=str, help="Make viewer") - # For coupled_global + # > For coupled_global parser.add_argument("--case_dir", type=str, help="Case directory") parser.add_argument("--experiment_name", type=str, help="Experiment name") parser.add_argument("--figstr", type=str, help="Figure string") @@ -40,29 +76,32 @@ def _get_args() -> Parameters: parser.add_argument("--ts_num_years", type=str, help="Time series number of years") parser.add_argument("--results_dir", type=str, help="Results directory") parser.add_argument("--regions", type=str, help="Regions") - # For both ocean_month and coupled_global + # > For both ocean_month and coupled_global parser.add_argument("--start_yr", type=str, help="Start year") parser.add_argument("--end_yr", type=str, help="End year") - # Classic PDF only - # For ocean_month + # For plots_original + # > For ocean_month parser.add_argument("--subsection", type=str, help="Subtask name") parser.add_argument("--use_ocn", type=str, help="Use ocean") parser.add_argument("--input", type=str, help="Input directory") parser.add_argument("--input_subdir", type=str, help="Input subdirectory") parser.add_argument("--moc_file", type=str, help="MOC file") - # For coupled_global + # > For coupled_global parser.add_argument("--plots_original", type=str, help="Plots original") - parser.add_argument("--nrows", type=str, help="Number of rows in pdf") - parser.add_argument("--ncols", type=str, help="Number of columns in pdf") - # Viewer only - # For coupled_global + # For plots_component + # > For coupled_global parser.add_argument("--plots_atm", type=str, help="Plots atmosphere") parser.add_argument("--plots_ice", type=str, help="Plots ice") parser.add_argument("--plots_lnd", type=str, help="Plots land") parser.add_argument("--plots_ocn", type=str, help="Plots ocean") + # Used for mode_pdf, regardless of plot type + # > For coupled_global + parser.add_argument("--ncols", type=str, help="Number of columns in pdf") + parser.add_argument("--nrows", type=str, help="Number of rows in pdf") + # Ignore the first arg # (zi-global-time-series) args: argparse.Namespace = parser.parse_args(sys.argv[1:]) diff --git a/zppy_interfaces/global_time_series/classic/driver.py b/zppy_interfaces/global_time_series/classic/driver.py deleted file mode 100644 index 83a0f03..0000000 --- a/zppy_interfaces/global_time_series/classic/driver.py +++ /dev/null @@ -1,59 +0,0 @@ -import os -import shutil - -from zppy_interfaces.global_time_series.classic.coupled_global import run_coupled_global -from zppy_interfaces.global_time_series.classic.ocean_month import ocean_month -from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger - -logger = _setup_custom_logger(__name__) - - -def run(parameters: Parameters): - # From zppy's default.ini: - # Remove the 3 ocean plots (change_ohc,max_moc,change_sea_level) if you don't have ocean data. - # plots_original = string(default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance") - if parameters.use_ocn: - if set(["change_ohc", "max_moc", "change_sea_level"]) & set( - parameters.plots_original - ): - logger.info("Create ocean time series") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Creates the directory post/ocn - os.makedirs( - f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr", - exist_ok=True, - ) - input_dir: str = f"{parameters.input}/{parameters.input_subdir}" - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Modifies post/ocn (which we just created in the first place) - ocean_month( - input_dir, - parameters.subtask_name, - parameters.case_dir, - parameters.year1, - parameters.year2, - int(parameters.ts_num_years_str), - ) - - src: str = ( - f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}" - ) - dst: str = ( - f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/" - ) - logger.info(f"Copy moc file from {src} to {dst}") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Copies files to post//ocn (which we just created in the first place) - shutil.copy( - src, - dst, - ) - else: - logger.info( - "use_ocn is set unecessarily. ocn plots have not been requested" - ) - logger.info("Update time series figures") - # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) - # Creates the directory parameters.results_dir - run_coupled_global(parameters) diff --git a/zppy_interfaces/global_time_series/classic/__init__.py b/zppy_interfaces/global_time_series/coupled_global/__init__.py similarity index 100% rename from zppy_interfaces/global_time_series/classic/__init__.py rename to zppy_interfaces/global_time_series/coupled_global/__init__.py diff --git a/zppy_interfaces/global_time_series/coupled_global/driver.py b/zppy_interfaces/global_time_series/coupled_global/driver.py new file mode 100644 index 0000000..dacb28d --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/driver.py @@ -0,0 +1,99 @@ +from typing import Any, Dict, List, Tuple + +from zppy_interfaces.global_time_series.coupled_global.mix_viewer_component import ( + produce_pngs_for_viewer, +) +from zppy_interfaces.global_time_series.coupled_global.mode_pdf import ( + assemble_cumulative_pdf, +) +from zppy_interfaces.global_time_series.coupled_global.mode_viewer import produce_viewer +from zppy_interfaces.global_time_series.coupled_global.plots_component import ( + process_data as component_process_data, +) +from zppy_interfaces.global_time_series.coupled_global.plots_original import ( + process_data as original_process_data, +) +from zppy_interfaces.global_time_series.coupled_global.utils import RequestedVariables +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + +# Main functionality ########################################################## + + +def run_coupled_global(parameters: Parameters) -> None: + requested_variables = RequestedVariables(parameters) + run(parameters, requested_variables) + if parameters.make_viewer: + produce_viewer(parameters, requested_variables) + + +def run(parameters: Parameters, requested_variables: RequestedVariables): + # Experiments + exps_original: List[Dict[str, Any]] = original_process_data( + parameters, requested_variables + ) + exps_component: List[Dict[str, Any]] = component_process_data( + parameters, requested_variables + ) + + xlim: List[float] = [float(parameters.year1), float(parameters.year2)] + + # Use list of tuples rather than a dict, to keep order + # Note: we use `parameters.plots_original` rather than `requested_variables.vars_original` + # because the "original" plots are expecting plot names that are not variable names. + # The model components however are expecting plot names to be variable names. + mapping: List[Tuple[str, List[str]]] = [ + ("original", parameters.plots_original), + ("atm", list(map(lambda v: v.variable_name, requested_variables.vars_atm))), + ("ice", list(map(lambda v: v.variable_name, requested_variables.vars_ice))), + ("lnd", list(map(lambda v: v.variable_name, requested_variables.vars_land))), + ("ocn", list(map(lambda v: v.variable_name, requested_variables.vars_ocn))), + ] + for rgn in parameters.regions: + valid_plots: List[str] = [] + invalid_plots: List[str] = [] + for component, plot_list in mapping: + exps: List[Dict[str, Any]] + if component == "original": + exps = exps_original + assemble_cumulative_pdf( + parameters, + rgn, + component, + xlim, + exps, + parameters.plots_original, + valid_plots, + invalid_plots, + ) + else: + exps = exps_component + if parameters.make_viewer: + produce_pngs_for_viewer( + parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, + ) + else: + assemble_cumulative_pdf( + parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, + ) + logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") + if invalid_plots: + logger.error( + f"These {rgn} region plots could not be generated successfully: {invalid_plots}" + ) diff --git a/zppy_interfaces/global_time_series/classic/coupled_global_plotting.py b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py similarity index 70% rename from zppy_interfaces/global_time_series/classic/coupled_global_plotting.py rename to zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py index 43a6cda..7094e7a 100644 --- a/zppy_interfaces/global_time_series/classic/coupled_global_plotting.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py @@ -1,103 +1,20 @@ import math -import os -import traceback +from typing import List import matplotlib as mpl -import matplotlib.backends.backend_pdf -import matplotlib.pyplot as plt import numpy as np -from zppy_interfaces.global_time_series.coupled_global_plotting import plot -from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.global_time_series.coupled_global.plotting import plot from zppy_interfaces.multi_utils.logger import _setup_custom_logger mpl.use("Agg") logger = _setup_custom_logger(__name__) +# This file is for making cumulative PDFs for the original plots. +# Hence, "mix_pdf_original" -# FIXME: C901 'make_plot_pdfs' is too complex (20) -def make_plot_pdfs( # noqa: C901 - parameters: Parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, -): - logger.info(f"Global Time Series Classic PDF: make_plot_pdfs for rgn={rgn}") - num_plots = len(plot_list) - if num_plots == 0: - return - - plots_per_page = parameters.nrows * parameters.ncols - num_pages = math.ceil(num_plots / plots_per_page) - - counter = 0 - os.makedirs(parameters.results_dir, exist_ok=True) - # https://stackoverflow.com/questions/58738992/save-multiple-figures-with-subplots-into-a-pdf-with-multiple-pages - pdf = matplotlib.backends.backend_pdf.PdfPages( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.pdf" - ) - for page in range(num_pages): - fig = plt.figure(1, figsize=[13.5, 16.5]) - logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") - fig.suptitle(f"{parameters.figstr}_{rgn}_{component}") - for j in range(plots_per_page): - logger.info( - f"Plotting plot {j} on page {page}. This is plot {counter} in total." - ) - # The final page doesn't need to be filled out with plots. - if counter >= num_plots: - break - ax = plt.subplot( - parameters.nrows, - parameters.ncols, - j + 1, - ) - plot_name = plot_list[counter] - try: - plot_function = PLOT_DICT[plot_name] - except KeyError: - raise KeyError(f"Invalid plot name: {plot_name}") - try: - plot_function(ax, xlim, exps, rgn) - valid_plots.append(plot_name) - except Exception: - traceback.print_exc() - required_vars = [] - if plot_name == "net_toa_flux_restom": - required_vars = ["RESTOM"] - elif plot_name == "net_atm_energy_imbalance": - required_vars = ["RESTOM", "RESSURF"] - elif plot_name == "global_surface_air_temperature": - required_vars = ["TREFHT"] - elif plot_name == "toa_radiation": - required_vars = ["FSNTOA", "FLUT"] - elif plot_name == "net_atm_water_imbalance": - required_vars = ["PRECC", "PRECL", "QFLX"] - logger.error( - f"Failed plot_function for {plot_name}. Check that {required_vars} are available." - ) - invalid_plots.append(plot_name) - counter += 1 - - fig.tight_layout() - pdf.savefig(1) - if num_pages > 1: - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{page}.png", - dpi=150, - ) - else: - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.png", - dpi=150, - ) - plt.close(fig) - pdf.close() +# Used by mode_pdf.assemble_cumulative_pdf #################################### # 1 @@ -366,3 +283,18 @@ def plot_net_atm_water_imbalance(ax, xlim, exps, rgn): "change_sea_level": plot_change_sea_level, # only glb "net_atm_water_imbalance": plot_net_atm_water_imbalance, } + + +def get_required_vars(plot_name: str) -> List[str]: + required_vars = [] + if plot_name == "net_toa_flux_restom": + required_vars = ["RESTOM"] + elif plot_name == "net_atm_energy_imbalance": + required_vars = ["RESTOM", "RESSURF"] + elif plot_name == "global_surface_air_temperature": + required_vars = ["TREFHT"] + elif plot_name == "toa_radiation": + required_vars = ["FSNTOA", "FLUT"] + elif plot_name == "net_atm_water_imbalance": + required_vars = ["PRECC", "PRECL", "QFLX"] + return required_vars diff --git a/zppy_interfaces/global_time_series/viewer/coupled_global_viewer.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py similarity index 50% rename from zppy_interfaces/global_time_series/viewer/coupled_global_viewer.py rename to zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py index 5c7323b..ce47363 100644 --- a/zppy_interfaces/global_time_series/viewer/coupled_global_viewer.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py @@ -1,16 +1,31 @@ -import importlib.resources as imp_res import os -from typing import List, Tuple - -from bs4 import BeautifulSoup - -from zppy_interfaces.global_time_series.coupled_global_utils import Metric, Variable +import traceback +from typing import List + +import matplotlib as mpl +import matplotlib.pyplot as plt + +from zppy_interfaces.global_time_series.coupled_global.plots_component import ( + plot_generic, +) +from zppy_interfaces.global_time_series.coupled_global.utils import ( + Metric, + RequestedVariables, + Variable, +) from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger from zppy_interfaces.multi_utils.viewer import OutputViewer +mpl.use("Agg") + logger = _setup_child_logger(__name__) +# This file is for making Viewers for the component plots. +# Hence, "mix_viewer_component" + +# Class ####################################################################### + class VariableGroup(object): def __init__(self, name: str, variables: List[Variable]): @@ -18,32 +33,89 @@ def __init__(self, name: str, variables: List[Variable]): self.variables = variables -def get_variable_groups(variables: List[Variable]) -> List[VariableGroup]: - group_names: List[str] = [] - groups: List[VariableGroup] = [] - for v in variables: - g: str = v.group - if g not in group_names: - # A new group! - group_names.append(g) - groups.append(VariableGroup(g, [v])) - else: - # Add a new variable to this existing group - for group in groups: - if g == group.group_name: - group.variables.append(v) - return groups - - -def create_viewer(parameters: Parameters, vars: List[Variable], component: str) -> str: +# Used by coupled_global.run ################################################## + + +def produce_pngs_for_viewer( + parameters: Parameters, + rgn: str, + component: str, + xlim, + exps, + plot_list: List[str], + valid_plots: List[str], + invalid_plots: List[str], +): + logger.info(f"Assembling Viewer for rgn={rgn}, component={component}") + num_plots = len(plot_list) + if num_plots == 0: + return + nrows = 1 + ncols = 1 + + os.makedirs(parameters.results_dir, exist_ok=True) + for i in range(num_plots): + fig = plt.figure(1, figsize=[13.5 / 2, 16.5 / 4]) + logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") + plot_name = plot_list[i] + fig.suptitle(plot_name) + ax = plt.subplot( + nrows, + ncols, + 1, + ) + try: + plot_generic(ax, xlim, exps, plot_name, rgn) + valid_plots.append(plot_name) + except Exception: + traceback.print_exc() + logger.error(f"plot_generic failed. Invalid plot={plot_name}, rgn={rgn}") + invalid_plots.append(plot_name) + + fig.tight_layout() + # Save individual PNGs + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{plot_name}.png", + dpi=150, + ) + plt.close(fig) + + +# Used by mode_viewer.produce_viewer ########################################## + + +def get_vars_component( + requested_variables: RequestedVariables, component: str +) -> List[Variable]: + vars: List[Variable] + if component == "atm": + vars = requested_variables.vars_atm + elif component == "ice": + vars = requested_variables.vars_ice + elif component == "lnd": + vars = requested_variables.vars_land + elif component == "ocn": + vars = requested_variables.vars_ocn + else: + raise ValueError(f"Invalid component={component}") + return vars + + +def create_viewer_for_component( + parameters: Parameters, vars: List[Variable], component: str +) -> str: logger.info(f"Creating viewer for {component}") + if not vars: + raise RuntimeError("No vars specified for viewer.") index_name = f"zppy global time-series plot: {parameters.experiment_name} {component} component ({parameters.year1}-{parameters.year2})" viewer = OutputViewer(path=parameters.results_dir, index_name=index_name) viewer.add_page(f"table_{component}", parameters.regions) - groups: List[VariableGroup] = get_variable_groups(vars) + groups: List[VariableGroup] = _get_variable_groups(vars) + if not groups: + raise RuntimeError("No groups specified for viewer.") for group in groups: logger.info(f"Adding group {group.group_name}") - # Only groups that have at least one variable will be returned by `get_variable_groups` + # Only groups that have at least one variable will be returned by `_get_variable_groups` # So, we know this group will be non-empty and should therefore be added to the viewer. viewer.add_group(group.group_name) for var in group.variables: @@ -78,63 +150,18 @@ def create_viewer(parameters: Parameters, vars: List[Variable], component: str) return url -# Copied from E3SM Diags and modified -def create_viewer_index( - root_dir: str, title_and_url_list: List[Tuple[str, str]] -) -> str: - """ - Creates the index page in root_dir which - joins the individual viewers. - Each tuple is on its own row. - """ - - logger.info("Creating viewer index") - - def insert_data_in_row(row_obj, name: str, url: str): - """ - Given a row object, insert the name and url. - """ - td = soup.new_tag("td") - a = soup.new_tag("a") - a["href"] = url - a.string = name - td.append(a) - row_obj.append(td) - - path: str = str( - imp_res.files("zppy_interfaces.global_time_series") / "index_template.html" - ) - output: str = os.path.join(root_dir, "index.html") - - soup = BeautifulSoup(open(path), "lxml") - - # If no one changes it, the template only has - # one element in the find command below. - table = soup.find_all("table", {"class": "table"})[0] - - # Adding the title. - tr = soup.new_tag("tr") - th = soup.new_tag("th") - th.string = "Output Sets" - tr.append(th) - - # Adding each of the rows. - for row in title_and_url_list: - tr = soup.new_tag("tr") - - if isinstance(row, list): - for elt in row: - name, url = elt - insert_data_in_row(tr, name, url) +def _get_variable_groups(variables: List[Variable]) -> List[VariableGroup]: + group_names: List[str] = [] + groups: List[VariableGroup] = [] + for v in variables: + g: str = v.group + if g not in group_names: + # A new group! + group_names.append(g) + groups.append(VariableGroup(g, [v])) else: - name, url = row - insert_data_in_row(tr, name, url) - - table.append(tr) - - html = soup.prettify("utf-8") - - with open(output, "wb") as f: - f.write(html) - - return output + # Add a new variable to this existing group + for group in groups: + if g == group.group_name: + group.variables.append(v) + return groups diff --git a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py new file mode 100644 index 0000000..ec87db8 --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py @@ -0,0 +1,40 @@ +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.viewer import OutputViewer + +logger = _setup_custom_logger(__name__) + +# This file is for making Viewers for the original plots. +# Hence, "mix_viewer_original" + +# Used by mode_viewer.produce_viewer ########################################## + + +def create_viewer_for_original(parameters: Parameters) -> str: + component: str = "original" + logger.info(f"Creating viewer for {component}") + index_name = f"zppy global time-series plot: {parameters.experiment_name} {component} component ({parameters.year1}-{parameters.year2})" + viewer = OutputViewer(path=parameters.results_dir, index_name=index_name) + viewer.add_page(f"table_{component}", parameters.regions) + viewer.add_group("Original Plots") + # PDFs -- these don't show up when clicked on, but they CAN be downloaded + row_title: str = "Original Plots, PDFs" + viewer.add_row(row_title) + for rgn in parameters.regions: + viewer.add_col( + f"{parameters.figstr}_{rgn}_{component}.pdf", + is_file=True, + title=f"{rgn}_{component}", + ) + # PNGs -- these show up when clicked on + row_title = "Original Plots, PNGs" + viewer.add_row(row_title) + for rgn in parameters.regions: + viewer.add_col( + f"{parameters.figstr}_{rgn}_{component}.png", + is_file=True, + title=f"{rgn}_{component}", + ) + url = viewer.generate_page() + viewer.generate_viewer() + return url diff --git a/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py b/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py new file mode 100644 index 0000000..970efdd --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py @@ -0,0 +1,112 @@ +import math +import os +import traceback +from typing import List + +import matplotlib as mpl +import matplotlib.backends.backend_pdf +import matplotlib.pyplot as plt + +from zppy_interfaces.global_time_series.coupled_global.mix_pdf_original import ( + PLOT_DICT, + get_required_vars, +) +from zppy_interfaces.global_time_series.coupled_global.plots_component import ( + plot_generic, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +mpl.use("Agg") + +logger = _setup_custom_logger(__name__) + +# This file is for making cumulative PDFs +# Hence, "mode_pdf" + +# Used by driver.run ########################################################## + + +def assemble_cumulative_pdf( + parameters: Parameters, + rgn: str, + component: str, + xlim, + exps, + plot_list: List[str], + valid_plots: List[str], + invalid_plots: List[str], +): + logger.info(f"Assembling Cumulative PDF for rgn={rgn}, component={component}") + num_plots = len(plot_list) + if num_plots == 0: + return + + plots_per_page = parameters.nrows * parameters.ncols + num_pages = math.ceil(num_plots / plots_per_page) + + counter = 0 + os.makedirs(parameters.results_dir, exist_ok=True) + # https://stackoverflow.com/questions/58738992/save-multiple-figures-with-subplots-into-a-pdf-with-multiple-pages + pdf = matplotlib.backends.backend_pdf.PdfPages( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.pdf" + ) + for page in range(num_pages): + fig = plt.figure(1, figsize=[13.5, 16.5]) + logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") + fig.suptitle(f"{parameters.figstr}_{rgn}_{component}") + for j in range(plots_per_page): + logger.info( + f"Plotting plot {j} on page {page}. This is plot {counter} in total." + ) + # The final page doesn't need to be filled out with plots. + if counter >= num_plots: + break + ax = plt.subplot( + parameters.nrows, + parameters.ncols, + j + 1, + ) + plot_name = plot_list[counter] + if component == "original": + try: + plot_function = PLOT_DICT[plot_name] + except KeyError: + raise KeyError(f"Invalid plot name: {plot_name}") + try: + plot_function(ax, xlim, exps, rgn) + valid_plots.append(plot_name) + except Exception: + traceback.print_exc() + required_vars: List[str] = get_required_vars(plot_name) + logger.error( + f"Failed plot_function for {plot_name}. Check that {required_vars} are available." + ) + invalid_plots.append(plot_name) + else: + try: + plot_generic(ax, xlim, exps, plot_name, rgn) + valid_plots.append(plot_name) + except Exception: + traceback.print_exc() + logger.error( + f"plot_generic failed. Invalid plot={plot_name}, rgn={rgn}" + ) + invalid_plots.append(plot_name) + counter += 1 + + fig.tight_layout() + pdf.savefig(1) + # Also save PNGs + if num_pages > 1: + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{page}.png", + dpi=150, + ) + else: + fig.savefig( + f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}.png", + dpi=150, + ) + plt.close(fig) + pdf.close() diff --git a/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py b/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py new file mode 100644 index 0000000..d5ef033 --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py @@ -0,0 +1,116 @@ +import importlib.resources as imp_res +import os +from typing import List, Tuple + +from bs4 import BeautifulSoup + +from zppy_interfaces.global_time_series.coupled_global.mix_viewer_component import ( + create_viewer_for_component, + get_vars_component, +) +from zppy_interfaces.global_time_series.coupled_global.mix_viewer_original import ( + create_viewer_for_original, +) +from zppy_interfaces.global_time_series.coupled_global.utils import ( + RequestedVariables, + Variable, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +logger = _setup_custom_logger(__name__) + +# This file is for making Viewers +# Hence, "mode_viewer" + +# Used by driver.run_coupled_global ########################################### + + +def produce_viewer(parameters: Parameters, requested_variables: RequestedVariables): + title_and_url_list: List[Tuple[str, str]] = [] + for component in [ + "original", + "atm", + "ice", + "lnd", + "ocn", + ]: + if component == "original": + if parameters.plots_original: + url = create_viewer_for_original(parameters) + else: + continue + else: + vars_list: List[Variable] = get_vars_component( + requested_variables, component + ) + if vars_list: + url = create_viewer_for_component(parameters, vars_list, component) + else: + continue + logger.info(f"Viewer URL for {component}: {url}") + title_and_url_list.append((component, url)) + index_url: str = _create_viewer_index(parameters.results_dir, title_and_url_list) + logger.info(f"Viewer index URL: {index_url}") + + +# Copied from E3SM Diags and modified +def _create_viewer_index( + root_dir: str, title_and_url_list: List[Tuple[str, str]] +) -> str: + """ + Creates the index page in root_dir which + joins the individual viewers. + Each tuple is on its own row. + """ + + logger.info("Creating viewer index") + + def insert_data_in_row(row_obj, name: str, url: str): + """ + Given a row object, insert the name and url. + """ + td = soup.new_tag("td") + a = soup.new_tag("a") + a["href"] = url + a.string = name + td.append(a) + row_obj.append(td) + + path: str = str( + imp_res.files("zppy_interfaces.global_time_series") / "index_template.html" + ) + output: str = os.path.join(root_dir, "index.html") + + soup = BeautifulSoup(open(path), "lxml") + + # If no one changes it, the template only has + # one element in the find command below. + table = soup.find_all("table", {"class": "table"})[0] + + # Adding the title. + tr = soup.new_tag("tr") + th = soup.new_tag("th") + th.string = "Output Sets" + tr.append(th) + + # Adding each of the rows. + for row in title_and_url_list: + tr = soup.new_tag("tr") + + if isinstance(row, list): + for elt in row: + name, url = elt + insert_data_in_row(tr, name, url) + else: + name, url = row + insert_data_in_row(tr, name, url) + + table.append(tr) + + html = soup.prettify("utf-8") + + with open(output, "wb") as f: + f.write(html) + + return output diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_component.py b/zppy_interfaces/global_time_series/coupled_global/plots_component.py new file mode 100644 index 0000000..ef037a3 --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/plots_component.py @@ -0,0 +1,125 @@ +from typing import Any, Dict, List + +import matplotlib as mpl +import numpy as np + +from zppy_interfaces.global_time_series.coupled_global.plotting import plot +from zppy_interfaces.global_time_series.coupled_global.utils import ( + RequestedVariables, + get_data_dir, + set_var, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_custom_logger + +mpl.use("Agg") + +logger = _setup_custom_logger(__name__) + +# This file is for handling the component plots +# Hence, "plots_component" + +# Used by driver.run ########################################################## + + +def process_data( + parameters: Parameters, requested_variables: RequestedVariables +) -> List[Dict[str, Any]]: + exps: List[Dict[str, Any]] = _get_exps(parameters) + valid_vars: List[str] = [] + invalid_vars: List[str] = [] + exp: Dict[str, Any] + for exp in exps: + exp["annual"] = {} + + requested_variables.vars_atm = set_var( + exp, + "atmos", + requested_variables.vars_atm, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ice = set_var( + exp, + "ice", + requested_variables.vars_ice, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_land = set_var( + exp, + "land", + requested_variables.vars_land, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ocn = set_var( + exp, + "ocean", + requested_variables.vars_ocn, + valid_vars, + invalid_vars, + parameters, + ) + + logger.info( + f"globalAnnual was computed successfully for these variables: {valid_vars}" + ) + if invalid_vars: + logger.error( + f"globalAnnual could not be computed for these variables: {invalid_vars}" + ) + return exps + + +def _get_exps(parameters: Parameters) -> List[Dict[str, Any]]: + # Experiments + exps: List[Dict[str, Any]] = [ + { + "atmos": get_data_dir(parameters, "atm", parameters.plots_atm != []), + "ice": get_data_dir(parameters, "ice", parameters.plots_ice != []), + "land": get_data_dir(parameters, "lnd", parameters.plots_lnd != []), + "ocean": get_data_dir(parameters, "ocn", parameters.plots_ocn != []), + "name": parameters.experiment_name, + "yoffset": 0.0, + "yr": ([parameters.year1, parameters.year2],), + "color": f"{parameters.color}", + } + ] + return exps + + +# Plotting #################################################################### +# Used by mix_viewer_component.produce_pngs_for_viewer +# Used by mode_pdf.assemble_cumulative_pdf + + +def plot_generic(ax, xlim, exps, var_name, rgn): + logger.info(f"plot_generic for {var_name}, rgn={rgn}") + param_dict = { + "2nd_var": False, + "axhline_y": 0, + "check_exp_ocean": False, + "check_exp_vol": False, + "check_exp_year": True, + "default_ylim": [], + "do_add_line": True, + "do_add_trend": True, + "format": "%4.2f", + "glb_only": False, + "lw": 1.0, + "ohc": False, + "set_axhline": False, + "set_legend": True, + "shorten_year": False, + "title": var_name, + "use_getmoc": False, + "var": lambda exp: np.array(exp["annual"][var_name][rgn][0]), + "verbose": False, + "vol": False, + "ylabel": lambda exp: np.array(exp["annual"][var_name][rgn][1]), + } + plot(ax, xlim, exps, param_dict, rgn) diff --git a/zppy_interfaces/global_time_series/classic/coupled_global.py b/zppy_interfaces/global_time_series/coupled_global/plots_original.py similarity index 53% rename from zppy_interfaces/global_time_series/classic/coupled_global.py rename to zppy_interfaces/global_time_series/coupled_global/plots_original.py index 5f28ef3..f93fc96 100644 --- a/zppy_interfaces/global_time_series/classic/coupled_global.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_original.py @@ -1,10 +1,10 @@ from typing import Any, Dict, List -from zppy_interfaces.global_time_series.classic.coupled_global_plotting import ( - make_plot_pdfs, -) -from zppy_interfaces.global_time_series.coupled_global_utils import ( +import matplotlib as mpl + +from zppy_interfaces.global_time_series.coupled_global.utils import ( DatasetWrapper, + RequestedVariables, Variable, get_data_dir, set_var, @@ -12,79 +12,20 @@ from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_custom_logger -logger = _setup_custom_logger(__name__) - - -# Classes ##################################################################### -class RequestedVariables(object): - def __init__(self, parameters: Parameters): - self.vars_original: List[Variable] = get_vars_original( - parameters.plots_original - ) - - -def get_vars_original(plots_original: List[str]) -> List[Variable]: - # NOTE: These are ALL atmosphere variables - vars_original: List[Variable] = [] - if ("net_toa_flux_restom" in plots_original) or ( - "net_atm_energy_imbalance" in plots_original - ): - vars_original.append(Variable("RESTOM")) - if "net_atm_energy_imbalance" in plots_original: - vars_original.append(Variable("RESSURF")) - if "global_surface_air_temperature" in plots_original: - vars_original.append(Variable("TREFHT")) - if "toa_radiation" in plots_original: - vars_original.append(Variable("FSNTOA")) - vars_original.append(Variable("FLUT")) - if "net_atm_water_imbalance" in plots_original: - vars_original.append(Variable("PRECC")) - vars_original.append(Variable("PRECL")) - vars_original.append(Variable("QFLX")) - return vars_original - +mpl.use("Agg") -# Main functionality ########################################################## - - -def run_coupled_global(parameters: Parameters) -> None: - requested_variables = RequestedVariables(parameters) - run(parameters, requested_variables) - - -def run(parameters: Parameters, requested_variables: RequestedVariables): - # Experiments - exps: List[Dict[str, Any]] = process_data(parameters, requested_variables) +logger = _setup_custom_logger(__name__) - xlim: List[float] = [float(parameters.year1), float(parameters.year2)] +# This file is for handling the original plots +# Hence, "plots_original" - # Note: we use `parameters.plots_original` rather than `requested_variables.vars_original` - # because the "original" plots are expecting plot names that are not variable names. - # The model components however are expecting plot names to be variable names. - for rgn in parameters.regions: - valid_plots: List[str] = [] - invalid_plots: List[str] = [] - make_plot_pdfs( - parameters, - rgn, - "original", - xlim, - exps, - parameters.plots_original, - valid_plots, - invalid_plots, - ) - logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") - if invalid_plots: - logger.error( - f"These {rgn} region plots could not be generated successfully: {invalid_plots}" - ) +# Used by driver.run ########################################################## def process_data( parameters: Parameters, requested_variables: RequestedVariables ) -> List[Dict[str, Any]]: - exps: List[Dict[str, Any]] = get_exps(parameters) + exps: List[Dict[str, Any]] = _get_exps(parameters) valid_vars: List[str] = [] invalid_vars: List[str] = [] exp: Dict[str, Any] @@ -130,7 +71,7 @@ def process_data( return exps -def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: +def _get_exps(parameters: Parameters) -> List[Dict[str, Any]]: # Experiments atm_set_intersection: set = set( [ diff --git a/zppy_interfaces/global_time_series/coupled_global_plotting.py b/zppy_interfaces/global_time_series/coupled_global/plotting.py similarity index 95% rename from zppy_interfaces/global_time_series/coupled_global_plotting.py rename to zppy_interfaces/global_time_series/coupled_global/plotting.py index dec4810..a6c94dd 100644 --- a/zppy_interfaces/global_time_series/coupled_global_plotting.py +++ b/zppy_interfaces/global_time_series/coupled_global/plotting.py @@ -36,7 +36,13 @@ def plot(ax, xlim, exps, param_dict, rgn): # noqa: C901 continue else: year = np.array(exp["annual"]["year"]) + exp["yoffset"] - var = param_dict["var"](exp) + try: + var = param_dict["var"](exp) + except KeyError as e: + # We have no direct way of getting `var_name`, but + # `title`` also happens to be set to `var_name` in plot_generic + logger.error(f"Key {param_dict['title']} not found in exp={exp}") + raise e extreme_values.append(np.amax(var)) extreme_values.append(np.amin(var)) if param_dict["shorten_year"]: diff --git a/zppy_interfaces/global_time_series/coupled_global_utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py similarity index 79% rename from zppy_interfaces/global_time_series/coupled_global_utils.py rename to zppy_interfaces/global_time_series/coupled_global/utils.py index c8587ff..c10da48 100644 --- a/zppy_interfaces/global_time_series/coupled_global_utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -1,3 +1,5 @@ +import csv +import importlib.resources as imp_res import os.path from enum import Enum from typing import Any, Dict, List, Optional, Tuple @@ -49,6 +51,107 @@ def __init__( self.long_name: str = long_name +class RequestedVariables(object): + def __init__(self, parameters: Parameters): + # Original plots + self.vars_original: List[Variable] = get_vars_original( + parameters.plots_original + ) + # Component plots + # > Land variables are constructed differently + self.vars_land: List[Variable] = construct_land_variables(parameters.plots_lnd) + # > Other variables use the generic constructor + self.vars_atm: List[Variable] = construct_generic_variables( + parameters.plots_atm + ) + self.vars_ice: List[Variable] = construct_generic_variables( + parameters.plots_ice + ) + self.vars_ocn: List[Variable] = construct_generic_variables( + parameters.plots_ocn + ) + + +# For plots_original + + +def get_vars_original(plots_original: List[str]) -> List[Variable]: + # NOTE: These are ALL atmosphere variables + vars_original: List[Variable] = [] + if ("net_toa_flux_restom" in plots_original) or ( + "net_atm_energy_imbalance" in plots_original + ): + vars_original.append(Variable("RESTOM")) + if "net_atm_energy_imbalance" in plots_original: + vars_original.append(Variable("RESSURF")) + if "global_surface_air_temperature" in plots_original: + vars_original.append(Variable("TREFHT")) + if "toa_radiation" in plots_original: + vars_original.append(Variable("FSNTOA")) + vars_original.append(Variable("FLUT")) + if "net_atm_water_imbalance" in plots_original: + vars_original.append(Variable("PRECC")) + vars_original.append(Variable("PRECL")) + vars_original.append(Variable("QFLX")) + return vars_original + + +# For plots_components + + +def construct_land_variables(requested_vars: List[str]) -> List[Variable]: + var_list: List[Variable] = [] + header = True + csv_filename = str( + imp_res.files("zppy_interfaces.global_time_series") / "zppy_land_fields.csv" + ) + with open(csv_filename, newline="") as csv_file: + logger.debug("Reading zppy_land_fields.csv") + var_reader = csv.reader(csv_file) + for row in var_reader: + # logger.debug(f"row={row}") + # Skip the header row + if header: + header = False + else: + # If set to "all" then we want all variables. + # Design note: we can't simply run all variables if requested_vars is empty because + # that would actually mean the user doesn't want to make *any* land plots. + if (requested_vars == ["all"]) or (row[0] in requested_vars): + row_elements_strip_whitespace: List[str] = list( + map(lambda x: x.strip(), row) + ) + var_list.append(_land_csv_row_to_var(row_elements_strip_whitespace)) + return var_list + + +def _land_csv_row_to_var(csv_row: List[str]) -> Variable: + # “A” or “T” for global average over land area or global total, respectively + metric: Metric + if csv_row[1] == "A": + metric = Metric.AVERAGE + elif csv_row[1] == "T": + metric = Metric.TOTAL + else: + raise ValueError(f"Invalid metric={csv_row[1]}") + return Variable( + variable_name=csv_row[0], + metric=metric, + scale_factor=float(csv_row[2]), + original_units=csv_row[3], + final_units=csv_row[4], + group=csv_row[5], + long_name=csv_row[6], + ) + + +def construct_generic_variables(requested_vars: List[str]) -> List[Variable]: + var_list: List[Variable] = [] + for var_name in requested_vars: + var_list.append(Variable(var_name)) + return var_list + + class DatasetWrapper(object): def __init__( self, diff --git a/zppy_interfaces/global_time_series/classic/ocean_month.py b/zppy_interfaces/global_time_series/create_ocean_ts.py similarity index 77% rename from zppy_interfaces/global_time_series/classic/ocean_month.py rename to zppy_interfaces/global_time_series/create_ocean_ts.py index 54279d0..3b52d06 100644 --- a/zppy_interfaces/global_time_series/classic/ocean_month.py +++ b/zppy_interfaces/global_time_series/create_ocean_ts.py @@ -1,17 +1,55 @@ # Compute time series of ocean heat content (ohc) using MPAS-O output import glob +import os +import shutil from datetime import datetime import numpy as np from mpas_tools.cime.constants import constants from netCDF4 import Dataset, chartostring, date2num +from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger logger = _setup_child_logger(__name__) +def create_ocean_ts(parameters: Parameters): + logger.info("Create ocean time series") + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Creates the directory post/ocn + os.makedirs( + f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr", + exist_ok=True, + ) + input_dir: str = f"{parameters.input}/{parameters.input_subdir}" + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Modifies post/ocn (which we just created in the first place) + ocean_month( + input_dir, + parameters.subtask_name, + parameters.case_dir, + parameters.year1, + parameters.year2, + int(parameters.ts_num_years_str), + ) + + src: str = ( + f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}" + ) + dst: str = ( + f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/" + ) + logger.info(f"Copy moc file from {src} to {dst}") + # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory + # Copies files to post//ocn (which we just created in the first place) + shutil.copy( + src, + dst, + ) + + def ocean_month( path_in: str, subtask_name: str, diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index b792ff5..52f46ea 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -51,40 +51,18 @@ def __init__(self, args: Dict[str, str]): self.plots_ocn: List[str] = param_get_list(args["plots_ocn"]) # Input validation - if self.make_viewer: - if self.plots_original: - logger.warning( - f"plots_original={self.plots_original} will not be plotted in Viewer mode." - ) - else: - if self.use_ocn and (not self.moc_file): - raise ValueError( - "moc_file must be set for ocean plots in Classic PDF mode." - ) - if self.nrows != 4: - logger.warning( - f"nrows={self.nrows} is DEPRECATED. It will be overridden as 4." - ) - if self.ncols != 2: - logger.warning( - f"ncols={self.ncols} is DEPRECATED. It will be overridden as 2." - ) - if self.plots_atm: - logger.warning( - f"plots_atm={self.plots_atm} will not be plotted in Classic PDF mode." - ) - if self.plots_ice: - logger.warning( - f"plots_ice={self.plots_ice} will not be plotted in Classic PDF mode." - ) - if self.plots_lnd: - logger.warning( - f"plots_lnd={self.plots_lnd} will not be plotted in Classic PDF mode." - ) - if self.plots_ocn: - logger.warning( - f"plots_ocn={self.plots_ocn} will not be plotted in Classic PDF mode." - ) + if self.plots_original and self.use_ocn and (not self.moc_file): + raise ValueError( + "moc_file must be set for ocean plots in the original 8-plot set." + ) + if not ( + self.plots_original + or self.plots_atm + or self.plots_ice + or self.plots_lnd + or self.plots_ocn + ): + raise ValueError("No plots are specified, so nothing will be generated.") def _str2bool(s: str) -> bool: diff --git a/zppy_interfaces/global_time_series/viewer/__init__.py b/zppy_interfaces/global_time_series/viewer/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/zppy_interfaces/global_time_series/viewer/coupled_global.py b/zppy_interfaces/global_time_series/viewer/coupled_global.py deleted file mode 100644 index fd54305..0000000 --- a/zppy_interfaces/global_time_series/viewer/coupled_global.py +++ /dev/null @@ -1,234 +0,0 @@ -import csv -import importlib.resources as imp_res -from typing import Any, Dict, List, Tuple - -from zppy_interfaces.global_time_series.coupled_global_utils import ( - Metric, - Variable, - get_data_dir, - set_var, -) -from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.global_time_series.viewer.coupled_global_plotting import ( - make_plot_pdfs, -) -from zppy_interfaces.global_time_series.viewer.coupled_global_viewer import ( - create_viewer, - create_viewer_index, -) -from zppy_interfaces.multi_utils.logger import _setup_custom_logger - -logger = _setup_custom_logger(__name__) - - -# Classes ##################################################################### -class RequestedVariables(object): - def __init__(self, parameters: Parameters): - # Land variables are constructed differently - self.vars_land: List[Variable] = construct_land_variables(parameters.plots_lnd) - - # Other variables use the generic constructor - self.vars_atm: List[Variable] = construct_generic_variables( - parameters.plots_atm - ) - self.vars_ice: List[Variable] = construct_generic_variables( - parameters.plots_ice - ) - self.vars_ocn: List[Variable] = construct_generic_variables( - parameters.plots_ocn - ) - - -def construct_land_variables(requested_vars: List[str]) -> List[Variable]: - var_list: List[Variable] = [] - header = True - csv_filename = str( - imp_res.files("zppy_interfaces.global_time_series") / "zppy_land_fields.csv" - ) - with open(csv_filename, newline="") as csv_file: - logger.debug("Reading zppy_land_fields.csv") - var_reader = csv.reader(csv_file) - for row in var_reader: - # logger.debug(f"row={row}") - # Skip the header row - if header: - header = False - else: - # If set to "all" then we want all variables. - # Design note: we can't simply run all variables if requested_vars is empty because - # that would actually mean the user doesn't want to make *any* land plots. - if (requested_vars == ["all"]) or (row[0] in requested_vars): - row_elements_strip_whitespace: List[str] = list( - map(lambda x: x.strip(), row) - ) - var_list.append(land_csv_row_to_var(row_elements_strip_whitespace)) - return var_list - - -def land_csv_row_to_var(csv_row: List[str]) -> Variable: - # “A” or “T” for global average over land area or global total, respectively - metric: Metric - if csv_row[1] == "A": - metric = Metric.AVERAGE - elif csv_row[1] == "T": - metric = Metric.TOTAL - else: - raise ValueError(f"Invalid metric={csv_row[1]}") - return Variable( - variable_name=csv_row[0], - metric=metric, - scale_factor=float(csv_row[2]), - original_units=csv_row[3], - final_units=csv_row[4], - group=csv_row[5], - long_name=csv_row[6], - ) - - -def construct_generic_variables(requested_vars: List[str]) -> List[Variable]: - var_list: List[Variable] = [] - for var_name in requested_vars: - var_list.append(Variable(var_name)) - return var_list - - -# Main functionality ########################################################## - - -def run_coupled_global(parameters: Parameters) -> None: - requested_variables = RequestedVariables(parameters) - run(parameters, requested_variables) - title_and_url_list: List[Tuple[str, str]] = [] - for component in [ - "atm", - "ice", - "lnd", - "ocn", - ]: # Don't create viewer for original component - vars_list: List[Variable] = get_vars(requested_variables, component) - if vars_list: - url = create_viewer(parameters, vars_list, component) - logger.info(f"Viewer URL for {component}: {url}") - title_and_url_list.append((component, url)) - - index_url: str = create_viewer_index(parameters.results_dir, title_and_url_list) - logger.info(f"Viewer index URL: {index_url}") - - -def run(parameters: Parameters, requested_variables: RequestedVariables): - # Experiments - exps: List[Dict[str, Any]] = process_data(parameters, requested_variables) - - xlim: List[float] = [float(parameters.year1), float(parameters.year2)] - - # Use list of tuples rather than a dict, to keep order - mapping: List[Tuple[str, List[str]]] = [ - ("atm", list(map(lambda v: v.variable_name, requested_variables.vars_atm))), - ("ice", list(map(lambda v: v.variable_name, requested_variables.vars_ice))), - ("lnd", list(map(lambda v: v.variable_name, requested_variables.vars_land))), - ("ocn", list(map(lambda v: v.variable_name, requested_variables.vars_ocn))), - ] - for rgn in parameters.regions: - valid_plots: List[str] = [] - invalid_plots: List[str] = [] - for component, plot_list in mapping: - make_plot_pdfs( - parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, - ) - logger.info(f"These {rgn} region plots generated successfully: {valid_plots}") - if invalid_plots: - logger.error( - f"These {rgn} region plots could not be generated successfully: {invalid_plots}" - ) - - -def get_vars(requested_variables: RequestedVariables, component: str) -> List[Variable]: - vars: List[Variable] - if component == "atm": - vars = requested_variables.vars_atm - elif component == "ice": - vars = requested_variables.vars_ice - elif component == "lnd": - vars = requested_variables.vars_land - elif component == "ocn": - vars = requested_variables.vars_ocn - else: - raise ValueError(f"Invalid component={component}") - return vars - - -def process_data( - parameters: Parameters, requested_variables: RequestedVariables -) -> List[Dict[str, Any]]: - exps: List[Dict[str, Any]] = get_exps(parameters) - valid_vars: List[str] = [] - invalid_vars: List[str] = [] - exp: Dict[str, Any] - for exp in exps: - exp["annual"] = {} - - requested_variables.vars_atm = set_var( - exp, - "atmos", - requested_variables.vars_atm, - valid_vars, - invalid_vars, - parameters, - ) - requested_variables.vars_ice = set_var( - exp, - "ice", - requested_variables.vars_ice, - valid_vars, - invalid_vars, - parameters, - ) - requested_variables.vars_land = set_var( - exp, - "land", - requested_variables.vars_land, - valid_vars, - invalid_vars, - parameters, - ) - requested_variables.vars_ocn = set_var( - exp, - "ocean", - requested_variables.vars_ocn, - valid_vars, - invalid_vars, - parameters, - ) - - logger.info( - f"globalAnnual was computed successfully for these variables: {valid_vars}" - ) - if invalid_vars: - logger.error( - f"globalAnnual could not be computed for these variables: {invalid_vars}" - ) - return exps - - -def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: - # Experiments - exps: List[Dict[str, Any]] = [ - { - "atmos": get_data_dir(parameters, "atm", parameters.plots_atm != []), - "ice": get_data_dir(parameters, "ice", parameters.plots_ice != []), - "land": get_data_dir(parameters, "lnd", parameters.plots_lnd != []), - "ocean": get_data_dir(parameters, "ocn", parameters.plots_ocn != []), - "name": parameters.experiment_name, - "yoffset": 0.0, - "yr": ([parameters.year1, parameters.year2],), - "color": f"{parameters.color}", - } - ] - return exps diff --git a/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py b/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py deleted file mode 100644 index e7fb03d..0000000 --- a/zppy_interfaces/global_time_series/viewer/coupled_global_plotting.py +++ /dev/null @@ -1,92 +0,0 @@ -import os -import traceback - -import matplotlib as mpl -import matplotlib.backends.backend_pdf -import matplotlib.pyplot as plt -import numpy as np - -from zppy_interfaces.global_time_series.coupled_global_plotting import plot -from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger - -mpl.use("Agg") - -logger = _setup_custom_logger(__name__) - - -# FIXME: C901 'make_plot_pdfs' is too complex (20) -def make_plot_pdfs( # noqa: C901 - parameters: Parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, -): - logger.info( - f"Global Time Series Viewer: make_plot_pdfs for rgn={rgn}, component={component}" - ) - num_plots = len(plot_list) - if num_plots == 0: - return - nrows = 1 - ncols = 1 - - os.makedirs(parameters.results_dir, exist_ok=True) - for i in range(num_plots): - fig = plt.figure(1, figsize=[13.5 / 2, 16.5 / 4]) - logger.info(f"Figure size={fig.get_size_inches() * fig.dpi}") - plot_name = plot_list[i] - fig.suptitle(plot_name) - ax = plt.subplot( - nrows, - ncols, - 1, - ) - try: - plot_generic(ax, xlim, exps, plot_name, rgn) - valid_plots.append(plot_name) - except Exception: - traceback.print_exc() - logger.error(f"plot_generic failed. Invalid plot={plot_name}, rgn={rgn}") - invalid_plots.append(plot_name) - - fig.tight_layout() - # Save individual PNGs - fig.savefig( - f"{parameters.results_dir}/{parameters.figstr}_{rgn}_{component}_{plot_name}.png", - dpi=150, - ) - plt.close(fig) - - -# Generic plot function -def plot_generic(ax, xlim, exps, var_name, rgn): - logger.info(f"plot_generic for {var_name}, rgn={rgn}") - param_dict = { - "2nd_var": False, - "axhline_y": 0, - "check_exp_ocean": False, - "check_exp_vol": False, - "check_exp_year": True, - "default_ylim": [], - "do_add_line": True, - "do_add_trend": True, - "format": "%4.2f", - "glb_only": False, - "lw": 1.0, - "ohc": False, - "set_axhline": False, - "set_legend": True, - "shorten_year": False, - "title": var_name, - "use_getmoc": False, - "var": lambda exp: np.array(exp["annual"][var_name][rgn][0]), - "verbose": False, - "vol": False, - "ylabel": lambda exp: np.array(exp["annual"][var_name][rgn][1]), - } - plot(ax, xlim, exps, param_dict, rgn) diff --git a/zppy_interfaces/global_time_series/viewer/driver.py b/zppy_interfaces/global_time_series/viewer/driver.py deleted file mode 100644 index e603a8a..0000000 --- a/zppy_interfaces/global_time_series/viewer/driver.py +++ /dev/null @@ -1,12 +0,0 @@ -from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.global_time_series.viewer.coupled_global import run_coupled_global -from zppy_interfaces.multi_utils.logger import _setup_custom_logger - -logger = _setup_custom_logger(__name__) - - -def run(parameters: Parameters): - logger.info("Update time series figures") - # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) - # Creates the directory parameters.results_dir - run_coupled_global(parameters) From 02556a6434ac9e5281f30dd74791cddb0d075f09 Mon Sep 17 00:00:00 2001 From: Ryan Forsyth Date: Wed, 3 Sep 2025 17:17:35 -0500 Subject: [PATCH 03/17] Update logger imports post-rebase --- zppy_interfaces/global_time_series/coupled_global/driver.py | 4 ++-- .../global_time_series/coupled_global/mix_pdf_original.py | 4 ++-- .../global_time_series/coupled_global/mix_viewer_original.py | 4 ++-- zppy_interfaces/global_time_series/coupled_global/mode_pdf.py | 4 ++-- .../global_time_series/coupled_global/mode_viewer.py | 4 ++-- .../global_time_series/coupled_global/plots_component.py | 4 ++-- .../global_time_series/coupled_global/plots_original.py | 4 ++-- zppy_interfaces/global_time_series/coupled_global/utils.py | 4 ++-- zppy_interfaces/global_time_series/utils.py | 4 ++-- 9 files changed, 18 insertions(+), 18 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/driver.py b/zppy_interfaces/global_time_series/coupled_global/driver.py index dacb28d..fd58987 100644 --- a/zppy_interfaces/global_time_series/coupled_global/driver.py +++ b/zppy_interfaces/global_time_series/coupled_global/driver.py @@ -15,9 +15,9 @@ ) from zppy_interfaces.global_time_series.coupled_global.utils import RequestedVariables from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # Main functionality ########################################################## diff --git a/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py index 7094e7a..4f29f8e 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py @@ -5,11 +5,11 @@ import numpy as np from zppy_interfaces.global_time_series.coupled_global.plotting import plot -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger mpl.use("Agg") -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for making cumulative PDFs for the original plots. # Hence, "mix_pdf_original" diff --git a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py index ec87db8..5adf936 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py @@ -1,8 +1,8 @@ from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger from zppy_interfaces.multi_utils.viewer import OutputViewer -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for making Viewers for the original plots. # Hence, "mix_viewer_original" diff --git a/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py b/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py index 970efdd..11ec0dc 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py +++ b/zppy_interfaces/global_time_series/coupled_global/mode_pdf.py @@ -15,11 +15,11 @@ plot_generic, ) from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger mpl.use("Agg") -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for making cumulative PDFs # Hence, "mode_pdf" diff --git a/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py b/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py index d5ef033..49e4acd 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py +++ b/zppy_interfaces/global_time_series/coupled_global/mode_viewer.py @@ -16,9 +16,9 @@ Variable, ) from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for making Viewers # Hence, "mode_viewer" diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_component.py b/zppy_interfaces/global_time_series/coupled_global/plots_component.py index ef037a3..79e3376 100644 --- a/zppy_interfaces/global_time_series/coupled_global/plots_component.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_component.py @@ -10,11 +10,11 @@ set_var, ) from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger mpl.use("Agg") -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for handling the component plots # Hence, "plots_component" diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_original.py b/zppy_interfaces/global_time_series/coupled_global/plots_original.py index f93fc96..c229686 100644 --- a/zppy_interfaces/global_time_series/coupled_global/plots_original.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_original.py @@ -10,11 +10,11 @@ set_var, ) from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger mpl.use("Agg") -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # This file is for handling the original plots # Hence, "plots_original" diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index c10da48..c401dd4 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -10,9 +10,9 @@ import xcdat from zppy_interfaces.global_time_series.utils import Parameters -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) class Metric(Enum): diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index 52f46ea..762bdb7 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -1,8 +1,8 @@ from typing import Dict, List -from zppy_interfaces.multi_utils.logger import _setup_custom_logger +from zppy_interfaces.multi_utils.logger import _setup_child_logger -logger = _setup_custom_logger(__name__) +logger = _setup_child_logger(__name__) # Parameters ################################################################## From df821cdffaeaf1e625514e5c041e9529cc3afeed Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 4 Sep 2025 15:37:53 -0500 Subject: [PATCH 04/17] avoid writing data into input directory/case_dir --- .../global_time_series/__main__.py | 2 ++ .../coupled_global/plots_original.py | 9 ++--- .../global_time_series/create_ocean_ts.py | 36 ++++++++----------- 3 files changed, 21 insertions(+), 26 deletions(-) diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 316c6ec..17f4cb2 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -50,6 +50,8 @@ def main(parameters=None): if set(["change_ohc", "max_moc", "change_sea_level"]) & set( parameters.plots_original ): + # NOTE: READS FROM case_dir (existing MPAS-Analysis results) and input (raw MPAS-O data) + # WRITES TO results_dir (new ocean time series output) create_ocean_ts(parameters) logger.info("Update time series figures") # NOTE: PRODUCES OUTPUT IN THE CURRENT DIRECTORY (not necessarily the case directory) diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_original.py b/zppy_interfaces/global_time_series/coupled_global/plots_original.py index c229686..90dbe89 100644 --- a/zppy_interfaces/global_time_series/coupled_global/plots_original.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_original.py @@ -87,10 +87,11 @@ def _get_exps(parameters: Parameters) -> List[Dict[str, Any]]: ["change_ohc", "max_moc", "change_sea_level"] ) & set(parameters.plots_original) ocean_dir = get_data_dir(parameters, "ocn", ocn_set_intersection != set()) - ocean_month_dir = get_data_dir( - parameters, - f"{parameters.subtask_name}/ocn", - ocn_set_intersection != set(), + # Ocean time series data is in results_dir with simplified structure + ocean_month_dir = ( + f"{parameters.results_dir}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/" + if ocn_set_intersection != set() + else "" ) exps: List[Dict[str, Any]] = [ { diff --git a/zppy_interfaces/global_time_series/create_ocean_ts.py b/zppy_interfaces/global_time_series/create_ocean_ts.py index 3b52d06..bee0307 100644 --- a/zppy_interfaces/global_time_series/create_ocean_ts.py +++ b/zppy_interfaces/global_time_series/create_ocean_ts.py @@ -17,48 +17,40 @@ def create_ocean_ts(parameters: Parameters): logger.info("Create ocean time series") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Creates the directory post/ocn - os.makedirs( - f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr", - exist_ok=True, - ) + # Create output directory in results_dir (simplified structure) + output_dir = f"{parameters.results_dir}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr" + os.makedirs(output_dir, exist_ok=True) + + # Input: Raw MPAS-O data input_dir: str = f"{parameters.input}/{parameters.input_subdir}" - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Modifies post/ocn (which we just created in the first place) + + # Generate ocean time series in results_dir ocean_month( input_dir, - parameters.subtask_name, - parameters.case_dir, + output_dir, parameters.year1, parameters.year2, int(parameters.ts_num_years_str), ) + # Input: MOC file from case_dir (existing MPAS-Analysis results) src: str = ( f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}" ) - dst: str = ( - f"{parameters.case_dir}/post/{parameters.subtask_name}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr/" - ) + # Output: Copy to results_dir (new ocean time series location) + dst: str = output_dir + "/" logger.info(f"Copy moc file from {src} to {dst}") - # NOTE: MODIFIES THE CASE DIRECTORY (parameters.case_dir) post subdirectory - # Copies files to post//ocn (which we just created in the first place) - shutil.copy( - src, - dst, - ) + shutil.copy(src, dst) def ocean_month( path_in: str, - subtask_name: str, - case_dir: str, + path_out: str, start_yr: int, end_yr: int, ts_num_years: int, ): - path_out = f"{case_dir}/post/{subtask_name}/ocn/glb/ts/monthly/{ts_num_years}yr" + # path_out is now directly provided as the output directory # Ocean constants # specific heat [J/(kg*degC)] From 369f1cfeddfbbdb48c825887da793c2ec7547e2d Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 4 Sep 2025 15:42:17 -0500 Subject: [PATCH 05/17] remove irrelavent parameters subsection/subtask_name --- tests/unit/global_time_series/test_global_time_series.py | 2 -- zppy_interfaces/global_time_series/__main__.py | 1 - zppy_interfaces/global_time_series/utils.py | 1 - 3 files changed, 4 deletions(-) diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index 45155d8..a72ba3b 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -83,7 +83,6 @@ def test_Parameters_and_related_functions(): "regions": "glb,n,s", "start_yr": "1985", "end_yr": "1989", - "subsection": "subsection_name", "use_ocn": "True", "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", "input_subdir": "archive/atm/hist", @@ -111,7 +110,6 @@ def test_Parameters_and_related_functions(): assert parameters.regions == ["glb", "n", "s"] assert parameters.year1 == 1985 assert parameters.year2 == 1989 - assert parameters.subtask_name == "global_time_series_subsection_name" assert parameters.plots_original == [] assert parameters.nrows == 1 assert parameters.ncols == 1 diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 17f4cb2..48e7ba3 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -84,7 +84,6 @@ def _get_args() -> Parameters: # For plots_original # > For ocean_month - parser.add_argument("--subsection", type=str, help="Subtask name") parser.add_argument("--use_ocn", type=str, help="Use ocean") parser.add_argument("--input", type=str, help="Input directory") parser.add_argument("--input_subdir", type=str, help="Input subdirectory") diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index 762bdb7..19f4244 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -29,7 +29,6 @@ def __init__(self, args: Dict[str, str]): # Used by Classic PDF only # For ocean_month - self.subtask_name: str = f"global_time_series_{args['subsection']}" self.use_ocn: bool = _str2bool(args["use_ocn"]) self.input: str = args["input"] self.input_subdir: str = args["input_subdir"] From 06435ad4a6f95dcdef96ed9d509276a95a7c0b7c Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 4 Sep 2025 15:53:02 -0500 Subject: [PATCH 06/17] fix pre-committing; and viewer row texts for clarity --- .../coupled_global/mix_viewer_original.py | 4 ++-- zppy_interfaces/global_time_series/create_ocean_ts.py | 10 ++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py index 5adf936..a8b25b4 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_original.py @@ -18,7 +18,7 @@ def create_viewer_for_original(parameters: Parameters) -> str: viewer.add_page(f"table_{component}", parameters.regions) viewer.add_group("Original Plots") # PDFs -- these don't show up when clicked on, but they CAN be downloaded - row_title: str = "Original Plots, PDFs" + row_title: str = "Original Plots, PDFs (click for download)" viewer.add_row(row_title) for rgn in parameters.regions: viewer.add_col( @@ -27,7 +27,7 @@ def create_viewer_for_original(parameters: Parameters) -> str: title=f"{rgn}_{component}", ) # PNGs -- these show up when clicked on - row_title = "Original Plots, PNGs" + row_title = "Original Plots, PNGs (click to view)" viewer.add_row(row_title) for rgn in parameters.regions: viewer.add_col( diff --git a/zppy_interfaces/global_time_series/create_ocean_ts.py b/zppy_interfaces/global_time_series/create_ocean_ts.py index bee0307..790228e 100644 --- a/zppy_interfaces/global_time_series/create_ocean_ts.py +++ b/zppy_interfaces/global_time_series/create_ocean_ts.py @@ -18,12 +18,14 @@ def create_ocean_ts(parameters: Parameters): logger.info("Create ocean time series") # Create output directory in results_dir (simplified structure) - output_dir = f"{parameters.results_dir}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr" + output_dir = ( + f"{parameters.results_dir}/ocn/glb/ts/monthly/{parameters.ts_num_years_str}yr" + ) os.makedirs(output_dir, exist_ok=True) - + # Input: Raw MPAS-O data input_dir: str = f"{parameters.input}/{parameters.input_subdir}" - + # Generate ocean time series in results_dir ocean_month( input_dir, @@ -33,7 +35,7 @@ def create_ocean_ts(parameters: Parameters): int(parameters.ts_num_years_str), ) - # Input: MOC file from case_dir (existing MPAS-Analysis results) + # Input: MOC file from case_dir (existing MPAS-Analysis results) src: str = ( f"{parameters.case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc/{parameters.moc_file}" ) From e8b871ec25abaedabb90e7b5c32e79e1b792aeaf Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 4 Sep 2025 16:16:13 -0500 Subject: [PATCH 07/17] parameter alignment with zppy; fix xarray furture warning --- .../test_global_time_series.py | 4 ++- .../global_time_series/__main__.py | 30 +++++++++---------- .../coupled_global/utils.py | 3 ++ zppy_interfaces/global_time_series/utils.py | 3 ++ 4 files changed, 24 insertions(+), 16 deletions(-) diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index a72ba3b..ba9fbee 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -83,7 +83,7 @@ def test_Parameters_and_related_functions(): "regions": "glb,n,s", "start_yr": "1985", "end_yr": "1989", - "use_ocn": "True", + "use_ocn": "False", "input": "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", "input_subdir": "archive/atm/hist", "moc_file": "mocTimeSeries_1985-1995.nc", @@ -143,6 +143,7 @@ def test_Parameters_and_related_functions(): "HR", ] assert parameters.plots_ocn == [] + assert parameters.use_ocn == False # test_get_data_dir assert ( @@ -201,6 +202,7 @@ def test_Parameters_and_related_functions(): assert exps[0] == expected + def test_Variable(): v = Variable( "var_name", diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 48e7ba3..365bc51 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -69,39 +69,39 @@ def _get_args() -> Parameters: # Used in all cases # > For determining which output type to produce - parser.add_argument("--make_viewer", type=str, help="Make viewer") + parser.add_argument("--make_viewer", type=str, default="False", help="Make viewer") # > For coupled_global parser.add_argument("--case_dir", type=str, help="Case directory") parser.add_argument("--experiment_name", type=str, help="Experiment name") - parser.add_argument("--figstr", type=str, help="Figure string") - parser.add_argument("--color", type=str, help="Color") - parser.add_argument("--ts_num_years", type=str, help="Time series number of years") + parser.add_argument("--figstr", type=str, default="", help="Figure string") + parser.add_argument("--color", type=str, default="Blue", help="Color") + parser.add_argument("--ts_num_years", type=str, default="5", help="Time series number of years") parser.add_argument("--results_dir", type=str, help="Results directory") - parser.add_argument("--regions", type=str, help="Regions") + parser.add_argument("--regions", type=str, default="glb,n,s", help="Regions") # > For both ocean_month and coupled_global parser.add_argument("--start_yr", type=str, help="Start year") parser.add_argument("--end_yr", type=str, help="End year") # For plots_original # > For ocean_month - parser.add_argument("--use_ocn", type=str, help="Use ocean") + parser.add_argument("--use_ocn", type=str, default="False", help="Use ocean (should match ocean plots in plots_original/plots_ocn)") parser.add_argument("--input", type=str, help="Input directory") - parser.add_argument("--input_subdir", type=str, help="Input subdirectory") - parser.add_argument("--moc_file", type=str, help="MOC file") + parser.add_argument("--input_subdir", type=str, default="archive/ocn/hist", help="Input subdirectory") + parser.add_argument("--moc_file", type=str, default="None", help="MOC file") # > For coupled_global - parser.add_argument("--plots_original", type=str, help="Plots original") + parser.add_argument("--plots_original", type=str, default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", help="Plots original") # For plots_component # > For coupled_global - parser.add_argument("--plots_atm", type=str, help="Plots atmosphere") - parser.add_argument("--plots_ice", type=str, help="Plots ice") - parser.add_argument("--plots_lnd", type=str, help="Plots land") - parser.add_argument("--plots_ocn", type=str, help="Plots ocean") + parser.add_argument("--plots_atm", type=str, default="None", help="Plots atmosphere") + parser.add_argument("--plots_ice", type=str, default="None", help="Plots ice") + parser.add_argument("--plots_lnd", type=str, default="None", help="Plots land") + parser.add_argument("--plots_ocn", type=str, default="None", help="Plots ocean") # Used for mode_pdf, regardless of plot type # > For coupled_global - parser.add_argument("--ncols", type=str, help="Number of columns in pdf") - parser.add_argument("--nrows", type=str, help="Number of rows in pdf") + parser.add_argument("--ncols", type=str, default="2", help="Number of columns in pdf") + parser.add_argument("--nrows", type=str, default="4", help="Number of rows in pdf") # Ignore the first arg # (zi-global-time-series) diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index c401dd4..57fadc3 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -9,6 +9,9 @@ import xarray import xcdat +# Set xarray options to suppress FutureWarning about compat parameter defaults +xarray.set_options(use_new_combine_kwarg_defaults=True) + from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index 19f4244..ced6967 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -50,6 +50,9 @@ def __init__(self, args: Dict[str, str]): self.plots_ocn: List[str] = param_get_list(args["plots_ocn"]) # Input validation + # Note: use_ocn should be True if ocean plots are requested in plots_original + # (change_ohc, max_moc, change_sea_level) or plots_ocn is non-empty. + # This follows zppy's logic where use_ocn is auto-determined from plot content. if self.plots_original and self.use_ocn and (not self.moc_file): raise ValueError( "moc_file must be set for ocean plots in the original 8-plot set." From 9dd615f6cc14e1f8c9288322a0842023006234f2 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 4 Sep 2025 16:50:50 -0500 Subject: [PATCH 08/17] parameter alignment with zppy; fix xarray furture warning --- .../test_global_time_series.py | 3 +- .../global_time_series/__main__.py | 33 +++++++++++++++---- .../coupled_global/utils.py | 6 ++-- 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index ba9fbee..b999eb8 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -143,7 +143,7 @@ def test_Parameters_and_related_functions(): "HR", ] assert parameters.plots_ocn == [] - assert parameters.use_ocn == False + assert parameters.use_ocn is False # test_get_data_dir assert ( @@ -202,7 +202,6 @@ def test_Parameters_and_related_functions(): assert exps[0] == expected - def test_Variable(): v = Variable( "var_name", diff --git a/zppy_interfaces/global_time_series/__main__.py b/zppy_interfaces/global_time_series/__main__.py index 365bc51..ffa934e 100644 --- a/zppy_interfaces/global_time_series/__main__.py +++ b/zppy_interfaces/global_time_series/__main__.py @@ -75,7 +75,9 @@ def _get_args() -> Parameters: parser.add_argument("--experiment_name", type=str, help="Experiment name") parser.add_argument("--figstr", type=str, default="", help="Figure string") parser.add_argument("--color", type=str, default="Blue", help="Color") - parser.add_argument("--ts_num_years", type=str, default="5", help="Time series number of years") + parser.add_argument( + "--ts_num_years", type=str, default="5", help="Time series number of years" + ) parser.add_argument("--results_dir", type=str, help="Results directory") parser.add_argument("--regions", type=str, default="glb,n,s", help="Regions") # > For both ocean_month and coupled_global @@ -84,23 +86,42 @@ def _get_args() -> Parameters: # For plots_original # > For ocean_month - parser.add_argument("--use_ocn", type=str, default="False", help="Use ocean (should match ocean plots in plots_original/plots_ocn)") + parser.add_argument( + "--use_ocn", + type=str, + default="False", + help="Use ocean (should match ocean plots in plots_original/plots_ocn)", + ) parser.add_argument("--input", type=str, help="Input directory") - parser.add_argument("--input_subdir", type=str, default="archive/ocn/hist", help="Input subdirectory") + parser.add_argument( + "--input_subdir", + type=str, + default="archive/ocn/hist", + help="Input subdirectory", + ) parser.add_argument("--moc_file", type=str, default="None", help="MOC file") # > For coupled_global - parser.add_argument("--plots_original", type=str, default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", help="Plots original") + parser.add_argument( + "--plots_original", + type=str, + default="net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", + help="Plots original", + ) # For plots_component # > For coupled_global - parser.add_argument("--plots_atm", type=str, default="None", help="Plots atmosphere") + parser.add_argument( + "--plots_atm", type=str, default="None", help="Plots atmosphere" + ) parser.add_argument("--plots_ice", type=str, default="None", help="Plots ice") parser.add_argument("--plots_lnd", type=str, default="None", help="Plots land") parser.add_argument("--plots_ocn", type=str, default="None", help="Plots ocean") # Used for mode_pdf, regardless of plot type # > For coupled_global - parser.add_argument("--ncols", type=str, default="2", help="Number of columns in pdf") + parser.add_argument( + "--ncols", type=str, default="2", help="Number of columns in pdf" + ) parser.add_argument("--nrows", type=str, default="4", help="Number of rows in pdf") # Ignore the first arg diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 57fadc3..eea170f 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -9,12 +9,12 @@ import xarray import xcdat -# Set xarray options to suppress FutureWarning about compat parameter defaults -xarray.set_options(use_new_combine_kwarg_defaults=True) - from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger +# Set xarray options to suppress FutureWarning about compat parameter defaults +xarray.set_options(use_new_combine_kwarg_defaults=True) + logger = _setup_child_logger(__name__) From e5dd503b40c1417fe93a53c6317b39b751178b68 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Fri, 5 Sep 2025 17:01:28 -0500 Subject: [PATCH 09/17] simplify logic for component plots; staging for memory/performance improvement --- .../coupled_global/utils.py | 146 ++++++++++++++---- 1 file changed, 119 insertions(+), 27 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index eea170f..271afe3 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -102,6 +102,108 @@ def get_vars_original(plots_original: List[str]) -> List[Variable]: # For plots_components +def get_variable_files( + var_name: str, directory: str, parameters: Parameters +) -> List[str]: + """Get list of NetCDF files for a single variable.""" + file_path_list: List[str] = [] + num_years: int = int(parameters.ts_num_years_str) + y1: int = parameters.year1 + y2: int = parameters.year1 + num_years - 1 + + while y2 <= parameters.year2: + file_path: str = f"{directory}{var_name}_{y1:04d}01_{y2:04d}12.nc" + if os.path.exists(file_path): + file_path_list.append(file_path) + y1 += num_years + y2 += num_years + + return file_path_list + + +def apply_scaling( + data_array: xarray.core.dataarray.DataArray, + metric: Metric, + dataset: xarray.core.dataset.Dataset, +) -> xarray.core.dataarray.DataArray: + """Apply area scaling for TOTAL metrics.""" + if metric != Metric.TOTAL: + return data_array + + # Calculate land areas for scaling + keys = list(dataset.keys()) + if "valid_area_per_gridcell" in keys: + land_area_per_gridcell = dataset["valid_area_per_gridcell"] + total_land_area = land_area_per_gridcell.sum() + north_land_area = land_area_per_gridcell.where( + land_area_per_gridcell.lat >= 0 + ).sum() + south_land_area = land_area_per_gridcell.where( + land_area_per_gridcell.lat < 0 + ).sum() + else: + area = dataset["area"] + landfrac = dataset["landfrac"] + total_land_area = (area * landfrac).sum() + north_area = area.where(area.lat >= 0) + north_landfrac = landfrac.where(landfrac.lat >= 0) + north_land_area = (north_area * north_landfrac).sum() + south_area = area.where(area.lat < 0) + south_landfrac = landfrac.where(landfrac.lat < 0) + south_land_area = (south_area * south_landfrac).sum() + + # Apply scaling + data_array[:, 0] *= total_land_area + data_array[:, 1] *= north_land_area + data_array[:, 2] *= south_land_area + + return data_array + + +def process_variable( + var: Variable, directory: str, parameters: Parameters +) -> Tuple[xarray.core.dataarray.DataArray, str]: + """Process a single variable independently - load, compute, cleanup.""" + try: + # 1. Get file paths for this variable + file_paths = get_variable_files(var.variable_name, directory, parameters) + if not file_paths: + raise ValueError(f"No data files found for variable {var.variable_name}") + + # 2. Load only this variable's data + dataset = xcdat.open_mfdataset(file_paths, center_times=True) + + try: + # 3. Compute annual average + annual_dataset = dataset.temporal.group_average(var.variable_name, "year") + data_array = annual_dataset.data_vars[var.variable_name] + + # 4. Apply area scaling if needed + data_array = apply_scaling(data_array, var.metric, dataset) + + # 5. Apply unit scaling + units = data_array.units + if ( + (units != "1") + and (var.original_units != "") + and var.original_units != units + ): + raise ValueError(f"Units don't match: {units} vs {var.original_units}") + if (var.scale_factor != 1) and (var.final_units != ""): + data_array *= var.scale_factor + units = var.final_units + + return data_array, units + + finally: + # 6. Always cleanup immediately + dataset.close() + + except Exception as e: + logger.error(f"Failed to process variable {var.variable_name}: {e}") + raise + + def construct_land_variables(requested_vars: List[str]) -> List[Variable]: var_list: List[Variable] = [] header = True @@ -392,43 +494,33 @@ def set_var( ) -> List[Variable]: new_var_list: List[Variable] = [] if parameters and (var_list == []): - # If we want to load specific variables, - # but none are specified, - # then we can just immediately return. return new_var_list + if exp[exp_key] != "": - try: - dataset_wrapper: DatasetWrapper - if parameters: - # If this is passed in, then we want to load specific vars. - dataset_wrapper = DatasetWrapper(exp[exp_key], var_list, parameters) - else: - dataset_wrapper = DatasetWrapper(exp[exp_key]) - except Exception as e: - logger.critical(e) - logger.critical( - f"DatasetWrapper object could not be created for {exp_key}={exp[exp_key]}" - ) - raise e + directory = exp[exp_key] + for var in var_list: var_str: str = var.variable_name try: - data_array: xarray.core.dataarray.DataArray - units: str - data_array, units = dataset_wrapper.globalAnnual(var) - valid_vars.append(str(var_str)) # Append the name - new_var_list.append(var) # Append the variable itself + if parameters: + # Use simplified approach with lazy loading + data_array, units = process_variable(var, directory, parameters) + else: + # Legacy fallback: use DatasetWrapper for backward compatibility + dataset_wrapper = DatasetWrapper(directory) + data_array, units = dataset_wrapper.globalAnnual(var) + del dataset_wrapper + + valid_vars.append(str(var_str)) + new_var_list.append(var) except Exception as e: logger.error(e) - logger.error(f"globalAnnual failed for {var_str}") + logger.error(f"Processing failed for {var_str}") invalid_vars.append(str(var_str)) continue + exp["annual"][var_str] = {"glb": (data_array.isel(rgn=0), units)} if data_array.sizes["rgn"] > 1: - # data_array.shape => number of years x 3 regions - # 3 regions = global, northern hemisphere, southern hemisphere - # We get here if we used the updated `ts` task - # (using `rgn_avg` rather than `glb_avg`). exp["annual"][var_str]["n"] = (data_array.isel(rgn=1), units) exp["annual"][var_str]["s"] = (data_array.isel(rgn=2), units) if "year" not in exp["annual"]: @@ -436,5 +528,5 @@ def set_var( "time" ].values exp["annual"]["year"] = [x.year for x in years] - del dataset_wrapper + return new_var_list From a5ceb5ddc6942f38b0993649671db0271d911e7d Mon Sep 17 00:00:00 2001 From: Ryan Forsyth Date: Thu, 4 Sep 2025 18:47:32 -0500 Subject: [PATCH 10/17] Fix unit test --- tests/unit/global_time_series/test_global_time_series.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/global_time_series/test_global_time_series.py b/tests/unit/global_time_series/test_global_time_series.py index b999eb8..87a339e 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -192,8 +192,8 @@ def test_Parameters_and_related_functions(): expected = { "atmos": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/atm/glb/ts/monthly/5yr/", "ocean": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/ocn/glb/ts/monthly/5yr/", - "moc": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/global_time_series_subsection_name/ocn/glb/ts/monthly/5yr/", - "vol": "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/global_time_series_subsection_name/ocn/glb/ts/monthly/5yr/", + "moc": "results/ocn/glb/ts/monthly/5yr/", + "vol": "results/ocn/glb/ts/monthly/5yr/", "name": "v3.LR.historical_0051", "yoffset": 0.0, "yr": ([1985, 1989],), From 9ea2a0ab47c0d54afc6f8dea69c74f2b355e3755 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Fri, 5 Sep 2025 17:43:08 -0500 Subject: [PATCH 11/17] use multiprocessing for component processing --- .../coupled_global/plots_component.py | 3 +- .../coupled_global/utils.py | 134 ++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_component.py b/zppy_interfaces/global_time_series/coupled_global/plots_component.py index 79e3376..c220025 100644 --- a/zppy_interfaces/global_time_series/coupled_global/plots_component.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_component.py @@ -8,6 +8,7 @@ RequestedVariables, get_data_dir, set_var, + set_var_parallel, ) from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger @@ -48,7 +49,7 @@ def process_data( invalid_vars, parameters, ) - requested_variables.vars_land = set_var( + requested_variables.vars_land = set_var_parallel( exp, "land", requested_variables.vars_land, diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 271afe3..f965402 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -1,5 +1,6 @@ import csv import importlib.resources as imp_res +import multiprocessing as mp import os.path from enum import Enum from typing import Any, Dict, List, Optional, Tuple @@ -204,6 +205,90 @@ def process_variable( raise +def process_variable_worker(args): + """ + Worker function for multiprocessing - unpacks arguments and processes single variable. + + Args: + args: Tuple of (var, directory, parameters) + + Returns: + Tuple of (var_name, data_array, units, success_flag, error_msg) + """ + var, directory, parameters = args + var_name = var.variable_name + + try: + data_array, units = process_variable(var, directory, parameters) + return (var_name, data_array, units, True, None) + except Exception as e: + return (var_name, None, None, False, str(e)) + + +def process_variables_parallel( + var_list: List[Variable], + directory: str, + parameters: Parameters, + num_processes: Optional[int] = None, +) -> Dict[str, Tuple[xarray.core.dataarray.DataArray, str]]: + """ + Process multiple variables in parallel using multiprocessing.Pool. + + Args: + var_list: List of variables to process + directory: Data directory path + parameters: Processing parameters + num_processes: Number of parallel processes (default: CPU count) + + Returns: + Dictionary mapping variable names to (data_array, units) tuples + + Raises: + Exception: If no variables processed successfully + """ + if num_processes is None: + num_processes = min(mp.cpu_count(), len(var_list)) + + logger.info( + f"Starting parallel processing of {len(var_list)} variables with {num_processes} processes" + ) + + # Prepare arguments for worker processes + worker_args = [(var, directory, parameters) for var in var_list] + + results = {} + failed_vars = [] + + try: + with mp.Pool(processes=num_processes) as pool: + # Process all variables in parallel + worker_results = pool.map(process_variable_worker, worker_args) + + # Collect results + for var_name, data_array, units, success, error_msg in worker_results: + if success: + results[var_name] = (data_array, units) + logger.info(f"✓ Completed processing variable: {var_name}") + else: + failed_vars.append(var_name) + logger.error(f"✗ Failed processing variable {var_name}: {error_msg}") + + except Exception as e: + logger.error(f"Parallel processing failed: {e}") + raise + + if failed_vars: + logger.warning(f"Failed to process {len(failed_vars)} variables: {failed_vars}") + + if not results: + raise Exception("No variables processed successfully") + + logger.info( + f"Parallel processing complete. {len(results)}/{len(var_list)} variables processed successfully" + ) + return results + + def construct_land_variables(requested_vars: List[str]) -> List[Variable]: var_list: List[Variable] = [] header = True @@ -484,6 +569,55 @@ def get_data_dir(parameters: Parameters, component: str, conditional: bool) -> s ) +def set_var_parallel( + exp: Dict[str, Any], + exp_key: str, + var_list: List[Variable], + valid_vars: List[str], + invalid_vars: List[str], + parameters: Parameters, + num_processes: Optional[int] = None, +) -> List[Variable]: + """Parallel version of set_var for component plots.""" + new_var_list: List[Variable] = [] + if var_list == []: + return new_var_list + + if exp[exp_key] != "": + directory = exp[exp_key] + + if len(var_list) > 1: + # Use parallel processing + results = process_variables_parallel( + var_list, directory, parameters, num_processes + ) + + # Process results + for var in var_list: + var_str = var.variable_name + if var_str in results: + data_array, units = results[var_str] + valid_vars.append(var_str) + new_var_list.append(var) + + exp["annual"][var_str] = {"glb": (data_array.isel(rgn=0), units)} + if data_array.sizes["rgn"] > 1: + exp["annual"][var_str]["n"] = (data_array.isel(rgn=1), units) + exp["annual"][var_str]["s"] = (data_array.isel(rgn=2), units) + if "year" not in exp["annual"]: + years: np.ndarray[cftime.DatetimeNoLeap] = data_array.coords[ + "time" + ].values + exp["annual"]["year"] = [x.year for x in years] + else: + invalid_vars.append(var_str) + else: + # Single variable - use sequential + return set_var(exp, exp_key, var_list, valid_vars, invalid_vars, parameters) + + return new_var_list + + def set_var( exp: Dict[str, Any], exp_key: str, From 63b0c9389a880b1e93dfeaf07426943e18b41bcb Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Fri, 5 Sep 2025 22:49:38 -0500 Subject: [PATCH 12/17] make plotting more performant and memo efficient by combining in processing --- .../coupled_global/driver.py | 28 ++- .../coupled_global/plots_component.py | 11 +- .../coupled_global/utils.py | 234 ++++++++++++++++++ 3 files changed, 257 insertions(+), 16 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/driver.py b/zppy_interfaces/global_time_series/coupled_global/driver.py index fd58987..7a1174f 100644 --- a/zppy_interfaces/global_time_series/coupled_global/driver.py +++ b/zppy_interfaces/global_time_series/coupled_global/driver.py @@ -71,16 +71,24 @@ def run(parameters: Parameters, requested_variables: RequestedVariables): else: exps = exps_component if parameters.make_viewer: - produce_pngs_for_viewer( - parameters, - rgn, - component, - xlim, - exps, - plot_list, - valid_plots, - invalid_plots, - ) + # Skip original plotting for all components since combined approach already generated plots + if component in ["atm", "ice", "lnd", "ocn"]: + # For all components: plots already generated by combined processing+plotting + valid_plots.extend(plot_list) + logger.info( + f"Skipping original plotting for {component} - plots already generated during processing" + ) + else: + produce_pngs_for_viewer( + parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, + ) else: assemble_cumulative_pdf( parameters, diff --git a/zppy_interfaces/global_time_series/coupled_global/plots_component.py b/zppy_interfaces/global_time_series/coupled_global/plots_component.py index c220025..7e1177d 100644 --- a/zppy_interfaces/global_time_series/coupled_global/plots_component.py +++ b/zppy_interfaces/global_time_series/coupled_global/plots_component.py @@ -7,8 +7,7 @@ from zppy_interfaces.global_time_series.coupled_global.utils import ( RequestedVariables, get_data_dir, - set_var, - set_var_parallel, + set_var_parallel_with_plots, ) from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger @@ -33,7 +32,7 @@ def process_data( for exp in exps: exp["annual"] = {} - requested_variables.vars_atm = set_var( + requested_variables.vars_atm = set_var_parallel_with_plots( exp, "atmos", requested_variables.vars_atm, @@ -41,7 +40,7 @@ def process_data( invalid_vars, parameters, ) - requested_variables.vars_ice = set_var( + requested_variables.vars_ice = set_var_parallel_with_plots( exp, "ice", requested_variables.vars_ice, @@ -49,7 +48,7 @@ def process_data( invalid_vars, parameters, ) - requested_variables.vars_land = set_var_parallel( + requested_variables.vars_land = set_var_parallel_with_plots( exp, "land", requested_variables.vars_land, @@ -57,7 +56,7 @@ def process_data( invalid_vars, parameters, ) - requested_variables.vars_ocn = set_var( + requested_variables.vars_ocn = set_var_parallel_with_plots( exp, "ocean", requested_variables.vars_ocn, diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index f965402..4947414 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -225,6 +225,134 @@ def process_variable_worker(args): return (var_name, None, None, False, str(e)) +def process_and_plot_worker(args): + """ + Combined worker: process variable and generate plots immediately. + + Args: + args: Tuple of (var, directory, parameters, plot_config) + + Returns: + Tuple of (var_name, success_flag, error_msg, plot_info, data_array, units) + """ + var, directory, parameters, plot_config = args + var_name = var.variable_name + + try: + # Process the variable first + data_array, units = process_variable(var, directory, parameters) + + # Check if we got valid data + if data_array is None: + return ( + var_name, + False, + "No data returned from processing", + None, + None, + None, + ) + + # Generate plots only if processing succeeded + component_name = plot_config.get("component", "lnd") + plot_info = generate_variable_plots( + var_name, data_array, units, parameters, plot_config, component_name + ) + + # Return data for populating exp["annual"] - don't delete here + return (var_name, True, None, plot_info, data_array, units) + + except Exception as e: + # Processing failed - don't attempt plotting + return (var_name, False, str(e), None, None, None) + + +def generate_variable_plots( + var_name: str, + data_array: xarray.core.dataarray.DataArray, + units: str, + parameters: Parameters, + plot_config: Dict[str, Any], + component: str, +) -> Dict[str, Any]: + """ + Generate PNG plots for a single variable immediately after processing. + + Args: + var_name: Variable name + data_array: Processed data array + units: Data units + parameters: Processing parameters + plot_config: Plotting configuration + + Returns: + Dictionary with plot file paths and metadata + """ + import matplotlib.pyplot as plt + + plot_info: Dict[str, Any] = {"var_name": var_name, "plots": []} + + # Validate input data + if data_array is None: + logger.error(f"Cannot plot {var_name}: data_array is None") + return plot_info + + if data_array.size == 0: + logger.error(f"Cannot plot {var_name}: data_array is empty") + return plot_info + + # Create temporary exp structure for plotting compatibility + temp_exp = { + "annual": {var_name: {"glb": (data_array.isel(rgn=0), units)}}, + "color": plot_config.get("color", "blue"), + "name": plot_config.get("name", "data"), + "yoffset": plot_config.get("yoffset", 0), + "yr": ([parameters.year1, parameters.year2],), + } + + if data_array.sizes["rgn"] > 1: + temp_exp["annual"][var_name]["n"] = (data_array.isel(rgn=1), units) + temp_exp["annual"][var_name]["s"] = (data_array.isel(rgn=2), units) + + # Add year data + years = data_array.coords["time"].values + temp_exp["annual"]["year"] = [x.year for x in years] + + # Generate plots for each region + regions = ["glb"] + if data_array.sizes["rgn"] > 1: + regions.extend(["n", "s"]) + + for rgn in regions: + try: + fig, ax = plt.subplots(figsize=(10, 6)) + + # Use existing plot_generic function with proper xlim + xlim = [parameters.year1, parameters.year2] + from zppy_interfaces.global_time_series.coupled_global.plots_component import ( + plot_generic, + ) + + plot_generic(ax, xlim, [temp_exp], var_name, rgn) + + # Save plot + plot_filename = f"{parameters.figstr}_{rgn}_{component}_{var_name}.png" + plot_path = f"{parameters.results_dir}/{plot_filename}" + fig.savefig(plot_path, dpi=150, bbox_inches="tight") + plt.close(fig) + + plot_info["plots"].append( + {"region": rgn, "filename": plot_filename, "path": plot_path} + ) + + logger.debug(f"Generated plot: {plot_filename}") + + except Exception as e: + logger.error(f"Failed to generate plot for {var_name}_{rgn}: {e}") + + return plot_info + + def process_variables_parallel( var_list: List[Variable], directory: str, @@ -569,6 +697,112 @@ def get_data_dir(parameters: Parameters, component: str, conditional: bool) -> s ) +def set_var_parallel_with_plots( + exp: Dict[str, Any], + exp_key: str, + var_list: List[Variable], + valid_vars: List[str], + invalid_vars: List[str], + parameters: Parameters, + num_processes: Optional[int] = None, +) -> List[Variable]: + """Combined parallel processing + plotting version.""" + new_var_list: List[Variable] = [] + if var_list == []: + return new_var_list + + if exp[exp_key] != "": + directory = exp[exp_key] + + if len(var_list) > 1: + # Map exp_key to component name for filename + component_map = { + "atmos": "atm", + "ice": "ice", + "land": "lnd", + "ocean": "ocn", + } + component_name = component_map.get(exp_key, exp_key) + + # Combined processing + plotting + plot_config = { + "color": exp.get("color", "blue"), + "name": exp.get("name", "data"), + "yoffset": exp.get("yoffset", 0), + "component": component_name, + } + + # Prepare arguments for combined workers + worker_args = [ + (var, directory, parameters, plot_config) for var in var_list + ] + + if num_processes is None: + num_processes = min(mp.cpu_count(), len(var_list)) + + logger.info( + f"Starting combined processing+plotting of {len(var_list)} variables with {num_processes} processes" + ) + + try: + with mp.Pool(processes=num_processes) as pool: + worker_results = pool.map(process_and_plot_worker, worker_args) + + # Process results + for ( + var_name, + success, + error_msg, + plot_info, + data_array, + units, + ) in worker_results: + if success: + valid_vars.append(var_name) + # Find the variable object + var_obj = next( + v for v in var_list if v.variable_name == var_name + ) + new_var_list.append(var_obj) + + # Populate exp["annual"] with processed data + exp["annual"][var_name] = { + "glb": (data_array.isel(rgn=0), units) + } + if data_array.sizes["rgn"] > 1: + exp["annual"][var_name]["n"] = ( + data_array.isel(rgn=1), + units, + ) + exp["annual"][var_name]["s"] = ( + data_array.isel(rgn=2), + units, + ) + if "year" not in exp["annual"]: + years = data_array.coords["time"].values + exp["annual"]["year"] = [x.year for x in years] + + # Clean up data_array after storing + del data_array + + logger.info(f"Completed processing+plotting for: {var_name}") + if plot_info: + logger.debug( + f"Generated {len(plot_info['plots'])} plots for {var_name}" + ) + else: + invalid_vars.append(var_name) + + except Exception as e: + logger.error(f"Combined parallel processing+plotting failed: {e}") + raise + else: + # Single variable - use sequential + return set_var(exp, exp_key, var_list, valid_vars, invalid_vars, parameters) + + return new_var_list + + def set_var_parallel( exp: Dict[str, Any], exp_key: str, From 2333e8d335731ea2b5ef05663472b7ab799b7eab Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 18 Sep 2025 18:39:31 -0500 Subject: [PATCH 13/17] Switch from multiprocessing to sequential processing for variable handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multiprocessing with 128+ processes was causing deadlocks and memory issues when processing 300+ NetCDF variables. Sequential processing provides better reliability and memory management for this I/O-intensive workload. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- .../coupled_global/utils.py | 108 +++++++++--------- 1 file changed, 51 insertions(+), 57 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 4947414..2cb9edf 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -9,6 +9,9 @@ import numpy as np import xarray import xcdat +import matplotlib +matplotlib.use('Agg') # Use non-interactive backend +import matplotlib.pyplot as plt from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger @@ -288,8 +291,6 @@ def generate_variable_plots( Returns: Dictionary with plot file paths and metadata """ - import matplotlib.pyplot as plt - plot_info: Dict[str, Any] = {"var_name": var_name, "plots": []} # Validate input data @@ -375,7 +376,7 @@ def process_variables_parallel( Exception: If no variables processed successfully """ if num_processes is None: - num_processes = min(mp.cpu_count(), len(var_list)) + num_processes = min(16, len(var_list)) logger.info( f"Starting parallel processing of {len(var_list)} variables with {num_processes} processes" @@ -738,64 +739,57 @@ def set_var_parallel_with_plots( ] if num_processes is None: - num_processes = min(mp.cpu_count(), len(var_list)) + num_processes = min(16, len(var_list)) - logger.info( - f"Starting combined processing+plotting of {len(var_list)} variables with {num_processes} processes" - ) + logger.info(f"Processing {len(var_list)} variables") - try: - with mp.Pool(processes=num_processes) as pool: - worker_results = pool.map(process_and_plot_worker, worker_args) - - # Process results - for ( - var_name, - success, - error_msg, - plot_info, - data_array, - units, - ) in worker_results: - if success: - valid_vars.append(var_name) - # Find the variable object - var_obj = next( - v for v in var_list if v.variable_name == var_name + worker_results = [] + for i, var in enumerate(var_list): + logger.info(f"Processing {i+1}/{len(var_list)}: {var.variable_name}") + try: + result = process_and_plot_worker((var, directory, parameters, plot_config)) + worker_results.append(result) + if not result[1]: + logger.error(f"Failed {var.variable_name}: {result[2]}") + except Exception as e: + logger.error(f"Exception processing {var.variable_name}: {e}") + worker_results.append((var.variable_name, False, str(e), None, None, None)) + + # Process results + for ( + var_name, + success, + error_msg, + plot_info, + data_array, + units, + ) in worker_results: + if success: + valid_vars.append(var_name) + var_obj = next( + v for v in var_list if v.variable_name == var_name + ) + new_var_list.append(var_obj) + + exp["annual"][var_name] = { + "glb": (data_array.isel(rgn=0), units) + } + if data_array.sizes["rgn"] > 1: + exp["annual"][var_name]["n"] = ( + data_array.isel(rgn=1), + units, ) - new_var_list.append(var_obj) - - # Populate exp["annual"] with processed data - exp["annual"][var_name] = { - "glb": (data_array.isel(rgn=0), units) - } - if data_array.sizes["rgn"] > 1: - exp["annual"][var_name]["n"] = ( - data_array.isel(rgn=1), - units, - ) - exp["annual"][var_name]["s"] = ( - data_array.isel(rgn=2), - units, - ) - if "year" not in exp["annual"]: - years = data_array.coords["time"].values - exp["annual"]["year"] = [x.year for x in years] - - # Clean up data_array after storing - del data_array - - logger.info(f"Completed processing+plotting for: {var_name}") - if plot_info: - logger.debug( - f"Generated {len(plot_info['plots'])} plots for {var_name}" - ) - else: - invalid_vars.append(var_name) + exp["annual"][var_name]["s"] = ( + data_array.isel(rgn=2), + units, + ) + if "year" not in exp["annual"]: + years = data_array.coords["time"].values + exp["annual"]["year"] = [x.year for x in years] - except Exception as e: - logger.error(f"Combined parallel processing+plotting failed: {e}") - raise + del data_array + else: + invalid_vars.append(var_name) else: # Single variable - use sequential return set_var(exp, exp_key, var_list, valid_vars, invalid_vars, parameters) From 5de22baf76950795d4baab0d501d55e38ba07f30 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 18 Sep 2025 18:42:34 -0500 Subject: [PATCH 14/17] Add example script for global time series processing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds run_zi_global-time-series.py to demonstrate usage of the global time series functionality with real data paths and performance timing. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- examples/run_zi_global-time-series.py | 35 +++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 examples/run_zi_global-time-series.py diff --git a/examples/run_zi_global-time-series.py b/examples/run_zi_global-time-series.py new file mode 100644 index 0000000..d83da60 --- /dev/null +++ b/examples/run_zi_global-time-series.py @@ -0,0 +1,35 @@ +from zppy_interfaces.global_time_series.__main__ import main +import sys +import time + +sys.argv.extend([ + "--use_ocn", "True", + "--input", "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", + "--input_subdir", "archive/ocn/hist", + "--moc_file", "mocTimeSeries_1985-1995.nc", + "--case_dir", "/lcrc/group/e3sm/ac.forsyth2/zppy_weekly_comprehensive_v3_output/test_issue-23-rebased-20250903/v3.LR.historical_0051", + "--experiment_name", "v3.LR.historical_0051", + "--figstr", "v3.LR.historical_0051", + "--color", "Blue", + "--ts_num_years", "5", + "--plots_original","net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", + #"--plots_original","None", + "--plots_atm", "TREFHT", + "--plots_ice", "None", + "--plots_lnd", "all", + #"--plots_lnd", "FSH", + "--plots_ocn", "None", + "--nrows", "4", + "--ncols", "2", + "--results_dir", "/lcrc/group/e3sm/public_html/diagnostic_output/ac.zhang40/tests/zi", + "--regions", "glb,n,s", + "--make_viewer", "True", + "--start_yr", "1985", + "--end_yr", "1995" +]) + +start_time = time.time() +main() +end_time = time.time() + +print(f"Execution time: {end_time - start_time:.2f} seconds") From 27729f07bdba2fc530ee4cb02c70b1d611980dbc Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Thu, 18 Sep 2025 18:51:11 -0500 Subject: [PATCH 15/17] fix pre-committing --- examples/run_zi_global-time-series.py | 76 ++++++++++++------- .../coupled_global/utils.py | 31 ++++---- 2 files changed, 63 insertions(+), 44 deletions(-) diff --git a/examples/run_zi_global-time-series.py b/examples/run_zi_global-time-series.py index d83da60..ba8fff9 100644 --- a/examples/run_zi_global-time-series.py +++ b/examples/run_zi_global-time-series.py @@ -1,32 +1,56 @@ -from zppy_interfaces.global_time_series.__main__ import main import sys import time -sys.argv.extend([ - "--use_ocn", "True", - "--input", "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", - "--input_subdir", "archive/ocn/hist", - "--moc_file", "mocTimeSeries_1985-1995.nc", - "--case_dir", "/lcrc/group/e3sm/ac.forsyth2/zppy_weekly_comprehensive_v3_output/test_issue-23-rebased-20250903/v3.LR.historical_0051", - "--experiment_name", "v3.LR.historical_0051", - "--figstr", "v3.LR.historical_0051", - "--color", "Blue", - "--ts_num_years", "5", - "--plots_original","net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", - #"--plots_original","None", - "--plots_atm", "TREFHT", - "--plots_ice", "None", - "--plots_lnd", "all", - #"--plots_lnd", "FSH", - "--plots_ocn", "None", - "--nrows", "4", - "--ncols", "2", - "--results_dir", "/lcrc/group/e3sm/public_html/diagnostic_output/ac.zhang40/tests/zi", - "--regions", "glb,n,s", - "--make_viewer", "True", - "--start_yr", "1985", - "--end_yr", "1995" -]) +from zppy_interfaces.global_time_series.__main__ import main + +sys.argv.extend( + [ + "--use_ocn", + "True", + "--input", + "/lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051", + "--input_subdir", + "archive/ocn/hist", + "--moc_file", + "mocTimeSeries_1985-1995.nc", + "--case_dir", + "/lcrc/group/e3sm/ac.forsyth2/zppy_weekly_comprehensive_v3_output/test_issue-23-rebased-20250903/v3.LR.historical_0051", + "--experiment_name", + "v3.LR.historical_0051", + "--figstr", + "v3.LR.historical_0051", + "--color", + "Blue", + "--ts_num_years", + "5", + "--plots_original", + "net_toa_flux_restom,global_surface_air_temperature,toa_radiation,net_atm_energy_imbalance,change_ohc,max_moc,change_sea_level,net_atm_water_imbalance", + # "--plots_original","None", + "--plots_atm", + "TREFHT", + "--plots_ice", + "None", + "--plots_lnd", + "all", + # "--plots_lnd", "FSH", + "--plots_ocn", + "None", + "--nrows", + "4", + "--ncols", + "2", + "--results_dir", + "/lcrc/group/e3sm/public_html/diagnostic_output/ac.zhang40/tests/zi", + "--regions", + "glb,n,s", + "--make_viewer", + "True", + "--start_yr", + "1985", + "--end_yr", + "1995", + ] +) start_time = time.time() main() diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 2cb9edf..80d67f0 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -6,17 +6,17 @@ from typing import Any, Dict, List, Optional, Tuple import cftime +import matplotlib +import matplotlib.pyplot as plt import numpy as np import xarray import xcdat -import matplotlib -matplotlib.use('Agg') # Use non-interactive backend -import matplotlib.pyplot as plt from zppy_interfaces.global_time_series.utils import Parameters from zppy_interfaces.multi_utils.logger import _setup_child_logger -# Set xarray options to suppress FutureWarning about compat parameter defaults +# Set matplotlib backend and xarray options +matplotlib.use("Agg") # Use non-interactive backend xarray.set_options(use_new_combine_kwarg_defaults=True) logger = _setup_child_logger(__name__) @@ -733,11 +733,6 @@ def set_var_parallel_with_plots( "component": component_name, } - # Prepare arguments for combined workers - worker_args = [ - (var, directory, parameters, plot_config) for var in var_list - ] - if num_processes is None: num_processes = min(16, len(var_list)) @@ -745,15 +740,19 @@ def set_var_parallel_with_plots( worker_results = [] for i, var in enumerate(var_list): - logger.info(f"Processing {i+1}/{len(var_list)}: {var.variable_name}") + logger.info(f"Processing {i + 1}/{len(var_list)}: {var.variable_name}") try: - result = process_and_plot_worker((var, directory, parameters, plot_config)) + result = process_and_plot_worker( + (var, directory, parameters, plot_config) + ) worker_results.append(result) if not result[1]: logger.error(f"Failed {var.variable_name}: {result[2]}") except Exception as e: logger.error(f"Exception processing {var.variable_name}: {e}") - worker_results.append((var.variable_name, False, str(e), None, None, None)) + worker_results.append( + (var.variable_name, False, str(e), None, None, None) + ) # Process results for ( @@ -766,14 +765,10 @@ def set_var_parallel_with_plots( ) in worker_results: if success: valid_vars.append(var_name) - var_obj = next( - v for v in var_list if v.variable_name == var_name - ) + var_obj = next(v for v in var_list if v.variable_name == var_name) new_var_list.append(var_obj) - exp["annual"][var_name] = { - "glb": (data_array.isel(rgn=0), units) - } + exp["annual"][var_name] = {"glb": (data_array.isel(rgn=0), units)} if data_array.sizes["rgn"] > 1: exp["annual"][var_name]["n"] = ( data_array.isel(rgn=1), From 6fab8a39730c005de9df58a8c8f72c2c9b9327a5 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Tue, 30 Sep 2025 15:30:30 -0500 Subject: [PATCH 16/17] fix directory generation for component only config --- .../global_time_series/coupled_global/mix_viewer_component.py | 4 ++++ zppy_interfaces/global_time_series/coupled_global/utils.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py index ce47363..3856e8f 100644 --- a/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py @@ -107,6 +107,10 @@ def create_viewer_for_component( logger.info(f"Creating viewer for {component}") if not vars: raise RuntimeError("No vars specified for viewer.") + + # Ensure results directory exists before creating viewer + os.makedirs(parameters.results_dir, exist_ok=True) + index_name = f"zppy global time-series plot: {parameters.experiment_name} {component} component ({parameters.year1}-{parameters.year2})" viewer = OutputViewer(path=parameters.results_dir, index_name=index_name) viewer.add_page(f"table_{component}", parameters.regions) diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 80d67f0..77fe545 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -1,6 +1,7 @@ import csv import importlib.resources as imp_res import multiprocessing as mp +import os import os.path from enum import Enum from typing import Any, Dict, List, Optional, Tuple @@ -336,6 +337,9 @@ def generate_variable_plots( plot_generic(ax, xlim, [temp_exp], var_name, rgn) + # Ensure results directory exists before saving plot + os.makedirs(parameters.results_dir, exist_ok=True) + # Save plot plot_filename = f"{parameters.figstr}_{rgn}_{component}_{var_name}.png" plot_path = f"{parameters.results_dir}/{plot_filename}" From 3909dfae8722a083baaf53190210605f4b16cbf6 Mon Sep 17 00:00:00 2001 From: chengzhuzhang Date: Wed, 1 Oct 2025 15:42:27 -0500 Subject: [PATCH 17/17] fix for single variable --- .../coupled_global/utils.py | 136 +++++++++--------- 1 file changed, 66 insertions(+), 70 deletions(-) diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py index 77fe545..664fa4b 100644 --- a/zppy_interfaces/global_time_series/coupled_global/utils.py +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -719,79 +719,75 @@ def set_var_parallel_with_plots( if exp[exp_key] != "": directory = exp[exp_key] - if len(var_list) > 1: - # Map exp_key to component name for filename - component_map = { - "atmos": "atm", - "ice": "ice", - "land": "lnd", - "ocean": "ocn", - } - component_name = component_map.get(exp_key, exp_key) - - # Combined processing + plotting - plot_config = { - "color": exp.get("color", "blue"), - "name": exp.get("name", "data"), - "yoffset": exp.get("yoffset", 0), - "component": component_name, - } - - if num_processes is None: - num_processes = min(16, len(var_list)) - - logger.info(f"Processing {len(var_list)} variables") - - worker_results = [] - for i, var in enumerate(var_list): - logger.info(f"Processing {i + 1}/{len(var_list)}: {var.variable_name}") - try: - result = process_and_plot_worker( - (var, directory, parameters, plot_config) + # Map exp_key to component name for filename + component_map = { + "atmos": "atm", + "ice": "ice", + "land": "lnd", + "ocean": "ocn", + } + component_name = component_map.get(exp_key, exp_key) + + # Combined processing + plotting + plot_config = { + "color": exp.get("color", "blue"), + "name": exp.get("name", "data"), + "yoffset": exp.get("yoffset", 0), + "component": component_name, + } + + if num_processes is None: + num_processes = min(16, len(var_list)) + + logger.info(f"Processing {len(var_list)} variables") + + worker_results = [] + for i, var in enumerate(var_list): + logger.info(f"Processing {i + 1}/{len(var_list)}: {var.variable_name}") + try: + result = process_and_plot_worker( + (var, directory, parameters, plot_config) + ) + worker_results.append(result) + if not result[1]: + logger.error(f"Failed {var.variable_name}: {result[2]}") + except Exception as e: + logger.error(f"Exception processing {var.variable_name}: {e}") + worker_results.append( + (var.variable_name, False, str(e), None, None, None) + ) + + # Process results + for ( + var_name, + success, + error_msg, + plot_info, + data_array, + units, + ) in worker_results: + if success: + valid_vars.append(var_name) + var_obj = next(v for v in var_list if v.variable_name == var_name) + new_var_list.append(var_obj) + + exp["annual"][var_name] = {"glb": (data_array.isel(rgn=0), units)} + if data_array.sizes["rgn"] > 1: + exp["annual"][var_name]["n"] = ( + data_array.isel(rgn=1), + units, ) - worker_results.append(result) - if not result[1]: - logger.error(f"Failed {var.variable_name}: {result[2]}") - except Exception as e: - logger.error(f"Exception processing {var.variable_name}: {e}") - worker_results.append( - (var.variable_name, False, str(e), None, None, None) + exp["annual"][var_name]["s"] = ( + data_array.isel(rgn=2), + units, ) + if "year" not in exp["annual"]: + years = data_array.coords["time"].values + exp["annual"]["year"] = [x.year for x in years] - # Process results - for ( - var_name, - success, - error_msg, - plot_info, - data_array, - units, - ) in worker_results: - if success: - valid_vars.append(var_name) - var_obj = next(v for v in var_list if v.variable_name == var_name) - new_var_list.append(var_obj) - - exp["annual"][var_name] = {"glb": (data_array.isel(rgn=0), units)} - if data_array.sizes["rgn"] > 1: - exp["annual"][var_name]["n"] = ( - data_array.isel(rgn=1), - units, - ) - exp["annual"][var_name]["s"] = ( - data_array.isel(rgn=2), - units, - ) - if "year" not in exp["annual"]: - years = data_array.coords["time"].values - exp["annual"]["year"] = [x.year for x in years] - - del data_array - else: - invalid_vars.append(var_name) - else: - # Single variable - use sequential - return set_var(exp, exp_key, var_list, valid_vars, invalid_vars, parameters) + del data_array + else: + invalid_vars.append(var_name) return new_var_list