Làm cách nào để tạo biểu đồ phân tán được tô màu theo mật độ trong matplotlib?

Tôi muốn tạo một biểu đồ phân tán trong đó mỗi điểm được tô màu theo mật độ không gian của các điểm lân cận.

Tôi đã gặp một câu hỏi tương tự, cho thấy một ví dụ về điều này bằng cách sử dụng R:

R Scatter Plot: màu biểu tượng thể hiện số điểm chồng chéo

Cách tốt nhất để thực hiện điều gì đó tương tự trong python bằng matplotlib là gì?

Ngoài hist2dhoặc hexbinnhư @askewchan đã đề xuất, bạn có thể sử dụng cùng một phương pháp mà câu trả lời được chấp nhận trong câu hỏi mà bạn đã liên kết sử dụng.

Nếu bạn muốn làm điều đó:

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

Nếu bạn muốn các điểm được vẽ theo thứ tự mật độ để các điểm dày đặc nhất luôn ở trên cùng (tương tự như ví dụ được liên kết), chỉ cần sắp xếp chúng theo giá trị z. Tôi cũng sẽ sử dụng kích thước điểm đánh dấu nhỏ hơn ở đây vì nó trông đẹp hơn một chút:

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

Bạn có thể tạo một biểu đồ:

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

Ngoài ra, nếu số điểm làm cho phép tính KDE quá chậm, màu có thể được nội suy trong np.histogram2d [Cập nhật theo nhận xét: Nếu bạn muốn hiển thị thanh màu, hãy sử dụng plt.scatter () thay vì ax.scatter () theo sau bởi 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] )

Mưu đồ> 100k điểm dữ liệu?

Các câu trả lời được chấp nhận , sử dụng gaussian_kde () sẽ mất rất nhiều thời gian. Trên máy của tôi, 100k hàng mất khoảng 11 phút . Ở đây tôi sẽ thêm hai phương pháp thay thế ( mpl-scatter-mật độ và datashader ) và so sánh các câu trả lời đã cho với cùng một tập dữ liệu.

Trong phần sau, tôi đã sử dụng tập dữ liệu thử nghiệm gồm 100 nghìn hàng:

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)

So sánh đầu ra và thời gian tính toán

Dưới đây là so sánh các phương pháp khác nhau.

1: mpl-scatter-density

Cài đặt

pip install mpl-scatter-density

Mã mẫu

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

Vẽ cái này mất 0,05 giây:

Và phóng to trông khá đẹp:

2: datashader

• Datashader là một dự án thú vị. Tuy nhiên, hỗ trợ cho matplotlib là WIP kể từ tháng 9 năm 2020. Tôi đã cài đặt nhánh mpl từ bản sao của nvictus :

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

Code (nguồn dsshow đây ):

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

• Phải mất 0,83 giây để vẽ cái này:

và hình ảnh được phóng to trông rất tuyệt!

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='')

• Phải mất 11 phút để vẽ cái này:

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)

• Mất 0,021 giây để vẽ thùng này = (50,50):
• Mất 0,173 giây để vẽ thùng này = (1000,1000):
• Nhược điểm: Dữ liệu được phóng to trông không đẹp bằng với mật độ phân tán mpl hoặc datashader. Ngoài ra bạn phải tự xác định số lượng thùng.

5: density_scatter

• Mã như trong câu trả lời của Guillaume .
• Mất 0,073 giây để vẽ cái này với thùng = (50,50):
• Mất 0,368 giây để vẽ cái này với thùng = (1000,1000):