def make_filled_contour_plot(x,
y,
z,
ax=None,
n_pixels=50,
interp_func='Rbf',
contour_levels=[
1.0, 5.0, 10.0, 20.0, 30.0, 50.0, 60.0, 80.0,
90.0, 95.0, 99.99
],
cmap=None,
add_colorbar=True,
interp_func_args={},
interp_func_default_args={
'Rbf': {
'function': 'quintic',
'smooth': 0.01
},
'griddata': {
'method': 'cubic',
'fill_value': 0
},
'SmoothBivariateSpline': {
'kx': 5,
'ky': 5,
's': 0.00002
}
}):
'''
Makes filled contours for Z data on X-Y axes (where these 3 are the
first three arguments)
Input:
------
x, y, z : 1-D scalar arrays
ax : matplotlib.Figure.axes object on which to draw the contours.
If not provided, make a default figure and add an axis.
n_pixels : number of pixels or grains along both X-Y axes. default: 50
interp_func : Interpolation method to use to convert the 1-D input data
to 2-D gridded data that can be plotted as contours
interp_func_args : Arguments to be passed to the interpolation method
interp_func_default_args : Default arguments to be passed to the
interpolation method
contour_levels : Percentage levels of Z at which to draw contours.
Default: many contours from 1% to 99%.
Output:
-------
If an axis is passed in:
im, ax : (contour object, matplotlib.Figure.axes object)
If no axis is passed in:
im, ax, fig : (contour object, matplotlib.Figure.axes object,
matplotlib.Figure object)
'''
if interp_func not in interp_func_default_args:
raise IOError(
"Interpolation function {} not supported. Use one of {}.".format(
interp_func, list(interp_func_default_args.keys())))
interp_kwargs = {a: interp_func_args[a] for a in interp_func_args}
for a in interp_func_default_args[interp_func]:
if a not in interp_kwargs:
interp_kwargs[a] = interp_func_default_args[interp_func][a]
x1vals = np.linspace(min(x), max(x), n_pixels)
x2vals = np.linspace(min(y), max(y), n_pixels)
q, w = np.meshgrid(x1vals, x2vals)
if interp_func == 'Rbf':
from scipy.interpolate import Rbf
z_int = Rbf(x, y, z, **interp_kwargs)
r1 = z_int(q, w)
elif interp_func == 'griddata':
from scipy.interpolate import griddata
points = np.column_stack(
[np.array(x).flatten(),
np.array(y).flatten()])
r1 = griddata(points, z, (q, w), **interp_kwargs)
elif interp_func == 'SmoothBivariateSpline':
from scipy.interpolate import SmoothBivariateSpline
z_int = SmoothBivariateSpline(x, y, z, **interp_kwargs)
r1 = z_int.ev(q, w)
else:
raise RuntimeError("How did we even get here?")
# Create axis (and figure) if not provided
fig = None
if ax is None:
fig = plt.figure()
ax = fig.add_subplot(111)
# Create contour levels
zlevels = [
z.min() + lev * (z.max() - z.min()) / 100.0 for lev in contour_levels
]
# Make contour plot
im = ax.contourf(x1vals, x2vals, r1, zlevels, cmap=cmap)
# Beautify
ax.grid(True)
if add_colorbar:
cb = plt.colorbar(im, ax=ax)
if fig == None:
return im, ax
else:
return im, ax, fig
