blob: 3de540b7e8f982b5ab6182c0bb3c9394250ce853 [file] [log] [blame]
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from tempfile import TemporaryFile
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
from matplotlib import rcParams, cm
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from mpl_toolkits.basemap import Basemap
from mpl_toolkits.axes_grid1 import ImageGrid
import scipy.stats.mstats as mstats
import numpy as np
import numpy.ma as ma
import ocw.utils as utils
# Set the default colormap to coolwarm
mpl.rc('image', cmap='coolwarm')
def set_cmap(name):
'''
Sets the default colormap (eg when setting cmap=None in a function)
See: http://matplotlib.org/examples/pylab_examples/show_colormaps.html
for a list of possible colormaps.
Appending '_r' to a matplotlib colormap name will give you a reversed
version of it.
:param name: The name of the colormap.
:type name: :mod:`string`
'''
# The first line is redundant but it prevents the user from setting
# the cmap rc value improperly
cmap = plt.get_cmap(name)
mpl.rc('image', cmap=cmap.name)
def _nice_intervals(data, nlevs):
'''
Purpose::
Calculates nice intervals between each color level for colorbars
and contour plots. The target minimum and maximum color levels are
calculated by taking the minimum and maximum of the distribution
after cutting off the tails to remove outliers.
Input::
data - an array of data to be plotted
nlevs - an int giving the target number of intervals
Output::
clevs - A list of floats for the resultant colorbar levels
'''
# Find the min and max levels by cutting off the tails of the distribution
# This mitigates the influence of outliers
data = data.ravel()
mn = mstats.scoreatpercentile(data, 5)
mx = mstats.scoreatpercentile(data, 95)
# if min less than 0 and or max more than 0 put 0 in center of color bar
if mn < 0 and mx > 0:
level = max(abs(mn), abs(mx))
mnlvl = -1 * level
mxlvl = level
# if min is larger than 0 then have color bar between min and max
else:
mnlvl = mn
mxlvl = mx
# hack to make generated intervals from mpl the same for all versions
autolimit_mode = mpl.rcParams.get('axes.autolimit_mode')
if autolimit_mode:
mpl.rc('axes', autolimit_mode='round_numbers')
locator = mpl.ticker.MaxNLocator(nlevs)
clevs = locator.tick_values(mnlvl, mxlvl)
if autolimit_mode:
mpl.rc('axes', autolimit_mode=autolimit_mode)
# Make sure the bounds of clevs are reasonable since sometimes
# MaxNLocator gives values outside the domain of the input data
clevs = clevs[(clevs >= mnlvl) & (clevs <= mxlvl)]
return clevs
def _best_grid_shape(nplots, oldshape):
'''
Purpose::
Calculate a better grid shape in case the user enters more columns
and rows than needed to fit a given number of subplots.
Input::
nplots - an int giving the number of plots that will be made
oldshape - a tuple denoting the desired grid shape (nrows, ncols) for arranging
the subplots originally requested by the user.
Output::
newshape - the smallest possible subplot grid shape needed to fit nplots
'''
nrows, ncols = oldshape
size = nrows * ncols
diff = size - nplots
if diff < 0:
raise ValueError(
'gridshape=(%d, %d): Cannot fit enough subplots for data' % (nrows, ncols))
else:
# If the user enters an excessively large number of
# rows and columns for gridshape, automatically
# correct it so that it fits only as many plots
# as needed
while diff >= nrows:
ncols -= 1
size = nrows * ncols
diff = size - nplots
# Don't forget to remove unnecessary columns too
if ncols == 1:
nrows = nplots
newshape = nrows, ncols
return newshape
def _fig_size(gridshape, aspect=None):
'''
Purpose::
Calculates the figure dimensions from a subplot gridshape
Input::
gridshape - Tuple denoting the subplot gridshape
aspect - Float denoting approximate aspect ratio of each subplot
(width / height). Default is 8.5 / 5.5
Output::
width - float for width of the figure in inches
height - float for height of the figure in inches
'''
if aspect is None:
aspect = 8.5 / 5.5
# Default figure size is 8.5" x 5.5" if nrows == ncols
# and then adjusted by given aspect ratio
nrows, ncols = gridshape
if nrows >= ncols:
# If more rows keep width constant
width, height = (aspect * 5.5), 5.5 * (nrows // ncols)
else:
# If more columns keep height constant
width, height = (aspect * 5.5) * (ncols // nrows), 5.5
return width, height
def draw_taylor_diagram(results, names, refname, fname, fmt='png',
gridshape=(1, 1), ptitle='', subtitles=None,
pos='upper right', frameon=True, radmax=1.5, legend_size=13):
''' Draw a Taylor diagram.
:param results: An Nx2 array containing normalized standard deviations,
correlation coefficients, and names of evaluation results.
:type results: :class:`numpy.ndarray`
:param names: A list of names for each evaluated dataset
:type names: :class:`list` of :mod:`string`
:param refname: The name of the reference dataset.
:type refname: :mod:`string`
:param fname: The filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output plot.
:type fmt: :mod:`string`
:param gridshape: (Optional) Tuple denoting the desired grid shape
(num_rows, num_cols) for arranging the subplots.
:type gridshape: A :class:`tuple` of the form (num_rows, num_cols)
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param subtitles: (Optional) list of strings specifying the title for each
subplot.
:type subtitles: :class:`list` of :mod:`string`
:param pos: (Optional) string or tuple of floats used to set the position
of the legend. Check the `Matplotlib docs <http://matplotlib.org/api/legend_api.html#matplotlib.legend.Legend>`_
for additional information.
:type pos: :mod:`string` or :func:`tuple` of :class:`float`
:param frameon: (Optional) boolean specifying whether to draw a frame
around the legend box.
:type frameon: :class:`bool`
:param radmax: (Optional) float to adjust the extent of the axes in terms of
standard deviation.
:type radmax: :class:`float`
:param legend_size: (Optional) float to control the font size of the legend
:type legend_size: :class:`float`
'''
# Handle the single plot case.
if results.ndim == 2:
results = results.reshape(1, *results.shape)
# Make sure gridshape is compatible with input data
nplots = results.shape[0]
gridshape = _best_grid_shape(nplots, gridshape)
# Set up the figure
fig = plt.figure()
fig.set_size_inches((8.5, 11))
fig.dpi = 300
for i, data in enumerate(results):
rect = gridshape + (i + 1,)
# Convert rect to string form as expected by TaylorDiagram constructor
rect = str(rect[0]) + str(rect[1]) + str(rect[2])
# Create Taylor Diagram object
dia = TaylorDiagram(1, fig=fig, rect=rect,
label=refname, radmax=radmax)
for i, (stddev, corrcoef) in enumerate(data):
dia.add_sample(stddev, corrcoef, marker='$%d$' % (i + 1), ms=6,
label=names[i])
if subtitles is not None:
dia._ax.set_title(subtitles[i])
# Add legend
legend = fig.legend(dia.samplePoints,
[p.get_label() for p in dia.samplePoints],
handlelength=0., prop={'size': legend_size}, numpoints=1,
loc=pos)
legend.draw_frame(frameon)
plt.subplots_adjust(wspace=0)
# Add title and save the figure
fig.suptitle(ptitle)
plt.tight_layout(.05, .05)
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_subregions(subregions, lats, lons, fname, fmt='png', ptitle='',
parallels=None, meridians=None, subregion_masks=None):
''' Draw subregion domain(s) on a map.
:param subregions: The subregion objects to plot on the map.
:type subregions: :class:`list` of subregion objects (Bounds objects)
:param lats: Array of latitudes values.
:type lats: :class:`numpy.ndarray`
:param lons: Array of longitudes values.
:type lons: :class:`numpy.ndarray`
:param fname: The filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param parallels: (Optional) :class:`list` of :class:`int` or :class:`float` for the parallels to
be drawn. See the `Basemap documentation <http://matplotlib.org/basemap/users/graticule.html>`_
for additional information.
:type parallels: :class:`list` of :class:`int` or :class:`float`
:param meridians: (Optional) :class:`list` of :class:`int` or :class:`float` for the meridians to
be drawn. See the `Basemap documentation <http://matplotlib.org/basemap/users/graticule.html>`_
for additional information.
:type meridians: :class:`list` of :class:`int` or :class:`float`
:param subregion_masks: (Optional) :class:`dict` of :class:`bool` arrays for each
subregion for giving finer control of the domain to be drawn, by default
the entire domain is drawn.
:type subregion_masks: :class:`dict` of :class:`bool` arrays
'''
# Set up the figure
fig = plt.figure()
fig.set_size_inches((8.5, 11.))
fig.dpi = 300
ax = fig.add_subplot(111)
# Determine the map boundaries and construct a Basemap object
lonmin = lons.min()
lonmax = lons.max()
latmin = lats.min()
latmax = lats.max()
m = Basemap(projection='cyl', llcrnrlat=latmin, urcrnrlat=latmax,
llcrnrlon=lonmin, urcrnrlon=lonmax, resolution='l', ax=ax)
# Draw the borders for coastlines and countries
m.drawcoastlines(linewidth=1)
m.drawcountries(linewidth=.75)
m.drawstates()
# Create default meridians and parallels. The interval between
# them should be 1, 5, 10, 20, 30, or 40 depending on the size
# of the domain
length = max((latmax - latmin), (lonmax - lonmin)) / 5
if length <= 1:
dlatlon = 1
elif length <= 5:
dlatlon = 5
else:
dlatlon = np.round(length, decimals=-1)
if meridians is None:
meridians = np.r_[
np.arange(0, -180, -dlatlon)[::-1], np.arange(0, 180, dlatlon)]
if parallels is None:
parallels = np.r_[np.arange(0, -90, -dlatlon)
[::-1], np.arange(0, 90, dlatlon)]
# Draw parallels / meridians
m.drawmeridians(meridians, labels=[0, 0, 0, 1], linewidth=.75, fontsize=10)
m.drawparallels(parallels, labels=[1, 0, 0, 1], linewidth=.75, fontsize=10)
# Set up the color scaling
cmap = plt.cm.rainbow
norm = mpl.colors.BoundaryNorm(np.arange(1, len(subregions) + 3), cmap.N)
# Process the subregions
for i, reg in enumerate(subregions):
if subregion_masks is not None and reg.name in subregion_masks.keys():
domain = (i + 1) * subregion_masks[reg.name]
else:
domain = (i + 1) * np.ones((2, 2))
nlats, nlons = domain.shape
domain = ma.masked_equal(domain, 0)
reglats = np.linspace(reg.lat_min, reg.lat_max, nlats)
reglons = np.linspace(reg.lon_min, reg.lon_max, nlons)
reglons, reglats = np.meshgrid(reglons, reglats)
# Convert to to projection coordinates. Not really necessary
# for cylindrical projections but keeping it here in case we need
# support for other projections.
x, y = m(reglons, reglats)
# Draw the subregion domain
m.pcolormesh(x, y, domain, cmap=cmap, norm=norm, alpha=.5)
# Label the subregion
xm, ym = x.mean(), y.mean()
m.plot(xm, ym, marker='$%s$' %
("R" + str(i + 1)), markersize=12, color='k')
# Add the title
ax.set_title(ptitle)
# Save the figure
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def _get_colors(num_colors):
"""
matplotlib will recycle colors after a certain number. This can make
line type charts confusing as colors will be reused. This function
provides a distribution of colors across the default color map
to better approximate uniqueness.
:param num_colors: The number of unique colors to generate.
:return: A color map with num_colors.
"""
cmap = plt.get_cmap()
return [cmap(1. * i / num_colors) for i in range(num_colors)]
def draw_time_series(results, times, labels, fname, fmt='png', gridshape=(1, 1),
xlabel='', ylabel='', ptitle='', subtitles=None,
label_month=False, yscale='linear', aspect=None,
cycle_colors=True, cmap=None):
''' Draw a time series plot.
:param results: 3D array of time series data.
:type results: :class:`numpy.ndarray`
:param times: List of Python datetime objects used by Matplotlib to handle
axis formatting.
:type times: :class:`list` of :class:`datetime.datetime`
:param labels: List of names for each data being plotted.
:type labels: :class:`list` of :mod:`string`
:param fname: Filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param gridshape: (Optional) tuple denoting the desired grid shape
(num_rows, num_cols) for arranging the subplots.
:type gridshape: :func:`tuple` of the form (num_rows, num_cols)
:param xlabel: (Optional) x-axis title.
:type xlabel: :mod:`string`
:param ylabel: (Optional) y-axis title.
:type ylabel: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param subtitles: (Optional) list of titles for each subplot.
:type subtitles: :class:`list` of :mod:`string`
:param label_month: (Optional) flag to toggle drawing month labels on the
x-axis.
:type label_month: :class:`bool`
:param yscale: (Optional) y-axis scale value, 'linear' for linear and 'log'
for log base 10.
:type yscale: :mod:`string`
:param aspect: (Optional) approximate aspect ratio of each subplot
(width / height). Default is 8.5 / 5.5
:type aspect: :class:`float`
:param cycle_colors: (Optional) flag to toggle whether to allow matlibplot
to re-use colors when plotting or force an evenly distributed range.
:type cycle_colors: :class:`bool`
:param cmap: (Optional) string or :class:`matplotlib.colors.LinearSegmentedColormap`
instance denoting the colormap. This must be able to be recognized by
`Matplotlib's get_cmap function <http://matplotlib.org/api/cm_api.html#matplotlib.cm.get_cmap>`_.
Maps like rainbow and spectral with wide spectrum of colors are nice choices when used with
the cycle_colors option. tab20, tab20b, and tab20c are good if the plot has less than 20 datasets.
:type cmap: :mod:`string` or :class:`matplotlib.colors.LinearSegmentedColormap`
'''
if cmap is not None:
set_cmap(cmap)
# Handle the single plot case.
if results.ndim == 2:
results = results.reshape(1, *results.shape)
# Make sure gridshape is compatible with input data
nplots = results.shape[0]
gridshape = _best_grid_shape(nplots, gridshape)
# Set up the figure
width, height = _fig_size(gridshape)
fig = plt.figure()
fig.set_size_inches((width, height))
fig.dpi = 300
# Make the subplot grid
grid = ImageGrid(fig, 111,
nrows_ncols=gridshape,
axes_pad=0.3,
share_all=True,
add_all=True,
ngrids=nplots,
label_mode='L',
aspect=False,
cbar_mode='single',
cbar_location='bottom',
cbar_size=.05,
cbar_pad=.20
)
# Make the plots
for i, ax in enumerate(grid):
data = results[i]
if not cycle_colors:
ax.set_prop_cycle('color', _get_colors(data.shape[0]))
if label_month:
xfmt = mpl.dates.DateFormatter('%b')
xloc = mpl.dates.MonthLocator()
ax.xaxis.set_major_formatter(xfmt)
ax.xaxis.set_major_locator(xloc)
# Set the y-axis scale
ax.set_yscale(yscale)
# Set up list of lines for legend
lines = []
ymin, ymax = 0, 0
# Plot each line
for tSeries in data:
line = ax.plot_date(times, tSeries, '')
lines.extend(line)
cmin, cmax = tSeries.min(), tSeries.max()
ymin = min(ymin, cmin)
ymax = max(ymax, cmax)
# Add a bit of padding so lines don't touch bottom and top of the plot
ymin = ymin - ((ymax - ymin) * 0.1)
ymax = ymax + ((ymax - ymin) * 0.1)
ax.set_ylim((ymin, ymax))
# Set the subplot title if desired
if subtitles is not None:
ax.set_title(subtitles[i], fontsize='small')
# Create a master axes rectangle for figure wide labels
fax = fig.add_subplot(111, frameon=False)
fax.tick_params(labelcolor='none', top='off',
bottom='off', left='off', right='off')
fax.set_ylabel(ylabel)
fax.set_title(ptitle, fontsize=16)
fax.title.set_y(1.04)
# Create the legend using a 'fake' colorbar axes. This lets us have a nice
# legend that is in sync with the subplot grid
cax = ax.cax
cax.set_frame_on(False)
cax.set_xticks([])
cax.set_yticks([])
cax.legend((lines), labels, loc='upper center', ncol=10, fontsize='small',
mode='expand', frameon=False)
# Note that due to weird behavior by axes_grid, it is more convenient to
# place the x-axis label relative to the colorbar axes instead of the
# master axes rectangle.
cax.set_title(xlabel, fontsize=12)
cax.title.set_y(-1.5)
# Rotate the x-axis tick labels
for ax in grid:
for xtick in ax.get_xticklabels():
xtick.set_ha('right')
xtick.set_rotation(30)
# Save the figure
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_barchart(results, yvalues, fname, ptitle='', fmt='png',
xlabel='', ylabel=''):
''' Draw a barchart.
:param results: 1D array of data.
:type results: :class:`numpy.ndarray`
:param yvalues: List of y-axis labels
:type times: :class:`list`
:param fname: Filename of the plot.
:type fname: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param xlabel: (Optional) x-axis title.
:type xlabel: :mod:`string`
:param ylabel: (Optional) y-axis title.
:type ylabel: :mod:`string`
'''
y_pos = list(range(len(yvalues)))
fig = plt.figure()
fig.set_size_inches((11., 8.5))
fig.dpi = 300
ax = plt.subplot(111)
plt.barh(y_pos, results, align="center", height=0.8, linewidth=0)
plt.yticks(y_pos, yvalues)
plt.tick_params(axis="both", which="both", bottom="on", top="off",
labelbottom="on", left="off", right="off", labelleft="on")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ymin = min(y_pos)
ymax = max(y_pos)
ymin = min((ymin - ((ymax - ymin) * 0.1) / 2), 0.5)
ymax = ymax + ((ymax - ymin) * 0.1)
ax.set_ylim((ymin, ymax))
plt.xlabel(xlabel)
plt.tight_layout()
# Save the figure
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_marker_on_map(lat, lon, fname, fmt='png', location_name=' ', gridshape=(1, 1)):
'''Draw a marker on a map.
:param lat: Latitude for plotting a marker.
:type lat: :class:`float`
:param lon: Longitude for plotting a marker.
:type lon: :class:`float`
:param fname: The filename of the plot.
:type fname: :class:`string`
:param fmt: (Optional) Filetype for the output.
:type fmt: :class:`string`
:param location_name: (Optional) A label for the map marker.
:type location_name: :class:`string`
'''
fig = plt.figure()
fig.dpi = 300
ax = fig.add_subplot(111)
m = Basemap(projection='cyl', resolution='c', llcrnrlat=lat -
30, urcrnrlat=lat + 30, llcrnrlon=lon - 60, urcrnrlon=lon + 60)
m.drawcoastlines(linewidth=1)
m.drawcountries(linewidth=1)
m.drawmapboundary(fill_color='aqua')
m.fillcontinents(color='coral', lake_color='aqua')
m.ax = ax
xpt, ypt = m(lon, lat)
m.plot(xpt, ypt, 'bo') # plot a blue dot there
# put some text next to the dot, offset a little bit
# (the offset is in map projection coordinates)
plt.text(xpt + 0.5, ypt + 1.5, location_name +
'\n(lon: %5.1f, lat: %3.1f)' % (lon, lat))
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_contour_map(dataset, lats, lons, fname, fmt='png', gridshape=(1, 1),
clabel='', ptitle='', subtitles=None, cmap=None,
clevs=None, nlevs=10, parallels=None, meridians=None,
extend='neither', aspect=8.5 / 2.5):
''' Draw a multiple panel contour map plot.
:param dataset: 3D array of data to be plotted with shape (nT, nLat, nLon).
:type dataset: :class:`numpy.ndarray`
:param lats: Array of latitudes values.
:type lats: :class:`numpy.ndarray`
:param lons: Array of longitudes
:type lons: :class:`numpy.ndarray`
:param fname: The filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param gridshape: (Optional) tuple denoting the desired grid shape
(num_rows, num_cols) for arranging the subplots.
:type gridshape: :func:`tuple` of the form (num_rows, num_cols)
:param clabel: (Optional) colorbar title.
:type clabel: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param subtitles: (Optional) list of titles for each subplot.
:type subtitles: :class:`list` of :mod:`string`
:param cmap: (Optional) string or :class:`matplotlib.colors.LinearSegmentedColormap`
instance denoting the colormap. This must be able to be recognized by
`Matplotlib's get_cmap function <http://matplotlib.org/api/cm_api.html#matplotlib.cm.get_cmap>`_.
:type cmap: :mod:`string` or :class:`matplotlib.colors.LinearSegmentedColormap`
:param clevs: (Optional) contour levels values.
:type clevs: :class:`list` of :class:`int` or :class:`float`
:param nlevs: (Optional) target number of contour levels if clevs is None.
:type nlevs: :class:`int`
:param parallels: (Optional) list of ints or floats for the parallels to
be drawn. See the `Basemap documentation <http://matplotlib.org/basemap/users/graticule.html>`_
for additional information.
:type parallels: :class:`list` of :class:`int` or :class:`float`
:param meridians: (Optional) list of ints or floats for the meridians to
be drawn. See the `Basemap documentation <http://matplotlib.org/basemap/users/graticule.html>`_
for additional information.
:type meridians: :class:`list` of :class:`int` or :class:`float`
:param extend: (Optional) flag to toggle whether to place arrows at the colorbar
boundaries. Default is 'neither', but can also be 'min', 'max', or
'both'. Will be automatically set to 'both' if clevs is None.
:type extend: :mod:`string`
'''
# Handle the single plot case. Meridians and Parallels are not labeled for
# multiple plots to save space.
if dataset.ndim == 2 or (dataset.ndim == 3 and dataset.shape[0] == 1):
if dataset.ndim == 2:
dataset = dataset.reshape(1, *dataset.shape)
mlabels = [0, 0, 0, 1]
plabels = [1, 0, 0, 1]
else:
mlabels = [0, 0, 0, 0]
plabels = [0, 0, 0, 0]
# Make sure gridshape is compatible with input data
nplots = dataset.shape[0]
gridshape = _best_grid_shape(nplots, gridshape)
# Set up the figure
fig = plt.figure()
fig.set_size_inches((8.5, 11.))
fig.dpi = 300
# Make the subplot grid
grid = ImageGrid(fig, 111,
nrows_ncols=gridshape,
axes_pad=0.3,
share_all=True,
add_all=True,
ngrids=nplots,
label_mode='L',
cbar_mode='single',
cbar_location='bottom',
cbar_size=.15,
cbar_pad='0%'
)
# Determine the map boundaries and construct a Basemap object
lonmin = lons.min()
lonmax = lons.max()
latmin = lats.min()
latmax = lats.max()
m = Basemap(projection='cyl', llcrnrlat=latmin, urcrnrlat=latmax,
llcrnrlon=lonmin, urcrnrlon=lonmax, resolution='l')
# Convert lats and lons to projection coordinates
if lats.ndim == 1 and lons.ndim == 1:
lons, lats = np.meshgrid(lons, lats)
# Calculate contour levels if not given
if clevs is None:
# Cut off the tails of the distribution
# for more representative contour levels
clevs = _nice_intervals(dataset, nlevs)
extend = 'both'
cmap = plt.get_cmap(cmap)
# Create default meridians and parallels. The interval between
# them should be 1, 5, 10, 20, 30, or 40 depending on the size
# of the domain
length = max((latmax - latmin), (lonmax - lonmin)) / 5
if length <= 1:
dlatlon = 1
elif length <= 5:
dlatlon = 5
else:
dlatlon = np.round(length, decimals=-1)
if meridians is None:
meridians = np.r_[
np.arange(0, -180, -dlatlon)[::-1], np.arange(0, 180, dlatlon)]
if parallels is None:
parallels = np.r_[np.arange(0, -90, -dlatlon)
[::-1], np.arange(0, 90, dlatlon)]
x, y = m(lons, lats)
for i, ax in enumerate(grid):
# Load the data to be plotted
data = dataset[i]
m.ax = ax
# Draw the borders for coastlines and countries
m.drawcoastlines(linewidth=1)
m.drawcountries(linewidth=.75)
# Draw parallels / meridians
m.drawmeridians(meridians, labels=mlabels, linewidth=.75, fontsize=10)
m.drawparallels(parallels, labels=plabels, linewidth=.75, fontsize=10)
# Draw filled contours
cs = m.contourf(x, y, data, cmap=cmap, levels=clevs, extend=extend)
# Add title
if subtitles is not None:
ax.set_title(subtitles[i], fontsize='small')
# Add colorbar
cbar = fig.colorbar(cs, cax=ax.cax, drawedges=True,
orientation='horizontal', extendfrac='auto')
cbar.set_label(clabel)
cbar.set_ticks(clevs)
cbar.ax.tick_params(labelsize=6)
cbar.ax.xaxis.set_ticks_position('none')
cbar.ax.yaxis.set_ticks_position('none')
# This is an ugly hack to make the title show up at the correct height.
# Basically save the figure once to achieve tight layout and calculate
# the adjusted heights of the axes, then draw the title slightly above
# that height and save the figure again
fig.savefig(TemporaryFile(), bbox_inches='tight', dpi=fig.dpi)
ymax = 0
for ax in grid:
bbox = ax.get_position()
ymax = max(ymax, bbox.ymax)
# Add figure title
fig.suptitle(ptitle, y=ymax + .06, fontsize=16)
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_portrait_diagram(results, rowlabels, collabels, fname, fmt='png',
gridshape=(1, 1), xlabel='', ylabel='', clabel='',
ptitle='', subtitles=None, cmap=None, clevs=None,
nlevs=10, extend='neither', aspect=None):
''' Draw a portrait diagram plot.
:param results: 3D array of the fields to be plotted. The second dimension
should correspond to the number of rows in the diagram and the
third should correspond to the number of columns.
:type results: :class:`numpy.ndarray`
:param rowlabels: Labels for each row.
:type rowlabels: :class:`list` of :mod:`string`
:param collabels: Labels for each row.
:type collabels: :class:`list` of :mod:`string`
:param fname: Filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param gridshape: (Optional) tuple denoting the desired grid shape
(num_rows, num_cols) for arranging the subplots.
:type gridshape: :func:`tuple` of the form (num_rows, num_cols)
:param xlabel: (Optional) x-axis title.
:type xlabel: :mod:`string`
:param ylabel: (Optional) y-ayis title.
:type ylabel: :mod:`string`
:param clabel: (Optional) colorbar title.
:type clabel: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param subtitles: (Optional) list of titles for each subplot.
:type subtitles: :class:`list` of :mod:`string`
:param cmap: (Optional) string or :class:`matplotlib.colors.LinearSegmentedColormap`
instance denoting the colormap. This must be able to be recognized by
`Matplotlib's get_cmap function <http://matplotlib.org/api/cm_api.html#matplotlib.cm.get_cmap>`_.
:type cmap: :mod:`string` or :class:`matplotlib.colors.LinearSegmentedColormap`
:param clevs: (Optional) contour levels values.
:type clevs: :class:`list` of :class:`int` or :class:`float`
:param nlevs: Optional target number of contour levels if clevs is None.
:type nlevs: :class:`int`
:param extend: (Optional) flag to toggle whether to place arrows at the colorbar
boundaries. Default is 'neither', but can also be 'min', 'max', or
'both'. Will be automatically set to 'both' if clevs is None.
:type extend: :mod:`string`
:param aspect: (Optional) approximate aspect ratio of each subplot
(width / height). Default is 8.5 / 5.5
:type aspect: :class:`float`
'''
# Handle the single plot case.
if results.ndim == 2:
results = results.reshape(1, *results.shape)
nplots = results.shape[0]
# Make sure gridshape is compatible with input data
gridshape = _best_grid_shape(nplots, gridshape)
# Row and Column labels must be consistent with the shape of
# the input data too
prows, pcols = results.shape[1:]
if len(rowlabels) != prows or len(collabels) != pcols:
raise ValueError(
'rowlabels and collabels must have %d and %d elements respectively' % (prows, pcols))
# Set up the figure
width, height = _fig_size(gridshape)
fig = plt.figure()
fig.set_size_inches((width, height))
fig.dpi = 300
# Make the subplot grid
grid = ImageGrid(fig, 111,
nrows_ncols=gridshape,
axes_pad=0.4,
share_all=True,
aspect=False,
add_all=True,
ngrids=nplots,
label_mode='all',
cbar_mode='single',
cbar_location='bottom',
cbar_size=.15,
cbar_pad='3%'
)
# Calculate colorbar levels if not given
if clevs is None:
# Cut off the tails of the distribution
# for more representative colorbar levels
clevs = _nice_intervals(results, nlevs)
extend = 'both'
cmap = plt.get_cmap(cmap)
norm = mpl.colors.BoundaryNorm(clevs, cmap.N)
# Do the plotting
for i, ax in enumerate(grid):
data = results[i]
cs = ax.matshow(data, cmap=cmap, aspect='auto',
origin='lower', norm=norm)
# Add grid lines
ax.xaxis.set_ticks(np.arange(data.shape[1] + 1))
ax.yaxis.set_ticks(np.arange(data.shape[0] + 1))
x = (ax.xaxis.get_majorticklocs() - .5)
y = (ax.yaxis.get_majorticklocs() - .5)
ax.vlines(x, y.min(), y.max())
ax.hlines(y, x.min(), x.max())
# Configure ticks
ax.xaxis.tick_bottom()
ax.xaxis.set_ticks_position('none')
ax.yaxis.set_ticks_position('none')
ax.set_xticklabels(collabels, fontsize='xx-small')
ax.set_yticklabels(rowlabels, fontsize='xx-small')
# Add axes title
if subtitles is not None:
ax.text(0.5, 1.04, subtitles[i], va='center', ha='center',
transform=ax.transAxes, fontsize='small')
# Create a master axes rectangle for figure wide labels
fax = fig.add_subplot(111, frameon=False)
fax.tick_params(labelcolor='none', top='off',
bottom='off', left='off', right='off')
fax.set_ylabel(ylabel)
fax.set_title(ptitle, fontsize=16)
fax.title.set_y(1.04)
# Add colorbar
cax = ax.cax
cbar = fig.colorbar(cs, cax=cax, norm=norm, boundaries=clevs, drawedges=True,
extend=extend, orientation='horizontal', extendfrac='auto')
cbar.set_label(clabel)
cbar.set_ticks(clevs)
cbar.ax.tick_params(labelsize=6)
cbar.ax.xaxis.set_ticks_position('none')
cbar.ax.yaxis.set_ticks_position('none')
# Note that due to weird behavior by axes_grid, it is more convenient to
# place the x-axis label relative to the colorbar axes instead of the
# master axes rectangle.
cax.set_title(xlabel, fontsize=12)
cax.title.set_y(1.5)
# Save the figure
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
class TaylorDiagram(object):
""" Taylor diagram helper class
Plot model standard deviation and correlation to reference (data)
sample in a single-quadrant polar plot, with r=stddev and
theta=arccos(correlation).
This class was released as public domain by the original author
Yannick Copin. You can find the original Gist where it was
released at: https://gist.github.com/ycopin/3342888
"""
def __init__(self, refstd, radmax=1.5, fig=None, rect=111, label='_'):
"""Set up Taylor diagram axes, i.e. single quadrant polar
plot, using mpl_toolkits.axisartist.floating_axes. refstd is
the reference standard deviation to be compared to.
"""
from matplotlib.projections import PolarAxes
import mpl_toolkits.axisartist.floating_axes as FA
import mpl_toolkits.axisartist.grid_finder as GF
self.refstd = refstd # Reference standard deviation
tr = PolarAxes.PolarTransform()
# Correlation labels
rlocs = np.concatenate((np.arange(10) / 10., [0.95, 0.99]))
tlocs = np.arccos(rlocs) # Conversion to polar angles
gl1 = GF.FixedLocator(tlocs) # Positions
tf1 = GF.DictFormatter(dict(zip(tlocs, map(str, rlocs))))
# Standard deviation axis extent
self.smin = 0
self.smax = radmax * self.refstd
ghelper = FA.GridHelperCurveLinear(tr,
extremes=(0, np.pi / 2, # 1st quadrant
self.smin, self.smax),
grid_locator1=gl1,
tick_formatter1=tf1,
)
if fig is None:
fig = plt.figure()
ax = FA.FloatingSubplot(fig, rect, grid_helper=ghelper)
fig.add_subplot(ax)
# Adjust axes
ax.axis["top"].set_axis_direction("bottom") # "Angle axis"
ax.axis["top"].toggle(ticklabels=True, label=True)
ax.axis["top"].major_ticklabels.set_axis_direction("top")
ax.axis["top"].label.set_axis_direction("top")
ax.axis["top"].label.set_text("Correlation")
ax.axis["left"].set_axis_direction("bottom") # "X axis"
ax.axis["left"].label.set_text("Standard deviation")
ax.axis["right"].set_axis_direction("top") # "Y axis"
ax.axis["right"].toggle(ticklabels=True)
ax.axis["right"].major_ticklabels.set_axis_direction("left")
ax.axis["bottom"].set_visible(False) # Useless
# Contours along standard deviations
ax.grid(False)
self._ax = ax # Graphical axes
self.ax = ax.get_aux_axes(tr) # Polar coordinates
# Add reference point and stddev contour
# print "Reference std:", self.refstd
l, = self.ax.plot([0], self.refstd, 'k*',
ls='', ms=10, label=label)
t = np.linspace(0, np.pi / 2)
r = np.zeros_like(t) + self.refstd
self.ax.plot(t, r, 'k--', label='_')
# Collect sample points for latter use (e.g. legend)
self.samplePoints = [l]
def add_sample(self, stddev, corrcoef, *args, **kwargs):
"""Add sample (stddev,corrcoeff) to the Taylor diagram. args
and kwargs are directly propagated to the Figure.plot
command."""
l, = self.ax.plot(np.arccos(corrcoef), stddev,
*args, **kwargs) # (theta,radius)
self.samplePoints.append(l)
return l
def add_rms_contours(self, levels=5, **kwargs):
"""Add constant centered RMS difference contours."""
rs, ts = np.meshgrid(np.linspace(self.smin, self.smax),
np.linspace(0, np.pi / 2))
# Compute centered RMS difference
rms = np.sqrt(self.refstd**2 + rs**2 - 2 *
self.refstd * rs * np.cos(ts))
contours = self.ax.contour(ts, rs, rms, levels, **kwargs)
def add_stddev_contours(self, std, corr1, corr2, **kwargs):
"""Add a curved line with a radius of std between two points
[std, corr1] and [std, corr2]"""
t = np.linspace(np.arccos(corr1), np.arccos(corr2))
r = np.zeros_like(t) + std
return self.ax.plot(t, r, 'red', linewidth=2)
def add_contours(self, std1, corr1, std2, corr2, **kwargs):
"""Add a line between two points
[std1, corr1] and [std2, corr2]"""
t = np.linspace(np.arccos(corr1), np.arccos(corr2))
r = np.linspace(std1, std2)
return self.ax.plot(t, r, 'red', linewidth=2)
def draw_histogram(dataset_array, data_names, fname, fmt='png', nbins=10):
'''
Purpose:: Draw a histogram for the input dataset.
:param dataset_array: A list of data values [data1, data2, ....].
:type dataset_array: :class:`list` of :class:`float`
:param data_names: A list of data names ['name1','name2',....].
:type data_names: :class:`list` of :class:`string`
:param fname: The filename of the plot.
:type fname: :class:`string`
:param fmt: (Optional) Filetype for the output.
:type fmt: :class:`string`
:param bins: (Optional) Number of bins.
:type bins: :class:`integer`
'''
fig = plt.figure()
fig.dpi = 300
ndata = len(dataset_array)
data_min = 500.
data_max = 0.
for data in dataset_array:
data_min = np.min([data_min, data.min()])
data_max = np.max([data_max, data.max()])
bins = np.linspace(np.round(data_min), np.round(data_max + 1), nbins)
for idata, data in enumerate(dataset_array):
ax = fig.add_subplot(ndata, 1, idata + 1)
ax.hist(data, bins, alpha=0.5, label=data_names[idata], normed=True)
leg = ax.legend()
leg.get_frame().set_alpha(0.5)
ax.set_xlim([data_min - (data_max - data_min) * 0.15,
data_max + (data_max - data_min) * 0.15])
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
def fill_US_states_with_color(regions, fname, fmt='png', ptitle='',
colors=False, values=None, region_names=None):
''' Fill the States over the contiguous US with colors
:param regions: The list of subregions(lists of US States) to be filled
with different colors.
:type regions: :class:`list`
:param fname: The filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param colors: (Optional) : If True, each region will be filled
with different colors without using values
:type colors: :class:`bool`
:param values: (Optional) : If colors==False, color for each region scales
an associated element in values
:type values: :class:`numpy.ndarray`
'''
nregion = len(regions)
if colors:
cmap = plt.cm.rainbow
if not (values is None):
cmap = plt.cm.seismic
max_abs = np.abs(values).max()
# Set up the figure
fig = plt.figure()
fig.set_size_inches((8.5, 11.))
fig.dpi = 300
ax = fig.add_subplot(111)
# create the map
m = Basemap(llcrnrlon=-127,llcrnrlat=22,urcrnrlon=-65,urcrnrlat=52,
ax=ax)
for iregion, region in enumerate(regions):
shapes = utils.shapefile_boundary('us_states', region)
patches=[]
lats=np.empty(0)
lons=np.empty(0)
for shape in shapes:
patches.append(Polygon(np.array(shape), True))
lons = np.append(lons, shape[:,0])
lats = np.append(lats, shape[:,1])
if colors:
color_to_fill=cmap((iregion+0.5)/nregion)
if not (values is None):
value = values[iregion]
color_to_fill = cmap(0.5+np.sign(value)*abs(value)/max_abs*0.45)
ax.add_collection(PatchCollection(patches, facecolor=color_to_fill))
if region_names:
ax.text(lons.mean(), lats.mean(), region_names[iregion],
ha='center', va='center', fontsize=10)
m.drawcountries(linewidth=0.)
# Add the title
ax.set_title(ptitle)
# Save the figure
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight', dpi=fig.dpi)
fig.clf()
def draw_plot_to_compare_trends(obs_data, ens_data, model_data,
fname, fmt='png', ptitle='', data_labels=None,
xlabel='', ylabel=''):
''' Fill the States over the contiguous US with colors
:param obs_data: An array of observed trend and standard errors for regions
:type obs_data: :class:'numpy.ndarray'
:param ens_data: An array of trend and standard errors from a multi-model ensemble for regions
:type ens_data: : class:'numpy.ndarray'
:param model_data: An array of trends from models for regions
:type model_data: : class:'numpy.ndarray'
:param fname: The filename of the plot.
:type fname: :mod:`string`
:param fmt: (Optional) filetype for the output.
:type fmt: :mod:`string`
:param ptitle: (Optional) plot title.
:type ptitle: :mod:`string`
:param data_labels: (Optional) names of the regions
:type data_labels: :mod:`list`
:param xlabel: (Optional) a label for x-axis
:type xlabel: :mod:`string`
:param ylabel: (Optional) a label for y-axis
:type ylabel: :mod:`string`
'''
nregions = obs_data.shape[1]
# Set up the figure
fig = plt.figure()
fig.set_size_inches((8.5, 11.))
fig.dpi = 300
ax = fig.add_subplot(111)
b_plot = ax.boxplot(model_data, widths=np.repeat(0.2, nregions), positions=np.arange(nregions)+1.3)
plt.setp(b_plot['medians'], color='black')
plt.setp(b_plot['whiskers'], color='black')
plt.setp(b_plot['boxes'], color='black')
plt.setp(b_plot['fliers'], color='black')
ax.errorbar(np.arange(nregions)+0.8, obs_data[0,:], yerr=obs_data[1,:],
fmt='o', color='r', ecolor='r')
ax.errorbar(np.arange(nregions)+1., ens_data[0,:], yerr=ens_data[1,:],
fmt='o', color='b', ecolor='b')
ax.set_xticks(np.arange(nregions)+1)
ax.set_xlim([0, nregions+1])
if data_labels:
ax.set_xticklabels(data_labels)
fig.savefig('%s.%s' % (fname, fmt), bbox_inches='tight')
def draw_precipitation_JPDF (plot_data, plot_level, x_ticks, x_names,y_ticks,y_names,
output_file, title=None, diff_plot=False, cmap = cm.BrBG,
cbar_ticks=[0.01, 0.10, 0.5, 2, 5, 25],
cbar_label=['0.01', '0.10', '0.5', '2', '5', '25']):
'''
:param plot_data: a numpy array of data to plot (dimY, dimX)
:type plot_data: :class:'numpy.ndarray'
:param plot_level: levels to plot
:type plot_level: :class:'numpy.ndarray'
:param x_ticks: x values where tick makrs are located
:type x_ticks: :class:'numpy.ndarray'
:param x_names: labels for the ticks on x-axis (dimX)
:type x_names: :class:'list'
:param y_ticks: y values where tick makrs are located
:type y_ticks: :class:'numpy.ndarray'
:param y_names: labels for the ticks on y-axis (dimY)
:type y_names: :class:'list'
:param output_file: name of output png file
:type output_file: :mod:'string'
:param title: (Optional) title of the plot
:type title: :mod:'string'
:param diff_plot: (Optional) if true, a difference plot will be generated
:type diff_plot: :mod:'bool'
:param cbar_ticks: (Optional) tick marks for the color bar
:type cbar_ticks: :class:'list'
:param cbar_label: (Optional) lables for the tick marks of the color bar
:type cbar_label: :class:'list'
'''
if diff_plot:
cmap = cm.RdBu_r
fig = plt.figure()
sb = fig.add_subplot(111)
dimY, dimX = plot_data.shape
plot_data2 = np.zeros([dimY,dimX]) # sectioned array for plotting
# fontsize
rcParams['axes.labelsize'] = 8
rcParams['xtick.labelsize'] = 8
rcParams['ytick.labelsize'] = 8
# assign the values in plot_level to plot_data
for iy in range(dimY):
for ix in range(dimX):
if plot_data[iy,ix] <= np.min(plot_level):
plot_data2[iy,ix] = -1.
else:
plot_data2[iy,ix] = plot_level[np.where(plot_level <= plot_data[iy,ix])].max()
sb.set_xticks(x_ticks)
sb.set_xticklabels(x_names)
sb.set_yticks(y_ticks)
sb.set_yticklabels(y_names)
norm = BoundaryNorm(plot_level[1:], cmap.N)
a=sb.pcolor(plot_data2, edgecolors = 'k', linewidths=0.5, cmap = cmap, norm = norm)
a.cmap.set_under('w')
sb.set_xlabel('Spell duration [hrs]')
sb.set_ylabel('Peak rainfall [mm/hr]')
cax = fig.add_axes([0.2, -0.06, 0.6, 0.02])
cbar = plt.colorbar(a, cax=cax, orientation = 'horizontal', extend='both')
cbar.set_ticks(cbar_ticks)
cbar.set_ticklabels(cbar_label)
if title:
fig.suptitle(title)
fig.savefig(output_file, dpi=600,bbox_inches='tight')