[FIXED] How to add color and legend by points' label one by one in python?

Issue

I want to divide and color points,val_lab(611,3) by their labels,pred_LAB(611,)

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

fig = plt.figure()
ax = plt.axes(projection = '3d')
ax.set_xlabel('L')
ax.set_ylabel('A')
ax.set_zlabel('B')
for i in range(0, len(val_lab)):
    ax.scatter3D(
        val_lab[i,0], 
        val_lab[i,1],
        val_lab[i,2],
        s = 8,
        marker='o',
        c = pred_LAB
        #cmap = 'rainbow'
    )
#ax.legend(*points.legend_elements(), title = 'clusters')  
plt.show()

The problem is it shows error,

c' argument has 611 elements, which is not acceptable for use with 'x' with size 1, 'y' with size 1.

However, if the dataset only have ten points,it can show the figure correctly, I don't know how to solve this problem, besides, how to add legend of this figure?

Solution

In your solution you would want to replace c = pred_LAB with c = pred_LAB[i]. But you do not have to use a for loop to plot the data. You can just use the following:

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

# generate random input data
val_lab = np.random.randint(0,10,(611,3))
pred_LAB = np.random.uniform(0,1, (611,))

# plot data
fig = plt.figure()
ax = plt.axes(projection = '3d')
ax.set_xlabel('L')
ax.set_ylabel('A')
ax.set_zlabel('B')

# create one 3D scatter plot - no for loop
ax.scatter3D(
    val_lab[:,0], 
    val_lab[:,1],
    val_lab[:,2],
    s = 8,
    marker='o',
    c = pred_LAB,
    cmap = 'rainbow',
    label='my points'
)
# add legend
plt.legend()
plt.show()

Answered By - sehan2