首页 > 解决方案 > Changing 3D scatter plot color based on specific column

问题描述

I have some example of information as shown below and I want to make 3D scatter plot with different color of the scatter based on the "clusters" (e.g. 0,1,2)

ID    TP    ALB   BUN   clusters
1     153   101  698    1
2     100   90   400    0
3     50    199  500    1
4     113   102  340    2

Currently I have tried:

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(10, 8))
ax1 = fig.add_subplot(111,projection='3d')

for i in range(len(df_tr)):
    x, y, z = df_tr.iloc[i]['BUN'], df_tr.iloc[i]['ALB'], df_tr.iloc[i]['TP']
    ax1.scatter(x, y, z, c=['blue'])
    ax1.text(x, y, z, '{0}'.format(df_tr.iloc[i] 
    ['clusters']), size=12)

ax1.set_xlabel('BUN')
ax1.set_ylabel('ALB')
ax1.set_zlabel('TP')   

ax1.legend('012')
plt.show()

Getting the result of the scatter has the cluster information (0,1, and 2), but is there anyway to change the scatter color based on the ifnormation on specific column using Axes3D?

Current result

标签: pythonmatplotlibcluster-analysis

解决方案


根据您期望的不同颜色的预期数量clusters(只要不是很高)制作一个颜色列表,例如

colors = ['blue', 'green', 'red']

然后只需使用该clusters值作为列表的索引来获取颜色;

colors = ['blue', 'green', 'red']
for i in range(len(df_tr)):
    x, y, z = df_tr.iloc[i]['BUN'], df_tr.iloc[i]['ALB'], df_tr.iloc[i]['TP']
    ax1.scatter(x, y, z, c=colors[int(df_tr.iloc[i]['clusters'])])

推荐阅读