Как настроить цветовую легенду при использовании для цикла в Matplotlib, Scatter

Question

Я хочу нарисовать 3D-разброс, в котором данные раскрашены по группам.Вот пример данных:

aa=pd.DataFrame({'a':[1,2,3,4,5],
                 'b':[2,3,4,5,6],
                 'c':[1,3,4,6,9],
                 'd':[0,0,1,2,3],
                 'e':['abc','sdf','ert','hgf','nhkm']})

Здесь a, b, c - оси x, y, z.е - текст, показанный в разбросе.Мне нужно d, чтобы сгруппировать данные и показать разные цвета.

Вот мой код:

fig = plt.figure()
ax = fig.gca(projection='3d')
zdirs = aa.loc[:,'e'].__array__()
xs = aa.loc[:,'a'].__array__()
ys = aa.loc[:,'b'].__array__()
zs = aa.loc[:,'c'].__array__()
colors = aa.loc[:,'d'].__array__()
colors1=np.where(colors==0,'grey',
                 np.where(colors==1,'yellow',
                          np.where(colors==2,'green',
                                   np.where(colors==3,'pink','red'))))
for i in range(len(zdirs)): #plot each point + it's index as text above
    ax.scatter(xs[i],ys[i],zs[i],color=colors1[i])
    ax.text(xs[i],ys[i],zs[i],  '%s' % (str(zdirs[i])), size=10, zorder=1, color='k')
ax.set_xlabel('a')
ax.set_ylabel('b')
ax.set_zlabel('c')
plt.show()

Но я не знаю, как поставить легенду на сюжет.Я надеюсь, что моя легенда выглядит так:

Цвета и цифры должны совпадать и заказываться.

Может ли кто-нибудь помочь мне скак настроить цветовую полосу?

Answer 1

Прежде всего, я позволил себе немного уменьшить ваш код:

• Я бы предложил создать ListedColormap для отображения целого числа -> цвета, который позволяет вам передаватьколонка цвета через c=aa['d'] (обратите внимание, что c=, а не color=!)
• вам не нужно использовать __array__() здесь, в приведенном ниже коде вы можете напрямую использовать aa['a']
• наконец, вы можете добавить пустой график рассеяния для каждого цвета в ListedColormap, и тогда это может быть правильно отрисовано с помощью ax.legend()

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

from matplotlib.colors import ListedColormap
import matplotlib.patches as mpatches

aa=pd.DataFrame({'a':[1,2,3,4,5],
                 'b':[2,3,4,5,6],
                 'c':[1,3,4,6,9],
                 'd':[0,0,1,2,3],
                 'e':['abc','sdf','ert','hgf','nhkm']})

fig = plt.figure()
ax = fig.gca(projection='3d')

cmap = ListedColormap(['grey', 'yellow', 'green', 'pink','red'])
ax.scatter(aa['a'],aa['b'],aa['c'],c=aa['d'],cmap=cmap)

for x,y,z,label in zip(aa['a'],aa['b'],aa['c'],aa['e']):
    ax.text(x,y,z,label,size=10,zorder=1)

# Create a legend through an *empty* scatter plot
[ax.scatter([], [], c=cmap(i), label=str(i)) for i in range(len(aa))]
ax.legend()

ax.set_xlabel('a')
ax.set_ylabel('b')
ax.set_zlabel('c')

plt.show()