public
Last active

Make a 2D density contour plot with matplotlib

  • Download Gist
density_contour.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51
import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as so
 
def find_confidence_interval(x, pdf, confidence_level):
return pdf[pdf > x].sum() - confidence_level
 
def density_contour(xdata, ydata, nbins_x, nbins_y, ax=None, **contour_kwargs):
""" Create a density contour plot.
 
Parameters
----------
xdata : numpy.ndarray
ydata : numpy.ndarray
nbins_x : int
Number of bins along x dimension
nbins_y : int
Number of bins along y dimension
ax : matplotlib.Axes (optional)
If supplied, plot the contour to this axis. Otherwise, open a new figure
contour_kwargs : dict
kwargs to be passed to pyplot.contour()
"""
 
H, xedges, yedges = np.histogram2d(xdata, ydata, bins=(nbins_x,nbins_y), normed=True)
x_bin_sizes = (xedges[1:] - xedges[:-1]).reshape((1,nbins_x))
y_bin_sizes = (yedges[1:] - yedges[:-1]).reshape((nbins_y,1))
 
pdf = (H*(x_bin_sizes*y_bin_sizes))
 
one_sigma = so.brentq(find_confidence_interval, 0., 1., args=(pdf, 0.68))
two_sigma = so.brentq(find_confidence_interval, 0., 1., args=(pdf, 0.95))
three_sigma = so.brentq(find_confidence_interval, 0., 1., args=(pdf, 0.99))
levels = [one_sigma, two_sigma, three_sigma]
 
X, Y = 0.5*(xedges[1:]+xedges[:-1]), 0.5*(yedges[1:]+yedges[:-1])
Z = pdf.T
 
if ax == None:
contour = plt.contour(X, Y, Z, levels=levels, origin="lower", **contour_kwargs)
else:
contour = ax.contour(X, Y, Z, levels=levels, origin="lower", **contour_kwargs)
 
return contour
 
def test_density_contour():
norm = np.random.normal(10., 15., size=(12540035, 2))
density_contour(norm[:,0], norm[:,1], 100, 100)
plt.show()
 
test_density_contour()

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.