def pcolor_latlon(data, lat=None, lon=None, m=None, cmap='RdBu_r',
axlims=None, fancy=True, cb_kwargs={}, **kwargs):
"""Create a pseudo-color plot of geo data.
Parameters
----------
data : ndarray or xray.DataArray
Data to be plotted.
lat, lon : ndarray, optional
Latitude and longitude arrays. Only used if data is an ndarray.
If data is an xray.DataArray then lat = data['lat'] and
lon = data['lon']
m : Basemap object, optional
Basemap to plot on. If omitted, then a map is created with
init_latlon().
cmap : string or colormap object, optional
Colormap to use.
axlims : 4-tuple of ints or floats, optional
Lat-lon limits for map (lat1, lat2, lon1, lon2). If None, then
data range is used.
fancy : bool, optional
If True, init_latlon will label axes with fancy lat-lon labels.
cb_kwargs : dict, optional
Keyword arguments to plt.colorbar().
**kwargs : keyword arguments, optional
Additional keyword arguments to plt.pcolormesh().
Returns
-------
m : Basemap object
pc : plt.pcolormesh object
cb : plt.colorbar object
"""
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
lon = dat.get_coord(data, 'lon')
vals = np.squeeze(data.values)
else:
vals = np.squeeze(data)
# Lat-lon ranges
if axlims is None:
lat1, lat2 = np.floor(lat.min()), np.ceil(lat.max())
lon1, lon2 = np.floor(lon.min()), np.ceil(lon.max())
else:
lat1, lat2, lon1, lon2 = axlims
# Use a masked array so that pcolormesh displays NaNs properly
vals_plot = np.ma.array(vals, mask=np.isnan(vals))
x, y = np.meshgrid(lon, lat)
if m is None:
m = init_latlon(lat1, lat2, lon1, lon2, fancy)
pc = m.pcolormesh(x, y, vals_plot, cmap=cmap, latlon=True, **kwargs)
cb = m.colorbar(**cb_kwargs)
plt.draw()
return m, pc, cb
def stipple_pts(pts_mask, xname, yname, xsample=1, ysample=1, ax=None,
marker='+', color='k', alpha=0.25, markersize=6,
markeredgewidth=1.5, **kwargs):
"""Plot points to stipple a figure.
Parameters
----------
pts_mask: xray.DataArray
Masking array of points to exclude from stippling.
xname, yname : str
Name of x and y dimensions in pts_mask.
xsample, ysample : int
Sub-sampling of x- and y- dimensions.
ax : plt.axes object
Axes to plot on. If None, then use current axes.
Remaining parameters are keyword arguments to plt.plot() to specify
format for plotting the stipple points.
"""
# Get grid points
x = dat.get_coord(pts_mask, xname)
y = dat.get_coord(pts_mask, yname)
xgrid, ygrid = np.meshgrid(x, y)
# Mask out points to be excluded
if not xgrid.shape == pts_mask.shape:
pts_mask = pts_mask.T
xpts = np.ma.masked_array(xgrid, mask=pts_mask)
ypts = np.ma.masked_array(ygrid, mask=pts_mask)
# Sub-sample
xpts = xpts[::ysample, ::xsample]
ypts = ypts[::ysample, ::xsample]
# Plot stippling
if ax is not None:
plt.sca(ax)
plt.plot(xpts, ypts, marker=marker, color=color, alpha=alpha,
markersize=markersize, markeredgewidth=markeredgewidth, **kwargs)
return None
def testing(data, lat1, lat2, lon1, lon2, t, k):
latname = get_coord(data, 'lat', 'name')
lonname = get_coord(data, 'lon', 'name')
data_sub = dat.subset(data, latname, lat1, lat2, lonname, lon1, lon2)
plt.figure()
ap.pcolor_latlon(data_sub[t,k], axlims=(lat1,lat2,lon1,lon2), cmap='jet')
avg0 = data_sub.mean(axis=-1).mean(axis=-1)
avg1 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=False)
avg2 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=True)
avg3 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=True,
land_only=True)
print(avg0[t, k].values)
print(avg1[t, k].values)
print(avg2[t, k].values)
print(avg3[t, k].values)
import atmos.plots as ap
import atmos.data as dat
import atmos.utils as utils
from atmos.data import get_coord, mean_over_geobox
# ----------------------------------------------------------------------
# Read data
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)
T = ds['T']
ps = ds['PS']
q = ds['QV']
lat = get_coord(ps, 'lat')
lon = get_coord(ps, 'lon')
plev = get_coord(T, 'plev')
# ----------------------------------------------------------------------
# Averaging over box - constant array
lon1, lon2 = 20, 80
lat1, lat2 = 65, 85
data = 2.5 * np.ones(ps.shape)
avg1 = mean_over_geobox(data,
lat1,
lat2,
lon1,
from atmos.utils import print_if
from atmos.constants import const as constants
from atmos.data import get_coord
# ----------------------------------------------------------------------
# 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)
T = ds['T']
ps = ds['PS']
q = ds['QV']
lat = get_coord(T, 'lat')
lon = get_coord(T, 'lon')
plev = get_coord(T, 'plev')
pname = get_coord(T, 'plev', 'name')
units_in = dat.pres_units(T[pname].units)
plev = dat.pres_convert(plev, units_in, 'Pa')
p0 = 1e5
print('Calculating potential temperature')
theta = av.potential_temp(T, plev, p0)
print('Calculating equivalent potential temperature')
theta_e = av.equiv_potential_temp(T, plev, q, p0)
# t, k = 0, 6
# t, k = 0, 22
t, k = 0, 14
def streamfunction(v,
lat=None,
pres=None,
pdim=None,
scale=1e-9,
sector_scale=None):
"""Return the Eulerian mass streamfunction.
Parameters
----------
v : ndarray or xray.DataArray
Meridional wind speed in m/s.
lat : ndarray, optional
Array of latitude in degrees. If omitted, lat is extracted
from xray.DataArray input v.
pres : ndarray, optional
Pressures in Pa. Can be a vector for pressure-level data, or
a grid of pressures of the same shape as v. If omitted, pres
is extracted from xray.DataArray input v.
pdim : {-3, -2}, optional
Dimension of v corresponding to vertical levels. Can be either
the second-last or third-last dimension. If omitted, pdim is
extracted from xray.DataArray v.
scale : float, optional
Scale factor for output, e.g. 1e-9 to output streamfunction in
10^9 kg/s.
sector_scale : float, optional
Scaling factor for the streamfunction over a sector lon1-lon2
(i.e. should be set to (lon2 - lon1)/360.0 or None if full longitude
range is going to be included).
Returns
-------
psi : ndarray or xray.DataArray
Eulerian mass stream function in units of 1/scale kg/s.
"""
R = constants.radius_earth.values
g = constants.g.values
if isinstance(v, xray.DataArray):
i_DataArray = True
name, attrs, coords, _ = xr.meta(v)
if lat is None:
lat = dat.get_coord(v, 'lat')
if pres is None:
pres = dat.get_coord(v, 'plev')
pname = dat.get_coord(v, 'plev', 'name')
pres = dat.pres_convert(pres, v[pname].units, 'Pa')
if pdim is None:
pdim = dat.get_coord(v, 'plev', 'dim') - v.ndim
v = v.values.copy()
else:
i_DataArray = False
if lat is None or pres is None or pdim is None:
raise ValueError('Inputs lat, pres and pdim must be provided when '
'v is an ndarray.')
# Standardize the shape of v
pdim_in = pdim
if pdim == -2:
v = np.expand_dims(v, axis=-1)
pdim = -3
elif pdim != -3:
raise ValueError('Invalid pdim %d. Must be -2 or -3.' % pdim)
dims = list(v.shape)
nlevel = dims[pdim]
dims[pdim] += 1
nlat = dims[-2]
# Tile the pressure levels to make a grid
pmid = np.zeros(dims, dtype=float)
pmid[..., :-1, :, :] = dat.biggify(pres, v, tile=True)
# Cosine-weighted meridional velocity
coslat = np.cos(np.radians(lat))
vcos = np.zeros(v.shape, dtype=float)
for j in range(nlat):
vcos[..., j, :] = v[..., j, :] * coslat[j]
# Compute streamfunction, integrating from top of atmosphere down
sfctn = np.zeros(dims, dtype=float)
for k in range(nlevel - 1, -1, -1):
dp = pmid[..., k + 1, :, :] - pmid[..., k, :, :]
sfctn[..., k, :, :] = sfctn[..., k + 1, :, :] + vcos[..., k, :, :] * dp
# Scale the output and remove the added dimension(s)
sfctn *= scale * 2 * np.pi * R / g
psi = sfctn[..., :-1, :, :]
if pdim_in == -2:
psi = psi[..., 0]
if sector_scale is not None:
psi = psi * sector_scale
if i_DataArray:
psi = xray.DataArray(psi,
name='Eulerian mass streamfunction',
coords=coords)
psi.attrs['units'] = '%.1e kg/s' % (1 / scale)
return psi
def divergence_spherical_2d(Fx, Fy, lat=None, lon=None):
"""Return the 2-D spherical divergence.
Parameters
----------
Fx, Fy : ndarrays or xray.DataArrays
Longitude, latitude components of a vector function in
spherical coordinates. Latitude and longitude should be the
second-last and last dimensions, respectively, of Fx and Fy.
Maximum size 5-D.
lat, lon : ndarrays, optional
Longitude and latitude in degrees. If these are omitted, then
Fx and Fy must be xray.DataArrays with latitude and longitude
in degrees within the coordinates.
Returns
-------
d, d1, d2 : ndarrays or xray.DataArrays
d1 = dFx/dx, d2 = dFy/dy, and d = d1 + d2.
Reference
---------
Atmospheric and Oceanic Fluid Dynamics: Fundamentals and
Large-Scale Circulation, by Geoffrey K. Vallis, Cambridge
University Press, 2006 -- Equation 2.30.
"""
nmax = 5
ndim = Fx.ndim
if ndim > nmax:
raise ValueError('Input data has too many dimensions. Max 5-D.')
if isinstance(Fx, xray.DataArray):
i_DataArray = True
name, attrs, coords, _ = xr.meta(Fx)
if lat is None:
lat = get_coord(Fx, 'lat')
if lon is None:
lon = get_coord(Fx, 'lon')
else:
i_DataArray = False
if lat is None or lon is None:
raise ValueError('Lat/lon inputs must be provided when input '
'data is an ndarray.')
R = constants.radius_earth.values
lon_rad = np.radians(lon)
lat_rad = np.radians(lat)
# Add singleton dimensions for looping, if necessary
for i in range(ndim, nmax):
Fx = np.expand_dims(Fx, axis=0)
Fy = np.expand_dims(Fy, axis=0)
dims = Fx.shape
nlon = dims[-1]
nlat = dims[-2]
d1 = np.zeros(dims, dtype=float)
d2 = np.zeros(dims, dtype=float)
for i in range(nlat):
dx = np.gradient(lon_rad)
coslat = np.cos(lat_rad[i])
for k1 in range(dims[0]):
for k2 in range(dims[1]):
for k3 in range(dims[2]):
sub = Fx[k1, k2, k3, i, :]
d1[k1, k2, k3, i, :] = np.gradient(sub, dx) / (R * coslat)
for j in range(nlon):
dy = np.gradient(lat_rad)
coslat = np.cos(lat_rad)
# Set to NaN at poles to keep from blowing up
coslat[abs(lat) > 89] = np.nan
for k1 in range(dims[0]):
for k2 in range(dims[1]):
for k3 in range(dims[2]):
sub = Fy[k1, k2, k3, :, j] * coslat
d2[k1, k2, k3, :, j] = np.gradient(sub, dy) / (R * coslat)
# Collapse any additional dimensions that were added
for i in range(ndim, d1.ndim):
d1, d2 = d1[0], d2[0]
d = d1 + d2
if i_DataArray:
d = xray.DataArray(d, coords=coords)
d1 = xray.DataArray(d1, coords=coords)
d2 = xray.DataArray(d2, coords=coords)
return d, d1, d2
import atmos.plots as ap import atmos.data as dat from atmos.data import get_coord, interp_latlon # ---------------------------------------------------------------------- # NCEP2 # ---------------------------------------------------------------------- filename = 'data/more/ncep2_climatology_monthly.nc' ds = dat.ncload(filename) mon = ds['mon'].values u = ds['u'] v = ds['v'] T = ds['T'] ps = ds['ps'] lat = get_coord(ps, 'lat') lon = get_coord(ps, 'lon') plev = get_coord(u, 'plev') # ---------------------------------------------------------------------- # interp_latlon lat_new = np.arange(-90, 90, 1.) lon_new = np.arange(0, 360, 1.) # DataArray input ps_i = interp_latlon(ps, lat_new, lon_new) # ndarray input ps_i2 = interp_latlon(ps.values, lat_new, lon_new, lat, lon)
def contourf_latpres(data, lat=None, plev=None, clev=None, init=True,
cmap='RdBu_r', symmetric=True, topo=None, topo_clr='black',
p_units='hPa', axlims=(-90, 90, 0, 1000),
lattick_width=None, ptick_width=None):
"""Plot filled contours of data in latitude-pressure plane.
Parameters
----------
data : ndarray or xray.DataArray
Data to be contoured.
lat : ndarray, optional
Latitude (degrees). If data is an xray.DataArray, lat is
extracted from data['lat'].
plev : ndarray, optional
Pressure levels. If data is an xray.DataArray, plev is
extracted from data['plev'].
clev : float or ndarray, optional
Contour levels (ndarray) or spacing interval (float)
init : bool, optional
If True, initialize plot axes and topography with init_latpres().
cmap : string or colormap object, optional
Colormap to use.
symmetric : bool, optional
Set contour levels to be symmetric about zero.
topo : ndarray or xray.DataArray, optional
Topography to shade (average surface pressure in units of plev).
If topo is an ndarray, it must be on the same latitude grid as
data. If topo is an xray.DataArray, its latitude grid is
extracted from topo['lat']. Only used if init is True.
topo_clr : string or mpl_color, optional
Color to fill topographic profile. Only used if init is True.
p_units : string, optional
Units for pressure axis label. Only used if init is True.
axlims : 4-tuple of floats or ints
Axis limits (latmin, latmax, pmin, pmax). Only used if init is True.
lattick_width, ptick_width : int or float, optional
Spacing for latitude and pressure ticks. Only used if init is True.
"""
# Data to be contoured
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
plev = dat.get_coord(data, 'plev')
# Contour levels
if isinstance(clev, float) or isinstance(clev, int):
# Define contour levels from selected interval spacing
clev = clevels(data, clev, symmetric=symmetric)
# Plot contours
y, z = np.meshgrid(lat, plev)
if clev is None:
plt.contourf(y, z, np.squeeze(data), cmap=cmap)
else:
plt.contourf(y, z, np.squeeze(data), clev, cmap=cmap)
plt.colorbar()
# Initialize plot
if init:
# Topography data
if isinstance(topo, np.ndarray) or isinstance(topo, list):
topo_ps, topo_lat = topo, lat
elif isinstance(topo, xray.DataArray):
topo_ps, topo_lat = topo.values, topo['lat']
else:
topo_ps, topo_lat = None, None
# Initialize axes and plot topography
latmin, latmax, pmin, pmax = axlims
init_latpres(latmin, latmax, pmin, pmax, topo_ps=topo_ps,
topo_lat=topo_lat, topo_clr=topo_clr, p_units=p_units,
lattick_width=lattick_width, ptick_width=ptick_width)
plt.draw()
def contour_latpres(data, lat=None, plev=None, clev=None, init=True,
colors='black', topo=None, topo_clr='black', p_units='hPa',
axlims=(-90, 90, 0, 1000), lattick_width=None,
ptick_width=None, omitzero=False, zerolinewidth=2,
contour_kw={}):
"""
Plot contour lines in latitude-pressure plane
Parameters
----------
data : ndarray or xray.DataArray
Data to be contoured.
lat : ndarray, optional
Latitude (degrees). If data is an xray.DataArray, lat is
extracted from data['lat'].
plev : ndarray, optional
Pressure levels. If data is an xray.DataArray, plev is
extracted from data['plev'].
clev : float or ndarray, optional
Contour levels (ndarray) or spacing interval (float)
init : bool, optional
If True, initialize plot axes and topography with init_latpres().
colors: string or mpl_color, optional
Contour line color.
topo : ndarray or xray.DataArray, optional
Topography to shade (average surface pressure in units of plev).
If topo is an ndarray, it must be on the same latitude grid as
data. If topo is an xray.DataArray, its latitude grid is
extracted from topo['lat']. Only used if init is True.
topo_clr : string or mpl_color, optional
Color to fill topographic profile. Only used if init is True.
p_units : string, optional
Units for pressure axis label. Only used if init is True.
axlims : 4-tuple of floats or ints
Axis limits (latmin, latmax, pmin, pmax). Only used if init is True.
lattick_width, ptick_width : int or float, optional
Spacing for latitude and pressure ticks. Only used if init is True.
omitzero : bool, optional
If True, omit zero contour.
zerolinewidth : int or float, optional
Include zero contour with specified line width.
contour_kw : dict, optional
Dict of additional keyword arguments to plt.contour().
Returns
-------
cs : plt.contour object (non-zero contours only)
"""
# Data to be contoured
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
plev = dat.get_coord(data, 'plev')
# Initialize plot
if init:
# Topography data
if isinstance(topo, np.ndarray) or isinstance(topo, list):
topo_ps, topo_lat = topo, lat
elif isinstance(topo, xray.DataArray):
topo_ps, topo_lat = topo.values, topo['lat']
else:
topo_ps, topo_lat = None, None
# Initialize axes and plot topography
latmin, latmax, pmin, pmax = axlims
init_latpres(latmin, latmax, pmin, pmax, topo_ps=topo_ps,
topo_lat=topo_lat, topo_clr=topo_clr, p_units=p_units,
lattick_width=lattick_width, ptick_width=ptick_width)
# Contour levels
if isinstance(clev, float) or isinstance(clev, int):
# Define contour levels from selected interval spacing
clev = clevels(data, clev, omitzero=omitzero)
# Plot contours
y, z = np.meshgrid(lat, plev)
if clev is None:
cs = plt.contour(y, z, np.squeeze(data), colors=colors, **contour_kw)
else:
cs = plt.contour(y, z, np.squeeze(data), clev, colors=colors, **contour_kw)
# Zero contour
if not omitzero and zerolinewidth > 0:
plt.contour(y, z, np.squeeze(data), 0, colors=colors,
linewidths=zerolinewidth, **contour_kw)
plt.draw()
return cs
def pcolor_latpres(data, lat=None, plev=None, init=True, cmap='RdBu_r',
topo=None, topo_clr='black', p_units='hPa',
axlims=(-90, 90, 0, 1000), lattick_width=None,
ptick_width=None):
"""Create pseudo-color plot of data in latitude-pressure plane.
Parameters
----------
data : ndarray or xray.DataArray
Data to be contoured.
lat : ndarray, optional
Latitude (degrees). If data is an xray.DataArray, lat is
extracted from data['lat'].
plev : ndarray, optional
Pressure levels. If data is an xray.DataArray, plev is
extracted from data['plev'].
init : bool, optional
If True, initialize plot axes and topography with init_latpres().
cmap : string or colormap object, optional
Colormap to use.
topo : ndarray or xray.DataArray, optional
Topography to shade (average surface pressure in units of plev).
If topo is an ndarray, it must be on the same latitude grid as
data. If topo is an xray.DataArray, its latitude grid is
extracted from topo['lat']. Only used if init is True.
topo_clr : string or mpl_color, optional
Color to fill topographic profile. Only used if init is True.
p_units : string, optional
Units for pressure axis label. Only used if init is True.
axlims : 4-tuple of floats or ints
Axis limits (latmin, latmax, pmin, pmax). Only used if init is True.
lattick_width, ptick_width : int or float, optional
Spacing for latitude and pressure ticks. Only used if init is True.
Returns
-------
pc : plt.pcolormesh object
"""
# Data to be plotted
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
plev = dat.get_coord(data, 'plev')
vals = np.squeeze(data.values)
else:
vals = np.squeeze(data)
# Use a masked array so that pcolormesh displays NaNs properly
vals_plot = np.ma.array(vals, mask=np.isnan(vals))
# Pseudo-color plot of data
y, z = np.meshgrid(lat, plev)
pc = plt.pcolormesh(y, z, vals_plot, cmap=cmap)
plt.colorbar()
plt.draw()
# Initialize plot
if init:
# Topography data
if isinstance(topo, np.ndarray) or isinstance(topo, list):
topo_ps, topo_lat = topo, lat
elif isinstance(topo, xray.DataArray):
topo_ps, topo_lat = topo.values, topo['lat']
else:
topo_ps, topo_lat = None, None
# Initialize axes and plot topography
latmin, latmax, pmin, pmax = axlims
init_latpres(latmin, latmax, pmin, pmax, topo_ps=topo_ps,
topo_lat=topo_lat, topo_clr=topo_clr, p_units=p_units,
lattick_width=lattick_width, ptick_width=ptick_width)
return pc
def contour_latlon(data, lat=None, lon=None, clev=None, m=None, colors='black',
linewidths=2.0, linestyles=None, axlims=None, fancy=True):
"""Create a contour line plot of geo data.
Parameters
----------
data : ndarray or xray.DataArray
Data to be plotted.
lat, lon : ndarray, optional
Latitude and longitude arrays. Only used if data is an ndarray.
If data is an xray.DataArray then lat = data['lat'] and
lon = data['lon']
clev : scalar or array of ints or floats, optional
Contour spacing (scalar) or contour levels (array).
m : Basemap object, optional
Basemap to plot on. If omitted, then a map is created with
init_latlon().
colors : string or mpl_color, optional
Contour line color(s).
linewidths : int or float, optional
Line width for contour lines
linestyles : {None, 'solid', 'dashed', 'dashdot', 'dotted'}, optional
Line style for contour lines.
axlims : 4-tuple of ints or floats, optional
Lat-lon limits for map (lat1, lat2, lon1, lon2). If None, then
data range is used.
fancy : bool, optional
If True, init_latlon will label axes with fancy lat-lon labels.
Returns
-------
m : Basemap object
cs : plt.contour object
"""
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
lon = dat.get_coord(data, 'lon')
if isinstance(clev, float) or isinstance(clev, int):
# Define contour levels from selected interval spacing
clev = clevels(data, clev)
# Lat-lon ranges
if axlims is None:
lat1, lat2 = np.floor(lat.min()), np.ceil(lat.max())
lon1, lon2 = np.floor(lon.min()), np.ceil(lon.max())
else:
lat1, lat2, lon1, lon2 = axlims
x, y = np.meshgrid(lon, lat)
if m is None:
m = init_latlon(lat1, lat2, lon1, lon2, fancy)
if clev is None:
cs = m.contour(x, y, np.squeeze(data), colors=colors,
linewidths=linewidths, linestyles=linestyles,
latlon=True)
else:
cs = m.contour(x, y, np.squeeze(data), clev, colors=colors,
linewidths=linewidths, linestyles=linestyles,
latlon=True)
plt.draw()
return m, cs
def contourf_latlon(data, lat=None, lon=None, clev=None, m=None, cmap='RdBu_r',
symmetric=True, axlims=None, fancy=True, colorbar=True,
**kwargs):
"""Create a filled contour plot of geo data.
Parameters
----------
data : ndarray or xray.DataArray
Data to be plotted.
lat, lon : ndarray, optional
Latitude and longitude arrays. Only used if data is an ndarray.
If data is an xray.DataArray then lat = data['lat'] and
lon = data['lon']
clev : scalar or array of ints or floats, optional
Contour spacing (scalar) or contour levels (array).
m : Basemap object, optional
Basemap to plot on. If omitted, then a map is created with
init_latlon().
cmap : string or colormap object, optional
Colormap to use.
symmetric : bool, optional
Set contour levels to be symmetric about zero.
axlims : 4-tuple of ints or floats, optional
Lat-lon limits for map (lat1, lat2, lon1, lon2). If None, then
data range is used.
fancy : bool, optional
If True, init_latlon will label axes with fancy lat-lon labels.
colorbar : bool, optional
If True, include a colorbar.
**kwargs : keyword arguments, optional
Additional keyword arguments to plt.contourf().
Returns
-------
m : Basemap object
"""
if isinstance(data, xray.DataArray):
lat = dat.get_coord(data, 'lat')
lon = dat.get_coord(data, 'lon')
if isinstance(clev, float) or isinstance(clev, int):
# Define contour levels from selected interval spacing
clev = clevels(data, clev, symmetric=symmetric)
# Lat-lon ranges
if axlims is None:
lat1, lat2 = np.floor(lat.min()), np.ceil(lat.max())
lon1, lon2 = np.floor(lon.min()), np.ceil(lon.max())
else:
lat1, lat2, lon1, lon2 = axlims
x, y = np.meshgrid(lon, lat)
if m is None:
m = init_latlon(lat1, lat2, lon1, lon2, fancy)
if clev is None:
m.contourf(x, y, np.squeeze(data), cmap=cmap, latlon=True, **kwargs)
else:
m.contourf(x, y, np.squeeze(data), clev, cmap=cmap, latlon=True, **kwargs)
if colorbar:
m.colorbar()
plt.draw()
return m
from atmos.utils import print_if
from atmos.constants import const as constants
from atmos.data import get_coord
# ----------------------------------------------------------------------
# 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)
T = ds['T']
ps = ds['PS']
q = ds['QV']
lat = get_coord(T, 'lat')
lon = get_coord(T, 'lon')
plev = get_coord(T, 'plev')
pname = get_coord(T, 'plev', 'name')
units_in = dat.pres_units(T[pname].units)
plev = dat.pres_convert(plev, units_in, 'Pa')
p0 = 1e5
print('Calculating potential temperature')
theta = av.potential_temp(T, plev, p0)
print('Calculating equivalent potential temperature')
theta_e = av.equiv_potential_temp(T, plev, q, p0)
# t, k = 0, 6
# t, k = 0, 22
t, k = 0, 14
import atmos.utils as utils import atmos.plots as ap import atmos.data as dat from atmos.data import get_coord, interp_plevels # ---------------------------------------------------------------------- # Read monthly mean climatologies and do some test calcs # ---------------------------------------------------------------------- filename = 'data/more/ncep2_climatology_monthly.nc' ds = dat.ncload(filename) u = ds['u'] v = ds['v'] T = ds['T'] ps = ds['ps'] lat = get_coord(u, 'lat') lon = get_coord(u, 'lon') plev = get_coord(u, 'plev') mon = ds['mon'].values topo = dat.get_ps_clim(lat, lon) / 100 topo.units = 'hPa' # ---------------------------------------------------------------------- # Correct for topography u_orig = u u = dat.correct_for_topography(u_orig, topo) # ---------------------------------------------------------------------- # Interpolate onto new pressure grid
import atmos.plots as ap
import atmos.data as dat
from atmos.data import get_coord, latlon_equal, lon_convention, set_lon
# ----------------------------------------------------------------------
# NCEP2
# ----------------------------------------------------------------------
filename = 'data/more/ncep2_climatology_monthly.nc'
ds = dat.ncload(filename)
mon = ds['mon'].values
u = ds['u']
v = ds['v']
T = ds['T']
ps = ds['ps']
lat = get_coord(ps, 'lat')
lon = get_coord(ps, 'lon')
plev = get_coord(u, 'plev')
topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'
# ----------------------------------------------------------------------
# MERRA
# ----------------------------------------------------------------------
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_m = xray.open_dataset(url)
T_m = ds_m['T']
import atmos.plots as ap
import atmos.data as dat
from atmos.data import get_coord, latlon_equal, lon_convention, set_lon
# ----------------------------------------------------------------------
# NCEP2
# ----------------------------------------------------------------------
filename = 'data/more/ncep2_climatology_monthly.nc'
ds = dat.ncload(filename)
mon = ds['mon'].values
u = ds['u']
v = ds['v']
T = ds['T']
ps = ds['ps']
lat = get_coord(ps, 'lat')
lon = get_coord(ps, 'lon')
plev = get_coord(u, 'plev')
topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'
# ----------------------------------------------------------------------
# MERRA
# ----------------------------------------------------------------------
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_m = xray.open_dataset(url)
T_m = ds_m['T']
from atmos.data import get_coord
from atmos.variables import vorticity, rossby_num
# ----------------------------------------------------------------------
# Read data
url = ('http://goldsmr3.sci.gsfc.nasa.gov/opendap/MERRA_MONTHLY/'
'MAIMCPASM.5.2.0/1979/MERRA100.prod.assim.instM_3d_asm_Cp.197907.hdf')
ds = xray.open_dataset(url)
#T = ds['T']
#ps = ds['PS']
u = ds['U']
v = ds['V']
#q = ds['QV']
lat = get_coord(v, 'lat')
lon = get_coord(v, 'lon')
plev = get_coord(v, 'plev')
topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'
# ----------------------------------------------------------------------
# Relative and absolute vorticity
# DataArray
rel_vort, abs_vort, f = vorticity(u, v)
# ndarray
rel_vort2, abs_vort2, f2 = vorticity(u.values, v.values, lat, lon)
# Multiple linear regression
nobs = 100
x = np.random.random((nobs, 2))
x = sm.add_constant(x)
beta = [1, .1, .5]
e = np.random.random(nobs)
y = np.dot(x, beta) + e
results = sm.OLS(y, x).fit()
print results.summary()
# ----------------------------------------------------------------------
# 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)
T = ds['T']
ps = ds['PS']
q = ds['QV']
plev = get_coord(T, 'plev')
lat = get_coord(T, 'lat')
lon = get_coord(T, 'lon')
# t, k = 0, 6
# t, k = 0, 22
t, k = 0, 14
pstr = '%d hPa' % (plev[k]/100)
# ----------------------------------------------------------------------
