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

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áfica 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 hist2d o hexbin como sugirió @askewchan, puede usar el mismo método que usa la respuesta aceptada en la pregunta a la que se 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)
plt.show()

Si desea que los puntos se representen 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)
plt.show()

¿Trazando >100k puntos de datos?

La respuesta aceptada, usando gaussian_kde() llevará mucho tiempo. En mi máquina, 100k filas tomaron alrededor 11 minutos. Aquí agregaré dos métodos alternativos (mpl-dispersión-densidad y sombreador de datos) y compare 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 salida y tiempo de 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 acercamiento se ve bastante bien:

2: datashader

• sombreador de datos es un proyecto interesante. Tiene Se agregó soporte para matplotlib en datashader 0.12.

Instalación

pip install datashader

Código (listado de fuentes y parámetros para mostrar):

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


def using_datashader(ax, x, y):

    df = pd.DataFrame(dict(x=x, y=y))
    dsartist = dsshow(
        df,
        ds.Point("x", "y"),
        ds.count(),
        vmin=0,
        vmax=35,
        norm="linear",
        aspect="auto",
        ax=ax,
    )

    plt.colorbar(dsartist)


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

• Tomó 0.83 s dibujar esto:

• También hay posibilidad de colorear por tercera variable. El tercer parámetro de dsshow controla la coloración. Ver más ejemplos aquí y la fuente de dsshow aquí.

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)

• Tomó 0.021 s dibujar estos contenedores = (50,50):
  
• Tomó 0.173 s dibujar estos contenedores = (1000,1000):
  
• Contras: los datos ampliados no se ven tan bien como con mpl-scatter-density o datashader. También debe determinar el número de contenedores usted mismo.

5: density_scatter

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

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 [Update in response to comments: If you wish to show the colorbar, use plt.scatter() instead of ax.scatter() followed by 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] )

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()