Matplotlib 3D散点图含NaN时颜色映射报错问题求助

当你将0替换为NaN后，bm.values中包含NaN值，而matplotlib的3D散点图函数scatter无法处理颜色参数c中的NaN值，直接传入二维数组时，NaN点会干扰颜色映射逻辑，导致绘图失败。

解决方案 #

核心思路是先过滤掉所有含NaN的数据点，只保留有效数据用于绘图，具体修改后的代码如下：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

N = 100

bm = pd.DataFrame(
    index=pd.bdate_range(start='2012-01-01', periods=N, freq='B'), 
    data={x: np.random.randn(N) for x in range(1, 11)}
)

# Simulate some zeros
bm = pd.DataFrame(index=bm.index, columns=bm.columns, data=np.where(np.abs(bm.values) < 0.02, 0, bm.values))

# Set zeros to Nan so that I don't plot them
bm = bm.replace({0: np.nan})

z = bm.values
x = bm.columns.tolist()
y = bm.reset_index().index.tolist()
x, y = np.meshgrid(x, y)

# 扁平化数组并过滤NaN点
x_flat = x.flatten()
y_flat = y.flatten()
z_flat = z.flatten()
c_flat = bm.values.flatten()

# 创建掩码过滤NaN值
mask = ~np.isnan(z_flat)

# Set up plot
fig = plt.figure(figsize = (15,10))
ax = plt.axes(projection ='3d')

# 使用过滤后的数据绘图
ax.scatter(x_flat[mask], y_flat[mask], z_flat[mask], '.', c=c_flat[mask], cmap='Reds')

ax.xaxis.set_ticklabels(bm.columns);
ax.yaxis.set_ticklabels(bm.index.strftime('%Y-%m-%d'));

plt.show()

关键修改说明 #

• 将二维的x、y、z和颜色数组c全部扁平化，转为一维数组
• 通过~np.isnan(z_flat)创建掩码，筛选出所有非NaN的有效数据点
• 只将过滤后的有效数据传入scatter函数，彻底避免NaN对颜色映射的干扰

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