Как сгруппировать метки легенды в Python?

Question

У меня есть наборы данных разных размеров, которые я хочу включить в график, используя scatter с colorbar из matplotlib. Чтобы избежать слишком большого количества элементов в легенде (и не заканчиваться типами маркеров), я бы хотел разделить все элементы легенды на lab_a, lab_b и lab_c, встречающиеся только один раз каждый.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mplcol

data_ax = np.random.rand(10,4)
data_bx = np.random.rand(7, 1)
data_cx = [np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)]

data_ay = np.random.rand(10,4)
data_by = np.random.rand(7, 1)
data_cy = [np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)]

data_az = np.random.rand(10,4)
data_bz = np.random.rand(7, 1)
data_cz = [np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)]

lab_a = 'Data A'
lab_b = 'Data B'
lab_c = 'Data C'

mrk_a = 'o'
mrk_b = 's'
mrk_c = 'D'

cm = plt.cm.get_cmap('viridis')

fig, ax = plt.subplots()

for x, y, z, l, m in zip([data_ax, data_bx, data_cx], [data_ay, data_by, data_cy], [data_az, data_bz, data_cz], [lab_a, lab_b, lab_c], [mrk_a, mrk_b, mrk_c]):
    #sc = ax.scatter(x, y, c=z, label=l, marker=m, cmap=cm, norm=mplcol.Normalize(vmin=0, vmax=1))
    for x2, y2, z2 in zip(x, y, z):
        sc = ax.scatter(x2, y2, c=z2, label=l, marker=m, cmap=cm, norm=mplcol.Normalize(vmin=0, vmax=1))

ax1.legend()

cb = plt.colorbar(sc)

plt.show()

Когда я запускаю приведенный выше код, метки в легенде повторяются. Если я запускаю только один l oop (т.е. #sc), data_cx, data_cy и data_cz генерируют проблему, потому что форма не согласована и возникает ошибка:

ValueError: 'c' argument must either be valid as mpl color(s) or as numbers to be mapped to colors.

• Есть ли способ сгруппировать метки, чтобы я мог запускать оба цикла, и легенда содержит только единственное число lab_a, lab_b и lab_c?

• Возможно, есть лучший способ сгруппировать ярлыки, о которых я не знаю?

Answer 1

Один из подходов состоит в том, чтобы только дать метку для первого разброса в l oop:

for x, y, z, l, m in zip([data_ax, data_bx, data_cx], [data_ay, data_by, data_cy], [data_az, data_bz, data_cz], [lab_a, lab_b, lab_c], [mrk_a, mrk_b, mrk_c]):
    for ind, (x2, y2, z2) in enumerate(zip(x, y, z)):
        sc = ax.scatter(x2, y2, c=z2, label=l if ind == 0 else None,
                        marker=m, cmap=cm, norm=mplcol.Normalize(vmin=0, vmax=1))

Другой подход заключается в преобразовании всех данных в np.arrays (data_cx et c). теперь список np.arrays), а затем используйте ravel (или flatten) из numpy.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mplcol

data_ax = np.random.rand(10,4)
data_bx = np.random.rand(7, 1)
data_cx = np.concatenate([np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)])

data_ay = np.random.rand(10,4)
data_by = np.random.rand(7, 1)
data_cy = np.concatenate( [np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)])

data_az = np.random.rand(10,4)
data_bz = np.random.rand(7, 1)
data_cz = np.concatenate( [np.random.rand(6,1), np.random.rand(5,1), np.random.rand(8,1)])

lab_a = 'Data A'
lab_b = 'Data B'
lab_c = 'Data C'

mrk_a = 'o'
mrk_b = 's'
mrk_c = 'D'

cm = plt.cm.get_cmap('viridis')

fig, ax = plt.subplots()

for x, y, z, l, m in zip([data_ax, data_bx, data_cx], [data_ay, data_by, data_cy], [data_az, data_bz, data_cz], [lab_a, lab_b, lab_c], [mrk_a, mrk_b, mrk_c]):
    sc = ax.scatter(x.ravel(), y.ravel(), c=z.ravel(), label=l, marker=m, cmap=cm, norm=mplcol.Normalize(vmin=0, vmax=1))
ax.legend()

cb = plt.colorbar(sc)

plt.show()