从熊猫数据框中绘制3D散点图并按组绘制颜色

Question

我正在做PCA，有10个组件。我想根据组的类型来绘制前三个组件和颜色。

from mpl_toolkits.mplot3d import Axes3D

df=pd.DataFrame(np.random.rand(30,20))
grp=round(pd.DataFrame(np.random.rand(30)*10),0)
df['grp']=grp

fig = plt.figure(figsize=(12, 9))
ax = Axes3D(fig)
y = df.iloc[:,1]
x = df.iloc[:,0]
z = df.iloc[:,2]
c = df['grp']
ax.scatter(x,y,z, c=c, cmap='coolwarm')
plt.title('First 3 Principal Components')
ax.set_ylabel('PC2')
ax.set_xlabel('PC1')
ax.set_zlabel('PC3')
plt.legend()

这行得通，但不幸的是它没有显示图例，我也不相信所有可能的组。

任何建议

Answer 1

您可以尝试添加颜色栏：

# note the difference
cb = ax.scatter(x,y,z, c=c, cmap='coolwarm')

plt.title('First 3 Principal Components')
ax.set_ylabel('PC2')
ax.set_xlabel('PC1')
ax.set_zlabel('PC3')

# and here we add a colorbar
plt.colorbar(cb)

plt.legend()

Answer 2

检出大熊猫groupby，将数据按组分组并分别绘制组：

from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(12, 9))
ax = Axes3D(fig)

for grp_name, grp_idx in df.groupby('grp').groups.items():
    y = df.iloc[grp_idx,1]
    x = df.iloc[grp_idx,0]
    z = df.iloc[grp_idx,2]
    ax.scatter(x,y,z, label=grp_name)  # this way you can control color/marker/size of each group freely
    ax.scatter(*df.iloc[grp_idx, [0, 1, 2]].T.values, label=grp_name)  # if you want to do everything in one line, lol

ax.legend()