From 75c9757c4ee43b23f178c7fa000a8b1b2173b502 Mon Sep 17 00:00:00 2001 From: Ryan Forsyth Date: Mon, 21 Oct 2024 22:13:20 -0500 Subject: [PATCH] Refactoring and multi viewer support --- .../global_time_series/test_coupled_global.py | 356 +++++ ...e_global_time_series_viewers_chrysalis.cfg | 47 + ...e_min_case_global_time_series_viewers.cfg} | 5 +- tests/integration/utils.py | 10 +- zppy/global_time_series.py | 2 +- zppy/templates/coupled_global.py | 1272 ++++++++++------- zppy/templates/default.ini | 1 + zppy/templates/global_time_series.bash | 1 - zppy/templates/readTS.py | 78 - 9 files changed, 1156 insertions(+), 616 deletions(-) create mode 100644 tests/global_time_series/test_coupled_global.py create mode 100644 tests/integration/generated/test_min_case_global_time_series_viewers_chrysalis.cfg rename tests/integration/{template_min_case_global_time_series_single_plots.cfg => template_min_case_global_time_series_viewers.cfg} (85%) delete mode 100644 zppy/templates/readTS.py diff --git a/tests/global_time_series/test_coupled_global.py b/tests/global_time_series/test_coupled_global.py new file mode 100644 index 00000000..f01e7f4e --- /dev/null +++ b/tests/global_time_series/test_coupled_global.py @@ -0,0 +1,356 @@ +import unittest +from typing import Any, Dict, List + +from zppy.templates.coupled_global import ( + Metric, + Parameters, + Variable, + VariableGroup, + construct_generic_variables, + get_data_dir, + get_exps, + get_region, + get_variable_groups, + get_vars_original, + get_ylim, + land_csv_row_to_var, + param_get_list, +) + + +# Helper function +def get_var_names(vars: List[Variable]): + return list(map(lambda v: v.variable_name, vars)) + + +# Run this test suite in the environment the global_time_series task runs in. +# I.e., whatever `environment_commands` is set to for `[global_time_series]` +# NOT the zppy dev environment. +# Run: python -u -m unittest tests/global_time_series/test_coupled_global.py +class TestCoupledGlobal(unittest.TestCase): + # TODO: fill out test suite as much as possible. + + # Useful classes and their helper functions ############################### + def test_param_get_list(self): + self.assertEqual(param_get_list("None"), []) + + self.assertEqual(param_get_list("a"), ["a"]) + self.assertEqual(param_get_list("a,b,c"), ["a", "b", "c"]) + + self.assertEqual(param_get_list(""), [""]) + self.assertEqual(param_get_list("a,"), ["a", ""]) + self.assertEqual(param_get_list("a,b,c,"), ["a", "b", "c", ""]) + + def test_get_region(self): + self.assertEqual(get_region("glb"), "glb") + self.assertEqual(get_region("global"), "glb") + self.assertEqual(get_region("GLB"), "glb") + self.assertEqual(get_region("Global"), "glb") + + self.assertEqual(get_region("n"), "n") + self.assertEqual(get_region("north"), "n") + self.assertEqual(get_region("northern"), "n") + self.assertEqual(get_region("N"), "n") + self.assertEqual(get_region("North"), "n") + self.assertEqual(get_region("Northern"), "n") + + self.assertEqual(get_region("s"), "s") + self.assertEqual(get_region("south"), "s") + self.assertEqual(get_region("southern"), "s") + self.assertEqual(get_region("S"), "s") + self.assertEqual(get_region("South"), "s") + self.assertEqual(get_region("Southern"), "s") + + self.assertRaises(ValueError, get_region, "not-a-region") + + def test_Parameters_and_related_functions(self): + # Consider the following parameters given by a user. + args = [ + "coupled_global.py", + "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051", + "v3.LR.historical_0051", + "v3.LR.historical_0051", + "1985", + "1989", + "Blue", + "5", + "None", + "false", + "TREFHT", + "None", + "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", + "None", + "1", + "1", + "glb,n,s", + ] + # Then: + parameters: Parameters = Parameters(args) + self.assertEqual( + parameters.case_dir, + "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051", + ) + self.assertEqual(parameters.experiment_name, "v3.LR.historical_0051") + self.assertEqual(parameters.figstr, "v3.LR.historical_0051") + self.assertEqual(parameters.year1, 1985) + self.assertEqual(parameters.year2, 1989) + self.assertEqual(parameters.color, "Blue") + self.assertEqual(parameters.ts_num_years_str, "5") + self.assertEqual(parameters.plots_original, []) + self.assertEqual(parameters.atmosphere_only, False) + self.assertEqual(parameters.plots_atm, ["TREFHT"]) + self.assertEqual(parameters.plots_ice, []) + self.assertEqual( + parameters.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", + ], + ) + self.assertEqual(parameters.plots_ocn, []) + self.assertEqual(parameters.nrows, 1) + self.assertEqual(parameters.ncols, 1) + self.assertEqual(parameters.regions, ["glb", "n", "s"]) + + # test_get_data_dir + self.assertEqual( + get_data_dir(parameters, "atm", True), + "/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/", + ) + self.assertEqual(get_data_dir(parameters, "atm", False), "") + self.assertEqual( + get_data_dir(parameters, "ice", True), + "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051/post/ice/glb/ts/monthly/5yr/", + ) + self.assertEqual(get_data_dir(parameters, "ice", False), "") + self.assertEqual( + get_data_dir(parameters, "lnd", True), + "/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/", + ) + self.assertEqual(get_data_dir(parameters, "lnd", False), "") + self.assertEqual( + get_data_dir(parameters, "ocn", True), + "/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/", + ) + self.assertEqual(get_data_dir(parameters, "ocn", False), "") + + # test_get_exps + self.maxDiff = None + exps: List[Dict[str, Any]] = get_exps(parameters) + self.assertEqual(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],), + "color": "Blue", + } + self.assertEqual(exps[0], expected) + # Change up parameters + parameters.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".split( + "," + ) + parameters.plots_atm = [] + parameters.plots_lnd = [] + exps = get_exps(parameters) + self.assertEqual(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/", + "name": "v3.LR.historical_0051", + "yoffset": 0.0, + "yr": ([1985, 1989],), + "color": "Blue", + } + self.assertEqual(exps[0], expected) + # Change up parameters + parameters.atmosphere_only = True + exps = get_exps(parameters) + self.assertEqual(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": "", + "moc": "", + "vol": "", + "name": "v3.LR.historical_0051", + "yoffset": 0.0, + "yr": ([1985, 1989],), + "color": "Blue", + } + self.assertEqual(exps[0], expected) + + # Metric + + def test_Variable(self): + v = Variable( + "var_name", + original_units="units1", + final_units="units2", + group="group_name", + long_name="long name", + ) + self.assertEqual(v.variable_name, "var_name") + self.assertEqual(v.metric, Metric.AVERAGE) + self.assertEqual(v.scale_factor, 1.0) + self.assertEqual(v.original_units, "units1") + self.assertEqual(v.final_units, "units2") + self.assertEqual(v.group, "group_name") + self.assertEqual(v.long_name, "long name") + + def test_get_vars_original(self): + self.assertEqual( + get_var_names(get_vars_original(["net_toa_flux_restom"])), ["RESTOM"] + ) + self.assertEqual( + get_var_names(get_vars_original(["net_atm_energy_imbalance"])), + ["RESTOM", "RESSURF"], + ) + self.assertEqual( + get_var_names(get_vars_original(["global_surface_air_temperature"])), + ["TREFHT"], + ) + self.assertEqual( + get_var_names(get_vars_original(["toa_radiation"])), ["FSNTOA", "FLUT"] + ) + self.assertEqual( + get_var_names(get_vars_original(["net_atm_water_imbalance"])), + ["PRECC", "PRECL", "QFLX"], + ) + self.assertEqual( + get_var_names( + get_vars_original( + [ + "net_toa_flux_restom", + "net_atm_energy_imbalance", + "global_surface_air_temperature", + "toa_radiation", + "net_atm_water_imbalance", + ] + ) + ), + ["RESTOM", "RESSURF", "TREFHT", "FSNTOA", "FLUT", "PRECC", "PRECL", "QFLX"], + ) + self.assertEqual(get_var_names(get_vars_original(["invalid_plot"])), []) + + def test_land_csv_row_to_var(self): + # Test with first row of land csv, whitespace stripped + 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) + self.assertEqual(v.variable_name, "BCDEP") + self.assertEqual(v.metric, Metric.AVERAGE) + self.assertEqual(v.scale_factor, 1.0) + self.assertEqual(v.original_units, "kg/m^2/s") + self.assertEqual(v.final_units, "kg/m^2/s") + self.assertEqual(v.group, "Aerosol Flux") + self.assertEqual( + v.long_name, "total black carbon deposition (dry+wet) from atmosphere" + ) + + # construct_land_variables -- requires IO + + def test_construct_generic_variables(self): + vars: List[str] = ["a", "b", "c"] + self.assertEqual(get_var_names(construct_generic_variables(vars)), vars) + + # RequestedVariables -- requires IO + + def test_VariableGroup(self): + var_str_list: List[str] = ["a", "b", "c"] + vars: List[Variable] = construct_generic_variables(var_str_list) + g: VariableGroup = VariableGroup("MyGroup", vars) + self.assertEqual(g.group_name, "MyGroup") + self.assertEqual(get_var_names(g.variables), var_str_list) + + # TS -- requires IO + # OutputViewer -- requires IO + + # Setup ################################################################### + + # set_var -- requires IO + # process_data -- requires IO + + # Plotting #################################################################### + + def test_get_variable_groups(self): + a: Variable = Variable(variable_name="a", group="GroupA") + b: Variable = Variable(variable_name="b", group="GroupA") + x: Variable = Variable(variable_name="x", group="GroupX") + y: Variable = Variable(variable_name="y", group="GroupX") + + def get_group_names(groups: List[VariableGroup]) -> List[str]: + return list(map(lambda g: g.group_name, groups)) + + self.assertEqual( + get_group_names(get_variable_groups([a, b, x, y])), ["GroupA", "GroupX"] + ) + + # getmoc -- requires IO + # add_line -- requires IO + # add_trend -- requires IO + + def test_get_ylim(self): + # Min is equal, max is equal + self.assertEqual(get_ylim([-1, 1], [-1, 1]), [-1, 1]) + # Min is lower, max is equal + self.assertEqual(get_ylim([-1, 1], [-2, 1]), [-2, 1]) + # Min is equal, max is higher + self.assertEqual(get_ylim([-1, 1], [-1, 2]), [-1, 2]) + # Min is lower, max is higher + self.assertEqual(get_ylim([-1, 1], [-2, 2]), [-2, 2]) + # Min is lower, max is higher, multiple extreme_values + self.assertEqual(get_ylim([-1, 1], [-2, -0.5, 0.5, 2]), [-2, 2]) + # No standard range + self.assertEqual(get_ylim([], [-2, 2]), [-2, 2]) + # No extreme range + self.assertEqual(get_ylim([-1, 1], []), [-1, 1]) + + # Plotting functions -- require IO + # make_plot_pdfs -- requires IO + + # Run coupled_global ###################################################### + + # run -- requires IO + # get_vars -- requires IO + # create_viewer -- requires IO + # create_viewer_index -- requires IO + # run_by_region -- requires IO + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/generated/test_min_case_global_time_series_viewers_chrysalis.cfg b/tests/integration/generated/test_min_case_global_time_series_viewers_chrysalis.cfg new file mode 100644 index 00000000..06cdbc97 --- /dev/null +++ b/tests/integration/generated/test_min_case_global_time_series_viewers_chrysalis.cfg @@ -0,0 +1,47 @@ +[default] +case = "v3.LR.historical_0051" +constraint = "" +dry_run = "False" +environment_commands = "" +input = /lcrc/group/e3sm2/ac.wlin/E3SMv3/v3.LR.historical_0051 +input_subdir = archive/atm/hist +mapping_file = "map_ne30pg2_to_cmip6_180x360_aave.20200201.nc" +output = "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_viewers_output/unique_id/v3.LR.historical_0051" +partition = "debug" +qos = "regular" +www = "/lcrc/group/e3sm/public_html/diagnostic_output/ac.forsyth2/zppy_min_case_global_time_series_viewers_www/unique_id" + +[ts] +active = True +e3sm_to_cmip_environment_commands = "" +walltime = "00:30:00" +years = "1985:1995:5", + + [[ atm_monthly_glb ]] + # Note global average won't work for 3D variables. + frequency = "monthly" + input_files = "eam.h0" + input_subdir = "archive/atm/hist" + mapping_file = "glb" + vars = "TREFHT" + + [[ lnd_monthly_glb ]] + frequency = "monthly" + input_files = "elm.h0" + input_subdir = "archive/lnd/hist" + mapping_file = "glb" + vars = "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" + +[global_time_series] +active = True +experiment_name = "v3.LR.historical_0051" +figstr = "v3.LR.historical_0051" +ncols = 1 +nrows = 1 +plots_original="" +plots_atm = "TREFHT" +# TODO: S LAISUN not showing up +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" +ts_num_years = 5 +walltime = "00:30:00" +years = "1985-1995", diff --git a/tests/integration/template_min_case_global_time_series_single_plots.cfg b/tests/integration/template_min_case_global_time_series_viewers.cfg similarity index 85% rename from tests/integration/template_min_case_global_time_series_single_plots.cfg rename to tests/integration/template_min_case_global_time_series_viewers.cfg index c883d8f5..b0d449b5 100644 --- a/tests/integration/template_min_case_global_time_series_single_plots.cfg +++ b/tests/integration/template_min_case_global_time_series_viewers.cfg @@ -6,10 +6,10 @@ environment_commands = "#expand environment_commands#" input = /lcrc/group/e3sm2/ac.wlin/E3SMv3/#expand case_name# input_subdir = archive/atm/hist mapping_file = "map_ne30pg2_to_cmip6_180x360_aave.20200201.nc" -output = "#expand user_output#zppy_min_case_global_time_series_single_plots_output/#expand unique_id#/#expand case_name#" +output = "#expand user_output#zppy_min_case_global_time_series_viewers_output/#expand unique_id#/#expand case_name#" partition = "#expand partition_short#" qos = "#expand qos_short#" -www = "#expand user_www#zppy_min_case_global_time_series_single_plots_www/#expand unique_id#" +www = "#expand user_www#zppy_min_case_global_time_series_viewers_www/#expand unique_id#" [ts] active = True @@ -40,6 +40,7 @@ ncols = 1 nrows = 1 plots_original="" plots_atm = "TREFHT" +# TODO: S LAISUN not showing up 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" ts_num_years = 5 walltime = "00:30:00" diff --git a/tests/integration/utils.py b/tests/integration/utils.py index 59740266..e62b866f 100644 --- a/tests/integration/utils.py +++ b/tests/integration/utils.py @@ -7,6 +7,14 @@ from mache import MachineInfo from PIL import Image, ImageChops, ImageDraw +# v2 -- using generate_viewer(), atm+lnd +# v3 -- atm+lnd +# v4 -- lnd only, accidentally commented cp table to viewer +# v5 -- lnd only => viewer looks correct +# v6 -- using generate_viewer(), lnd only => viewer looks correct +# v7 -- copy tree instead of move, atm+lnd (forgot to cp viewer_* in bash) +# v8 -- copy tree instead of move, atm+lnd => viewer back to bad formatting +# v9 -- v6, using generate_viewer(), lnd only UNIQUE_ID = "unique_id" # Image checking ########################################################## @@ -297,7 +305,7 @@ def generate_cfgs(unified_testing=False, dry_run=False): "min_case_global_time_series_custom", "min_case_global_time_series_original_8_no_ocn", "min_case_global_time_series_original_8", - "min_case_global_time_series_single_plots", + "min_case_global_time_series_viewers", "min_case_ilamb_land_only", "min_case_ilamb", "min_case_mpas_analysis", diff --git a/zppy/global_time_series.py b/zppy/global_time_series.py index fdf8e4a3..dd74b58a 100644 --- a/zppy/global_time_series.py +++ b/zppy/global_time_series.py @@ -50,7 +50,7 @@ def global_time_series(config, script_dir, existing_bundles, job_ids_file): c["global_time_series_dir"] = os.path.join(script_dir, f"{prefix}_dir") if not os.path.exists(c["global_time_series_dir"]): os.mkdir(c["global_time_series_dir"]) - scripts = ["coupled_global.py", "readTS.py", "ocean_month.py"] + scripts = ["coupled_global.py", "ocean_month.py"] for script in scripts: script_template = template_env.get_template(script) script_file = os.path.join(c["global_time_series_dir"], script) diff --git a/zppy/templates/coupled_global.py b/zppy/templates/coupled_global.py index 6f599744..79e9e518 100644 --- a/zppy/templates/coupled_global.py +++ b/zppy/templates/coupled_global.py @@ -7,7 +7,8 @@ import stat import sys import traceback -from typing import Any, Dict, List, Optional, Tuple +from enum import Enum +from typing import Any, Dict, List, Tuple import cftime import matplotlib as mpl @@ -15,6 +16,8 @@ import matplotlib.pyplot as plt import numpy as np import xarray +import xcdat +from bs4 import BeautifulSoup from netCDF4 import Dataset from output_viewer.build import build_page, build_viewer from output_viewer.index import ( @@ -25,11 +28,584 @@ OutputRow, ) from output_viewer.utils import rechmod -from readTS import TS mpl.use("Agg") +# Useful classes and their helper functions ################################### +def param_get_list(param_value: str) -> List[str]: + if param_value == "None": + return [] + else: + return param_value.split(",") + + +def get_region(rgn: str) -> str: + if rgn.lower() in ["glb", "global"]: + rgn = "glb" + elif rgn.lower() in ["n", "north", "northern"]: + rgn = "n" + elif rgn.lower() in ["s", "south", "southern"]: + rgn = "s" + else: + raise ValueError(f"Invalid rgn={rgn}") + return rgn + + +class Parameters(object): + def __init__(self, parameters): + self.case_dir: str = parameters[1] + self.experiment_name: str = parameters[2] + self.figstr: str = parameters[3] + self.year1: int = int(parameters[4]) + self.year2: int = int(parameters[5]) + self.color: str = parameters[6] + self.ts_num_years_str: str = parameters[7] + self.plots_original: List[str] = param_get_list(parameters[8]) + self.atmosphere_only: bool = ( + False if (parameters[9].lower() == "false") else True + ) + self.plots_atm: List[str] = param_get_list(parameters[10]) + self.plots_ice: List[str] = param_get_list(parameters[11]) + self.plots_lnd: List[str] = param_get_list(parameters[12]) + self.plots_ocn: List[str] = param_get_list(parameters[13]) + self.nrows: int = int(parameters[14]) + self.ncols: int = int(parameters[15]) + # 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), parameters[16].split(",")) + ) + + +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="", + 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 + + +def get_vars_original(plots_original: List[str]) -> List[Variable]: + 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 + # If this file is being run stand-alone, then + # it will search the directory above the git directory + with open("zppy_land_fields.csv", newline="") as csv_file: + print("In File") + var_reader = csv.reader(csv_file) + for row in var_reader: + print(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) + 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 + ) + + +class VariableGroup(object): + def __init__(self, name: str, variables: List[Variable]): + self.group_name = name + self.variables = variables + + +class TS(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.f: xarray.core.dataset.Dataset = xcdat.open_mfdataset( + f"{directory}*.nc", center_times=True + ) + + def __del__(self): + + self.f.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": + print("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 + if (metric == Metric.AVERAGE) or (metric == Metric.TOTAL): + annual_average_dataset_for_var: xarray.core.dataset.Dataset = ( + self.f.temporal.group_average(var, "year") + ) + data_array = annual_average_dataset_for_var.data_vars[var] + # elif metric == Metric.TOTAL: + # # TODO: Implement this! + # raise NotImplementedError() + else: + # This shouldn't be possible + raise ValueError(f"Invalid Enum option for metric={metric}") + 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." + ) + 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, + ) + + +# Copied from e3sm_diags +class OutputViewer(object): + def __init__(self, path=".", index_name="Results"): + self.path = os.path.abspath(path) + self.index = OutputIndex(index_name) + self.cache = {} # dict of { OutputPage: { OutputGroup: [OutputRow] } } + self.page = None + self.group = None + self.row = None + + def add_page(self, page_title, *args, **kwargs): + """Add a page to the viewer's index""" + self.page = OutputPage(page_title, *args, **kwargs) + self.cache[self.page] = {} + self.index.addPage(self.page) + + def set_page(self, page_title): + """Sets the page with the title name as the current page""" + for output_page in self.cache: + if page_title == output_page.title: + self.page = output_page + return + raise RuntimeError("There is no page titled: %s" % page_title) + + def add_group(self, group_name): + """Add a group to the current page""" + if self.page is None: + raise RuntimeError("You must first insert a page with add_page()") + self.group = OutputGroup(group_name) + if self.group not in self.cache[self.page]: + self.cache[self.page][self.group] = [] # group doesn't have any rows yet + self.page.addGroup(self.group) + + def set_group(self, group_name): + """Sets the group with the title name as the current group""" + for output_group in self.cache[self.page]: + if group_name == output_group.title: + self.group = output_group + return + raise RuntimeError("There is no group titled: %s" % group_name) + + def add_row(self, row_name): + """Add a row with the title name to the current group""" + if self.group is None: + raise RuntimeError("You must first insert a group with add_group()") + self.row = OutputRow(row_name, []) + if self.row not in self.cache[self.page][self.group]: + self.cache[self.page][self.group].append(self.row) + self.page.addRow(self.row, len(self.page.groups) - 1) # type: ignore + + def set_row(self, row_name): + """Sets the row with the title name as the current row""" + for output_row in self.cache[self.page][self.group]: + if row_name == output_row.title: + self.row = output_row + return + raise RuntimeError("There is no row titled: %s" % row_name) + + def add_cols(self, cols): + """Add multiple string cols to the current row""" + self.row.columns.append(cols) # type: ignore + + def add_col(self, col, is_file=False, **kwargs): + """Add a single col to the current row. Set is_file to True if the col is a file path.""" + if is_file: + self.row.columns.append(OutputFile(col, **kwargs)) # type: ignore + else: + self.row.columns.append(col) # type: ignore + + def generate_page(self) -> str: + """ + Generate and return the location of the current HTML page. + """ + self.index.toJSON(os.path.join(self.path, "index.json")) + + default_mask = stat.S_IMODE(os.stat(self.path).st_mode) + rechmod(self.path, default_mask) + + if os.access(self.path, os.W_OK): + default_mask = stat.S_IMODE( + os.stat(self.path).st_mode + ) # mode of files to be included + url = build_page( + self.page, + os.path.join(self.path, "index.json"), + default_mask=default_mask, + ) + return url + + raise RuntimeError("Error geneating the page.") + + def generate_viewer(self): + """Generate the webpage""" + self.index.toJSON(os.path.join(self.path, "index.json")) + + default_mask = stat.S_IMODE(os.stat(self.path).st_mode) + rechmod(self.path, default_mask) + + if os.access(self.path, os.W_OK): + default_mask = stat.S_IMODE( + os.stat(self.path).st_mode + ) # mode of files to be included + build_viewer( + os.path.join(self.path, "index.json"), + diag_name="Global Time Series", + default_mask=default_mask, + ) + + +# 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 ( + (not parameters.atmosphere_only) and 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, +) -> None: + if exp[exp_key] != "": + ts: TS = TS(exp[exp_key]) + for var in var_list: + var_str: str = var.variable_name + try: + data_array: xarray.core.dataarray.DataArray + units: str + data_array, units = ts.globalAnnual(var) + valid_vars.append(str(var_str)) + except Exception as e: + print(e) + print(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 ts + + +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"] = {} + + set_var( + exp, + "atmos", + requested_variables.vars_original, + valid_vars, + invalid_vars, + rgn, + ) + set_var( + exp, "atmos", requested_variables.vars_atm, valid_vars, invalid_vars, rgn + ) + set_var(exp, "ice", requested_variables.vars_ice, valid_vars, invalid_vars, rgn) + set_var( + exp, "land", requested_variables.vars_land, valid_vars, invalid_vars, rgn + ) + set_var( + exp, "ocean", requested_variables.vars_ocn, valid_vars, invalid_vars, rgn + ) + + # Optionally read ohc + if exp["ocean"] != "": + ts = TS(exp["ocean"]) + exp["annual"]["ohc"], _ = ts.globalAnnual(Variable("ohc")) + # anomalies with respect to first year + exp["annual"]["ohc"][:] = exp["annual"]["ohc"][:] - exp["annual"]["ohc"][0] + + if exp["vol"] != "": + ts = TS(exp["vol"]) + exp["annual"]["volume"], _ = ts.globalAnnual(Variable("volume")) + # annomalies with respect to first year + exp["annual"]["volume"][:] = ( + exp["annual"]["volume"][:] - exp["annual"]["volume"][0] + ) + + print( + f"{rgn} region globalAnnual was computed successfully for these variables: {valid_vars}" + ) + print( + f"{rgn} region globalAnnual could not be computed for these variables: {invalid_vars}" + ) + return exps + + +# Plotting #################################################################### + + +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 + + # ---additional function to get moc time series def getmoc(dir_in): files = sorted(glob.glob(dir_in + "mocTimeSeries*.nc")) @@ -74,8 +650,6 @@ def add_line(year, var, year1, year2, ax, format="%4.2f", lw=1, color="b"): 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 @@ -463,13 +1037,13 @@ def plot(ax, xlim, exps, param_dict, rgn): # noqa: C901 extreme_values = [] for exp in exps: # Relevant to "Plot 5: plot_change_ohc" - if param_dict["check_exp_ocean"] and (exp["ocean"] is None): + 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"] is None): + if param_dict["check_exp_vol"] and (exp["vol"] == ""): continue # Relevant to "Plot 6: plot_max_moc" if param_dict["use_getmoc"]: @@ -558,64 +1132,9 @@ def plot(ax, xlim, exps, param_dict, rgn): # noqa: C901 } -def param_get_list(param_value): - if param_value == "None": - return [] - else: - return param_value.split(",") - - -def set_var( - exp: Dict[str, Any], - exp_key: str, - var_list: List[str], - valid_vars: List[str], - invalid_vars: List[str], - rgn: str, -) -> None: - if exp[exp_key] is not None: - ts = TS(exp[exp_key]) - for var in var_list: - try: - v: xarray.core.dataarray.DataArray - units: Optional[str] - v, units = ts.globalAnnual(var) - valid_vars.append(str(var)) - except Exception as e: - print(e) - print(f"globalAnnual failed. Invalid var = {var}") - invalid_vars.append(str(var)) - continue - if v.sizes["rgn"] > 1: - # number of years x 3 regions = v.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}") - v = v.isel(rgn=n) # Just use nth region - elif rgn != "glb": - # v only has one dimension -- glb. - # Therefore it is not possible to get n or s plots. - raise RuntimeError( - f"var={var} only has global data. Cannot process rgn={rgn}" - ) - exp["annual"][var] = (v, units) - if "year" not in exp["annual"]: - years: np.ndarray[cftime.DatetimeNoLeap] = v.coords["time"].values - exp["annual"]["year"] = [x.year for x in years] - del ts - - -# FIXME: C901 'main' is too complex (19) +# FIXME: C901 'make_plot_pdfs' is too complex (20) def make_plot_pdfs( # noqa: C901 - figstr, + parameters: Parameters, rgn, component, xlim, @@ -623,29 +1142,29 @@ def make_plot_pdfs( # noqa: C901 plot_list, valid_plots, invalid_plots, - nrows, - ncols, ): num_plots = len(plot_list) if num_plots == 0: return - plots_per_page = nrows * ncols + plots_per_page = parameters.nrows * parameters.ncols num_pages = math.ceil(num_plots / plots_per_page) counter = 0 # https://stackoverflow.com/questions/58738992/save-multiple-figures-with-subplots-into-a-pdf-with-multiple-pages - pdf = matplotlib.backends.backend_pdf.PdfPages(f"{figstr}_{rgn}_{component}.pdf") + pdf = matplotlib.backends.backend_pdf.PdfPages( + f"{parameters.figstr}_{rgn}_{component}.pdf" + ) for page in range(num_pages): if plots_per_page == 1: fig = plt.figure(1, figsize=[13.5 / 2, 16.5 / 4]) else: fig = plt.figure(1, figsize=[13.5, 16.5]) - fig.suptitle(f"{figstr}_{rgn}_{component}") + fig.suptitle(f"{parameters.figstr}_{rgn}_{component}") for j in range(plots_per_page): # The final page doesn't need to be filled out with plots. if counter >= num_plots: break - ax = plt.subplot(nrows, ncols, j + 1) + ax = plt.subplot(parameters.nrows, parameters.ncols, j + 1) if component == "original": try: plot_function = PLOT_DICT[plot_list[counter]] @@ -687,507 +1206,194 @@ def make_plot_pdfs( # noqa: C901 fig.tight_layout() pdf.savefig(1) if plots_per_page == 1: - fig.savefig(f"{figstr}_{rgn}_{component}_{plot_name}.png", dpi=150) + fig.savefig( + f"{parameters.figstr}_{rgn}_{component}_{plot_name}.png", dpi=150 + ) elif num_pages > 1: - fig.savefig(f"{figstr}_{rgn}_{component}_{page}.png", dpi=150) + fig.savefig(f"{parameters.figstr}_{rgn}_{component}_{page}.png", dpi=150) else: - fig.savefig(f"{figstr}_{rgn}_{component}.png", dpi=150) + fig.savefig(f"{parameters.figstr}_{rgn}_{component}.png", dpi=150) plt.clf() pdf.close() +# Run coupled_global ########################################################## # ----------------------------------------------------------------------------- -# FIXME: C901 'run' is too complex (19) -def run(parameters, rgn): # noqa: C901 - # These are the "Tableau 20" colors as RGB. - t20: List[Tuple[float, float, float]] = [ - (31, 119, 180), - (174, 199, 232), - (255, 127, 14), - (255, 187, 120), - (44, 160, 44), - (152, 223, 138), - (214, 39, 40), - (255, 152, 150), - (148, 103, 189), - (197, 176, 213), - (140, 86, 75), - (196, 156, 148), - (227, 119, 194), - (247, 182, 210), - (127, 127, 127), - (199, 199, 199), - (188, 189, 34), - (219, 219, 141), - (23, 190, 207), - (158, 218, 229), - ] - # Scale the RGB values to the [0, 1] range, which is the format matplotlib accepts. - for i in range(len(t20)): - r, g, b = t20[i] - t20[i] = (r / 255.0, g / 255.0, b / 255.0) - - # "Tableau 10" uses every other color - t10 = [] - for i in range(0, len(t20), 2): - t10.append(t20[i]) - - # ----------------------------------------------------------------------------- - # --- Atmos data --- - +def run(parameters: Parameters, requested_variables: RequestedVariables, rgn: str): # Experiments - case_dir = parameters[1] - experiment_name = parameters[2] - figstr = parameters[3] - year1 = int(parameters[4]) - year2 = int(parameters[5]) - color = parameters[6] - ts_num_years = parameters[7] - plots_original = param_get_list(parameters[8]) - if parameters[9].lower() == "false": - atmosphere_only = False - else: - atmosphere_only = True - plots_atm = param_get_list(parameters[10]) - plots_ice = param_get_list(parameters[11]) - plots_lnd = param_get_list(parameters[12]) - plots_ocn = param_get_list(parameters[13]) - nrows = int(parameters[14]) - ncols = int(parameters[15]) - vars_original = [] - if "net_toa_flux_restom" or "net_atm_energy_imbalance" in plots_original: - vars_original.append("RESTOM") - if "net_atm_energy_imbalance" in plots_original: - vars_original.append("RESSURF") - if "global_surface_air_temperature" in plots_original: - vars_original.append("TREFHT") - if "toa_radiation" in plots_original: - vars_original.append("FSNTOA") - vars_original.append("FLUT") - if "net_atm_water_imbalance" in plots_original: - vars_original.append("PRECC") - vars_original.append("PRECL") - vars_original.append("QFLX") - use_atmos = plots_atm or plots_original - has_original_ocn_plots = ( - ("change_ohc" in plots_original) - or ("max_moc" in plots_original) - or ("change_sea_level" in plots_original) - ) - use_ocn = plots_ocn or (not atmosphere_only and has_original_ocn_plots) - exps: List[Dict[str, Any]] = [ - { - "atmos": ( - f"{case_dir}/post/atm/glb/ts/monthly/{ts_num_years}yr/" - if use_atmos - else None - ), - "ice": ( - f"{case_dir}/post/ice/glb/ts/monthly/{ts_num_years}yr/" - if plots_ice - else None - ), - "land": ( - f"{case_dir}/post/lnd/glb/ts/monthly/{ts_num_years}yr/" - if plots_lnd - else None - ), - "ocean": ( - f"{case_dir}/post/ocn/glb/ts/monthly/{ts_num_years}yr/" - if use_ocn - else None - ), - "moc": ( - f"{case_dir}/post/ocn/glb/ts/monthly/{ts_num_years}yr/" - if use_ocn - else None - ), - "vol": ( - f"{case_dir}/post/ocn/glb/ts/monthly/{ts_num_years}yr/" - if use_ocn - else None - ), - "name": experiment_name, - "yoffset": 0.0, - "yr": ([year1, year2],), - "color": f"{color}", - } - ] - - valid_vars: List[str] = [] - invalid_vars: List[str] = [] - - # Read data - exp: Dict[str, Any] - for exp in exps: - exp["annual"] = {} - - # Use vars_original rather than plots_original, - # since the plots have different names than the variables - set_var(exp, "atmos", vars_original, valid_vars, invalid_vars, rgn) - set_var(exp, "atmos", plots_atm, valid_vars, invalid_vars, rgn) - set_var(exp, "ice", plots_ice, valid_vars, invalid_vars, rgn) - set_var(exp, "land", plots_lnd, valid_vars, invalid_vars, rgn) - set_var(exp, "ocean", plots_ocn, valid_vars, invalid_vars, rgn) - - # Optionally read ohc - if exp["ocean"] is not None: - ts = TS(exp["ocean"]) - exp["annual"]["ohc"], _ = ts.globalAnnual("ohc") - # annomalies with respect to first year - exp["annual"]["ohc"][:] = exp["annual"]["ohc"][:] - exp["annual"]["ohc"][0] - - if exp["vol"] is not None: - ts = TS(exp["vol"]) - exp["annual"]["volume"], _ = ts.globalAnnual("volume") - # annomalies with respect to first year - exp["annual"]["volume"][:] = ( - exp["annual"]["volume"][:] - exp["annual"]["volume"][0] - ) - - print( - f"{rgn} region globalAnnual was computed successfully for these variables: {valid_vars}" - ) - print( - f"{rgn} region globalAnnual could not be computed for these variables: {invalid_vars}" - ) - - # ----------------------------------------------------------------------------- - # --- Generate plots --- + exps: List[Dict[str, Any]] = process_data(parameters, requested_variables, rgn) - xlim = [float(year1), float(year2)] + xlim: List[float] = [float(parameters.year1), float(parameters.year2)] valid_plots: List[str] = [] invalid_plots: List[str] = [] - make_plot_pdfs( - figstr, - rgn, - "original", - xlim, - exps, - plots_original, - valid_plots, - invalid_plots, - nrows, - ncols, - ) - make_plot_pdfs( - figstr, - rgn, - "atm", - xlim, - exps, - plots_atm, - valid_plots, - invalid_plots, - nrows, - ncols, - ) - make_plot_pdfs( - figstr, - rgn, - "ice", - xlim, - exps, - plots_ice, - valid_plots, - invalid_plots, - nrows, - ncols, - ) - make_plot_pdfs( - figstr, - rgn, - "lnd", - xlim, - exps, - plots_lnd, - valid_plots, - invalid_plots, - nrows, - ncols, - ) - make_plot_pdfs( - figstr, - rgn, - "ocn", - xlim, - exps, - plots_ocn, - valid_plots, - invalid_plots, - nrows, - ncols, - ) - + # Use list of tuples rather than a dict, to keep order + mapping: List[Tuple[str, List[str]]] = [ + ("original", parameters.plots_original), + ("atm", parameters.plots_atm), + ("ice", parameters.plots_ice), + ("lnd", parameters.plots_lnd), + ("ocn", parameters.plots_ocn), + ] + for component, plot_list in mapping: + make_plot_pdfs( + parameters, + rgn, + component, + xlim, + exps, + plot_list, + valid_plots, + invalid_plots, + ) print(f"These {rgn} region plots generated successfully: {valid_plots}") print( - f"These {rgn} region plots could not be generated successfully: {invalid_plots}" + f"These {rgn} regions plots could not be generated successfully: {invalid_plots}" ) -# Copied from e3sm_diags -class OutputViewer(object): - def __init__(self, path=".", index_name="Results"): - self.path = os.path.abspath(path) - self.index = OutputIndex(index_name) - self.cache = {} # dict of { OutputPage: { OutputGroup: [OutputRow] } } - self.page = None - self.group = None - self.row = None - - def add_page(self, page_title, *args, **kwargs): - """Add a page to the viewer's index""" - self.page = OutputPage(page_title, *args, **kwargs) - self.cache[self.page] = {} - self.index.addPage(self.page) - - def set_page(self, page_title): - """Sets the page with the title name as the current page""" - for output_page in self.cache: - if page_title == output_page.title: - self.page = output_page - return - raise RuntimeError("There is no page titled: %s" % page_title) - - def add_group(self, group_name): - """Add a group to the current page""" - if self.page is None: - raise RuntimeError("You must first insert a page with add_page()") - self.group = OutputGroup(group_name) - if self.group not in self.cache[self.page]: - self.cache[self.page][self.group] = [] # group doesn't have any rows yet - self.page.addGroup(self.group) - - def set_group(self, group_name): - """Sets the group with the title name as the current group""" - for output_group in self.cache[self.page]: - if group_name == output_group.title: - self.group = output_group - return - raise RuntimeError("There is no group titled: %s" % group_name) - - def add_row(self, row_name): - """Add a row with the title name to the current group""" - if self.group is None: - raise RuntimeError("You must first insert a group with add_group()") - self.row = OutputRow(row_name, []) - if self.row not in self.cache[self.page][self.group]: - self.cache[self.page][self.group].append(self.row) - self.page.addRow(self.row, len(self.page.groups) - 1) # type: ignore - - def set_row(self, row_name): - """Sets the row with the title name as the current row""" - for output_row in self.cache[self.page][self.group]: - if row_name == output_row.title: - self.row = output_row - return - raise RuntimeError("There is no row titled: %s" % row_name) - - def add_cols(self, cols): - """Add multiple string cols to the current row""" - self.row.columns.append(cols) # type: ignore - - def add_col(self, col, is_file=False, **kwargs): - """Add a single col to the current row. Set is_file to True if the col is a file path.""" - if is_file: - self.row.columns.append(OutputFile(col, **kwargs)) # type: ignore - else: - self.row.columns.append(col) # type: ignore - - def generate_page(self): - """ - Generate and return the location of the current HTML page. - """ - self.index.toJSON(os.path.join(self.path, "index.json")) - - default_mask = stat.S_IMODE(os.stat(self.path).st_mode) - rechmod(self.path, default_mask) - - if os.access(self.path, os.W_OK): - default_mask = stat.S_IMODE( - os.stat(self.path).st_mode - ) # mode of files to be included - url = build_page( - self.page, - os.path.join(self.path, "index.json"), - default_mask=default_mask, - ) - return url - - raise RuntimeError("Error geneating the page.") - - def generate_viewer(self): - """Generate the webpage""" - self.index.toJSON(os.path.join(self.path, "index.json")) - - default_mask = stat.S_IMODE(os.stat(self.path).st_mode) - rechmod(self.path, default_mask) - - if os.access(self.path, os.W_OK): - default_mask = stat.S_IMODE( - os.stat(self.path).st_mode - ) # mode of files to be included - build_viewer( - os.path.join(self.path, "index.json"), - diag_name="My Diagnostics", - default_mask=default_mask, - ) - - -class LandVariable(object): - def __init__(self, csv_row: List[str]): - # TODO: use these fields... - # the name of the ELM variable on the monthly h0 history file - self.variable_name = csv_row[0] - # “A” or “T” for global average over land area or global total, respectively - self.metric = csv_row[1] - # the factor that should convert from original units to final units, after standard processing with nco - self.scale_factor = csv_row[2] - # test string for the units as given on the history file (included here for possible testing) - self.original_units = csv_row[3] - # the units that should be reported in time series plots, based on A/T and Scale Factor - self.final_units = csv_row[4] - # a name used to cluster variables together, to be separated in groups within the output web pages - self.group = csv_row[5] - # Descriptive text to add to the plot page to help users identify the variable - self.long_name = csv_row[6] - - -def construct_land_variables() -> List[LandVariable]: - var_list: List[LandVariable] = [] - header = True - # If this file is being run stand-alone, then - # it will search the directory above the git directory - with open("zppy_land_fields.csv", newline="") as csv_file: - print("In File") - var_reader = csv.reader(csv_file) - for row in var_reader: - # Skip the header row - print(f"row={row}") - if header: - header = False - else: - var_list.append(LandVariable(row)) - return var_list - - -class LandGroup(object): - def __init__(self, name: str, land_variables: List[LandVariable]): - self.group_name = name - self.short_name = name.lower().replace(" ", "") - self.land_variables = land_variables - - -def get_land_groups(land_variables: List[LandVariable]) -> List[LandGroup]: - group_names: List[str] = [] - groups: List[LandGroup] = [] - for lv in land_variables: - g: str = lv.group - if g not in group_names: - # A new group! - group_names.append(g) - groups.append(LandGroup(g, [lv])) - else: - # Add a new variable to this existing group - for group in groups: - if g == group.group_name: - group.land_variables.append(lv) - return groups +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 create_viewer(figstr, regions, component, requested_plots): +def create_viewer(parameters: Parameters, vars: List[Variable], component: str) -> str: viewer = OutputViewer(path=".") - viewer.add_page("Table", regions) - land_variables: List[LandVariable] = construct_land_variables() - groups: List[LandGroup] = get_land_groups(land_variables) + viewer.add_page("Table", parameters.regions) + groups: List[VariableGroup] = get_variable_groups(vars) for group in groups: - vars_in_group = list(map(lambda lv: lv.variable_name, group.land_variables)) - for plot_name in requested_plots: - if plot_name in vars_in_group: - # There's at least one plot in this group. - # So, we want to add it to the viewer. - viewer.add_group(group.group_name) - break - for plot_name in vars_in_group: - # Only plot the requested plots - if plot_name in requested_plots: - viewer.add_row(plot_name) - for rgn in regions: - # v3.LR.historical_0051_glb_lnd_SOIL4C.png - # viewer/c-state/glb_lnd_soil4c.html - viewer.add_col( - f"{figstr}_{rgn}_{component}_{plot_name}.png", - is_file=True, - title=f"{rgn}_{component}_{plot_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.long_name != "": + row_title = f"{plot_name}: {var.long_name}" + else: + row_title = plot_name + viewer.add_row(row_title) + for rgn in parameters.regions: + # v3.LR.historical_0051_glb_lnd_SOIL4C.png + # viewer/c-state/glb_lnd_soil4c.html + 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 + viewer.generate_viewer() # Copy the contents of `table` into the `viewer` directory # (which initially only has `css` and `js` subdirectories) # Because `viewer` already exists, # `shutil.copytree` will not work. distutils.dir_util.copy_tree("table", "viewer") + print( + os.getcwd() + ) # /lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_viewers_output/test-pr616-20241022v2/v3.LR.historical_0051/post/scripts/global_time_series_1985-1995_dir + # shutil.rmtree("table") + # new_url = f"viewer_{component}" + # # shutil.move("viewer", new_url) + # distutils.dir_util.copy_tree("viewer", new_url) return url -def run_by_region(parameters): - regions = parameters[16].split(",") - for rgn in regions: - if rgn.lower() in ["glb", "global"]: - rgn = "glb" - elif rgn.lower() in ["n", "north", "northern"]: - rgn = "n" - elif rgn.lower() in ["s", "south", "southern"]: - rgn = "s" +# 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. + """ + + def insert_data_in_row(row_obj, name, url): + """ + 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) + + install_path = "" # TODO: figure this out + path = os.path.join(install_path, "viewer", "index_template.html") + output = 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: - raise RuntimeError(f"Invalid rgn={rgn}") - run(parameters, rgn) - figstr = parameters[3] - # plots_original = param_get_list(parameters[8]) - # if parameters[9].lower() == "false": - # atmosphere_only = False - # else: - # atmosphere_only = True - # plots_atm = param_get_list(parameters[10]) - # plots_ice = param_get_list(parameters[11]) - plots_lnd = param_get_list(parameters[12]) - # plots_ocn = param_get_list(parameters[13]) - nrows = int(parameters[14]) - ncols = int(parameters[15]) - plots_per_page = nrows * ncols + 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 + + +def run_by_region(command_line_arguments): + parameters: Parameters = Parameters(command_line_arguments) + requested_variables = RequestedVariables(parameters) + for rgn in parameters.regions: + run(parameters, requested_variables, rgn) + plots_per_page = parameters.nrows * parameters.ncols + # TODO: Is this how we want to determine when to make a viewer or should we have a `make_viewer` parameter in the cfg? if plots_per_page == 1: # In this case, we don't want the summary PDF. # Rather, we want to construct a viewer similar to E3SM Diags. - # TODO: make viewer home page to point to multiple viewers - url = create_viewer(figstr, regions, "lnd", plots_lnd) - print(url) + # TODO: determine directory paths for each viewer + # TODO: include "original"? + # for component in ["original", "atm", "ice", "lnd", "ocn"]: + title_and_url_list: List[Tuple[str, str]] = [] + for component in ["lnd"]: + vars = get_vars(requested_variables, component) + url = create_viewer(parameters, vars, component) + print(url) + title_and_url_list.append((component, url)) + # index_url: str = create_viewer_index(parameters.case_dir, title_and_url_list) + # print(f"Viewer index URL: {index_url}") if __name__ == "__main__": run_by_region(sys.argv) - """ - run_by_region( - [ - "coupled_global.py", - "/lcrc/group/e3sm/ac.forsyth2/zppy_min_case_global_time_series_single_plots_output/test-616-20240930/v3.LR.historical_0051", - "v3.LR.historical_0051", - "v3.LR.historical_0051", - "1985", - "1989", - "Blue", - "5", - "None", - "false", - "TREFHT", - "None", - "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", - "None", - "1", - "1", - "glb,n,s", - ] - ) - """ diff --git a/zppy/templates/default.ini b/zppy/templates/default.ini index 68a71677..33043bd9 100644 --- a/zppy/templates/default.ini +++ b/zppy/templates/default.ini @@ -323,6 +323,7 @@ plots_original = string(default="net_toa_flux_restom,global_surface_air_temperat # These should be a subset of the `vars` generated by the `ts` `glb` subtasks. plots_atm = string(default="") plots_ice = string(default="") +# Set `plots_lnd = "all"` to run every variable in the land csv file. plots_lnd = string(default="") plots_ocn = string(default="") # regions to plot: glb, n, s (global, northern hemisphere, southern hemisphere) diff --git a/zppy/templates/global_time_series.bash b/zppy/templates/global_time_series.bash index 7235a2a4..5b0bdcc8 100644 --- a/zppy/templates/global_time_series.bash +++ b/zppy/templates/global_time_series.bash @@ -83,7 +83,6 @@ mkdir -p ${results_dir_absolute_path} cp *.pdf ${results_dir_absolute_path} cp *.png ${results_dir_absolute_path} cp -r viewer ${results_dir_absolute_path} -cp -r table ${results_dir_absolute_path} ################################################################################ # Copy output to web server diff --git a/zppy/templates/readTS.py b/zppy/templates/readTS.py deleted file mode 100644 index 16a9031e..00000000 --- a/zppy/templates/readTS.py +++ /dev/null @@ -1,78 +0,0 @@ -from typing import Optional, Tuple - -import xarray -import xcdat # noqa: F401 - - -class TS(object): - def __init__(self, directory): - - self.directory: str = directory - - # `directory` will be of the form `{case_dir}/post//glb/ts/monthly/{ts_num_years}yr/` - self.f: xarray.core.dataset.Dataset = xcdat.open_mfdataset( - f"{directory}*.nc", center_times=True - ) - - def __del__(self): - - self.f.close() - - def globalAnnual( - self, var: str - ) -> Tuple[xarray.core.dataarray.DataArray, Optional[str]]: - - v: xarray.core.dataarray.DataArray - units: Optional[str] = None - - # Constants, from AMWG diagnostics - Lv = 2.501e6 - Lf = 3.337e5 - - # Is this a derived variable? - if var == "RESTOM": - - FSNT, _ = self.globalAnnual("FSNT") - FLNT, _ = self.globalAnnual("FLNT") - v = FSNT - FLNT - - elif var == "RESTOA": - - print("NOT READY") - FSNTOA, _ = self.globalAnnual("FSNTOA") - FLUT, _ = self.globalAnnual("FLUT") - v = FSNTOA - FLUT - - elif var == "LHFLX": - - QFLX, _ = self.globalAnnual("QFLX") - PRECC, _ = self.globalAnnual("PRECC") - PRECL, _ = self.globalAnnual("PRECL") - PRECSC, _ = self.globalAnnual("PRECSC") - PRECSL, _ = self.globalAnnual("PRECSL") - v = (Lv + Lf) * QFLX - Lf * 1.0e3 * (PRECC + PRECL - PRECSC - PRECSL) - - elif var == "RESSURF": - - FSNS, _ = self.globalAnnual("FSNS") - FLNS, _ = self.globalAnnual("FLNS") - SHFLX, _ = self.globalAnnual("SHFLX") - LHFLX, _ = self.globalAnnual("LHFLX") - v = FSNS - FLNS - SHFLX - LHFLX - - elif var == "PREC": - - PRECC, _ = self.globalAnnual("PRECC") - PRECL, _ = self.globalAnnual("PRECL") - v = 1.0e3 * (PRECC + PRECL) - - else: - # Non-derived variables - - annual_average_dataset_for_var: xarray.core.dataset.Dataset = ( - self.f.temporal.group_average(var, "year") - ) - v = annual_average_dataset_for_var.data_vars[var] - units = v.units - - return v, units