def latlon_plot(varnm, reg, day_or_season, coeff='m', stipple_kw={},
axlims=(-60, 60, 40, 120)):
regdata = reg[varnm + '_latlon']
keys = [key for key in regdata if key.endswith('_' + coeff)]
clim = atm.climits(regdata[keys].to_array(), symmetric=True,
percentile=99.9)
xname, yname = 'lon', 'lat'
lat = atm.get_coord(regdata, 'lat')
if max(np.diff(lat)) > 1:
xsample, ysample = 1, 1
else:
xsample, ysample = 2, 2
if isinstance(day_or_season, int):
key = varnm + '_DAILY_'
var = regdata[key + coeff].sel(dayrel=day_or_season)
p = regdata[key + 'p'].sel(dayrel=day_or_season)
titlestr = varnm + ' Day %d' % day_or_season
else:
key = varnm + '_' + day_or_season + '_'
var = regdata[key + coeff]
p = regdata[key + 'p']
titlestr = varnm + ' ' + day_or_season
pts_mask = stipple_mask(p)
atm.pcolor_latlon(var, axlims=axlims, fancy=False)
plt.clim(clim)
atm.stipple_pts(pts_mask, xname, yname, xsample, ysample, **stipple_kw)
plt.title(titlestr)
def plot_all_WLH(wlh, y=0, axlims=(0,50,50,180), cmap='jet', clines=True):
onset = wlh['onset']
Rsq = wlh['Rsq']
lat = wlh['lat']
lon = wlh['lon']
kmax = wlh['smoothing_kmax']
years=wlh['years']
if wlh['climatology']:
titlestr = 'CMAP %d-%d Climatology' % (years.min(), years.max())
else:
onset = onset[y]
Rsq = Rsq[y]
titlestr = 'CMAP %d' % years[y]
# Note: add 1 to pentad indices to index from 1-73 for comparison
# with Wang & LinHo
onset = onset + 1
# Calculate onset dates from pentads
nlat, nlon = onset.shape
onset_date = np.nan * np.ones((nlat, nlon))
for i in range(nlat):
for j in range(nlon):
jday = atm.pentad_to_jday(onset[i, j], pmin=1)
mon, day = atm.jday_to_mmdd(jday)
onset_date[i, j] = 100*mon + day
# -- Smooth with cubic spline
# lat_i = np.arange(-89.5, 90, 0.5)
# lon_i = np.arange(0, 360, 0.5)
lat_i = np.arange(-90, 90, 4.)
lon_i = np.arange(0, 360, 4.)
onset_date = atm.interp_latlon(onset_date, lat_i, lon_i, lat, lon, order=3)
clev_date = [501, 511, 521, 601, 611, 621, 701, 711]
clev_label = {}
for d in clev_date:
clev_label[d] = '%02d-%02d' % (d//100, d%100)
# Plot maps
manual = False
plt.figure(figsize=(14,10))
cmin, cmax = 20, 55
plt.subplot(211)
_, pc = atm.pcolor_latlon(onset, lat, lon, cmap=cmap, axlims=axlims)
pc.set_clim(cmin, cmax)
if clines:
_, cs = atm.contour_latlon(onset_date, lat_i, lon_i, clev=clev_date,
axlims=axlims)
plt.clabel(cs, clev_date, fmt=clev_label, manual=manual)
plt.title(titlestr + ' Onset Pentad & Day')
else:
plt.title(titlestr + ' Onset Pentad')
plt.subplot(212)
clev = np.arange(0, 1.1, 0.025)
_, pc = atm.pcolor_latlon(Rsq, lat, lon, cmap=cmap, axlims=axlims)
pc.set_clim(0., 1.)
plt.title('$R^2$ for kmax = %d' % kmax)
def animate(i):
plt.clf()
m, _ = atm.pcolor_latlon(animdata[i], axlims=axlims, cmap=cmap)
plt.clim(cmin, cmax)
day = animdata[daynm + 'rel'].values[i]
plt.title('%s %s RelDay %d' % (var, yearstr, day))
return m
import matplotlib.pyplot as plt
import xray
import atmos as atm
# ----------------------------------------------------------------------
# Read some data from OpenDAP url
url = ('http://goldsmr3.sci.gsfc.nasa.gov/opendap/MERRA_MONTHLY/'
'MAIMCPASM.5.2.0/1979/MERRA100.prod.assim.instM_3d_asm_Cp.197901.hdf')
ds = xray.open_dataset(url)
ps = ds['PS'] / 100
plt.figure()
atm.pcolor_latlon(ps, cmap='jet')
lon1, lon2 = 0, 100
lat1, lat2 = -45, 45
ps_sub = atm.subset(ps, 'lon', lon1, lon2, 'lat', lat1, lat2)
plt.figure()
atm.pcolor_latlon(ps_sub, cmap='jet')
import matplotlib.pyplot as plt
import xray
import atmos as atm
# ----------------------------------------------------------------------
# Read some data from OpenDAP url
url = ('http://goldsmr3.sci.gsfc.nasa.gov/opendap/MERRA_MONTHLY/'
'MAIMCPASM.5.2.0/1979/MERRA100.prod.assim.instM_3d_asm_Cp.197901.hdf')
ds = xray.open_dataset(url)
ps = ds['PS'] / 100
plt.figure()
atm.pcolor_latlon(ps, cmap='jet')
lon1, lon2 = 0, 100
lat1, lat2 = -45, 45
ps_sub = atm.subset(ps, 'lon', lon1, lon2, 'lat', lat1, lat2)
plt.figure()
atm.pcolor_latlon(ps_sub, cmap='jet')
import matplotlib.pyplot as plt import merra import atmos as atm from merra import load_daily_season datadir = '/home/jennifer/datastore/merra/daily/' #year = 1995 year = 1979 ustr = datadir + 'merra_u200_' vstr = datadir + 'merra_v200_' season = 'jan' u = load_daily_season(ustr, year, season, 'U') plt.figure() atm.pcolor_latlon(u.mean(dim='TIME')) season = 'jja' lon1, lon2 = 0, 100 lat1, lat2 = -20, 50 u = load_daily_season(ustr, year, season, 'U', lat1, lat2, lon1, lon2) plt.figure() atm.pcolor_latlon(u.mean(dim='TIME')) ds = load_daily_season(ustr, year, season, None, lat1, lat2, lon1, lon2) u2 = ds['U'] print((u == u2).any()) season = 'ann' lon1, lon2 = 20, 120 lat1, lat2 = -60, 60
lat = atm.get_coord(u, 'lat') lon = atm.get_coord(u, 'lon') # Number of time points per day n = 8 # Daily mean u_split = atm.split_timedim(u, n, time0_name='day') u_new = daily_from_subdaily(u, n, dayvals=np.arange(1,32)) print(np.array_equal(u_new, u_split.mean(axis=1))) # ndarray version u_new2 = daily_from_subdaily(u.values, n) print(np.array_equal(u_new, u_new2)) # Sub-sample version i = 2 u_new3 = daily_from_subdaily(u, n, method=i) print(np.array_equal(u_split[:,i], u_new3)) # Plot data to check d = 5 plt.figure() plt.subplot(211) atm.pcolor_latlon(u_new[d], cmap='jet') plt.subplot(212) atm.pcolor_latlon(u_new.mean(axis=0), cmap='jet')
filename = datadir + 'merra/daily/merra_u200_198601.nc' ds = atm.ncload(filename) u = ds['U'] lat = atm.get_coord(u, 'lat') lon = atm.get_coord(u, 'lon') # Number of time points per day n = 8 # Daily mean u_split = atm.split_timedim(u, n, time0_name='day') u_new = daily_from_subdaily(u, n, dayvals=np.arange(1, 32)) print(np.array_equal(u_new, u_split.mean(axis=1))) # ndarray version u_new2 = daily_from_subdaily(u.values, n) print(np.array_equal(u_new, u_new2)) # Sub-sample version i = 2 u_new3 = daily_from_subdaily(u, n, method=i) print(np.array_equal(u_split[:, i], u_new3)) # Plot data to check d = 5 plt.figure() plt.subplot(211) atm.pcolor_latlon(u_new[d], cmap='jet') plt.subplot(212) atm.pcolor_latlon(u_new.mean(axis=0), cmap='jet')
tt = onset_TT(T, north=north, south=south)
# Some weirdness going on in 1991, for now just set to NaN
# Troubleshoot later
for nm in ['ttn', 'tts', 'tseries']:
vals = tt[nm].values
vals = np.ma.masked_array(vals, abs(vals) > 1e30).filled(np.nan)
tt[nm].values = vals
# Plot TTN and TTS
plot_tseries_together(tt[['ttn', 'tts']], tt['onset'].values,
suptitle=suptitle, standardize=False)
# Summary plot and timeseries in individual years
summarize_indices(years, tt['onset'])
plt.suptitle(suptitle)
plot_index_years(tt, suptitle=suptitle, yearnm='year', daynm='day')
if save:
atm.savefigs(varname + '_', 'png')
# Plot contour map of pressure-level data
p_plot = 400
T_plot = T_p[p_plot]
y, d = 0, 80
lat = atm.get_coord(T_plot, 'lat')
lon = atm.get_coord(T_plot, 'lon')
axlims = (lat.min(), lat.max(), lon.min(), lon.max())
plt.figure()
atm.pcolor_latlon(T_plot[y, d], axlims=axlims)
# Days for pre/onset/post composites
compdays = comp_days_centered(5)
#compdays = comp_days_centered(1, 9)
# Calculate composites
comp = composite(data, compdays, daynm=daynm + 'rel')
compbar = collections.OrderedDict()
for key in comp:
compbar[key] = comp[key].mean(dim=yearnm)
# Plot pre/onset/post composites - climatology
cmin, cmax = 0, 20
plt.figure(figsize=(12,10))
for i, key in enumerate(comp):
plt.subplot(2, 2, i+1)
atm.pcolor_latlon(compbar[key], axlims=axlims, cmap=cmap)
plt.clim(cmin, cmax)
plt.title(var + ' ' + key)
# Plot pre/onset/post composites - individual years
cmin, cmax = 0, 30
nrow, ncol = 3, 4
figsize = (14, 10)
for key in comp.keys():
suptitle = var + ' ' + key
for y, year in enumerate(years):
if y % (nrow * ncol) == 0:
fig, axes = plt.subplots(nrow, ncol, figsize=figsize, sharex=True)
plt.subplots_adjust(left=0.08, right=0.95, wspace=0, hspace=0.2)
plt.suptitle(suptitle)
yplot = 1
unew = split_timedim(u, 12, time0_name='year', time0_vals=years,
time1_name='month', time1_vals=months)
unew2 = split_timedim(u, 12, slowfast=False)
# Check that reshaping worked properly
m, y, k = 11, 20, 10
data1 = u[y*12 + m, k]
data2 = unew[y, m, k]
data3 = unew2[m, y, k]
print(np.array_equal(data1,data2))
print(np.array_equal(data2, data3))
# Plot data to check
plt.figure()
plt.subplot(211)
atm.pcolor_latlon(data1, cmap='jet')
plt.subplot(212)
atm.pcolor_latlon(data2, cmap='jet')
# ----------------------------------------------------------------------
# MERRA Daily
filename = datadir + 'merra/daily/merra_u200_198601.nc'
ds = atm.ncload(filename)
u = ds['U']
lat = atm.get_coord(u, 'lat')
lon = atm.get_coord(u, 'lon')
# Number of time points per day
n = 8
days = np.arange(1, 32)
hrs = np.arange(0, 24, 3)
plt.title(pre_days + ' VIMT Climatology')
plt.subplot(223)
m = atm.init_latlon(lat1, lat2, lon1, lon2)
m.quiver(x, y, ds['uq_bar_post'], ds['vq_bar_post'])
plt.title(post_days + ' VIMT Climatology')
# Plot difference between pre- and post- composites
plt.subplot(222)
m = atm.init_latlon(lat1, lat2, lon1, lon2)
#m, _ = atm.pcolor_latlon(ds['vimt_bar_diff'], axlims=axlims, cmap='hot_r')
m.quiver(x, y, ds['uq_bar_diff'], ds['vq_bar_diff'])
plt.title(post_days + ' minus ' + pre_days + ' VIMT Climatology')
# Top N difference vectors
plt.subplot(224)
m, _ = atm.pcolor_latlon(ds['vimt_bar_diff_masked'],axlims=axlims, cmap='hot_r')
plt.title('Magnitude of Top %d Difference Fluxes' % npts)
# Plot vector VIMT fluxes for a few individual years
ylist = [0, 1, 2, 3]
plt.figure(figsize=(12, 10))
for yr in ylist:
plt.subplot(2, 2, yr + 1)
m = atm.init_latlon(lat1, lat2, lon1, lon2)
m.quiver(x, y, ds['uq'][yr].mean(dim='day'), ds['vq'][yr].mean(dim='day'))
m.contour(x, y, ds['mask'].astype(float), [0.99], colors='red')
plt.title('%d May-Sep VIMT Fluxes' % ds.year[yr])
# ----------------------------------------------------------------------
# TIMESERIES
# ----------------------------------------------------------------------
varstr = varnm.upper()
dat = {key : atm.subset(comp[varnm][key], subset_dict)
for key in keys}
if anom_plot:
cmax = max([abs(dat[key]).max().values for key in keys])
cmin = -cmax
else:
cmin, cmax = climits[varnm][0], climits[varnm][1]
# Lat-lon maps of composites
for j, key in enumerate(keys):
grp.next()
if comp_attrs[key]['axis'] == 1:
cmap = get_colormap(varnm, anom_plot)
else:
cmap = 'RdBu_r'
atm.pcolor_latlon(dat[key], axlims=axlims, cmap=cmap, fancy=False)
plt.xticks(range(40, 121, 20))
if comp_attrs[key]['axis'] == 1:
plt.clim(cmin, cmax)
else:
symmetric = atm.symm_colors(dat[key])
if symmetric:
cmax = np.nanmax(abs(dat[key]))
plt.clim(-cmax, cmax)
plt.title(varstr + ' ' + key.upper(), fontsize=11)
if grp.col > 0:
plt.gca().set_yticklabels([])
if grp.row < nrow - 1:
plt.gca().set_xticklabels([])
# Line plots of sector averages
grp.next()
# ----------------------------------------------------------------------
# Plots
cmap = 'hot_r'
axlims = (-30, 30, 40, 120)
clim1, clim2 = 0, 80
figsize = (12, 9)
months = [4, 5, 6]
plt.figure(figsize=figsize)
plt.suptitle(suptitle)
# Lat-lon maps
for i, month in enumerate(months):
plt.subplot(2, 2, i + 1)
atm.pcolor_latlon(mld[month-1], axlims=axlims, cmap=cmap)
plt.clim(clim1, clim2)
plt.title(atm.month_str(month))
# Sector mean line plots
latmin, latmax = -30, 30
mldplot = atm.subset(mldbar, {'lat' : (latmin, latmax)})
lat = atm.get_coord(mldplot, 'lat')
plt.subplot(2, 2, 4)
for month in months:
plt.plot(lat, mldplot[month-1], label=atm.month_str(month))
plt.legend()
plt.grid()
plt.title('%d-%d E Mean' % (lon1, lon2))
plt.xlabel('Latitude')
plt.ylabel('MLD (m)')
import atmos as atm
datadir = "/home/jennifer/datastore/merra/daily/"
filestr = "merra_uv200_40E-120E_60S-60N_"
year = 1980
lon1, lon2 = 60, 100
lat1, lat2 = 10, 30
filename = "%s%s%d.nc" % (datadir, filestr, year)
ds = atm.ncload(filename)
iplot = {"U": 1, "V": 2, "rel_vort": 3, "Ro": 4}
plt.figure(figsize=(12, 9))
for var in ds.data_vars:
plt.subplot(2, 2, iplot[var])
atm.pcolor_latlon(ds[var].mean(axis=0))
plt.title(var)
dsbar = xray.Dataset()
for var in ds.data_vars:
print(var)
dsbar[var] = atm.mean_over_geobox(ds[var], lat1, lat2, lon1, lon2)
plt.figure(figsize=(12, 9))
for var in dsbar.data_vars:
plt.subplot(2, 2, iplot[var])
plt.plot(dsbar.Day, dsbar[var])
plt.title(var)
