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/examples/run_zi_global-time-series.py b/examples/run_zi_global-time-series.py new file mode 100644 index 0000000..ba8fff9 --- /dev/null +++ b/examples/run_zi_global-time-series.py @@ -0,0 +1,59 @@ +import sys +import time + +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() +end_time = time.time() + +print(f"Execution time: {end_time - start_time:.2f} seconds") 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..87a339e 100644 --- a/tests/unit/global_time_series/test_global_time_series.py +++ b/tests/unit/global_time_series/test_global_time_series.py @@ -2,18 +2,24 @@ import pytest -from zppy_interfaces.global_time_series.coupled_global import ( +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.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_exps, get_vars_original, - land_csv_row_to_var, -) -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.utils import ( Parameters, @@ -67,41 +73,46 @@ 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", + "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", "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.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 +143,7 @@ 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"] + assert parameters.use_ocn is False # test_get_data_dir assert ( @@ -159,15 +168,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]] = 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/", "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 +187,13 @@ def test_Parameters_and_related_functions(): ) parameters.plots_atm = [] parameters.plots_lnd = [] - exps = 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/", - "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": "results/ocn/glb/ts/monthly/5yr/", + "vol": "results/ocn/glb/ts/monthly/5yr/", "name": "v3.LR.historical_0051", "yoffset": 0.0, "yr": ([1985, 1989],), @@ -248,7 +253,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 @@ -280,7 +285,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 0cad995..ffa934e 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 +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 @@ -17,38 +15,49 @@ def main(parameters=None): if not parameters: parameters = _get_args() + """ + Determine if we want the Classic PDF or the Viewer + There are several cases to consider. In markdown table format: - 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/", - ) + | 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 + ): + # 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) # Creates the directory parameters.results_dir - coupled_global(parameters) + run_coupled_global(parameters) + # TODO: Add tests for all of the above cases on the zppy side def _get_args() -> Parameters: @@ -58,33 +67,63 @@ 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") - - # For coupled_global + # Used in all cases + # > For determining which output type to produce + 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("--plots_original", type=str, help="Plots original") - 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("--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("--make_viewer", type=str, help="Make viewer") - - # For both + 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, + 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("--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", + ) + + # For plots_component + # > For coupled_global + 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("--nrows", type=str, default="4", 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/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/__init__.py b/zppy_interfaces/global_time_series/coupled_global/__init__.py new file mode 100644 index 0000000..e69de29 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..7a1174f --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/driver.py @@ -0,0 +1,107 @@ +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_child_logger + +logger = _setup_child_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: + # 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, + 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/coupled_global/mix_pdf_original.py b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py new file mode 100644 index 0000000..4f29f8e --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/mix_pdf_original.py @@ -0,0 +1,300 @@ +import math +from typing import List + +import matplotlib as mpl +import numpy as np + +from zppy_interfaces.global_time_series.coupled_global.plotting import plot +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +mpl.use("Agg") + +logger = _setup_child_logger(__name__) + +# This file is for making cumulative PDFs for the original plots. +# Hence, "mix_pdf_original" + +# Used by mode_pdf.assemble_cumulative_pdf #################################### + + +# 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, +} + + +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/coupled_global/mix_viewer_component.py b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py new file mode 100644 index 0000000..3856e8f --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/mix_viewer_component.py @@ -0,0 +1,171 @@ +import os +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]): + self.group_name = name + self.variables = variables + + +# 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.") + + # 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) + 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` + # 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: + plot_name: str = var.variable_name + row_title: str + if var.metric == Metric.AVERAGE: + metric_name = "AVERAGE" + elif var.metric == Metric.TOTAL: + metric_name = "TOTAL" + else: + # This shouldn't be possible + raise ValueError(f"Invalid Enum option for metric={var.metric}") + if var.long_name != "": + row_title = f"{plot_name}: {var.long_name}, metric={metric_name}" + else: + row_title = f"{plot_name}, metric={metric_name}" + viewer.add_row(row_title) + for rgn in parameters.regions: + viewer.add_col( + f"{parameters.figstr}_{rgn}_{component}_{plot_name}.png", + is_file=True, + title=f"{rgn}_{component}_{plot_name}", + ) + + url = viewer.generate_page() + viewer.generate_viewer() + # Example links: + # Viewer is expecting the actual images to be in the directory above `table`. + # table/index.html links to previews with: ../v3.LR.historical_0051_glb_lnd_FSH.png + # Viewer is expecting individual image html pages to be under both group and var subdirectories. + # table/energy-flux/fsh-sensible-heat/glb_lnd_fsh.html links to: ../../../v3.LR.historical_0051_glb_lnd_FSH.png + return 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: + # 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..a8b25b4 --- /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_child_logger +from zppy_interfaces.multi_utils.viewer import OutputViewer + +logger = _setup_child_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 (click for download)" + 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 (click to view)" + 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..11ec0dc --- /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_child_logger + +mpl.use("Agg") + +logger = _setup_child_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..49e4acd --- /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_child_logger + +logger = _setup_child_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..7e1177d --- /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_parallel_with_plots, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +mpl.use("Agg") + +logger = _setup_child_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_parallel_with_plots( + exp, + "atmos", + requested_variables.vars_atm, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ice = set_var_parallel_with_plots( + exp, + "ice", + requested_variables.vars_ice, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_land = set_var_parallel_with_plots( + exp, + "land", + requested_variables.vars_land, + valid_vars, + invalid_vars, + parameters, + ) + requested_variables.vars_ocn = set_var_parallel_with_plots( + 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/coupled_global/plots_original.py b/zppy_interfaces/global_time_series/coupled_global/plots_original.py new file mode 100644 index 0000000..90dbe89 --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/plots_original.py @@ -0,0 +1,108 @@ +from typing import Any, Dict, List + +import matplotlib as mpl + +from zppy_interfaces.global_time_series.coupled_global.utils import ( + DatasetWrapper, + RequestedVariables, + Variable, + get_data_dir, + set_var, +) +from zppy_interfaces.global_time_series.utils import Parameters +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +mpl.use("Agg") + +logger = _setup_child_logger(__name__) + +# This file is for handling the original plots +# Hence, "plots_original" + +# 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] + + 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 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]] = [ + { + "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/coupled_global/plotting.py b/zppy_interfaces/global_time_series/coupled_global/plotting.py new file mode 100644 index 0000000..a6c94dd --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/plotting.py @@ -0,0 +1,237 @@ +import glob +import math + +import matplotlib as mpl +import numpy as np +from netCDF4 import Dataset + +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +mpl.use("Agg") + +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"] + 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"]: + 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")) + nfiles = len(files) + logger.info(f"{dir_in} {nfiles} moc files in total") + var = np.array([]) + time = np.array([]) + for i in range(nfiles): + # Open input file + fin = Dataset(files[i], "r") + time0 = fin["year"][:] + var0 = fin["mocAtlantic26"][:] + for iyear in range(int(time0[0]), int(time0[-1]) + 1): + if i > 0 and iyear <= time[-1]: + logger.info( + f"the amoc value for year {iyear} has been included in the moc time series from another moc file {files[i - 1]} {time[-1]} Skipping..." + ) + else: + imon = np.where(time0 == iyear)[0] + if len(imon) == 12: + var = np.append(var, np.mean(var0[imon])) + time = np.append(time, iyear) + else: + logger.error(f"error in input file : {files[i]}") + + 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"): + + i1 = (np.abs(year - year1)).argmin() + i2 = (np.abs(year - year2)).argmin() + + tmp = np.average(var[i1 : i2 + 1]) + ax.plot((year[i1], year[i2]), (tmp, tmp), lw=lw, color=color, label="average") + ax.text(ax.get_xlim()[1] + 1, tmp, format % tmp, va="center", color=color) + + return + + +# Function to add line showing linear trend over a specified period +def add_trend( + year, + var, + year1, + year2, + ax, + format="%4.2f", + lw=1, + color="b", + verbose=False, + ohc=False, + vol=False, +): + + i1 = (np.abs(year - year1)).argmin() + i2 = (np.abs(year - year2)).argmin() + x = year[i1 : i2 + 1] + y = var[i1 : i2 + 1] + + fit = np.polyfit(x, y, 1) + if verbose: + logger.info(fit) + fit_fn = np.poly1d(fit) + ax.plot(x, fit_fn(x), lw=lw, ls="--", c=color, label="trend") + if ohc: + # Earth radius 6371229. from MPAS-O output files + heat_uptake = fit[0] / (4.0 * math.pi * (6371229.0) ** 2 * 365.0 * 86400.0) + ax.text( + ax.get_xlim()[1] + 1, + fit_fn(x[-1]), + "%+4.2f W m$^{-2}$" % (heat_uptake), + color=color, + ) + if vol: + # Earth radius 6371229. from MPAS-O output files + # sea_lvl = fit[0] / ( 4.0*math.pi*(6371229.)**2*0.7) #for oceanic portion of the Earth surface + ax.text( + ax.get_xlim()[1] + 1, + fit_fn(x[-1]), + "%+5.4f mm yr$^{-1}$" % (fit[0]), + color=color, + ) + + return + + +# Function to get ylim +def get_ylim(standard_range, extreme_values): + if len(extreme_values) > 0: + has_extreme_values = True + extreme_min = np.amin(extreme_values) + extreme_max = np.amax(extreme_values) + else: + has_extreme_values = False + extreme_min = None + extreme_max = None + if len(standard_range) == 2: + has_standard_range = True + standard_min = standard_range[0] + standard_max = standard_range[1] + else: + has_standard_range = False + standard_min = None + standard_max = None + if has_extreme_values and has_standard_range: + # Use at least the standard range, + # perhaps a wider window to include extremes + if standard_min <= extreme_min: + ylim_min = standard_min + else: + ylim_min = extreme_min + if standard_max >= extreme_max: + ylim_max = standard_max + else: + ylim_max = extreme_max + elif has_extreme_values and not has_standard_range: + ylim_min = extreme_min + ylim_max = extreme_max + elif has_standard_range and not has_extreme_values: + ylim_min = standard_min + ylim_max = standard_max + else: + raise ValueError("Not enough range information supplied") + return [ylim_min, ylim_max] diff --git a/zppy_interfaces/global_time_series/coupled_global/utils.py b/zppy_interfaces/global_time_series/coupled_global/utils.py new file mode 100644 index 0000000..664fa4b --- /dev/null +++ b/zppy_interfaces/global_time_series/coupled_global/utils.py @@ -0,0 +1,889 @@ +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 + +import cftime +import matplotlib +import matplotlib.pyplot as plt +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_child_logger + +# 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__) + + +class Metric(Enum): + AVERAGE = 1 + TOTAL = 2 + + +class Variable(object): + def __init__( + self, + variable_name, + metric=Metric.AVERAGE, + scale_factor=1.0, + original_units="", + final_units="", + group="All Variables", + long_name="", + ): + # The name of the EAM/ELM/etc. variable on the monthly h0 history file + self.variable_name: str = variable_name + + # These fields are used for computation + # Global average over land area or global total + self.metric: Metric = metric + # The factor that should convert from original_units to final_units, after standard processing with nco + self.scale_factor: float = scale_factor + # Test string for the units as given on the history file (included here for possible testing) + self.original_units: str = original_units + # The units that should be reported in time series plots, based on metric and scale_factor + self.final_units: str = final_units + + # These fields are used for plotting + # A name used to cluster variables together, to be separated in groups within the output web pages + 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 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 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 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_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 + """ + 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) + + # 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}" + 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, + 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(16, 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 + 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, + 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_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] + + # 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, + ) + 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) + + return new_var_list + + +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, + 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 == []): + return new_var_list + + if exp[exp_key] != "": + directory = exp[exp_key] + + for var in var_list: + var_str: str = var.variable_name + try: + 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"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: + 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] + + return new_var_list 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 deleted file mode 100644 index 76516d4..0000000 --- a/zppy_interfaces/global_time_series/coupled_global_plotting.py +++ /dev/null @@ -1,658 +0,0 @@ -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") - -logger = _setup_child_logger(__name__) - - -# ---additional function to get moc time series -def getmoc(dir_in): - files = sorted(glob.glob(dir_in + "mocTimeSeries*.nc")) - nfiles = len(files) - logger.info(f"{dir_in} {nfiles} moc files in total") - var = np.array([]) - time = np.array([]) - for i in range(nfiles): - # Open input file - fin = Dataset(files[i], "r") - time0 = fin["year"][:] - var0 = fin["mocAtlantic26"][:] - for iyear in range(int(time0[0]), int(time0[-1]) + 1): - if i > 0 and iyear <= time[-1]: - logger.info( - f"the amoc value for year {iyear} has been included in the moc time series from another moc file {files[i - 1]} {time[-1]} Skipping..." - ) - else: - imon = np.where(time0 == iyear)[0] - if len(imon) == 12: - var = np.append(var, np.mean(var0[imon])) - time = np.append(time, iyear) - else: - logger.error(f"error in input file : {files[i]}") - - 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"): - - i1 = (np.abs(year - year1)).argmin() - i2 = (np.abs(year - year2)).argmin() - - tmp = np.average(var[i1 : i2 + 1]) - ax.plot((year[i1], year[i2]), (tmp, tmp), lw=lw, color=color, label="average") - ax.text(ax.get_xlim()[1] + 1, tmp, format % tmp, va="center", color=color) - - return - - -# ----------------------------------------------------------------------------- -# Function to add line showing linear trend over a specified period -def add_trend( - year, - var, - year1, - year2, - ax, - format="%4.2f", - lw=1, - color="b", - verbose=False, - ohc=False, - vol=False, -): - - i1 = (np.abs(year - year1)).argmin() - i2 = (np.abs(year - year2)).argmin() - x = year[i1 : i2 + 1] - y = var[i1 : i2 + 1] - - fit = np.polyfit(x, y, 1) - if verbose: - logger.info(fit) - fit_fn = np.poly1d(fit) - ax.plot(x, fit_fn(x), lw=lw, ls="--", c=color, label="trend") - if ohc: - # Earth radius 6371229. from MPAS-O output files - heat_uptake = fit[0] / (4.0 * math.pi * (6371229.0) ** 2 * 365.0 * 86400.0) - ax.text( - ax.get_xlim()[1] + 1, - fit_fn(x[-1]), - "%+4.2f W m$^{-2}$" % (heat_uptake), - color=color, - ) - if vol: - # Earth radius 6371229. from MPAS-O output files - # sea_lvl = fit[0] / ( 4.0*math.pi*(6371229.)**2*0.7) #for oceanic portion of the Earth surface - ax.text( - ax.get_xlim()[1] + 1, - fit_fn(x[-1]), - "%+5.4f mm yr$^{-1}$" % (fit[0]), - color=color, - ) - - return - - -# ----------------------------------------------------------------------------- -# Function to get ylim -def get_ylim(standard_range, extreme_values): - if len(extreme_values) > 0: - has_extreme_values = True - extreme_min = np.amin(extreme_values) - extreme_max = np.amax(extreme_values) - else: - has_extreme_values = False - extreme_min = None - extreme_max = None - if len(standard_range) == 2: - has_standard_range = True - standard_min = standard_range[0] - standard_max = standard_range[1] - else: - has_standard_range = False - standard_min = None - standard_max = None - if has_extreme_values and has_standard_range: - # Use at least the standard range, - # perhaps a wider window to include extremes - if standard_min <= extreme_min: - ylim_min = standard_min - else: - ylim_min = extreme_min - if standard_max >= extreme_max: - ylim_max = standard_max - else: - ylim_max = extreme_max - elif has_extreme_values and not has_standard_range: - ylim_min = extreme_min - ylim_max = extreme_max - elif has_standard_range and not has_extreme_values: - ylim_min = standard_min - ylim_max = standard_max - 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 deleted file mode 100644 index e32d496..0000000 --- a/zppy_interfaces/global_time_series/coupled_global_utils.py +++ /dev/null @@ -1,37 +0,0 @@ -from enum import Enum - - -class Metric(Enum): - AVERAGE = 1 - TOTAL = 2 - - -class Variable(object): - def __init__( - self, - variable_name, - metric=Metric.AVERAGE, - scale_factor=1.0, - original_units="", - final_units="", - group="All Variables", - long_name="", - ): - # The name of the EAM/ELM/etc. variable on the monthly h0 history file - self.variable_name: str = variable_name - - # These fields are used for computation - # Global average over land area or global total - self.metric: Metric = metric - # The factor that should convert from original_units to final_units, after standard processing with nco - self.scale_factor: float = scale_factor - # Test string for the units as given on the history file (included here for possible testing) - self.original_units: str = original_units - # The units that should be reported in time series plots, based on metric and scale_factor - self.final_units: str = final_units - - # These fields are used for plotting - # A name used to cluster variables together, to be separated in groups within the output web pages - self.group: str = group - # Descriptive text to add to the plot page to help users identify the variable - self.long_name: str = long_name diff --git a/zppy_interfaces/global_time_series/coupled_global_viewer.py b/zppy_interfaces/global_time_series/coupled_global_viewer.py deleted file mode 100644 index 5c7323b..0000000 --- a/zppy_interfaces/global_time_series/coupled_global_viewer.py +++ /dev/null @@ -1,140 +0,0 @@ -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 -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 - -logger = _setup_child_logger(__name__) - - -class VariableGroup(object): - def __init__(self, name: str, variables: List[Variable]): - self.group_name = name - 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: - 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) - groups: List[VariableGroup] = get_variable_groups(vars) - 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` - # 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: - plot_name: str = var.variable_name - row_title: str - if var.metric == Metric.AVERAGE: - metric_name = "AVERAGE" - elif var.metric == Metric.TOTAL: - metric_name = "TOTAL" - else: - # This shouldn't be possible - raise ValueError(f"Invalid Enum option for metric={var.metric}") - if var.long_name != "": - row_title = f"{plot_name}: {var.long_name}, metric={metric_name}" - else: - row_title = f"{plot_name}, metric={metric_name}" - viewer.add_row(row_title) - for rgn in parameters.regions: - viewer.add_col( - f"{parameters.figstr}_{rgn}_{component}_{plot_name}.png", - is_file=True, - title=f"{rgn}_{component}_{plot_name}", - ) - - url = viewer.generate_page() - viewer.generate_viewer() - # Example links: - # Viewer is expecting the actual images to be in the directory above `table`. - # table/index.html links to previews with: ../v3.LR.historical_0051_glb_lnd_FSH.png - # Viewer is expecting individual image html pages to be under both group and var subdirectories. - # table/energy-flux/fsh-sensible-heat/glb_lnd_fsh.html links to: ../../../v3.LR.historical_0051_glb_lnd_FSH.png - 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) - 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/ocean_month.py b/zppy_interfaces/global_time_series/create_ocean_ts.py similarity index 78% rename from zppy_interfaces/global_time_series/ocean_month.py rename to zppy_interfaces/global_time_series/create_ocean_ts.py index 4295203..790228e 100644 --- a/zppy_interfaces/global_time_series/ocean_month.py +++ b/zppy_interfaces/global_time_series/create_ocean_ts.py @@ -1,21 +1,58 @@ # 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") + # 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}" + + # Generate ocean time series in results_dir + ocean_month( + input_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}" + ) + # Output: Copy to results_dir (new ocean time series location) + dst: str = output_dir + "/" + logger.info(f"Copy moc file from {src} to {dst}") + shutil.copy(src, dst) + + def ocean_month( - path_in: str, case_dir: str, start_yr: int, end_yr: int, ts_num_years: int + path_in: str, + path_out: 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 is now directly provided as the output directory # Ocean constants # specific heat [J/(kg*degC)] diff --git a/zppy_interfaces/global_time_series/utils.py b/zppy_interfaces/global_time_series/utils.py index 105f4fe..ced6967 100644 --- a/zppy_interfaces/global_time_series/utils.py +++ b/zppy_interfaces/global_time_series/utils.py @@ -1,10 +1,33 @@ from typing import Dict, List +from zppy_interfaces.multi_utils.logger import _setup_child_logger + +logger = _setup_child_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.use_ocn: bool = _str2bool(args["use_ocn"]) self.input: str = args["input"] @@ -14,31 +37,34 @@ 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 + # 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." + ) + 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: @@ -62,6 +88,3 @@ def get_region(rgn: str) -> str: else: raise ValueError(f"Invalid rgn={rgn}") return rgn - - -###############################################################################