Working Total plots

E3SM-Project · Dec 4, 2024 · be48809 · be48809
1 parent 9b7430f
commit be48809
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 35 deletions.
diff --git a/tests/integration/global_time_series/cases_global_time_series.py b/tests/integration/global_time_series/cases_global_time_series.py
@@ -23,8 +23,7 @@
         "atmosphere_only": "False",
         "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_lnd": "AR",
+        "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",

diff --git a/zppy_interfaces/global_time_series/coupled_global.py b/zppy_interfaces/global_time_series/coupled_global.py
@@ -163,10 +163,13 @@ def set_var(
     valid_vars: List[str],
     invalid_vars: List[str],
     rgn: str,
+    extra_vars_dict=None,
 ) -> None:
     if exp[exp_key] != "":
         try:
-            ts_object: DatasetWrapper = DatasetWrapper(exp[exp_key])
+            dataset_wrapper: DatasetWrapper = DatasetWrapper(
+                exp[exp_key], extra_vars_dict
+            )
         except Exception as e:
             logger.critical(e)
             logger.critical(
@@ -178,7 +181,7 @@ def set_var(
             try:
                 data_array: xarray.core.dataarray.DataArray
                 units: str
-                data_array, units = ts_object.globalAnnual(var)
+                data_array, units = dataset_wrapper.globalAnnual(var)
                 valid_vars.append(str(var_str))
             except Exception as e:
                 logger.error(e)
@@ -211,7 +214,7 @@ def set_var(
                     "time"
                 ].values
                 exp["annual"]["year"] = [x.year for x in years]
-        del ts_object
+        del dataset_wrapper
 
 
 def process_data(
@@ -237,22 +240,30 @@ def process_data(
         )
         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
+            exp,
+            "land",
+            requested_variables.vars_land,
+            valid_vars,
+            invalid_vars,
+            rgn,
+            extra_vars_dict=exp["land"].replace("glb", "180x360_aave"),
         )
         set_var(
             exp, "ocean", requested_variables.vars_ocn, valid_vars, invalid_vars, rgn
         )
 
         # Optionally read ohc
         if exp["ocean"] != "":
-            ts = DatasetWrapper(exp["ocean"])
-            exp["annual"]["ohc"], _ = ts.globalAnnual(Variable("ohc"))
+            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"] != "":
-            ts = DatasetWrapper(exp["vol"])
-            exp["annual"]["volume"], _ = ts.globalAnnual(Variable("volume"))
+            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]

diff --git a/zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py b/zppy_interfaces/global_time_series/coupled_global_dataset_wrapper.py
@@ -10,7 +10,7 @@
 
 
 class DatasetWrapper(object):
-    def __init__(self, directory):
+    def __init__(self, directory, extra_vars_directory=None):
 
         self.directory: str = directory
 
@@ -19,11 +19,12 @@ def __init__(self, directory):
             f"{directory}*.nc", center_times=True
         )
 
-        # TODO: going to need to get the alternative mapping_file and pass in an extra_vars specific directory to this function.
-        self.extra_var_dataset: xarray.core.dataset.Dataset = xcdat.open_mfdataset(
-            "/lcrc/group/e3sm/ac.forsyth2/zi-test-input-data/post/lnd/180x360_aave/ts/monthly/5yr/*nc",
-            center_times=True,
-        )
+        self.extra_var_dataset: xarray.core.dataset.Dataset = None
+        if extra_vars_directory:
+            self.extra_var_dataset = xcdat.open_mfdataset(
+                f"{extra_vars_directory}*.nc",
+                center_times=True,
+            )
 
     def __del__(self):
 
@@ -111,29 +112,22 @@ def globalAnnualHelper(
             )
             data_array = annual_average_dataset_for_var.data_vars[var]
             if metric == Metric.TOTAL:
-                # ['AR', 'time_bounds', 'CWDC', 'FSH', 'GPP', 'H2OSNO', 'HR', 'LAISHA', 'LAISUN', 'NBP', 'QINTR', 'QOVER', 'QRUNOFF', 'QSOIL', 'QVEGE', 'QVEGT', 'RH2M', 'SOIL1C', 'SOIL2C', 'SOIL3C', 'SOIL4C', 'SOILWATER_10CM', 'TOTLITC', 'TOTVEGC', 'TSA', 'WOOD_HARVESTC', 'lon_bnds', 'lat_bnds']
-                # TODO: looks like we don't actually have area or landfrac in the dataset
                 logger.debug(
                     f"self.extra_var_dataset.keys()={list(self.extra_var_dataset.keys())}"
                 )
-                area = self.extra_var_dataset["area"]
-                landfrac = self.extra_var_dataset["landfrac"]
-                try:
-                    data_array *= area * landfrac
-                except ValueError as e:
-                    logger.error(f"area={area}")
-                    logger.error(f"landfrac={landfrac}")
-                    logger.error(f"area.shape={area.shape}")  # area.shape=(180, 360)
-                    logger.error(
-                        f"landfrac.shape={landfrac.shape}"
-                    )  # landfrac.shape=(180, 360)
-                    logger.error(
-                        f"data_array.shape={data_array.shape}"
-                    )  # data_array.shape=(10, 3)
-                    # e: dimensions cannot change for in-place operations
-                    raise e
+                area: xarray.core.dataarray.DataArray = self.extra_var_dataset["area"]
+                landfrac: xarray.core.dataarray.DataArray = self.extra_var_dataset[
+                    "landfrac"
+                ]
+                # area.shape() = (180, 360)
+                total_area = area.sum()  # Sum over all dimensions
+                # landfrac.shape=(180, 360)
+                average_landfrac = landfrac.mean()  # Mean over all dimensions
+                total_land_area = total_area * average_landfrac
+                # data_array.shape = (number of years, number of regions)
+                data_array *= total_land_area
                 logger.info(
-                    "for Metric.TOTAL, data_array has been scaled by area*landfrac"
+                    "for Metric.TOTAL, data_array has been scaled by total land area"
                 )
             units = data_array.units
             # `units` will be "1" if it's a dimensionless quantity