Plotting heatmaps under blobs in a scatter plot with matplotlib

Question

I have an Nx2 matrix X and an N-dim vector of labels y. For instance:

from sklearn.datasets.samples_generator import make_blobs
import matplotlib.pyplot as plt

X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=2)

plt.scatter(X[:, 0], X[:, 1], c=y, edgecolor='k')
plt.show()

In the background of this plot I want to plot two heatmaps, with a colormap that has the point's colour in the zones of high points density, so that the image looks like having a purple and a yellow cloud, each centered at the purple and yellow blobs.

This has been challenging for me. I tried creating a 2D histogram for each blob as shown in this answer, and also created a custom colormap so that the low density areas of the plot are white, and the high density areas are coloured with the blob's colour:

import numpy as np
import seaborn as sns
from matplotlib.colors import ListedColormap

palette_colors = sns.color_palette("deep")
palette = sns.light_palette(palette_colors[0], input="husl", n_colors=100)
my_cmap = ListedColormap(sns.color_palette(palette).as_hex())

whr1 = np.where(y==0)
whr2 = np.where(y==1)
x1 = X[whr1][:, 0]
y1 = X[whr1][:, 1]
x2 = X[whr2][:, 0]
y2 = X[whr2][:, 1]

heatmap1, xedges1, yedges1 = np.histogram2d(x1, y1, bins=50)
extent1 = [xedges1[0], xedges1[-1], yedges1[0], yedges1[-1]]
heatmap2, xedges2, yedges2 = np.histogram2d(x2, y2, bins=50)
extent2 = [xedges2[0], xedges2[-1], yedges2[0], yedges2[-1]]

But now I don't know how to plot those heatmaps using imshow. I also want to make sure that if the blobs overlap, so will the heatmaps so that one heatmap does not cover the other heatmap, but rather there is a combination of the heatmaps colours and intensities in the overlapping region.

I really appreciate your help!

Answer 1

You could use seaborn's kdeplot

x1,y1 = np.random.normal(loc=0.0, scale=1.0, size=(100,)), np.random.normal(loc=2.0, scale=1.0, size=(100,))
x2,y2 = np.random.normal(loc=2., scale=1.0, size=(100,)), np.random.normal(loc=0.0, scale=1.0, size=(100,))

fig, ax = plt.subplots()
sns.kdeplot(x1,y1, shade=True, shade_lowest=False, alpha=0.5, cbar=False, ax=ax, cmap="Blues")
sns.kdeplot(x2,y2, shade=True, shade_lowest=False, alpha=0.5, cbar=False, ax=ax, cmap="Oranges")
ax.scatter(x1,y1, color="C0")
ax.scatter(x2,y2, color="C1")