Skip to content

Instantly share code, notes, and snippets.

@lukauskas
Last active September 10, 2020 18:06
Show Gist options
  • Save lukauskas/f5fea069ff45e10e57c1b7ec0f5804c8 to your computer and use it in GitHub Desktop.
Save lukauskas/f5fea069ff45e10e57c1b7ec0f5804c8 to your computer and use it in GitHub Desktop.
Density scatterplots using Turbo colormap from google
from matplotlib import pyplot as plt
import numpy as np
from scipy.interpolate import interpn
# Density scatterplots, based on https://stackoverflow.com/a/53865762
def density_scatterplot(x, y, bins=100, cmap='turbo', drop_na=True, ax=None, **kwargs):
if ax is None:
ax = plt.gca()
x = np.asarray(x)
y = np.asarray(y)
if drop_na:
na_mask = np.isnan(x) | np.isnan(y)
x = x[~na_mask]
y = y[~na_mask]
# Based on https://stackoverflow.com/a/53865762
data , x_e, y_e = np.histogram2d(x, y, bins = bins, density = True )
z = interpn(
( 0.5*(x_e[1:] + x_e[:-1]),
0.5*(y_e[1:]+y_e[:-1]) ),
data,
np.vstack([x,y]).T,
method = "splinef2d",
bounds_error = False
)
# To be sure to plot all data
z[np.where(np.isnan(z))] = 0.0
# Sort the points by density, so that the densest points are plotted last
idx = z.argsort()
x, y, z = x[idx], y[idx], z[idx]
return ax.scatter(x, y, c=z, cmap=cmap, **kwargs)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment