def profile(ifile, varkey, options, before='', after=''):
import matplotlib.pyplot as plt
print(varkey, end='')
outpath = getattr(options, 'outpath', '.')
try:
vert = cu.getpresmid(ifile)
vunit = 'Pa'
except Exception:
vert = cu.getsigmamid(ifile)
vunit = r'\sigma'
var = ifile.variables[varkey]
dims = list(var.dimensions)
for knownvert in ['layer', 'LAY'] + ['layer%d' % i for i in range(72)]:
if knownvert in dims:
vidx = dims.index(knownvert)
break
else:
raise KeyError("No known vertical coordinate; got %s" % str(dims))
vert = vert[:var[:].shape[vidx]]
units = var.units.strip()
vals = np.rollaxis(var[:], vidx, start=0).view(
np.ma.MaskedArray).reshape(vert.size, -1)
ax = plt.gca()
minmaxmean(ax, vals, vert)
ax.set_xlabel(varkey + ' (' + units + ')')
ax.set_ylabel(vunit)
ax.set_ylim(vert.max(), vert.min())
if options.logscale:
ax.set_xscale('log')
fmt = 'png'
figpath = os.path.join(outpath + '_profile_' + varkey + '.' + fmt)
exec(after)
plt.savefig(figpath)
print('Saved fig', figpath)
return figpath
def plot(ifiles, args):
from PseudoNetCDF.coordutil import getsigmamid, getpresmid, gettimes
import pylab as pl
from pylab import figure, NullFormatter, close, rcParams
rcParams['text.usetex'] = False
from matplotlib.colors import LinearSegmentedColormap, BoundaryNorm, LogNorm
scale = args.scale;
minmax = eval(args.minmax)
minmaxq = eval(args.minmaxq)
sigma = args.sigma
maskzeros = args.maskzeros
outunit = args.outunit
tespaths = args.tespaths
omipaths = args.omipaths
edges = args.edges
try:
f, = ifiles
except:
raise ValueError('curtain plot expects one file when done. Try stack time --stack=time to concatenate')
# Add CF conventions if necessary
if 'latitude_bounds' not in f.variables.keys():
try:
from PseudoNetCDF import getvarpnc
from PseudoNetCDF.conventions.ioapi import add_cf_from_ioapi
f = getvarpnc(f, None)
add_cf_from_ioapi(f)
except:
pass
if sigma:
vertcrd = getsigmamid(f)
else:
vertcrd = getpresmid(f, pref = 101325., ptop = getattr(f, 'VGTOP', 10000))
if vertcrd.max() > 2000: vertcrd /= 100.
try:
lonb = f.variables['geos_longitude_bounds']
latb = f.variables['geos_latitude_bounds']
except:
lonb = f.variables['longitude_bounds']
latb = f.variables['latitude_bounds']
for var_name in args.variables:
temp = defaultdict(lambda: 1)
try:
eval(var_name, None, temp)
var = eval(var_name, None, f.variables)[:]
except:
temp[var_name]
var = f.variables[var_name][:]
if maskzeros: var = np.ma.masked_values(var, 0)
unit = f.variables[temp.keys()[0]].units.strip()
if unit in unitconvert:
var = unitconvert.get((unit, outunit), lambda x: x)(var)
else:
outunit = unit
bmap = None
vmin, vmax = np.percentile(np.ma.compressed(var).ravel(), list(minmaxq))
if minmax[0] is not None:
vmin = minmax[0]
if minmax[1] is not None:
vmax = minmax[1]
if edges:
fig = pl.figure(figsize = (16, 4))
offset = 0.05
ax = fig.add_axes([.1 - offset, .15, .22, .725])
ax = fig.add_axes([.325 - offset, .15, .22, .725])
ax = fig.add_axes([.55 - offset, .15, .22, .725])
ax = fig.add_axes([.775 - offset, .15, .22, .725])
ss = 0
se = ss + f.NCOLS + 1
es = se
ee = se + f.NROWS + 1
ns = ee
ne = ee + f.NCOLS + 1
ws = ne
we = ws + f.NROWS + 1
axs = fig.axes
for ax in fig.axes[1:]:
ax.yaxis.set_major_formatter(pl.NullFormatter())
vars = [var[:, :, ss:se], var[:, :, es:ee], var[:, :, ns:ne][:, :, ::-1], var[:, :, ws:we][:, :, ::-1]]
lonbss = [lonb[ss:se], lonb[es:ee], lonb[ns:ne][::-1], lonb[ws:we][::-1]]
latbss = [latb[ss:se], latb[es:ee], latb[ns:ne][::-1], latb[ws:we][::-1]]
else:
fig = pl.figure(figsize = (8, 4))
ax = fig.add_axes([.1, .15, .8, .725])
axs = fig.axes
vars = [var]
lonbss = [lonb[:]]
latbss = [latb[:]]
for ax, var, lonbs, latbs in zip(axs, vars, lonbss, latbss):
vals = var.swapaxes(0, 1).reshape(var.shape[1], -1)
ax.text(.05, .9, 'n = %d' % vals.shape[1], transform = ax.transAxes)
modl, modr = minmaxmean(ax, vals, vertcrd, facecolor = 'k', edgecolor = 'k', alpha = .2, zorder = 4, label = 'GC', ls = '-', lw = 2, color = 'k')
llines = [(modl, modr)]
ymin, ymax = vertcrd.min(), vertcrd.max()
ax.set_ylim(ymax, ymin)
ax.set_xscale(scale)
ax.set_xlim(vmin, vmax)
#if scale == 'log':
# ax.set_xticklabels(['%.1f' % (10**x) for x in ax.get_xticks()])
if 'TFLAG' in f.variables.keys():
SDATE = f.variables['TFLAG'][:][0, 0, 0]
EDATE = f.variables['TFLAG'][:][-1, 0, 0]
STIME = f.variables['TFLAG'][:][0, 0, 1]
ETIME = f.variables['TFLAG'][:][-1, 0, 1]
if SDATE == 0:
SDATE = 1900001
EDATE = 1900001
sdate = datetime.strptime('%07d %06d' % (SDATE, STIME), '%Y%j %H%M%S')
edate = datetime.strptime('%07d %06d' % (EDATE, ETIME), '%Y%j %H%M%S')
elif 'tau0' in f.variables.keys():
sdate = datetime(1985, 1, 1, 0) + timedelta(hours = f.variables['tau0'][0])
edate = datetime(1985, 1, 1, 0) + timedelta(hours = f.variables['tau1'][-1])
else:
times = gettimes(f)
sdate = times[0]
edate = times[-1]
if len(tespaths) > 0:
tesl, tesr = plot_tes(ax, lonbs, latbs, tespaths)
if not tesl is None:
llines.append((tesl, tesr))
if len(omipaths) > 0:
omil, omir = plot_omi(ax, lonbs, latbs, omipaths, airden = f.variables['AIRDEN'][:].mean(0).mean(1), airdenvert = vertcrd)
if not omil is None:
llines.append((omil, omir))
try:
title = '%s to %s' % (sdate.strftime('%Y-%m-%d'), edate.strftime('%Y-%m-%d'))
except:
title = var_name
if sigma:
axs[0].set_ylabel('sigma')
else:
axs[0].set_ylabel('pressure')
xmax = -np.inf
xmin = np.inf
for ax in fig.axes:
tmp_xmin, tmp_xmax = ax.get_xlim()
xmax = max(tmp_xmax, xmax)
xmin = min(tmp_xmin, xmin)
for ax in fig.axes:
ax.set_xlim(xmin, xmax)
if len(axs) == 1:
axs[0].set_xlabel('%s %s' % (var_name, outunit))
else:
axs[0].set_xlabel('South')
axs[1].set_xlabel('East')
axs[2].set_xlabel('North')
axs[3].set_xlabel('West')
fig.text(.5, .90, '%s %s' % (var_name, outunit), horizontalalignment = 'center', fontsize = 16)
nl = 0
for ax in axs:
if len(ax.get_lines()) > nl:
nl = len(ax.get_lines())
pl.sca(ax)
llabels = [l[0].get_label() for l in llines]
pl.legend(llines, llabels, bbox_to_anchor = (.1, 1), loc = 'upper left', bbox_transform = fig.transFigure, ncol = 6)
if edges:
fig.text(0.95, 0.975, title, horizontalalignment = 'right', verticalalignment = "top", fontsize = 16)
else:
fig.text(0.95, 0.025, title, horizontalalignment = 'right', verticalalignment = "bottom", fontsize = 16)
fig.savefig('%s_%s.%s' % (args.outpath, var_name, args.figformat))
pl.close(fig)
return fig