"""
Library of utilities that can be used for plotting
"""
import numpy as np
from matplotlib import cm
import matplotlib.colors as mcolors
from matplotlib import pyplot as plt
import cubedsphere.const as c
from .utils import _flatten_ds
[docs]def overplot_wind(ds_reg, U, V, stepsize=1, **kwargs):
"""
Quick and dirty function for overplotting wind of a regridded dataset
Parameters
----------
ds_reg: xarray DataSet
regridded dataset
stepsize: integer
specify the stepsize for which wind arrows should be plotted
"""
ax = kwargs.pop("ax", plt.gca())
if 'lat' in ds_reg:
y, x = ds_reg["lat"].values, ds_reg["lon"].values
else:
y, x = ds_reg["YC"].values, ds_reg["XC"].values
xmesh, ymesh = np.meshgrid(x, y)
ax.quiver(xmesh[::stepsize, ::stepsize], ymesh[::stepsize, ::stepsize], U[::stepsize, ::stepsize],
V[::stepsize, ::stepsize], **kwargs)
[docs]def plotCS(dr, ds, mask_size=None, **kwargs):
"""
A quick way to plot cubed-sphere data of resolution 6*N*N.
Wrapping plotCS_quick_raw to work with xarray objects
Parameters
----------
dr: xarray DataArray
The dimensions must be (y,x)
ds: xarray DataSet (the parent DataSet of dr)
Must contain XC, YC as coordinate variables.
mask_size: None or int
The overlap size of individual tiles. If None is chosen one might likely experience issues
**kwargs
Other keyword arguments such as cmap will be passed to plotCS_quick_raw() and subsequently to plt.pcolormesh()
Returns
-------
ph: list
List of mappabales
"""
# must convert xarray objects to raw numpy arrays
# otherwise numpy masking functions won't work
# transform = kwargs.pop("transform")
# if transform is not None:
# return _plot_cs_cartopy(dr, ds, transform, **kwargs)
x_dim, y_dim = c.lon, c.lat
if len(ds[c.lon].shape) > 2:
if ds[c.lon].shape[-1] == dr[c.lon].shape[-1]:
x_dim = c.lon
x = _flatten_ds(ds[x_dim]).values
data = _flatten_ds(dr).values
else:
raise IndexError("your dataset is incompatible")
if ds[c.lat].shape[-2] == dr[c.lat].shape[-2]:
y_dim = c.lat
y = _flatten_ds(ds[y_dim]).values
else:
raise IndexError("your dataset is incompatible")
else:
x = ds[x_dim].values
y = ds[y_dim].values
data = dr.values
if mask_size is not None:
try:
mask = np.abs(x - 180) < mask_size
data = np.ma.masked_where(mask, data)
except IndexError:
print("caution: No masking possible!")
try:
return _plot_cs_raw(x, y, data, **kwargs)
except IndexError:
return _plot_cs_raw(x.T, y.T, data.T, **kwargs)
def _plot_cs_raw(x, y, data, projection=None, vmin=None, vmax=None, **kwargs):
"""
Plots 2D scalar fields on the MITgcm cubed sphere grid with pcolormesh.
Adapted from MITgcmutils (https://github.com/MITgcm/MITgcm/tree/master/utils/python/MITgcmutils/MITgcmutils/cs)
Parameters
----------
x: array_like
'xg', that is, x coordinate of the points one half grid cell to the
left and bottom, that is vorticity points for tracers, etc.
y: array_like
'yg', that is, y coordinate of same points
data: array_like
scalar field at tracer points
projection: Basemap instance, optional
used to transform if present.
Unfortunatly, cylindrical and conic maps are limited to
the [-180 180] range.
projection = 'sphere' results in a 3D visualization on the sphere
without any specific projection. Good for debugging.
Returns
-------
ph: list
List of mappabales
"""
# pcol first divides the 2D cs-field(6*n,n) into six faces. Then for
# each face, an extra row and colum is added from the neighboring faces in
# order to fool pcolor into drawing the entire field and not just
# (n-1,m-1) data points. There are two corner points that have no explicit
# coordinates so that they have to be found by
# interpolation/averaging. Then each face is divided into 4 tiles,
# assuming cs-geometry, and each tile is plotted individually in
# order to avoid problems due to ambigous longitude values (the jump
# between -180 and 180, or 360 and 0 degrees). As long as the poles
# are at the centers of the north and south faces and the first tile is
# symmetric about its center this should work.
# get the figure handle
fig = plt.gcf()
ax = kwargs.pop("ax", plt.gca())
mapit = 0
if projection != None:
mp = projection
if mp == 'sphere':
mapit = -1
else:
mapit = 1
# convert to [-180 180[ representation
x = np.where(x > 180, x - 360., x)
ny, nx = data.shape
# determine range for color range
cax = [data.min(), data.max()]
if cax[1] - cax[0] == 0: cax = [cax[0] - 1, cax[1] + 1]
if vmin != None: cax[0] = vmin
if vmax != None: cax[1] = vmax
norm = kwargs.pop('norm', None)
if norm is None:
norm = mcolors.Normalize(cax[0], cax[1])
if mapit == -1:
# set up 3D plot
if len(fig.axes) > 0:
# if present, remove and replace the last axis of fig
geom = fig.axes[-1].get_geometry()
plt.delaxes(fig.axes[-1])
else:
# otherwise use full figure
geom = ((1, 1, 1))
ax = kwargs.pop("ax", fig.add_subplot(geom[0], geom[1], geom[2], projection='3d',
facecolor='None'))
# define color range
tmp = data - data.min()
N = tmp / tmp.max()
# use this colormap
colmap = cm.jet
colmap.set_bad('w', 1.0)
mycolmap = colmap(N) # cm.jet(N)
ph = np.array([])
jc = x.shape[0] // 2
xxf = np.empty((jc + 1, jc + 1, 4))
yyf = xxf
ffld = np.empty((jc, jc, 4))
xff = []
yff = []
fldf = []
for k in range(0, 6):
ix = np.arange(0, ny) + k * ny
xff.append(x[0:ny, ix])
yff.append(y[0:ny, ix])
fldf.append(data[0:ny, ix])
# find the missing corners by interpolation (one in the North Atlantic)
xfodd = (xff[0][-1, 0] + xff[2][-1, 0] + xff[4][-1, 0]) / 3.
yfodd = (yff[0][-1, 0] + yff[2][-1, 0] + yff[4][-1, 0]) / 3.
# and one south of Australia
xfeven = (xff[1][0, -1] + xff[3][0, -1] + xff[5][0, -1]) / 3.
yfeven = (yff[1][0, -1] + yff[3][0, -1] + yff[5][0, -1]) / 3.
# loop over tiles
for k in range(0, 6):
kodd = 2 * (k // 2)
kodd2 = kodd
if kodd == 4: kodd2 = kodd - 6
keven = 2 * (k // 2)
keven2 = keven
if keven == 4: keven2 = keven - 6
fld = fldf[k]
if np.mod(k + 1, 2):
xf = np.vstack([np.column_stack([xff[k], xff[1 + kodd][:, 0]]),
np.flipud(np.append(xff[2 + kodd2][:, 0], xfodd))])
yf = np.vstack([np.column_stack([yff[k], yff[1 + kodd][:, 0]]),
np.flipud(np.append(yff[2 + kodd2][:, 0], yfodd))])
else:
xf = np.column_stack([np.vstack([xff[k], xff[2 + keven2][0, :]]),
np.flipud(np.append(xff[3 + keven2][0, :],
xfeven))])
yf = np.column_stack([np.vstack([yff[k], yff[2 + keven2][0, :]]),
np.flipud(np.append(yff[3 + keven2][0, :],
yfeven))])
if mapit == -1:
ix = np.arange(0, ny) + k * ny
# no projection at all (projection argument is 'sphere'),
# just convert to cartesian coordinates and plot a 3D sphere
deg2rad = np.pi / 180.
xcart, ycart, zcart = _sph2cart(xf * deg2rad, yf * deg2rad)
ax.plot_surface(xcart, ycart, zcart, rstride=1, cstride=1,
facecolors=mycolmap[0:ny, ix],
linewidth=2, shade=False)
ph = np.append(ph, ax)
else:
# divide all faces into 4 because potential problems arise at
# the centers
for kf in range(0, 4):
if kf == 0:
i0, i1, j0, j1 = 0, jc + 1, 0, jc + 1
elif kf == 1:
i0, i1, j0, j1 = 0, jc + 1, jc, 2 * jc + 1
elif kf == 2:
i0, i1, j0, j1 = jc, 2 * jc + 1, 0, jc + 1
elif kf == 3:
i0, i1, j0, j1 = jc, 2 * jc + 1, jc, 2 * jc + 1
xx = xf[i0:i1, j0:j1]
yy = yf[i0:i1, j0:j1]
ff = fld[i0:i1 - 1, j0:j1 - 1]
if np.median(xx) < 0:
xx = np.where(xx >= 180, xx - 360., xx)
else:
xx = np.where(xx <= -180, xx + 360., xx)
# if provided use projection
if mapit == 1: xx, yy = mp(xx, yy)
# now finally plot 4x6 tiles
ph = np.append(ph, ax.pcolormesh(xx, yy, ff,
norm=norm,
**kwargs))
if mapit == -1:
# ax.axis('image')
ax.set_axis_off()
# ax.set_visible=False
# add a reasonable colormap
m = cm.ScalarMappable(cmap=colmap)
m.set_array(data)
plt.colorbar(m)
return ph
def _sph2cart(azim_sph_coord, elev_sph_coord):
r = np.cos(elev_sph_coord)
x = -r * np.sin(azim_sph_coord)
y = r * np.cos(azim_sph_coord)
z = np.sin(elev_sph_coord)
return x, y, z