Last active
September 10, 2020 18:06
-
-
Save lukauskas/f5fea069ff45e10e57c1b7ec0f5804c8 to your computer and use it in GitHub Desktop.
Density scatterplots using Turbo colormap from google
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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