¿Cómo puedo hacer un diagrama de dispersión coloreado por densidad en matplotlib?

Me gustaría hacer un diagrama de dispersión donde cada punto esté coloreado por la densidad espacial de los puntos cercanos.

Me encontré con una pregunta muy similar, que muestra un ejemplo de esto usando R:

Gráfico de dispersión R: el color del símbolo representa el número de puntos superpuestos

¿Cuál es la mejor manera de lograr algo similar en Python usando matplotlib?

Además de hist2do hexbincomo sugirió @askewchan, puede usar el mismo método que usa la respuesta aceptada en la pregunta que vinculó.

Si quieres hacer eso:

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

# Generate fake data
x = np.random.normal(size=1000)
y = x * 3 + np.random.normal(size=1000)

# Calculate the point density
xy = np.vstack([x,y])
z = gaussian_kde(xy)(xy)

fig, ax = plt.subplots()
ax.scatter(x, y, c=z, s=100, edgecolor='')
plt.show()

Si desea que los puntos se tracen en orden de densidad para que los puntos más densos estén siempre en la parte superior (similar al ejemplo vinculado), simplemente ordénelos por los valores z. También voy a usar un tamaño de marcador más pequeño aquí, ya que se ve un poco mejor:

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde

# Generate fake data
x = np.random.normal(size=1000)
y = x * 3 + np.random.normal(size=1000)

# Calculate the point density
xy = np.vstack([x,y])
z = gaussian_kde(xy)(xy)

# 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]

fig, ax = plt.subplots()
ax.scatter(x, y, c=z, s=50, edgecolor='')
plt.show()

Podrías hacer un histograma:

import numpy as np
import matplotlib.pyplot as plt

# fake data:
a = np.random.normal(size=1000)
b = a*3 + np.random.normal(size=1000)

plt.hist2d(a, b, (50, 50), cmap=plt.cm.jet)
plt.colorbar()

Además, si el número de puntos hace que el cálculo de KDE sea demasiado lento, el color se puede interpolar en np.histogram2d [Actualización en respuesta a los comentarios: si desea mostrar la barra de colores, use plt.scatter () en lugar de ax.scatter () seguido por plt.colorbar ()]:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize 
from scipy.interpolate import interpn

def density_scatter( x , y, ax = None, sort = True, bins = 20, **kwargs )   :
    """
    Scatter plot colored by 2d histogram
    """
    if ax is None :
        fig , ax = plt.subplots()
    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
    if sort :
        idx = z.argsort()
        x, y, z = x[idx], y[idx], z[idx]

    ax.scatter( x, y, c=z, **kwargs )

    norm = Normalize(vmin = np.min(z), vmax = np.max(z))
    cbar = fig.colorbar(cm.ScalarMappable(norm = norm), ax=ax)
    cbar.ax.set_ylabel('Density')

    return ax


if "__main__" == __name__ :

    x = np.random.normal(size=100000)
    y = x * 3 + np.random.normal(size=100000)
    density_scatter( x, y, bins = [30,30] )

¿Trazar> 100k puntos de datos?

La respuesta aceptada , usar gaussian_kde () llevará mucho tiempo. En mi máquina, 100 mil filas tomaron aproximadamente 11 minutos . Aquí agregaré dos métodos alternativos ( mpl-scatter-density y datashader ) y compararé las respuestas dadas con el mismo conjunto de datos.

A continuación, utilicé un conjunto de datos de prueba de 100k filas:

import matplotlib.pyplot as plt
import numpy as np

# Fake data for testing
x = np.random.normal(size=100000)
y = x * 3 + np.random.normal(size=100000)

Comparación de tiempo de salida y cálculo

A continuación se muestra una comparación de diferentes métodos.

1: mpl-scatter-density

Instalación

pip install mpl-scatter-density

Código de ejemplo

import mpl_scatter_density # adds projection='scatter_density'
from matplotlib.colors import LinearSegmentedColormap

# "Viridis-like" colormap with white background
white_viridis = LinearSegmentedColormap.from_list('white_viridis', [
    (0, '#ffffff'),
    (1e-20, '#440053'),
    (0.2, '#404388'),
    (0.4, '#2a788e'),
    (0.6, '#21a784'),
    (0.8, '#78d151'),
    (1, '#fde624'),
], N=256)

def using_mpl_scatter_density(fig, x, y):
    ax = fig.add_subplot(1, 1, 1, projection='scatter_density')
    density = ax.scatter_density(x, y, cmap=white_viridis)
    fig.colorbar(density, label='Number of points per pixel')

fig = plt.figure()
using_mpl_scatter_density(fig, x, y)
plt.show()

Dibujar esto tomó 0.05 segundos:

Y el zoom se ve bastante bien:

2: datashader

• Datashader es un proyecto interesante. Sin embargo, el soporte para matplotlib es WIP a partir de septiembre de 2020. Instalé la rama mpl del clon de nvictus :

pip install "git+https://github.com/nvictus/datashader.git@mpl"

Código (fuente de dsshow aquí ):

from functools import partial

import datashader as ds
from datashader.mpl_ext import dsshow
import pandas as pd

dyn = partial(ds.tf.dynspread, max_px=40, threshold=0.5)

def using_datashader(ax, x, y):

    df = pd.DataFrame(dict(x=x, y=y))
    da1 = dsshow(df, ds.Point('x', 'y'), spread_fn=dyn, aspect='auto', ax=ax)
    plt.colorbar(da1)

fig, ax = plt.subplots()
using_datashader(ax, x, y)
plt.show()

• Se necesitaron 0.83 s para dibujar esto:

¡y la imagen ampliada se ve genial!

3: scatter_with_gaussian_kde

def scatter_with_gaussian_kde(ax, x, y):
    # https://stackoverflow.com/a/20107592/3015186
    # Answer by Joel Kington

    xy = np.vstack([x, y])
    z = gaussian_kde(xy)(xy)

    ax.scatter(x, y, c=z, s=100, edgecolor='')

• Tomó 11 minutos dibujar esto:

4: using_hist2d

import matplotlib.pyplot as plt
def using_hist2d(ax, x, y, bins=(50, 50)):
    # https://stackoverflow.com/a/20105673/3015186
    # Answer by askewchan
    ax.hist2d(x, y, bins, cmap=plt.cm.jet)

• Se necesitaron 0.021 s para dibujar estos contenedores = (50,50):
• Se necesitaron 0.173 s para dibujar estos contenedores = (1000,1000):
• Desventajas: Los datos ampliados no se ven tan bien como con mpl-scatter-density o datashader. También debe determinar la cantidad de contenedores usted mismo.

5: density_scatter

• El código es como en la respuesta de Guillaume .
• Se necesitaron 0.073 s para dibujar esto con bins = (50,50):
• Se necesitaron 0.368 s para dibujar esto con bins = (1000,1000):