def ts_slice_plot (temp, salt, grid, lon0=None, lat0=None, point0=None, point1=None, hmin=None, hmax=None, zmin=None, zmax=None, tmin=None, tmax=None, smin=None, smax=None, tcontours=None, scontours=None, date_string=None, fig_name=None):
# Choose what the endpoints of the colourbars should do
extend = [get_extend(vmin=tmin, vmax=tmax), get_extend(vmin=smin, vmax=smax)]
# Build the temperature patches and get the bounds
if lon0 is not None or lat0 is not None:
# Lat-lon slices
# Get gridded data and axes just in case
patches, temp_values, loc0, hmin, hmax, zmin, zmax, tmin_tmp, tmax_tmp, left, right, below, above, temp_grid, haxis, zaxis = slice_patches(temp, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
# Get the salinity values on the same patches, and their colour bounds
salt_values, smin_tmp, smax_tmp, salt_grid = slice_values(salt, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0, return_gridded=True)
elif point0 is not None and point1 is not None:
# Transect
loc0 = None
patches, temp_values, hmin, hmax, zmin, zmax, tmin_tmp, tmax_tmp, left, right, below, above, temp_grid, haxis, zaxis = transect_patches(temp, grid, point0, point1, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
salt_values, smin_tmp, smax_tmp, salt_grid = transect_values(salt, grid, point0, point1, left, right, below, above, hmin, hmax, zmin, zmax, return_gridded=True)
else:
print 'Error (ts_slice_plot): must specify either lon0, lat0, or point0 and point1'
sys.exit()
# Update any colour bounds which aren't already set
if tmin is None:
tmin = tmin_tmp
if tmax is None:
tmax = tmax_tmp
if smin is None:
smin = smin_tmp
if smax is None:
smax = smax_tmp
# Make the plot
make_ts_slice_plot(patches, temp_values, salt_values, loc0, hmin, hmax, zmin, zmax, tmin, tmax, smin, smax, lon0=lon0, lat0=lat0, point0=point0, point1=point1, tcontours=tcontours, scontours=scontours, temp_grid=temp_grid, salt_grid=salt_grid, haxis=haxis, zaxis=zaxis, extend=extend, date_string=date_string, fig_name=fig_name)
def ts_slice_plot (temp, salt, grid, lon0=None, lat0=None, hmin=None, hmax=None, zmin=None, zmax=None, tmin=None, tmax=None, smin=None, smax=None, date_string=None, fig_name=None):
# Choose what the endpoints of the colourbars should do
extend = [get_extend(vmin=tmin, vmax=tmax), get_extend(vmin=smin, vmax=smax)]
# Build the temperature patches and get the bounds
patches, temp_values, loc0, hmin, hmax, zmin, zmax, tmin_tmp, tmax_tmp, left, right, below, above = slice_patches(temp, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_bdry=True)
# Get the salinity values on the same patches, and their colour bounds
salt_values, smin_tmp, smax_tmp = slice_values(salt, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0)
# Update any colour bounds which aren't already set
if tmin is None:
tmin = tmin_tmp
if tmax is None:
tmax = tmax_tmp
if smin is None:
smin = smin_tmp
if smax is None:
smax = smax_tmp
# 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 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]
cax = [cax_t, cax_s]
title = [r'Temperature ($^{\circ}$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])
# 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')
for label in cbar.ax.xaxis.get_ticklabels()[1::2]:
label.set_visible(False)
# 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 ts_slice_plot_diff (temp_1, temp_2, salt_1, salt_2, grid, lon0=None, lat0=None, point0=None, point1=None, hmin=None, hmax=None, zmin=None, zmax=None, tmin=None, tmax=None, smin=None, smax=None, tcontours=None, scontours=None, date_string=None, fig_name=None):
# Choose what the endpoints of the colourbars should do
extend = [get_extend(vmin=tmin, vmax=tmax), get_extend(vmin=smin, vmax=smax)]
# Build the patches and temperature values for the first simulation
# vmin and vmax don't matter, so just store them as temporary variables
if lon0 is not None or lat0 is not None:
# Lat-lon slices
# Get gridded data and axes just in case
patches, temp_values_1, loc0, hmin, hmax, zmin, zmax, tmp1, tmp2, left, right, below, above, temp_grid_1, haxis, zaxis = slice_patches(temp_1, grid, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
# Get the temperature values for the second simulation on the same patches
temp_values_2, tmp3, tmp4, temp_grid_2 = slice_values(temp_2, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0, return_gridded=True)
# Get the salinity values for both simulations
salt_values_1, tmp5, tmp6, salt_grid_1 = slice_values(salt_1, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0, return_gridded=True)
salt_values_2, tmp7, tmp8, salt_grid_2 = slice_values(salt_2, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0, return_gridded=True)
elif point0 is not None and point1 is not None:
# Transect
loc0 = None
patches, temp_values_1, hmin, hmax, zmin, zmax, tmp1, tmp2, left, right, below, above, temp_grid_1, haxis, zaxis = transect_patches(temp_1, grid, point0, point1, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
temp_values_2, tmp3, tmp4, temp_grid_2 = transect_values(temp_2, grid, point0, point1, left, right, below, above, hmin, hmax, zmin, zmax, return_gridded=True)
salt_values_1, tmp5, tmp6, salt_grid_1 = transect_values(salt_1, grid, point0, point1, left, right, below, above, hmin, hmax, zmin, zmax, return_gridded=True)
salt_values_2, tmp7, tmp8, salt_grid_2 = transect_values(salt_2, grid, point0, point1, left, right, below, above, hmin, hmax, zmin, zmax, return_gridded=True)
else:
print 'Error (ts_slice_plot_diff): must specify either lon0, lat0, or point0 and point1'
sys.exit()
# Calculate the differences
temp_values_diff = temp_values_2 - temp_values_1
salt_values_diff = salt_values_2 - salt_values_1
temp_grid = temp_grid_2 - temp_grid_1
salt_grid = salt_grid_2 - salt_grid_1
# Now figure out the colour bounds
tmin_tmp, tmax_tmp = get_slice_minmax(np.reshape(temp_values_diff, left.shape), left, right, below, above, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_spatial=False)
smin_tmp, smax_tmp = get_slice_minmax(np.reshape(salt_values_diff, left.shape), left, right, below, above, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_spatial=False)
# Update any colour bounds which aren't already set
if tmin is None:
tmin = tmin_tmp
if tmax is None:
tmax = tmax_tmp
if smin is None:
smin = smin_tmp
if smax is None:
smax = smax_tmp
# Plot
make_ts_slice_plot(patches, temp_values_diff, salt_values_diff, loc0, hmin, hmax, zmin, zmax, tmin, tmax, smin, smax, lon0=lon0, lat0=lat0, point0=point0, point1=point1, tcontours=tcontours, scontours=scontours, temp_grid=temp_grid, salt_grid=salt_grid, haxis=haxis, zaxis=zaxis, extend=extend, diff=True, date_string=date_string, fig_name=fig_name)
def slice_plot (data, grid, gtype='t', lon0=None, lat0=None, point0=None, point1=None, hmin=None, hmax=None, zmin=None, zmax=None, vmin=None, vmax=None, contours=None, ctype='basic', title='', 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
if lon0 is not None or lat0 is not None:
# Lat-lon slices
# Get gridded data and axes just in case
patches, values, loc0, hmin, hmax, zmin, zmax, vmin_tmp, vmax_tmp, data_grid, haxis, zaxis = slice_patches(data, grid, gtype=gtype, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_gridded=True)
elif point0 is not None and point1 is not None:
# Transect
loc0 = None
patches, values, hmin, hmax, zmin, zmax, vmin_tmp, vmax_tmp, data_grid, haxis, zaxis = transect_patches(data, grid, point0, point1, gtype=gtype, zmin=zmin, zmax=zmax, return_gridded=True)
else:
print 'Error (slice_plot): must specify either lon0, lat0, or point0 and point1'
sys.exit()
# Update any colour bounds which aren't already set
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Plot
make_slice_plot(patches, values, loc0, hmin, hmax, zmin, zmax, vmin, vmax, lon0=lon0, lat0=lat0, point0=point0, point1=point1, contours=contours, data_grid=data_grid, haxis=haxis, zaxis=zaxis, ctype=ctype, extend=extend, title=title, date_string=date_string, fig_name=fig_name)
def hovmoller_ts_plot(temp,
salt,
time,
grid,
tmin=None,
tmax=None,
smin=None,
smax=None,
zmin=None,
zmax=None,
monthly=True,
t_contours=None,
s_contours=None,
ctype='basic',
loc_string='',
fig_name=None,
figsize=(12, 7),
dpi=None):
# Set panels
fig, gs, cax_t, cax_s = set_panels('2x1C2')
# Wrap things up in lists for easier iteration
data = [temp, salt]
vmin = [tmin, smin]
vmax = [tmax, smax]
contours = [t_contours, s_contours]
title = ['Temperature (' + deg_string + 'C)', 'Salinity (psu)']
cax = [cax_t, cax_s]
for i in range(2):
ax = plt.subplot(gs[i, 0])
# Make the plot
img = hovmoller_plot(data[i],
time,
grid,
ax=ax,
make_cbar=False,
vmin=vmin[i],
vmax=vmax[i],
zmin=zmin,
zmax=zmax,
monthly=monthly,
contours=contours[i],
ctype=ctype,
title=title[i])
# Add a colourbar
extend = get_extend(vmin=vmin[i], vmax=vmax[i])
plt.colorbar(img, cax=cax[i], extend=extend)
if i == 0:
# Remove x-tick labels from top plot
ax.set_xticklabels([])
plt.suptitle(loc_string, fontsize=22)
finished_plot(fig, fig_name=fig_name, dpi=dpi)
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 slice_plot_diff (data_1, data_2, grid, gtype='t', lon0=None, lat0=None, point0=None, point1=None, hmin=None, hmax=None, zmin=None, zmax=None, vmin=None, vmax=None, contours=None, title='', 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 for the first array
# vmin and vmax don't matter, so just store them as temporary variables
# Then get the values for the second array on the same patches
if lon0 is not None or lat0 is not None:
# Lat-lon slices
# Get gridded data and axes just in case
patches, values_1, loc0, hmin, hmax, zmin, zmax, tmp1, tmp2, left, right, below, above, data_grid_1, haxis, zaxis = slice_patches(data_1, grid, gtype=gtype, lon0=lon0, lat0=lat0, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
values_2, tmp3, tmp4, data_grid_2 = slice_values(data_2, grid, left, right, below, above, hmin, hmax, zmin, zmax, lon0=lon0, lat0=lat0, return_gridded=True)
elif point0 is not None and point1 is not None:
# Transect
loc0 = None
patches, values_1, hmin, hmax, zmin, zmax, tmp1, tmp2, left, right, below, above, data_grid_1, haxis, zaxis = transect_patches(data_1, grid, point0, point1, gtype=gtype, zmin=zmin, zmax=zmax, return_bdry=True, return_gridded=True)
values_2, tmp3, tmp4, data_grid_2 = transect_values(data_2, grid, point0, point1, left, right, below, above, hmin, hmax, zmin, zmax, gtype=gtype, return_gridded=True)
else:
print 'Error (slice_plot_diff): must specify either lon0, lat0, or point0 and point1'
sys.exit()
# Calculate the difference
values_diff = values_2 - values_1
data_grid = data_grid_2 - data_grid_1
# Now figure out the colour bounds
# Note we need to reshape the values array to be 2D again
vmin_tmp, vmax_tmp = get_slice_minmax(np.reshape(values_diff, left.shape), left, right, below, above, hmin=hmin, hmax=hmax, zmin=zmin, zmax=zmax, return_spatial=False)
# Update any colour bounds which aren't already set
if vmin is None:
vmin = vmin_tmp
if vmax is None:
vmax = vmax_tmp
# Plot
make_slice_plot(patches, values_diff, loc0, hmin, hmax, zmin, zmax, vmin, vmax, lon0=lon0, lat0=lat0, point0=point0, point1=point1, contours=contours, data_grid=data_grid, haxis=haxis, zaxis=zaxis, ctype='plusminus', extend=extend, title=title, date_string=date_string, 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 hovmoller_ts_plot(temp,
salt,
time,
grid,
smooth=0,
split_year=None,
tmin=None,
tmax=None,
smin=None,
smax=None,
zmin=None,
zmax=None,
monthly=True,
t_contours=None,
s_contours=None,
title=None,
date_since_start=False,
start=0,
t0=None,
s0=None,
ctype='basic',
loc_string='',
fig_name=None,
figsize=(12, 7),
dpi=None,
return_fig=False,
ab_inside=False,
rasterized=False):
# Set panels
fig, gs, cax_t, cax_s = set_panels('2x1C2', figsize=figsize)
if split_year is not None:
if date_since_start:
first_year = time[0].year - time[start].year
last_year = time[-1].year - time[start].year
else:
first_year = time[0].year
last_year = time[-1].year
width1 = (split_year - first_year)
width2 = (last_year + 1 - split_year)
gs = plt.GridSpec(2, 2, width_ratios=[width1, width2])
gs.update(left=0.08,
right=0.9,
bottom=0.1,
top=0.88,
hspace=0.2,
wspace=0.01)
# Need to choose the correct colour bounds
if any([zmin, zmax]):
tmin_tmp, tmax_tmp = var_min_max_zt(temp,
grid,
zmin=zmin,
zmax=zmax)
smin_tmp, smax_tmp = var_min_max_zt(salt,
grid,
zmin=zmin,
zmax=zmax)
if tmin is None:
tmin = tmin_tmp
if tmax is None:
tmax = tmax_tmp
if smin is None:
smin = smin_tmp
if smax is None:
smax = smax_tmp
# Wrap things up in lists for easier iteration
data = [temp, salt]
vmin = [tmin, smin]
vmax = [tmax, smax]
val0 = [t0, s0]
contours = [t_contours, s_contours]
if ab_inside:
titles = ['Temperature (' + deg_string + 'C)', 'Salinity (psu)']
ab = ['a', 'b']
else:
titles = ['a) Temperature (' + deg_string + 'C)', 'b) Salinity (psu)']
cax = [cax_t, cax_s]
axs = []
for i in range(2):
ax = plt.subplot(gs[i, 0])
# Make the plot
img = hovmoller_plot(data[i],
time,
grid,
smooth=smooth,
ax=ax,
make_cbar=False,
vmin=vmin[i],
vmax=vmax[i],
zmin=zmin,
zmax=zmax,
monthly=monthly,
contours=contours[i],
ctype=ctype,
title=titles[i],
date_since_start=date_since_start,
start=start,
val0=val0[i],
end_t=split_year,
rasterized=rasterized)
if ab_inside:
plt.text(0.01,
0.98,
ab[i],
weight='bold',
ha='left',
va='top',
fontsize=16,
transform=ax.transAxes)
# Add a colourbar
extend = get_extend(vmin=vmin[i], vmax=vmax[i])
cbar = plt.colorbar(img, cax=cax[i], extend=extend)
reduce_cbar_labels(cbar)
if i == 0:
# Remove x-tick labels from top plot
ax.set_xticklabels([])
else:
ax.set_xlabel('Year', fontsize=14)
ax.set_ylabel('')
axs.append(ax)
if split_year is not None:
# Now make another plot beside
ax2 = plt.subplot(gs[i, 1])
img = hovmoller_plot(data[i],
time,
grid,
smooth=smooth,
ax=ax2,
make_cbar=False,
vmin=vmin[i],
vmax=vmax[i],
zmin=zmin,
zmax=zmax,
monthly=monthly,
contours=contours[i],
ctype=ctype,
title='',
date_since_start=date_since_start,
start=start,
val0=val0[i],
start_t=split_year,
rasterized=rasterized)
ax2.set_yticklabels([])
ax2.set_ylabel('')
if i == 0:
ax2.set_xticklabels([])
axs.append(ax2)
if title is None:
title = loc_string
plt.suptitle(title, fontsize=22)
if return_fig:
return fig, axs
else:
finished_plot(fig, fig_name=fig_name, dpi=dpi)
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)
def animate_latlon_coupled (var, output_dir='./', file_name='output.nc', segment_dir=None, vmin=None, vmax=None, change_points=None, mov_name=None, fig_name_beg=None, fig_name_end=None, figsize=(8,6)):
import matplotlib.animation as animation
output_dir = real_dir(output_dir)
segment_dir = check_segment_dir(output_dir, segment_dir)
file_paths = segment_file_paths(output_dir, segment_dir, file_name)
# Inner function to read and process data from a single file
def read_process_data (file_path, var_name, grid, mask_option='3d', gtype='t', lev_option=None, ismr=False, psi=False):
data = read_netcdf(file_path, var_name)
if mask_option == '3d':
data = mask_3d(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'except_ice':
data = mask_except_ice(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'land':
data = mask_land(data, grid, gtype=gtype, time_dependent=True)
elif mask_option == 'land_ice':
data = mask_land_ice(data, grid, gtype=gtype, time_dependent=True)
else:
print 'Error (read_process_data): invalid mask_option ' + mask_option
sys.exit()
if lev_option is not None:
if lev_option == 'top':
data = select_top(data)
elif lev_option == 'bottom':
data = select_bottom(data)
else:
print 'Error (read_process_data): invalid lev_option ' + lev_option
sys.exit()
if ismr:
data = convert_ismr(data)
if psi:
data = np.sum(data, axis=-3)*1e-6
return data
all_data = []
all_grids = []
all_dates = []
# Loop over segments
for file_path in file_paths:
print 'Processing ' + file_path
# Build the grid
grid = Grid(file_path)
# Read and process the variable we need
ctype = 'basic'
gtype = 't'
include_shelf = var not in ['aice', 'hice', 'mld']
if var == 'ismr':
data = read_process_data(file_path, 'SHIfwFlx', grid, mask_option='except_ice', ismr=True)
title = 'Ice shelf melt rate (m/y)'
ctype = 'ismr'
elif var == 'bwtemp':
data = read_process_data(file_path, 'THETA', grid, lev_option='bottom')
title = 'Bottom water temperature ('+deg_string+'C)'
elif var == 'bwsalt':
data = read_process_data(file_path, 'SALT', grid, lev_option='bottom')
title = 'Bottom water salinity (psu)'
elif var == 'draft':
data = mask_except_ice(grid.draft, grid)
title = 'Ice shelf draft (m)'
elif var == 'aice':
data = read_process_data(file_path, 'SIarea', grid, mask_option='land_ice')
title = 'Sea ice concentration'
elif var == 'hice':
data = read_process_data(file_path, 'SIheff', grid, mask_option='land_ice')
title = 'Sea ice thickness (m)'
elif var == 'mld':
data = read_process_data(file_path, 'MXLDEPTH', grid, mask_option='land_ice')
title = 'Mixed layer depth (m)'
elif var == 'eta':
data = read_process_data(file_path, 'ETAN', grid, mask_option='land')
title = 'Free surface (m)'
elif var == 'psi':
data = read_process_data(file_path, 'PsiVEL', grid, psi=True)
title = 'Vertically integrated streamfunction (Sv)'
ctype = 'plusminus'
else:
print 'Error (animate_latlon): invalid var ' + var
sys.exit()
# Loop over timesteps
if var == 'draft':
# Just one timestep
all_data.append(data)
all_grids.append(grid)
all_dates.append(parse_date(file_path=file_path, time_index=0))
else:
for t in range(data.shape[0]):
# Extract the data from this timestep
# Save it and the grid to the long lists
all_data.append(data[t,:])
all_grids.append(grid)
all_dates.append(parse_date(file_path=file_path, time_index=t))
extend = get_extend(vmin=vmin, vmax=vmax)
if vmin is None:
vmin = np.amax(data)
for elm in all_data:
vmin = min(vmin, np.amin(elm))
if vmax is None:
vmax = np.amin(data)
for elm in all_data:
vmax = max(vmax, np.amax(elm))
num_frames = len(all_data)
# Make the first and last frames as stills
tsteps = [0, -1]
fig_names = [fig_name_beg, fig_name_end]
for t in range(2):
latlon_plot(all_data[tsteps[t]], all_grids[tsteps[t]], gtype=gtype, ctype=ctype, vmin=vmin, vmax=vmax, change_points=change_points, title=title, date_string=all_dates[tsteps[t]], figsize=figsize, fig_name=fig_names[t])
# Now make the animation
fig, ax = plt.subplots(figsize=figsize)
# Inner function to plot a frame
def plot_one_frame (t):
img = latlon_plot(all_data[t], all_grids[t], ax=ax, gtype=gtype, ctype=ctype, vmin=vmin, vmax=vmax, change_points=change_points, title=title+'\n'+all_dates[t], make_cbar=False)
if t==0:
return img
# First frame
img = plot_one_frame(0)
plt.colorbar(img, extend=extend)
# Function to update figure with the given frame
def animate(t):
ax.cla()
plot_one_frame(t)
# Call this for each frame
anim = animation.FuncAnimation(fig, func=animate, frames=range(num_frames))
writer = animation.FFMpegWriter(bitrate=500, fps=10)
anim.save(mov_name, writer=writer)
if mov_name is not None:
print 'Saving ' + mov_name
anim.save(mov_name, writer=writer)
else:
plt.show()
