for nm in varnms:
print('Loading ' + relfiles[nm])
with xray.open_dataset(relfiles[nm]) as ds:
if nm == 'PSI':
data[nm] = atm.streamfunction(ds['V'])
psimid = atm.subset(data[nm], {'plev' : (pmid, pmid)},
squeeze=True)
psimid.name = 'PSI%d' % pmid
data['PSI%d' % pmid] = psimid
elif nm == 'VFLXLQV':
var = atm.dim_mean(ds['VFLXQV'], 'lon', lon1, lon2)
data[nm] = var * atm.constants.Lv.values
elif nm == theta_nm:
theta = ds[nm]
_, _, dtheta = atm.divergence_spherical_2d(theta, theta)
data[nm] = atm.dim_mean(ds[nm], 'lon', lon1, lon2)
data[dtheta_nm] = atm.dim_mean(dtheta, 'lon', lon1, lon2)
elif nm == dtheta_nm:
continue
else:
data[nm] = atm.dim_mean(ds[nm], 'lon', lon1, lon2)
databar = {}
for nm in data:
if 'year' in data[nm].dims:
databar[nm] = data[nm].mean(dim='year')
else:
databar[nm] = data[nm]
# ----------------------------------------------------------------------