lat = atm.get_coord(data, 'lat')
lon = atm.get_coord(data, 'lon')
psfile = atm.homedir() + 'dynamics/python/atmos-tools/data/topo/ncep2_ps.nc'
ps = atm.get_ps_clim(lat, lon, psfile)
ps = ps / 100
figsize = (7, 9)
omitzero = False
for ssn in ['ANN', 'DJF', 'JJA', 'MAR']:
for lonlims in [(0, 360), (60, 100)]:
lon1, lon2 = lonlims
lonstr = atm.latlon_str(lon1, lon2, 'lon')
suptitle = ssn + ' ' + lonstr
months = atm.season_months(ssn)
v = data['V'].sel(month=months)
if (lon2 - lon1) < 360:
v = atm.subset(v, {'lon' : (lon1, lon2)})
sector_scale = (lon2 - lon1) / 360.0
psbar = atm.dim_mean(ps, 'lon', lon1, lon2)
clev = 10
else:
sector_scale = None
psbar = atm.dim_mean(ps, 'lon')
clev = 20
vssn = v.mean(dim='month')
vssn_bar = atm.dim_mean(vssn, 'lon')
psi1 = atm.streamfunction(vssn, sector_scale=sector_scale)
psi1 = atm.dim_mean(psi1, 'lon')
psi2 = atm.streamfunction(vssn_bar, sector_scale=sector_scale)
keys_dict = {'PRECTOT' : 'PRECTOT', 'CMAP' : 'precip', 'GPCP' : 'PREC',
'U200' : 'U', 'U850' : 'U', 'V200' : 'V', 'V850' : 'V'}
data = {}
for nm in datafiles:
print('Loading ' + datafiles[nm])
with xray.open_dataset(datafiles[nm]) as ds:
if 'year' in ds.dims:
ds = ds.mean(dim='year')
data[nm] = ds[keys_dict[nm]].load()
# ENSO indices
enso = pd.read_csv(ensofile, index_col=0)
enso = enso.loc[years]
for key in enso_keys:
if key not in enso.columns:
months = atm.season_months(key)
month_names = [(atm.month_str(m)).capitalize() for m in months]
enso[key] = enso[month_names].mean(axis=1)
enso = enso[enso_keys]
col_names = [enso_nm + ' ' + nm for nm in enso.columns]
enso.columns = col_names
# ----------------------------------------------------------------------
# Daily timeseries
ts = xray.Dataset()
for nm in ['GPCP', 'PRECTOT']:
ts[nm] = atm.mean_over_geobox(data[nm], lat1, lat2, lon1, lon2)
ts['MFC'] = utils.daily_rel2onset(index_all['CHP_MFC']['daily_ts'],
index[ind_nm], npre, npost)
ts['CMFC'] = utils.daily_rel2onset(index_all['CHP_MFC']['tseries'],
topo = atm.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'
# ----------------------------------------------------------------------
# Correct for topography
u_orig = u
u = atm.correct_for_topography(u_orig, topo)
# ----------------------------------------------------------------------
# Zonal mean zonal wind
season = 'jjas'
lon1, lon2 = 60, 100
cint = 5
months = atm.season_months(season)
uplot = atm.subset(u, 'lon', lon1, lon2, 'mon', months)
uplot = uplot.mean(['lon', 'mon'])
ps_plot = atm.subset(topo, 'lon', lon1, lon2)
ps_plot = ps_plot.mean('lon')
plt.figure()
cs = atm.contour_latpres(uplot, clev=cint, topo=ps_plot)
clev = atm.clevels(uplot, cint, omitzero=True)
plt.clabel(cs, clev[::2], fmt='%02d')
plt.figure()
atm.contourf_latpres(uplot, clev=cint, topo=ps_plot)
def load_daily_season(pathstr, year, season='ann', var_ids=None,
lat1=-90, lat2=90, lon1=0, lon2=360,
verbose=True, concat_dim=None):
"""Return daily data for a selected year, season and lat-lon subset.
Loads daily data from locally saved files and concatenates it into
a single DataArray or Dataset for that year and season.
Parameters
----------
pathstr : str
Beginning of path for each data file, where each file name is in
the format *yyyymm.nc.
e.g. pathstr = '~/datastore/merra/daily/u200_'
year : int
Year to load.
season : str, optional
Season to load. Valid values are as listed in atm.season_months()
e.g. 'jul', 'jja', 'ann'
Default is entire year ('ann')
var_ids : str or list of str, optional
Variable(s) to extract. If omitted, all variables in the data are
included and the output is a Dataset.
lat1, lat2, lon1, lon2 : floats, optional
Lat-lon subset to extract.
concat_dim : str, optional
Name of time dimension for concatenation. If None, then
atm.get_coord() is called to get the name from the data file.
verbose : bool, optional
If True, print updates while processing files.
Returns
-------
data : xray.DataArray or xray.Dataset
"""
months = atm.season_months(season)
paths = []
for m in months:
datestr = '%d%02d' % (year, m)
paths.append(pathstr + datestr + '.nc')
# Make sure longitude range is consistent with data
with xray.open_dataset(paths[0]) as ds:
lonmax = atm.lon_convention(atm.get_coord(ds, 'lon'))
if concat_dim is None:
concat_dim = atm.get_coord(ds, 'time', 'name')
if lon2 - lon1 == 360:
if lonmax < lon2:
offset = -180
elif lonmax > lon2:
offset = 180
else:
offset = 0
lon1, lon2 = lon1 + offset, lon2 + offset
print(lon1, lon2, lonmax)
# Load daily data
if var_ids is None:
var_nms = None
else:
var_nms = [get_varname(var_id) for var_id in atm.makelist(var_ids)]
subset_dict = {'lat' : (lat1, lat2), 'lon' : (lon1, lon2)}
data = atm.load_concat(paths, var_nms, concat_dim, subset_dict, verbose)
return data
