Simple example of 2D density plots in python | by Madalina Ciortan | Towards Data Science

This post will show you how to:
Use a Gaussian Kernel to estimate the PDF of 2 distributions
Use Matplotlib to represent the PDF with labelled contour lines around density plots
How to extract the contour lines
How to plot in 3D the above Gaussian kernel
How to use 2D histograms to plot the same PDF
Let’s start by generating an input dataset consisting of 3 blobs:
import numpy as npimport matplotlib.pyplot as pltimport scipy.stats as stfrom sklearn.datasets.samples_generator import make_blobsn_components = 3X, truth = make_blobs(n_samples=300, centers=n_components,                       cluster_std = [2, 1.5, 1],                       random_state=42)plt.scatter(X[:, 0], X[:, 1], s=50, c = truth)plt.title(f"Example of a mixture of {n_components} distributions")plt.xlabel("x")plt.ylabel("y");
For fitting the gaussian kernel, we specify a meshgrid which will use 100 points interpolation on each axis (e.g. mgrid(xmin:xmax:100j)):
# Extract x and yx = X[:, 0]y = X[:, 1]# Define the bordersdeltaX = (max(x) - min(x))/10deltaY = (max(y) - min(y))/10xmin = min(x) - deltaXxmax = max(x) + deltaXymin = min(y) - deltaYymax = max(y) + deltaYprint(xmin, xmax, ymin, ymax)# Create meshgridxx, yy = np.mgrid[xmin:xmax:100j, ymin:ymax:100j]
We will fit a gaussian kernel using the scipy’s gaussian_kde method:
positions = np.vstack([xx.ravel(), yy.ravel()])values = np.vstack([x, y])kernel = st.gaussian_kde(values)f = np.reshape(kernel(positions).T, xx.shape)
fig = plt.figure(figsize=(8,8))ax = fig.gca()ax.set_xlim(xmin, xmax)ax.set_ylim(ymin, ymax)cfset = ax.contourf(xx, yy, f, cmap='coolwarm')ax.imshow(np.rot90(f), cmap='coolwarm', extent=[xmin, xmax, ymin, ymax])cset = ax.contour(xx, yy, f, colors='k')ax.clabel(cset, inline=1, fontsize=10)ax.set_xlabel('X')ax.set_ylabel('Y')plt.title('2D Gaussian Kernel density estimation')
The matplotlib object doing the entire magic is called QuadContour set (cset in the code). We can programatically access the contour lines by iterating through allsegs object. The calculated labels are accessible from labelTexts.
plt.figure(figsize=(8,8))for j in range(len(cset.allsegs)):    for ii, seg in enumerate(cset.allsegs[j]):        plt.plot(seg[:,0], seg[:,1], '.-', label=f'Cluster{j}, level{ii}')plt.legend()
We will use matplotlib’s axes3d from mplot3d. We can plot the density as a surface:
fig = plt.figure(figsize=(13, 7))ax = plt.axes(projection='3d')surf = ax.plot_surface(xx, yy, f, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none')ax.set_xlabel('x')ax.set_ylabel('y')ax.set_zlabel('PDF')ax.set_title('Surface plot of Gaussian 2D KDE')fig.colorbar(surf, shrink=0.5, aspect=5) # add color bar indicating the PDFax.view_init(60, 35)
Or as a wireframe:
fig = plt.figure(figsize=(13, 7))ax = plt.axes(projection='3d')w = ax.plot_wireframe(xx, yy, f)ax.set_xlabel('x')ax.set_ylabel('y')ax.set_zlabel('PDF')ax.set_title('Wireframe plot of Gaussian 2D KDE');
Representation using 2D histograms
Another way to present the same information is by using 2D histograms. Setting the parameter normed to False returns actual frequencies while a True returns the PDF.
h =plt.hist2d(x, y)plt.colorbar(h[3])
The entire code is available on Github.