根据数据框绘制 3D 散点图并按组着色

Question

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

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

查看 pandas groupby，按组对数据进行分组并单独绘制各组的图：

测试于python 3.11.2, pandas 2.0.1,matplotlib 3.7.1

fig = plt.figure(figsize=(12, 9))
ax = fig.add_subplot(projection='3d')

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(bbox_to_anchor=(1, 0.5), loc='center left', frameon=False)
plt.show()