diff --git a/tests/global_time_series/test_coupled_global.py b/tests/global_time_series/test_coupled_global.py index eb1a8ffb..f01e7f4e 100644 --- a/tests/global_time_series/test_coupled_global.py +++ b/tests/global_time_series/test_coupled_global.py @@ -1,17 +1,355 @@ import unittest +from typing import Any, Dict, List -from zppy.templates.coupled_global import get_ylim +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): - actual = get_ylim([0, 1], [-1, 2]) - self.assertEqual(actual, [-1, 2]) + # 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__": diff --git a/zppy/templates/coupled_global.py b/zppy/templates/coupled_global.py index be1a1b51..fc314bfa 100644 --- a/zppy/templates/coupled_global.py +++ b/zppy/templates/coupled_global.py @@ -17,6 +17,7 @@ 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 ( @@ -61,14 +62,16 @@ def __init__(self, parameters): 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 = False if (parameters[9].lower() == "false") else True + 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 list elements are used often as strings, + # 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(",")) @@ -91,9 +94,10 @@ def __init__( group="", long_name="", ): - # TODO: Make use of all these fields # 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 @@ -102,6 +106,8 @@ def __init__( 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 @@ -110,7 +116,9 @@ def __init__( def get_vars_original(plots_original: List[str]) -> List[Variable]: vars_original: List[Variable] = [] - if "net_toa_flux_restom" or "net_atm_energy_imbalance" in plots_original: + 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")) @@ -160,11 +168,14 @@ def construct_land_variables(requested_vars: List[str]) -> List[Variable]: if header: header = False else: - # If no land variables are requested, assume user wants all - # TODO: this is not a great design. What if a user doesn't want land component plotted? - # Maybe have them literally write "all" as the var? - if (requested_vars == []) or (row[0] in requested_vars): - var_list.append(land_csv_row_to_var(row)) + # 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 @@ -195,7 +206,6 @@ def __init__(self, parameters: Parameters): class VariableGroup(object): def __init__(self, name: str, variables: List[Variable]): self.group_name = name - self.short_name = name.lower().replace(" ", "") self.variables = variables @@ -449,9 +459,9 @@ def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: set_intersection: set = set(["change_ohc", "max_moc", "change_sea_level"]) & set( parameters.plots_original ) - has_original_ocn_plots: bool = set_intersection == set() + has_original_ocn_plots: bool = set_intersection != set() use_ocn: bool = (parameters.plots_ocn != []) or ( - not parameters.atmosphere_only and has_original_ocn_plots + (not parameters.atmosphere_only) and has_original_ocn_plots ) ocean_dir = get_data_dir(parameters, "ocn", use_ocn) exps: List[Dict[str, Any]] = [ @@ -474,9 +484,7 @@ def get_exps(parameters: Parameters) -> List[Dict[str, Any]]: def set_var( exp: Dict[str, Any], exp_key: str, - var_list: List[ - Variable - ], # TODO: make changes to have List[Variable] instead of List[str] + var_list: List[Variable], valid_vars: List[str], invalid_vars: List[str], rgn: str, @@ -1208,6 +1216,7 @@ def make_plot_pdfs( # noqa: C901 pdf.close() +# Run coupled_global ########################################################## # ----------------------------------------------------------------------------- def run(parameters: Parameters, requested_variables: RequestedVariables, rgn: str): # Experiments @@ -1265,12 +1274,17 @@ def create_viewer(parameters: Parameters, vars: List[Variable], component: str) viewer.add_page("Table", parameters.regions) groups: List[VariableGroup] = get_variable_groups(vars) for group in groups: - vars_in_group: List[str] = list(map(lambda v: v.variable_name, group.variables)) # 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 plot_name in vars_in_group: - viewer.add_row(plot_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 @@ -1290,6 +1304,66 @@ def create_viewer(parameters: Parameters, vars: List[Variable], component: str) return url +# Copied from E3SM Diags and modified +def create_viewer_index( + root_dir: str, title_and_url_list: List[Tuple[str, str]] +) -> str: + """ + Creates the index page in root_dir which + joins the individual viewers. + + Each tuple is on its own row. + """ + + 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: + 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) @@ -1301,37 +1375,17 @@ def run_by_region(command_line_arguments): # In this case, we don't want the summary PDF. # Rather, we want to construct a viewer similar to E3SM Diags. # TODO: determine directory paths for each viewer - # TODO: make viewer home page to point to multiple viewers # 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)