Add legend to mplot3d

For the 3D plot given here (thanks to Redox for his elegant solution posted here):

how to go about adding the legend? I tried ax.legend() but got the message No artists with labels found to put in legend. Note that artists whose label start with an underscore are ignored when legend() is called with no argument.

Can you assist me in adding the legend to this 3D plot.

Code:

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

colors = dict(zip(['2', '1', '0'], sns.color_palette('tab10_r', 3))) #Your code
new_colors = list(colors.values()) #Get the rgb values for 3 colors into a list
del new_colors[0] #Remove first color - Grey

fig = plt.figure()
plt.figure(figsize=(15, 12))
ax = plt.axes(projection="3d")

x= x_pca[:,0]
y= x_pca[:,1]
z = x_pca[:,2]


y_train_new = y_train_orig.argmax(1)

ax.scatter3D(x, y, z, c = y_train_new, marker = 'o', alpha=0.6, s=55, cmap=mycmap, edgecolor='k')
plt.title("3D Scatterplot: 95% of the variability captured", pad = 15)
ax.set_xlabel('First principal component')
ax.set_ylabel('Second principal component')
ax.set_zlabel('Third principal component')

Data:

You can add these lines at the end of your code to add and adjust the legend. Do take a look if the 0 and 1 match to the purple and green as I have added them manually. You can adjust the position of the box by using bbox_to_anchor, the font and text of the title and legend items, size of the dots, etc. by using this. The colors are picked up from the new_colors that was created earlier, but you can change it to anything you want. Let me know in case of any questions...

from matplotlib.lines import Line2D
myHandle = [Line2D([], [], marker='.', color=new_colors[0], markersize=20, linestyle='None'),
          Line2D([], [], marker='.', color=new_colors[1], markersize=20, linestyle='None')]
legend = ax.legend(handles = myHandle, labels=['1', '0'], loc="upper right", title="My Title", bbox_to_anchor=(1.0, 1.0), title_fontsize=20, prop={'size': 10})