pyplot画3D散点图并根据标签设置颜色

import matplotlib.pyplot as plt
%matplotlib notebook
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
 
# 创建一个3D坐标轴
ax = fig.add_subplot(111, projection='3d')
# ax.scatter(x, y, z)

df_1 = pd.DataFrame()
df_1['dim0'] = X_test[:,0]
df_1['dim1'] = X_test[:,1]
df_1['dim2'] = X_test[:,2]
df_1['labels'] = labels_pred
colors = ['red', 'purple', 'orange', 'green', 'blue']
for label in  set(labels_pred):
    tmp_df = df_1[df_1['labels'] == label]
    tmp_labels = tmp_df['labels']
    ax.scatter(tmp_df['dim0'], tmp_df['dim1'],  tmp_df['dim2'],label="{}".format(label), s=1,c=colors[label],alpha=0.8)
ax.legend()
 
# 设置坐标轴标签
ax.set_xlabel('X Axis')
ax.set_ylabel('Y Axis')
ax.set_zlabel('Z Axis')

# 显示图像
plt.show()