def make_ts_slice_plot (patches, temp_values, salt_values, loc0, hmin, hmax, zmin, zmax, tmin, tmax, smin, smax, lon0=None, lat0=None, point0=None, point1=None, tcontours=None, scontours=None, temp_grid=None, salt_grid=None, haxis=None, zaxis=None, extend=['neither', 'neither'], diff=False, date_string=None, fig_name=None):
# Set colour map
if diff:
ctype = 'plusminus'
else:
ctype = 'basic'
cmap_t, tmin, tmax = set_colours(temp_values, ctype=ctype, vmin=tmin, vmax=tmax)
cmap_s, smin, smax = set_colours(salt_values, ctype=ctype, vmin=smin, vmax=smax)
# Figure out orientation and format slice location
h_axis, loc_string = get_loc(loc0, lon0=lon0, lat0=lat0, point0=point0, point1=point1)
# Set panels
fig, gs, cax_t, cax_s = set_panels('1x2C2')
# Wrap some things up in lists for easier iteration
values = [temp_values, salt_values]
vmin = [tmin, smin]
vmax = [tmax, smax]
cmap = [cmap_t, cmap_s]
cax = [cax_t, cax_s]
contours = [tcontours, scontours]
data_grid = [temp_grid, salt_grid]
if diff:
title = ['Change in temperature ('+deg_string+'C)', 'Change in salinity (psu)']
else:
title = ['Temperature ('+deg_string+'C)', 'Salinity (psu)']
for i in range(2):
ax = plt.subplot(gs[0,i])
# Plot patches
img = plot_slice_patches(ax, patches, values[i], hmin, hmax, zmin, zmax, vmin[i], vmax[i], cmap=cmap[i])
if contours[i] is not None:
# Overlay contours
if None in [data_grid[i], haxis, zaxis]:
print 'Error (make_ts_slice_plot): need to specify temp_grid/salt_grid, haxis, and zaxis to do tcontours/scontours'
sys.exit()
plt.contour(haxis, zaxis, data_grid[i], levels=contours[i], colors='black', linestyles='solid')
# Nice axes
slice_axes(ax, h_axis=h_axis)
if i == 1:
# Don't need depth labels a second time
ax.set_yticklabels([])
ax.set_ylabel('')
# Add a colourbar and hide every second label so they're not squished
cbar = plt.colorbar(img, cax=cax[i], extend=extend[i], orientation='horizontal')
reduce_cbar_labels(cbar)
# Variable title
plt.title(title[i], fontsize=18)
if date_string is not None:
# Add date to main title
loc_string += ', ' + date_string
# Main title
plt.suptitle(loc_string, fontsize=20)
finished_plot(fig, fig_name=fig_name)
def make_slice_plot (patches, values, loc0, hmin, hmax, zmin, zmax, vmin, vmax, lon0=None, lat0=None, point0=None, point1=None, contours=None, data_grid=None, haxis=None, zaxis=None, ctype='basic', extend='neither', title='', date_string=None, fig_name=None):
# Set colour map
cmap, vmin, vmax = set_colours(values, ctype=ctype, vmin=vmin, vmax=vmax)
# Figure out orientation and format slice location
h_axis, loc_string = get_loc(loc0, lon0=lon0, lat0=lat0, point0=point0, point1=point1)
# Set up the title
if h_axis in ['lat', 'lon']:
title += ' at ' + loc_string
elif h_axis == 'trans':
title += ' from ' + loc_string
# Plot
fig, ax = plt.subplots()
# Add patches
img = plot_slice_patches(ax, patches, values, hmin, hmax, zmin, zmax, vmin, vmax, cmap=cmap)
if contours is not None:
# Overlay contours
if None in [data_grid, haxis, zaxis]:
print 'Error (make_slice_plot): need to specify data_grid, haxis, and zaxis to do contours'
sys.exit()
plt.contour(haxis, zaxis, data_grid, levels=contours, colors='black', linestyles='solid')
# Make nice axis labels
slice_axes(ax, h_axis=h_axis)
# Add a colourbar
plt.colorbar(img, extend=extend)
# Add a title
plt.title(title, fontsize=18)
if date_string is not None:
# Add the date in the bottom right corner
plt.text(.99, .01, date_string, fontsize=14, ha='right', va='bottom', transform=fig.transFigure)
finished_plot(fig, fig_name=fig_name)
def slice_plot (data, grid, gtype='t', lon0=None, lat0=None, hmin=None, hmax=None, zmin=None, zmax=None, vmin=None, vmax=None, ctype='basic', title=None, date_string=None, fig_name=None):
# Choose what the endpoints of the colourbar should do
extend = get_extend(vmin=vmin, vmax=vmax)
# Build the patches and get the bounds
patches, values, loc0, hmin, hmax, zmin, zmax, vmin_tmp, vmax_tmp = slice_patches(data, grid, gtype=gtype, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax)
# Update any colour bounds which aren't already set
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Set colour map
cmap, vmin, vmax = set_colours(values, ctype=ctype, vmin=vmin, vmax=vmax)
# Figure out orientation and format slice location
if lon0 is not None:
h_axis = 'lat'
loc_string = lon_label(loc0, 3)
elif lat0 is not None:
h_axis = 'lon'
loc_string = lat_label(loc0, 3)
# Set up the title
if title is None:
title = ''
title += ' at ' + loc_string
# Plot
fig, ax = plt.subplots()
# Add patches
img = plot_slice_patches(ax, patches, values, hmin, hmax, zmin, zmax, vmin, vmax, cmap=cmap)
# Make nice axis labels
slice_axes(ax, h_axis=h_axis)
# Add a colourbar
plt.colorbar(img, extend=extend)
# Add a title
plt.title(title, fontsize=18)
if date_string is not None:
# Add the date in the bottom right corner
plt.text(.99, .01, date_string, fontsize=14, ha='right', va='bottom', transform=fig.transFigure)
finished_plot(fig, fig_name=fig_name)
def hovmoller_plot(data,
time,
grid,
ax=None,
make_cbar=True,
ctype='basic',
vmin=None,
vmax=None,
zmin=None,
zmax=None,
monthly=True,
contours=None,
title=None,
titlesize=18,
return_fig=False,
fig_name=None,
extend=None,
figsize=(14, 5),
dpi=None):
# Choose what the endpoints of the colourbar should do
if extend is None:
extend = get_extend(vmin=vmin, vmax=vmax)
# If we're zooming, we need to choose the correct colour bounds
if any([zmin, zmax]):
vmin_tmp, vmax_tmp = var_min_max_zt(data, grid, zmin=zmin, zmax=zmax)
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Get colourmap
cmap, vmin, vmax = set_colours(data, ctype=ctype, vmin=vmin, vmax=vmax)
if monthly:
# As in netcdf_time, the time axis will have been corrected so it is
# marked with the beginning of each month. So to get the boundaries of
# each time index, we just need to add one month to the end.
if time[-1].month == 12:
end_time = datetime.datetime(time[-1].year + 1, 1, 1)
else:
end_time = datetime.datetime(time[-1].year, time[-1].month + 1, 1)
time_edges = np.concatenate((time, [end_time]))
else:
# Following MITgcm convention, the time axis will be stamped with the
# first day of the next averaging period. So to get the boundaries of
# each time index, we just need to extrapolate to the beginning,
# assuming regularly spaced time intervals.
dt = time[1] - time[0]
start_time = time[0] - dt
time_edges = np.concatenate(([start_time], time))
# Update for versions of pcolormesh that don't support date axis:
time_flt = [(t - time_edges[0]).total_seconds() for t in time_edges]
# Ticks at each year
xtick_years = np.sort(list(set([t.year for t in time_edges])))
xtick_loc = [
(cftime.real_datetime(year, 1, 1) - time_edges[0]).total_seconds()
for year in xtick_years
]
xtick_labels = [str(year) for year in xtick_years]
time_edges = np.array(time_flt)
# Make the figure and axes, if needed
existing_ax = ax is not None
if not existing_ax:
fig, ax = plt.subplots(figsize=figsize)
# Plot the data
img = ax.pcolormesh(time_edges,
grid.z_edges,
np.transpose(data),
cmap=cmap,
vmin=vmin,
vmax=vmax)
ax.set_xticks(xtick_loc)
ax.set_xticklabels(xtick_labels)
if contours is not None:
# Overlay contours
# Need time at the centres of each index
# Have to do this with a loop unfortunately
time_centres = []
for t in range(time_edges.size - 1):
dt = (time_edges[t + 1] - time_edges[t]) / 2
time_centres.append(time_edges[t] + dt)
plt.contour(time_centres,
grid.z,
np.transpose(data),
levels=contours,
colors='black',
linestyles='solid')
# Set depth limits
if zmin is None:
# Index of last masked cell
k_bottom = np.argwhere(np.invert(data[0, :].mask))[-1][0]
zmin = grid.z_edges[k_bottom + 1]
if zmax is None:
# Index of first unmasked cell
k_top = np.argwhere(np.invert(data[0, :].mask))[0][0]
zmax = grid.z_edges[k_top]
ax.set_ylim([zmin, zmax])
# Make nice axes labels
depth_axis(ax)
if make_cbar:
# Add a colourbar
plt.colorbar(img, extend=extend)
if title is not None:
# Add a title
plt.title(title, fontsize=titlesize)
if return_fig:
return fig, ax
elif existing_ax:
return img
else:
finished_plot(fig, fig_name=fig_name, dpi=dpi)
def ua_plot(option,
data,
x,
y,
connectivity=None,
xGL=None,
yGL=None,
x_bdry=None,
y_bdry=None,
ax=None,
make_cbar=True,
ctype='basic',
vmin=None,
vmax=None,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
zoom_fris=False,
title=None,
titlesize=18,
return_fig=False,
fig_name=None,
extend=None,
figsize=None,
dpi=None,
rasterized=False):
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
if option == 'tri' and connectivity is None:
print 'Error (ua_plot): Need to provide connectivity'
sys.exit()
if figsize is None:
if zoom_fris:
figsize = (8, 6)
else:
figsize = (10, 6)
# Choose what the endpoints of the colourbar should do
if extend is None:
extend = get_extend(vmin=vmin, vmax=vmax)
# If we're zooming, choose the correct colour bounds
zoom = zoom_fris or any([xmin, xmax, ymin, ymax])
if zoom:
vmin_tmp, vmax_tmp = var_min_max(data, [x, y],
pster=True,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
ua=True)
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Get colourmap
cmap, vmin, vmax = set_colours(data, ctype=ctype, vmin=vmin, vmax=vmax)
levels = np.linspace(vmin, vmax, num=26)
# Figure out if we need to mask outside the model bounds
clip = option == 'reg' and x_bdry is not None and y_bdry is not None
if clip:
xy_bdry = np.stack((x_bdry, y_bdry), axis=-1)
if len(x.shape) == 2:
x_2d = x
y_2d = y
else:
x_2d, y_2d = np.meshgrid(x, y)
xy_points = np.stack((x_2d.ravel(), y_2d.ravel()), axis=-1)
bdry_path = matplotlib.path.Path(xy_bdry)
inside = bdry_path.contains_points(xy_points).reshape(data.shape)
data[~inside] = np.ma.masked
bdry = matplotlib.patches.Polygon(xy_bdry,
facecolor='none',
edgecolor='black',
linestyle='dotted',
linewidth=2)
# Make the figure and axes, if needed
existing_ax = ax is not None
if not existing_ax:
fig, ax = plt.subplots(figsize=figsize)
ax.axis('equal')
# Plot the data
if option == 'tri':
img = ax.tricontourf(x,
y,
connectivity,
data,
levels,
cmap=cmap,
vmin=vmin,
vmax=vmax,
extend=extend)
elif option == 'reg':
if clip:
# Draw the outline of the domain
ax.add_patch(bdry)
img = ax.pcolormesh(x,
y,
data,
cmap=cmap,
vmin=vmin,
vmax=vmax,
rasterized=rasterized)
if make_cbar:
# Add a colourbar
if option == 'tri':
plt.colorbar(img)
elif option == 'reg':
plt.colorbar(img, extend=extend)
if xGL is not None and yGL is not None:
ax.plot(xGL, yGL, color='black')
# Set axes limits etc.
latlon_axes(ax,
x,
y,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
pster=True,
ua=True)
if title is not None:
# Add a title
plt.title(title, fontsize=titlesize)
if return_fig:
return fig, ax
elif existing_ax:
return img
else:
finished_plot(fig, fig_name=fig_name, dpi=dpi)
def ua_tri_plot(data,
x,
y,
connectivity,
ax=None,
make_cbar=True,
ctype='basic',
vmin=None,
vmax=None,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
zoom_fris=False,
title=None,
titlesize=18,
return_fig=False,
fig_name=None,
extend=None,
figsize=(8, 6)):
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
# Choose what the endpoints of the colourbar should do
if extend is None:
extend = get_extend(vmin=vmin, vmax=vmax)
# If we're zooming, choose the correct colour bounds
zoom = zoom_fris or any([xmin, xmax, ymin, ymax])
if zoom:
vmin_tmp, vmax_tmp = var_min_max(data, [x, y],
pster=True,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
ua=True)
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Get colourmap
cmap, vmin, vmax = set_colours(data, ctype=ctype, vmin=vmin, vmax=vmax)
levels = np.linspace(vmin, vmax, num=26)
# Make the figure and axes, if needed
existing_ax = ax is not None
if not existing_ax:
fig, ax = plt.subplots(figsize=figsize)
# Plot the data
img = ax.tricontourf(x,
y,
connectivity,
data,
levels,
cmap=cmap,
vmin=vmin,
vmax=vmax,
extend=extend)
if make_cbar:
# Add a colourbar
plt.colorbar(img)
# Set axes limits etc.
latlon_axes(ax,
x,
y,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
pster=True)
if title is not None:
# Add a title
plt.title(title, fontsize=titlesize)
if return_fig:
return fig, ax
elif existing_ax:
return img
else:
finished_plot(fig, fig_name=fig_name)
def latlon_plot(data,
grid,
gtype='t',
include_shelf=True,
ctype='basic',
vmin=None,
vmax=None,
zoom_fris=False,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
date_string=None,
title=None,
return_fig=False,
fig_name=None,
change_points=None):
# Choose what the endpoints of the colourbar should do
extend = get_extend(vmin=vmin, vmax=vmax)
# If we're zooming, we need to choose the correct colour bounds
if zoom_fris or any([xmin, xmax, ymin, ymax]):
vmin_tmp, vmax_tmp = var_min_max(data,
grid,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax,
gtype=gtype)
# Don't override manually set bounds
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Get colourmap
cmap, vmin, vmax = set_colours(data,
ctype=ctype,
vmin=vmin,
vmax=vmax,
change_points=change_points)
# Prepare quadrilateral patches
lon, lat, data_plot = cell_boundaries(data, grid, gtype=gtype)
fig, ax = plt.subplots()
if include_shelf:
# Shade land in grey
shade_land(ax, grid, gtype=gtype)
else:
# Shade land and ice shelves in grey
shade_land_zice(ax, grid, gtype=gtype)
# Plot the data
img = ax.pcolormesh(lon, lat, data_plot, cmap=cmap, vmin=vmin, vmax=vmax)
if include_shelf:
# Contour ice shelf front
contour_iceshelf_front(ax, grid)
# Add a colourbar
plt.colorbar(img, extend=extend)
# Make nice axes
latlon_axes(ax,
lon,
lat,
zoom_fris=zoom_fris,
xmin=xmin,
xmax=xmax,
ymin=ymin,
ymax=ymax)
if date_string is not None:
# Add the date in the bottom right corner
plt.text(.99,
.01,
date_string,
fontsize=14,
ha='right',
va='bottom',
transform=fig.transFigure)
if title is not None:
# Add a title
plt.title(title, fontsize=18)
if return_fig:
return fig, ax
else:
finished_plot(fig, fig_name=fig_name)