Skip to content

Commit

Permalink
Replace CDAT (#519)
Browse files Browse the repository at this point in the history
* Replace CDAT

* types without xml

* glb ts without cdat

* Clean up code
  • Loading branch information
forsyth2 authored May 23, 2024
1 parent 905b755 commit 0883805
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 80 deletions.
74 changes: 37 additions & 37 deletions zppy/templates/coupled_global.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -549,19 +551,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
Expand All @@ -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


Expand Down Expand Up @@ -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],),
Expand All @@ -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"] = {}

Expand Down
31 changes: 0 additions & 31 deletions zppy/templates/global_time_series.bash
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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/
Expand Down
29 changes: 17 additions & 12 deletions zppy/templates/readTS.py
Original file line number Diff line number Diff line change
@@ -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/<componen>/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
Expand Down Expand Up @@ -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

0 comments on commit 0883805

Please sign in to comment.