Search code examples
pythonmatplotlibcluster-analysis

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

Solution

  • Make a list of colors based on the expected number of different clusters you expect (as long as is not very high), for example

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

    Then just use the clusters value as index of the list to get the color;

    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'])])