Modified Taylor diagrams

There is a type of diagram summarizing how well predictions from numerical models fit expectations; one obvious use case is comparing machine-learning regression models. Modified Taylor diagrams are described in this paper.

I'd like to know if my implementation follows best practices for using matplotlib, and how I might incorporate better testing.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.projections import PolarAxes
import mpl_toolkits.axisartist.floating_axes as fa
import mpl_toolkits.axisartist.grid_finder as gf

class TaylorDiagramPoint(object):
"""
A single point on a Modified Taylor Diagram.
How well do the values predicted match the values expected

* do the means match
* do the standard deviations match
* are they correlated
* what is the normalized error standard deviation
* what is the bias?

Notation:

* s_ := sample standard deviation
* m_ := sample mean
* nesd := normalized error standard deviation;
> nesd**2 = s_predicted**2 + s_expected**2 -
2 * s_predicted * s_expected * corcoeff
"""
def __init__(self, expected, predicted, pred_name, point_id):

self.pred = predicted
self.expd = expected
self.s_pred = np.std(self.pred)
self.s_expd = np.std(self.expd)
self.s_normd = self.s_pred / self.s_expd
self.bias = (np.mean(self.pred) - np.mean(self.expd)) / self.s_expd
self.corrcoef = np.corrcoef(self.pred, self.expd)[0, 1]
self.corrcoef = min([self.corrcoef, 1.0])
self.nesd = np.sqrt(self.s_pred**2 + self.s_expd**2 -
2 * self.s_pred * self.s_expd * self.corrcoef)
self.name = pred_name
self.point_id = point_id

class ModTaylorDiagram(object):
"""
Given predictions and expected numerical data
plot the standard deviation of the differences and correlation between
expected and predicted in a single-quadrant polar plot, with
r=stddev and theta=arccos(correlation).
"""
def __init__(self, fig=None, label='expected'):
"""
Set up Taylor diagram axes.
"""

self.title_polar = r'Correlation'
self.title_xy = r'Normalized Standard Deviation'
self.title_expected = r'Expected'
self.max_normed_std = 1.55
self.s_min = 0

# Correlation labels
corln_r = np.append(np.linspace(0.0, 0.9, 10), 0.95)
corln_ang = np.arccos(corln_r)      # Conversion to polar angles
grid_loc1 = gf.FixedLocator(corln_ang)    # Positions
tick_fmttr1 = gf.DictFormatter(dict(zip(corln_ang, map(str, corln_r))))

# Normalized standard deviation axis
tr = PolarAxes.PolarTransform()
grid_helper = fa.GridHelperCurveLinear(tr,
extremes=(0, np.pi/2, # 1st quadrant
self.s_min, self.max_normed_std),
grid_locator1=grid_loc1,
tick_formatter1=tick_fmttr1)
self.fig = fig
if self.fig is None:
self.fig = plt.figure()

# setup axes
ax = fa.FloatingSubplot(self.fig, 111, grid_helper=grid_helper)
# make the axis (polar ax child used for plotting)
# hide base-axis labels etc
self.ax.axis['bottom'].set_visible(False)
self._setup_axes()

# attach the ploar axes
self.polar_ax = self.ax.get_aux_axes(tr)

# Add norm error stddev and nesd==1 contours
self._plot_req1_cont(label)
self._plot_nesd_cont(levels=np.arange(0.0, 1.75, 0.25))
self.points = []

def add_prediction(self, expected, predicted, predictor_name,
plot_pt_id):
"""
Add a prediction/model to the diagram
"""
this_point = TaylorDiagramPoint(expected, predicted,
predictor_name, plot_pt_id)
self.points.append(this_point)

def plot(self):
"""
Place all the loaded points onto the figure
"""
rs = []
thetas = []
biases = []
names = []
point_tags = []
for point in self.points:
rs.append(point.s_expd)
thetas.append(np.arccos(point.corrcoef))
biases.append(point.bias)
names.append(point.name)
point_tags.append(point.point_id)

sc = self.polar_ax.scatter(thetas, rs, c=biases,
s=35, cmap=plt.cm.jet)
for i, tag in enumerate(point_tags):
self.polar_ax.text(thetas[i], rs[i], tag,
horizontalalignment='center',
verticalalignment='bottom')

cbaxes = self.fig.add_axes([0.238, 0.9, 0.55, 0.03])
cbar = plt.colorbar(sc, cax=cbaxes, orientation='horizontal',
format='%.2f')
cbaxes.set_xlabel('Normalized bias')
cbaxes.xaxis.set_ticks_position('top')
cbaxes.xaxis.set_label_position('top')
self.show_key()

def show_key(self):
"""
add annotation key for model IDs and normalization factors
"""
textstr = ''
for i, p in enumerate(self.points):
if i > 0:
textstr += '\n'

textstr += r'{0}$\rightarrow${1}: std/{2:.3f}'.format(p.point_id,
p.name, p.s_expd)
props = dict(boxstyle='round', facecolor='white', alpha=0.75)
# place a text box in upper left in axes coords
self.ax.text(0.75, 0.98, textstr, transform=self.ax.transAxes, fontsize=11,
verticalalignment='top', bbox=props)

def show_norm_factor(self):
"""
"""
n_fact = self.points[0]
out_str = r'Norm Factor {:.2f}'.format(n_fact)
x = 0.95 * self.max_normed_std
y = 0.95 * self.max_normed_std
self.ax.text(x, y,
out_str,
horizontalalignment='right',
verticalalignment='top',
bbox={'edgecolor': 'black', 'facecolor':'None'})

def _plot_req1_cont(self, label):
"""
plot the normalized standard deviation = 1 contour and label
"""
my_purple = [0.414, 0.254, 0.609]
t = np.linspace(0, np.pi/2)
r = np.ones_like(t)
self.polar_ax.plot(t, r, '--', color=my_purple, label=label)
self.polar_ax.text(0, 1,
self.title_expected,
color=my_purple,
horizontalalignment='center',
verticalalignment='center')

def _plot_nesd_cont(self, levels=6):
"""
plot the normalized error standard deviation contours
"""
my_blue = [0.171875, 0.39453125, 0.63671875]
rs, ts = np.meshgrid(np.linspace(self.s_min, self.max_normed_std),
np.linspace(0, np.pi/2))

nesd = np.sqrt(1.0 + rs**2 - 2 * rs * np.cos(ts))
contours = self.polar_ax.contour(ts, rs, nesd, levels,
colors=my_blue, linestyles='dotted')

self.polar_ax.clabel(contours, inline=1, fontsize=10)

def _setup_angle_axis(self):
"""
set the ticks labels etc for the angle axis
"""
loc = 'top'
self.ax.axis[loc].set_axis_direction('bottom')
self.ax.axis[loc].toggle(ticklabels=True, label=True)
self.ax.axis[loc].major_ticklabels.set_axis_direction('top')
self.ax.axis[loc].label.set_axis_direction('top')
self.ax.axis[loc].label.set_text(self.title_polar)

def _setup_x_axis(self):
"""
set the ticks labels etc for the x axis
"""
loc = 'left'
self.ax.axis[loc].set_axis_direction('bottom')
self.ax.axis[loc].label.set_text(self.title_xy)

def _setup_y_axis(self):
"""
set the ticks labels etc for the y axis
"""
loc = 'right'
self.ax.axis[loc].set_axis_direction('top')
self.ax.axis[loc].toggle(ticklabels=True)
self.ax.axis[loc].major_ticklabels.set_axis_direction('left')
self.ax.axis[loc].label.set_text(self.title_xy)

def _setup_axes(self):
"""
set the ticks labels etc for the angle x and y axes
"""
self._setup_angle_axis()
self._setup_x_axis()
self._setup_y_axis()

if '__main__'== __name__:
mtd = ModTaylorDiagram()
x = np.linspace(0.0, 4.0*np.pi, 100)
observed = np.sin(x)
# Models
pred_0 = observed + 0.2*np.random.randn(len(x))
pred_1 = 0.8*observed + 2*np.random.randn(len(x))
pred_2 = np.sin(x - np.pi/10)  - 0.5*np.random.randn(len(x))
mods = [pred_0, pred_1, pred_2]
mod_names = [r'Model 0', r'Model 1', r'Model 2']
mod_ids = [r'a', r'$\beta$', r'$\spadesuit$']

for i, model in enumerate(mods):
mtd.add_prediction(observed, model, mod_names[i], mod_ids[i])

mtd.plot()


I like how you split your functions. It's straightforward and maintainable, except the large amount of magic numbers. I have a personal dislike for numerals in function names and variables (for example: grid_locator1) and constructs like the following make me think you got naming problems:
grid_locator1=grid_loc1,

Basically, you got the usage of matplotlib down quite nicely. The rest could use some work.