diff --git a/zppy/templates/coupled_global.py b/zppy/templates/coupled_global.py index 4957ba55..762e7724 100644 --- a/zppy/templates/coupled_global.py +++ b/zppy/templates/coupled_global.py @@ -3,12 +3,14 @@ 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 from netCDF4 import Dataset from readTS import TS @@ -549,11 +551,20 @@ 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: @@ -561,7 +572,7 @@ def set_var(exp, exp_key, var_list, valid_vars, invalid_vars, rgn): 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 +585,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 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 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,36 +746,26 @@ 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 - else "{}/post/atm/glb/ts/monthly/{}yr/glb.xml".format( - case_dir, ts_num_years - ), - "ice": None - if not plots_ice - else "{}/post/ice/glb/ts/monthly/{}yr/glb.xml".format( - case_dir, ts_num_years - ), - "land": None - if not plots_lnd - else "{}/post/lnd/glb/ts/monthly/{}yr/glb.xml".format( - case_dir, ts_num_years - ), - "ocean": None - if not use_ocn - else "{}/post/ocn/glb/ts/monthly/{}yr/glb.xml".format( - case_dir, ts_num_years - ), - "moc": None - if not use_ocn - else "{}/post/ocn/glb/ts/monthly/{}yr/".format(case_dir, ts_num_years), - "vol": None - if not use_ocn - else "{}/post/ocn/glb/ts/monthly/{}yr/glb.xml".format( - case_dir, ts_num_years - ), + "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],), @@ -777,7 +777,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"] = {} diff --git a/zppy/templates/global_time_series.bash b/zppy/templates/global_time_series.bash index 2154e66b..22dee4d7 100644 --- a/zppy/templates/global_time_series.bash +++ b/zppy/templates/global_time_series.bash @@ -39,28 +39,7 @@ global_ts_dir={{ global_time_series_dir }} export CDMS_NO_MPI=true use_atm={{ use_atm }} -if [[ ${use_atm,,} == "true" ]]; then - echo 'Create xml files for atm' - cd ${case_dir}/post/atm/glb/ts/monthly/${ts_num_years}yr - cdscan -x glb.xml *.nc - if [ $? != 0 ]; then - cd {{ scriptDir }} - echo 'ERROR (1)' > {{ prefix }}.status - exit 1 - fi -fi - use_lnd={{ use_lnd }} -if [[ ${use_lnd,,} == "true" ]]; then - echo 'Create xml files for lnd' - cd ${case_dir}/post/lnd/glb/ts/monthly/${ts_num_years}yr - cdscan -x glb.xml *.nc - if [ $? != 0 ]; then - cd {{ scriptDir }} - echo 'ERROR (2)' > {{ prefix }}.status - exit 2 - fi -fi use_ocn={{ use_ocn }} if [[ ${use_ocn,,} == "true" ]]; then @@ -75,16 +54,6 @@ if [[ ${use_ocn,,} == "true" ]]; then exit 3 fi - echo 'Create xml for for ocn' - export CDMS_NO_MPI=true - cd ${case_dir}/post/ocn/glb/ts/monthly/${ts_num_years}yr - cdscan -x glb.xml mpaso.glb*.nc - if [ $? != 0 ]; then - cd {{ scriptDir }} - echo 'ERROR (4)' > {{ prefix }}.status - exit 4 - fi - echo 'Copy moc file' cd ${case_dir}/post/analysis/mpas_analysis/cache/timeseries/moc cp ${moc_file} ../../../../../ocn/glb/ts/monthly/${ts_num_years}yr/ diff --git a/zppy/templates/readTS.py b/zppy/templates/readTS.py index 58f4c712..312e604d 100644 --- a/zppy/templates/readTS.py +++ b/zppy/templates/readTS.py @@ -1,21 +1,27 @@ -import cdms2 -import cdutil +from typing import Optional, Tuple + +import xarray +import xcdat # noqa: F401 class TS(object): - def __init__(self, filename): + def __init__(self, directory): - self.filename = filename + self.directory: str = directory - self.f = cdms2.open(filename) + # `directory` will be of the form `{case_dir}/post//glb/ts/monthly/{ts_num_years}yr/` + self.f: xarray.core.dataset.Dataset = xarray.open_mfdataset(f"{directory}*.nc") def __del__(self): self.f.close() - def globalAnnual(self, var): + def globalAnnual( + self, var: str + ) -> Tuple[xarray.core.dataarray.DataArray, Optional[str]]: - units = None + v: xarray.core.dataarray.DataArray + units: Optional[str] = None # Constants, from AMWG diagnostics Lv = 2.501e6 @@ -61,11 +67,10 @@ def globalAnnual(self, var): else: # Non-derived variables - # Read variable - v = self.f(var) + 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 - v = cdutil.YEAR(v) - return v, units