#!/usr/bin/env python
########################################################################
# name: plTools.py
"""
Nice predefined plots using matplotlib
"""
#
# 2014, G. SERAZIN, PhD student at MEOM/LGGE, Grenoble
########################################################################
#=======================================================================
# Main module importation
#=======================================================================
try:
import pylab as plt
import matplotlib.mlab as mlab
from matplotlib.ticker import MultipleLocator
import matplotlib
import matplotlib.animation as manimation
except:
print 'Package pylab (matplotlib) is not available...'
exit()
try:
import numpy as np
except:
print 'Package numpy is not available...'
exit()
[docs]class subplot():
"""
Sublplot class containing axes matplotlib instances
"""
def __init__(self, *args, **kwargs):
self.ax = plt.subplot(*args, **kwargs)
[docs] def spectrum_plot(self, x, y, **kwargs):
"""
Define a nice spectrum with twin x-axis, one with frequencies, the
other one with periods, on a predefined axis object
** kwargs : optional keyword arguments
See the plot method in matplotlib documentation
"""
if not 'xlog' in kwargs:
xlog = False
else:
xlog = kwargs['xlog']
del kwargs['xlog']
if not 'ylog' in kwargs:
ylog = False
else:
ylog = kwargs['ylog']
del kwargs['ylog']
if not 'xlim' in kwargs:
xlim = None
else:
xlim = kwargs['xlim']
del kwargs['xlim']
if not 'ylim' in kwargs:
ylim = None
else:
ylim = kwargs['ylim']
del kwargs['ylim']
self.ax.plot(x, y, **kwargs)
if xlog:
self.ax.set_xscale('log', nonposx='clip')
try:
xmin = np.ceil(np.log10(np.abs(xlim[0])))-1
xmax = np.ceil(np.log10(np.abs(xlim[1])))
except:
xmin = np.ceil(np.log10(np.abs(x[1, ])))-1
xmax = np.ceil(np.log10(np.abs(x[-1, ])))
self.ax.set_xlim((10 ** xmin, 10 ** xmax))
else:
try:
self.ax.set_xlim(xlim)
except:
self.ax.set_xlim(np.min(x), np.max(x))
try:
self.ax.set_ylim(ylim)
except:
pass
if ylog:
self.ax.set_yscale('log', nonposx='clip')
self.twiny = self.ax.twiny()
if xlog:
self.twiny.set_xscale('log', nonposx='clip')
self.twiny.set_xlim((10 ** xmin, 10 ** xmax))
new_major_ticks = 10 ** np.arange(xmin + 1, xmax, 1)
new_major_ticklabels = 1. / new_major_ticks
new_major_ticklabels = \
["%.3g" % i for i in new_major_ticklabels]
self.twiny.set_xticks(new_major_ticks)
self.twiny.set_xticklabels(new_major_ticklabels, rotation=60,
fontsize=12)
A = np.arange(2, 10, 2)[np.newaxis]
B = 10 ** (np.arange(-xmax, -xmin, 1)[np.newaxis])
C = np.dot(B.transpose(), A)
new_minor_ticklabels = C.flatten()
new_minor_ticks = 1. / new_minor_ticklabels
new_minor_ticklabels = \
["%.3g" % i for i in new_minor_ticklabels]
self.twiny.set_xticks(new_minor_ticks, minor=True)
self.twiny.set_xticklabels(new_minor_ticklabels, minor=True,
rotation=60, fontsize=12)
self.ax.grid(True, which='both')
[docs] def spectrum2d_plot(self, x, y, z, xlog=False, ylog=False,
zlog=False, **kwargs):
"""
Define a nice spectrum with twin x-axis and twin y-axis, one with
frequencies, the other one with periods, on a predefined axis
object.
Parameters
----------
x,y : array_like
1D array defining the coordinates
z : array_like
2D array
xlog, ylog, zlog : bool, optional
Define if the x-axis, y-axis and z-axis are plotted with a
log scale
** kwargs : optional keyword arguments
See matplotlib.axes.Axes.contourf method in matplotlib
documentation
"""
if not 'xlim' in kwargs:
xlim = None
else:
xlim = kwargs['xlim']
del kwargs['xlim']
if not 'ylim' in kwargs:
ylim = None
else:
ylim = kwargs['ylim']
del kwargs['ylim']
if not 'zlim' in kwargs:
zlim = None
else:
zlim = kwargs['zlim']
del kwargs['zlim']
n_lev = 40
#if symmetric:
#lim = max(np.max(z), abs(np.min(z)))
#lev = np.hstack((np.linspace(- lim, 0, n_lev / 2 + 1),
#np.linspace(0, lim, n_lev / 2)[1:]))
#
#else:
#lev = np.linspace(np.min(z), np.max(z), n_lev / 2 + 1)
if zlog:
plot = self.ax.contourf(x, y, np.log10(z), n_lev,
**kwargs)
else:
plot = self.ax.contourf(x, y, z, levels=lev, **kwargs)
# X limits
if xlog:
self.ax.set_xscale('symlog', nonposx='clip')
xmin = np.ceil(np.log10(x[1, ]))-1
xmax = np.ceil(np.log10(x[-1, ]))
self.ax.set_xlim((10 ** xmin,10 ** xmax))
else:
try:
self.ax.set_xlim(xlim)
except:
self.ax.set_xlim(np.min(x), np.max(x))
# Y limits
if ylog:
self.ax.set_yscale('symlog', nonposx='clip')
ymin = np.ceil(np.log10(x[1, ]))-1
ymax = np.ceil(np.log10(x[-1, ]))
self.ax.set_ylim((-10 ** ymin,10 ** ymax))
else:
try:
self.ax.set_ylim(ylim)
except:
self.ax.set_ylim(np.min(y), np.max(y))
axtwiny = self.ax.twiny()
if xlog:
axtwiny.set_xscale('symlog', nonposx='clip')
axtwiny.set_xlim((-10 ** xmin, 10 ** xmax))
A = np.arange(2,10,2)[np.newaxis]
B = 10 ** (np.arange(-xmax, -xmin, 1)[np.newaxis])
C = np.dot(B.transpose(), A)
new_major_ticks = 10 ** np.arange(xmin + 1, xmax, 1)
new_minor_ticklabels = C.flatten()
new_minor_ticklabels = new_minor_ticklabels.astype(int)
new_minor_ticks = 1. / new_minor_ticklabels
axtwiny.set_xticks(new_minor_ticks, minor=True)
axtwiny.set_xticklabels(new_minor_ticklabels, minor=True,
rotation=30)
new_major_ticklabels = 1. / new_major_ticks
new_major_ticklabels = new_major_ticklabels.astype(int)
axtwiny.set_xticks(new_major_ticks)
axtwiny.set_xticklabels(new_major_ticklabels, rotation=30)
axtwinx = self.ax.twinx()
if ylog:
axtwinx.set_yscale('symlog', nonposx='clip')
axtwinx.set_ylim(y[1], y[-1])
axtwinx.set_ylim((10 ** ymin, 10 ** ymax))
new_major_ticks = 10 ** np.arange(ymin + 1, ymax, 1)
new_major_ticklabels = 1. / new_major_ticks
new_major_ticklabels = new_major_ticklabels.astype(int)
axtwinx.set_yticks(new_major_ticks)
axtwinx.set_yticklabels(new_major_ticklabels)
A = np.arange(2,10,2)[np.newaxis]
B = 10 ** (np.arange(-ymax, -ymin, 1)[np.newaxis])
C = np.dot(B.transpose(), A)
new_minor_ticklabels = C.flatten()
new_minor_ticklabels = new_minor_ticklabels.astype(int)
new_minor_ticks = 1. / new_minor_ticklabels
axtwinx.set_yticks(new_minor_ticks, minor=True)
axtwinx.set_yticklabels(new_minor_ticklabels, minor=True)
self.ax.grid(True, which='both')
[docs] def plot(self, *args, **kwargs):
"""
See `matplotlib.axes.Axes.plot`_ method in matplotlib
documentation
.. _matplotlib.axes.Axes.plot: \
http://matplotlib.org/api/axes_api.html\
#matplotlib.axes.Axes.plot
"""
lines = self.ax.plot(*args, **kwargs)
return lines
def zonal_plot(self, *args, **kwargs):
lines = self.ax.plot(*args, **kwargs)
self.ax.grid(True)
self.nice_lat_ticklabels(args[1])
return lines
def meridional_plot(self, *args, **kwargs):
self.ax.plot(*args, **kwargs)
self.ax.grid(True)
self.nice_lon_ticklabels(args[0])
def fill_plot(self, *args, **kwargs):
if len(args) == 1:
y = args[0]
n = np.size(y)
x = np.arange(n)
elif len(args) == 2:
x = args[0]
y = args[1]
n = np.size(y)
else:
raise TypeError
try:
ymin = kwargs['ymin']
del kwargs['ymin']
except:
ymin = 0
try:
ymax = kwargs['ymax']
del kwargs['ymax']
except:
ymax = 0
self.ax.fill_between(x, ymin, y, where=y < ymin,
facecolor='blue', interpolate=True)
self.ax.fill_between(x, ymax, y, where=y > ymax,
facecolor='red', interpolate=True)
self.ax.plot(x, y, color='k', **kwargs)
self.ax.plot(x, ymin * np.ones(n), lw=0.5, color='k')
self.ax.plot(x, ymax * np.ones(n), lw=0.5, color='k')
self.set_xlim((np.min(x), np.max(x)))
#self.ax.grid(True)
[docs] def pdf_plot(self, x, bins=10, normed=False, weights=None,
cumulative=False, bottom=None, histtype=u'bar',
align=u'mid', orientation=u'vertical', rwidth=None,
log=False, color=None, stacked=False, **kwargs):
"""
Plot the probability density function associated with the input
variable
Parameters
----------
x : array_like
bins : int, optional
normed : bool, optional
weights : array_like, optional
cumulative: bool, optional
"""
n, bins, patches = \
self.ax.hist(x, bins=bins, normed=normed,weights=weights,
cumulative=cumulative, bottom=bottom,
histtype=histtype, align=align,
orientation=orientation, rwidth=rwidth,
log=log, color=color, stacked=stacked, **kwargs)
y = mlab.normpdf(bins, np.mean(x), np.std(x))
self.ax.plot(bins, y, linestyle='--', color='k')
self.set_ylabel('Probability')
[docs] def timelon_plot(self, lon, time, value, v_min=None, v_max=None,
nb_levels=10, **kwargs):
"""
Create a time-longitude plot (Hovmueller)
Parameters
----------
lon : 1d array
The longitude array
time : 1d array
The time array
value : 2d array
The variable to plot
** kwargs : optional keyword arguments
See the contourf method in matplotlib documentation
"""
# Check lon
delta_lon = np.max(lon)-np.min(lon)
if delta_lon > 350:
indx = np.where(lon < 0)
if not indx:
indx = np.where(lon > 180)
lon[indx] -= 360
else:
lon[indx] += 360
if v_min == None:
v_min = np.min(value)
if v_max == None:
v_max = np.max(value)
levels = np.linspace(v_min, v_max, nb_levels)
self.ax.contourf(lon, time, value, levels=levels, **kwargs)
self.grid(True)
self.set_ylabel('Time')
self.nice_lon_ticklabels(lon)
def vline_plot(self, x, **kwargs):
ylim = self.get_ylim()
self.ax.plot([x, x], ylim, **kwargs)
def hline_plot(self, y, **kwargs):
xlim = self.get_xlim()
self.ax.plot(xlim, [y, y], **kwargs)
[docs] def scatter(self, x, y, s=20, c='b', marker='o', cmap=None,
norm=None, vmin=None, vmax=None, alpha=None,
linewidths=None, verts=None, **kwargs):
"""
Make a scatter plot of x vs y, where x and y are sequence like
objects of the same lengths.
See `matplotlib.axes.Axes.scatter`_ method in matplotlib
documentation.
.. _matplotlib.axes.Axes.scatter: \
http://matplotlib.org/api/axes_api.html\
#matplotlib.axes.Axes.scatter
"""
paths = self.ax.scatter(x, y, s=s, c=c, marker=marker,
cmap=cmap, norm=norm, vmin=vmin,
vmax=vmax, alpha=alpha,
linewidths=linewidths, verts=verts,
**kwargs)
return paths
[docs] def nice_lon_ticklabels(self, lon, axis='x', fontdict=None,
minor=False, **kwargs):
"""
Make nice ticklabels for the longitude axis
Parameters
----------
lon : 1darray
The origninal longitude array
axis : optional, {'x', 'y'}
The plot axis corresponding to the longitude array
**kwargs : optional, keyword arguments
"""
lon_min = min(lon)
lon_max = max(lon)
if axis == 'x':
self.set_xlim((lon_min, lon_max))
labels = self.get_xticks()
elif axis == 'y':
self.set_ylim((lon_min, lon_max))
labels = self.get_yticks()
new_labels = []
# Test if there is a cyclic point in the data
# TO DO: A test on the monotony would probably be better
for lab in labels:
if lab > 180:
lab = abs(lab - 360)
lab = "%.0fW" % lab
elif lab < 0:
lab = abs(lab)
lab = "%.0fW" % lab
else:
lab = "%.0fE" % lab
new_labels.append(lab)
if axis == 'x':
self.set_xticklabels(new_labels, fontdict=fontdict,
minor=minor, **kwargs)
self.set_xlabel('Longitude')
elif axis == 'y':
self.set_yticklabels(new_labels, fontdict=fontdict,
minor=minor, **kwargs)
self.set_ylabel('Longitude')
[docs] def nice_lat_ticklabels(self, lat, axis='y', fontdict=None,
minor=False, **kwargs):
"""
Make nice ticklabels for the latitude axis
Parameters
----------
lat : 1darray
The origninal longitude array
axis : optional, {'x', 'y'}
The plot axis corresponding to the longitude array
**kwargs : optional, keyword arguments
"""
lat_min = min(lat)
lat_max = max(lat)
if axis == 'x':
self.set_xlim((lat_min, lat_max))
labels = self.get_xticks()
elif axis == 'y':
self.set_ylim((lat_min, lat_max))
labels = self.get_yticks()
new_labels = []
# Test if there is a cyclic point in the data
# TO DO: A test on the monotony would probably be better
for lab in labels:
if lab > 0:
lab = "%.0fN" % lab
elif lab < 0:
lab = abs(lab)
lab = "%.0fS" % lab
else:
lab = "%.0f" % lab
new_labels.append(lab)
if axis == 'x':
self.set_xticklabels(new_labels, fontdict=fontdict,
minor=minor, **kwargs)
self.set_xlabel('Latitude')
elif axis == 'y':
self.set_yticklabels(new_labels, fontdict=fontdict,
minor=minor, **kwargs)
self.set_ylabel('Latitude')
#-------------------------------------------------------------------
# Axes and title modifications from matplotlib Axes
#-------------------------------------------------------------------
def text(self, x, y, s, fontdict=None, withdash=False, **kwargs):
self.ax.text(x, y, s, fontdict=None, withdash=False, **kwargs)
[docs] def set_title(self, label, fontdict=None, loc=u'center', **kwargs):
"""
Set the title of the subplot.
See `matplotlib.axes.Axes.set_title`_ method in matplotlib
documentation.
.. _matplotlib.axes.Axes.set_title: \
http://matplotlib.org/api/axes_api.html\
#matplotlib.axes.Axes.set_title
"""
text = self.ax.set_title(label, fontdict=None, loc=u'center',
**kwargs)
return text
def set_xlim(self, *args, **kwargs):
self.ax.set_xlim(*args, **kwargs)
def set_ylim(self, *args, **kwargs):
self.ax.set_ylim(*args, **kwargs)
[docs] def set_xlabel(self, xlabel, fontdict=None, labelpad=None,
**kwargs):
"""
See matplotlib.axes.Axes.set_xlabel method in matplotlib
documentation
"""
self.ax.set_xlabel(xlabel, fontdict=None, labelpad=None,
**kwargs)
def set_ylabel(self, ylabel, fontdict=None, labelpad=None,
**kwargs):
self.ax.set_ylabel(ylabel, fontdict=None, labelpad=None,
**kwargs)
def set_twinxlabel(self, xlabel, fontdict=None, labelpad=None,
**kwargs):
try:
self.twiny.set_xlabel(xlabel, fontdict=None, labelpad=None,
**kwargs)
except:
pass
def set_twinylabel(self, ylabel, fontdict=None, labelpad=None,
**kwargs):
try:
self.twinx.set_ylabel(ylabel, fontdict=None, labelpad=None,
**kwargs)
except:
pass
[docs] def set_xticks(self, ticks, minor=False):
"""
See matplotlib.axes.Axes.set_xticks method in matplotlib
documentation
"""
self.ax.set_xticks(ticks, minor=minor)
[docs] def set_yticks(self, ticks, minor=False):
"""
See matplotlib.axes.Axes.set_yticks method in matplotlib
documentation
"""
self.ax.set_yticks(ticks, minor=minor)
[docs] def set_xticklabels(self, labels, fontdict=None, minor=False,
**kwargs):
"""
See `matplotlib.axes.Axes.set_xticklabels`_ method in matplotlib
documentation
.. _matplotlib.axes.Axes.set_xticklabels: \
http://matplotlib.org/1.3.0/api/axes_api.html\
#matplotlib.axes.Axes.set_xticklabels
"""
self.ax.set_xticklabels(labels, fontdict=fontdict, minor=minor,
**kwargs)
[docs] def set_yticklabels(self, labels, fontdict=None, minor=False,
**kwargs):
"""
See `matplotlib.axes.Axes.set_yticklabels`_ method in matplotlib
documentation
.. _matplotlib.axes.Axes.set_yticklabels: \
http://matplotlib.org/1.3.0/api/axes_api.html\
#matplotlib.axes.Axes.set_yticklabels
"""
self.ax.set_yticklabels(labels, fontdict=fontdict, minor=minor,
**kwargs)
def set_xmajorlocator(self, base=1.0, vmin=None, vmax=None):
locator = MultipleLocator(base)
self.ax.xaxis.set_major_locator(locator)
def set_xminorlocator(self, base=1.0):
locator = MultipleLocator(base)
self.ax.xaxis.set_minor_locator(locator)
def set_ymajorlocator(self, base=1.0):
locator = MultipleLocator(base)
self.ax.yaxis.set_major_locator(locator)
def set_yminorlocator(self, base=1.0):
locator = MultipleLocator(base)
self.ax.yaxis.set_minor_locator(locator)
def get_xticks(self, minor=False):
return self.ax.get_xticks(minor=minor)
def get_yticks(self, minor=False):
return self.ax.get_yticks(minor=minor)
def get_xticklabels(self):
return self.ax.get_xticklabels()
def get_yticklabels(self):
return self.ax.get_xticklabels()
def get_xlim(self):
return self.ax.get_xlim()
def get_ylim(self):
return self.ax.get_ylim()
def get_transAxes(self):
return self.ax.transAxes
def set_tickparams(self, axis=u'both', **kwargs):
try:
self.twinx.tick_params(axis=axis, **kwargs)
except:
pass
try:
self.twiny.tick_params(axis=axis, **kwargs)
except:
pass
self.ax.tick_params(axis=axis, **kwargs)
[docs] def grid(self, b=None, which=u'major', axis=u'both', **kwargs):
"""
See `matplotlib.axes.Axes.grid`_ method in matplotlib
documentation
.. _matplotlib.axes.Axes.grid: \
http://matplotlib.org/api/axes_api.html\
#matplotlib.axes.Axes.grid
"""
self.ax.grid(b=None, which=u'major', axis=u'both', **kwargs)
[docs] def legend(self, *args, **kwargs):
"""
See `matplotlib.axes.Axes.legend`_ method in matplotlib
documentation
.. _matplotlib.axes.Axes.legend: \
http://matplotlib.org/api/axes_api.html\
#matplotlib.axes.Axes.legend
"""
self.ax.legend(*args, **kwargs)
#class colormap(matplotlib.colors.Colormap):
#"""
#A class herited from matplotlib Colormap
#"""
[docs]class movie():
"""
A class for making movie using matplotlib
"""
def __init__(self, title='movie', fps=15):
"""
Animation
"""
FFMpegWriter = manimation.writers['ffmpeg']
self.title = title
metadata = dict(title=title)
self.writer = FFMpegWriter(fps=fps, metadata=metadata)
[docs] def make_zonal_movie(self, lat, var, **kwargs):
"""
"""
xmin = np.min(var)
xmax = np.max(var)
fig = plt.figure(figsize=(8.27, 11.69))
ax = subplot(111)
n_t, _ = np.shape(var)
l, = ax.zonal_plot(var[0, :], lat, **kwargs)
ax.set_xlim((xmin, xmax))
with self.writer.saving(fig, self.title + '.mp4', 200):
for i in range(n_t):
l.set_data(var[i, :], lat)
self.writer.grab_frame()