Replace CDAT

zppy/templates/coupled_global.py

@@ -3,12 +3,15 @@
 import math
 import sys
 import traceback
-from typing import Any, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
+import cftime
 import matplotlib as mpl
 import matplotlib.backends.backend_pdf
 import matplotlib.pyplot as plt
 import numpy as np
+import xarray
+import xcdat  # noqa: F401
 from netCDF4 import Dataset
 from readTS import TS
 
@@ -549,19 +552,28 @@ def param_get_list(param_value):
         return param_value.split(",")
 
 
-def set_var(exp, exp_key, var_list, valid_vars, invalid_vars, rgn):
+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 len(v.shape) > 1:
+            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
@@ -574,18 +586,17 @@ def set_var(exp, exp_key, var_list, valid_vars, invalid_vars, rgn):
                     n = 2
                 else:
                     raise RuntimeError(f"Invalid rgn={rgn}")
-                v = v[:, n]  # Just use nth column
+                v = v.isel(rgn=n)  # Just use nth column
             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
             exp["annual"][var] = (v, units)
             if "year" not in exp["annual"]:
-                time = v.getTime()
-                exp["annual"]["year"] = [x.year for x in time.asComponentTime()]
+                years: np.ndarray[cftime.DatetimeNoLeap] = v.coords["time"].values
+                exp["annual"]["year"] = [x.year for x in years]
         del ts
 
 
@@ -736,7 +747,7 @@ def run(parameters, rgn):  # noqa: C901
         or ("change_sea_level" in plots_original)
     )
     use_ocn = plots_ocn or (not atmosphere_only and has_original_ocn_plots)
-    exps = [
+    exps: List[Dict[str, Any]] = [
         {
             "atmos": None
             if not use_atmos
@@ -767,7 +778,7 @@ def run(parameters, rgn):  # noqa: C901
     invalid_vars: List[str] = []
 
     # Read data
-    exp: Any
+    exp: Dict[str, Any]
     for exp in exps:
         exp["annual"] = {}
 
@@ -845,5 +856,27 @@ def run_by_region(parameters):
         run(parameters, rgn)
 
 
+def debug():
+    args = [
+        "coupled_global.py",
+        "/lcrc/group/e3sm/ac.forsyth2/zppy_test_debug_output/test-346/v2.LR.historical_0201",
+        "v2.LR.historical_0201",
+        "v2_historical_0201",
+        "1850",
+        "1860",
+        "Blue",
+        "5",
+        "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",
+        "False",
+        "None",
+        "None",
+        "LAISHA,LAISUN",
+        "None",
+        "glb,n,s",
+    ]
+    run_by_region(args)
+
+
 if __name__ == "__main__":
     run_by_region(sys.argv)
+    # debug()

zppy/templates/readTS.py

@@ -1,16 +1,17 @@
 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/<componen>/glb/ts/monthly/{ts_num_years}yr`
-        self.f: xarray.core.dataset.Dataset = xarray.open_mfdataset(f"{directory}/*.nc")
+        # `directory` will be of the form `{case_dir}/post/<componen>/glb/ts/monthly/{ts_num_years}yr`
         # Refactor note: `self.f = cdms2.open(filename)` gave `cdms2.dataset.Dataset`
+        self.f: xarray.core.dataset.Dataset = xarray.open_mfdataset(f"{directory}/*.nc")
 
     def __del__(self):
 
@@ -20,14 +21,13 @@ 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
 
-        v: xarray.core.dataarray.DataArray
-
         # Is this a derived variable?
         if var == "RESTOM":
 
@@ -68,16 +68,10 @@ def globalAnnual(
         else:
             # Non-derived variables
 
-            # Read variable
-            v = self.f.data_vars[var]
-            # Refactor note: `v = self.f(var)` gave `cdms2.tvariable.TransientVariable`
+            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
 
-            # Annual average
-
-            # Refactor note: `AttributeError: 'Dataset' object has no attribute 'temporal'` seems to always occur
-            # Regardless if using CDAT or not, if using as object or class method.
-            # v = self.f.temporal.group_average(v, "year")
-            v = xarray.Dataset.temporal.group_average(v, "year")
-
         return v, units