如何用mpl-scatter-density或datashader实现1D密度着色散点图

针对你提到的两个核心需求，以下是具体的实现方案：

一、用datashader实现大数据量下的1D密度着色散点图 #

datashader专门针对大规模数据的可视化做了性能优化，计算密度的速度远快于gaussian_kde，完全适配10k+数据量的场景。核心思路是对y轴数据计算1D核密度，再将密度值映射为散点颜色，最后按密度排序绘制散点（让高密度点显示在上方）。

代码示例 #

import numpy as np
import matplotlib.pyplot as plt
import datashader as ds
from matplotlib.colors import Normalize

# 生成测试数据
data_x = np.broadcast_to(np.array([[1], [2], [3]]), (2, 3, 1000))
data_y = data_x * np.random.normal(size=(2, 3, 1000))

# 创建子图
nrows = 1
ncols = 2
size = 5
fig, ax_array = plt.subplots(
    nrows, ncols,
    figsize=(16/9 * ncols * size, nrows * size),
    squeeze=False
)

# 定义颜色映射与归一化规则
cmap = plt.cm.get_cmap('Reds')
norm = Normalize(vmin=0, vmax=1)

for i, ax_row in enumerate(ax_array):
    for j, axes in enumerate(ax_row):
        index = nrows * i + j
        x = data_x[index, :, :]
        y = data_y[index, :, :]
        
        for x_values, y_values in zip(x, y):
            # 用datashader计算y轴方向的1D密度
            df = {'y': y_values}
            cvs = ds.Canvas(plot_width=100, plot_height=1)
            agg = cvs.line(df, 'y', agg=ds.count())
            density = agg.to_numpy()[0]
            
            # 归一化密度值并生成对应颜色
            normalized_density = norm(density / density.max())
            color = cmap(normalized_density)
            
            # 按密度排序，确保高密度点在图层上方
            idx = normalized_density.argsort()
            x_sorted, y_sorted, color_sorted = x_values[idx], y_values[idx], color[idx]
            
            axes.scatter(x_sorted, y_sorted, c=color_sorted, s=10)

plt.tight_layout()
plt.show()

二、用mpl-scatter-density实现并适配plt.subplots #

mpl-scatter-density默认通过fig.add_subplot(projection='scatter_density')创建专用轴，但也可以通过两种方式适配plt.subplots：

方法1：直接在plt.subplots中指定投影类型 #

通过subplot_kw参数统一设置所有子图的投影为scatter_density，无需逐个添加子图：

import numpy as np
import matplotlib.pyplot as plt
import mpl_scatter_density

# 生成测试数据
data_x = np.broadcast_to(np.array([[1], [2], [3]]), (2, 3, 1000))
data_y = data_x * np.random.normal(size=(2, 3, 1000))

# 创建带scatter_density投影的子图
nrows = 1
ncols = 2
size = 5
fig, ax_array = plt.subplots(
    nrows, ncols,
    figsize=(16/9 * ncols * size, nrows * size),
    squeeze=False,
    subplot_kw={'projection': 'scatter_density'}
)

cmap = plt.cm.get_cmap('Reds')

for i, ax_row in enumerate(ax_array):
    for j, axes in enumerate(ax_row):
        index = nrows * i + j
        x = data_x[index, :, :]
        y = data_y[index, :, :]
        
        for x_values, y_values in zip(x, y):
            # 绘制密度着色散点，自动计算密度并映射颜色
            density = axes.scatter_density(x_values, y_values, cmap=cmap)
            # 若需单独提取密度值，可通过axes.get_images()[0].get_array()获取

plt.colorbar(density, ax=ax_array)
plt.tight_layout()
plt.show()

方法2：手动给普通Axes添加ScatterDensityArtist #

如果不想修改子图投影，可直接创建普通Axes，再手动添加ScatterDensityArtist，关键是要正确绑定轴并设置参数：

import numpy as np
import matplotlib.pyplot as plt
from mpl_scatter_density import ScatterDensityArtist
from matplotlib.colors import LogNorm

# 生成测试数据
data_x = np.broadcast_to(np.array([[1], [2], [3]]), (2, 3, 1000))
data_y = data_x * np.random.normal(size=(2, 3, 1000))

# 创建普通子图
nrows = 1
ncols = 2
size = 5
fig, ax_array = plt.subplots(
    nrows, ncols,
    figsize=(16/9 * ncols * size, nrows * size),
    squeeze=False
)

cmap = plt.cm.get_cmap('Reds')

for i, ax_row in enumerate(ax_array):
    for j, axes in enumerate(ax_row):
        index = nrows * i + j
        x = data_x[index, :, :]
        y = data_y[index, :, :]
        
        for x_values, y_values in zip(x, y):
            # 创建ScatterDensityArtist并添加到轴上
            artist = ScatterDensityArtist(
                axes, x_values, y_values,
                cmap=cmap, norm=LogNorm(vmin=1, vmax=100)
            )
            axes.add_artist(artist)
            # 手动设置轴范围以适配数据
            axes.set_xlim(x_values.min(), x_values.max())
            axes.set_ylim(y_values.min(), y_values.max())
            
plt.colorbar(artist, ax=ax_array)
plt.tight_layout()
plt.show()

补充说明 #

• 数据量超10万时，优先选择datashader，其底层优化的密度计算逻辑能大幅缩短运行时间；mpl-scatter-density更贴近matplotlib生态，适合中等规模数据或需要与matplotlib其他功能结合的场景。
• 若需要完全复刻gaussian_kde的1D密度效果，可在datashader中调整核函数参数，或预计算密度后用datashader映射颜色（前者速度优势更明显）。

内容的提问来源于stack exchange，提问作者Rober