Legends not appearing in 3D plot-CodePudding

I'm using the following dummy code for generating a 3D plot.

import random 
from matplotlib import pyplot as plt

random.seed(0)
D = [[random.random() for x in range(3)] for y in range(1000)]
df = pd.DataFrame(D,columns=['Feature_1','Feature_2','Feature_3'])
predictions = [random.randint(0,4) for x in range(1000)]
df['predictions'] = predictions

plt.rcParams["figure.figsize"]=(10,10)
plt.rcParams['legend.fontsize'] = 10
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = Axes3D(fig)

ax.scatter(df['Feature_1'],df['Feature_2'],df['Feature_3'], c=df['predictions'], s =150,cmap='rainbow')
ax.legend(loc = 'upper left')

ax.set_xlabel('Feature_1',fontsize=20,labelpad=10)
ax.set_ylabel('Feature_2', fontsize=20, rotation=150,labelpad=10)
ax.set_zlabel('Feature_3', fontsize=20, rotation=60,labelpad=15)

plt.show()

I'm using as marker color the column predictions, and i would like for each element of that column to appear in the legend but it does not.

Here's a screenshot of the resulting plot

CodePudding user response：

You forgot to put a label handle into the scatter function. If you replace your scatter call with the following line, a legend will show up:

ax.scatter(
    df['Feature_1'], df['Feature_2'], df['Feature_3'],
    c=df['predictions'], s=150, cmap='rainbow', label='Dummy data'
)

Or to show the predictions classes as labels:

scatter = ax.scatter(df['Feature_1'], df['Feature_2'], df['Feature_3'],
                     c=df['predictions'], s=150, cmap='rainbow')
legend1 = ax.legend(*scatter.legend_elements(),
                    loc="upper left", title="Classes")
ax.add_artist(legend1)