Добавление ярлыков на легенду не работает, используя matplotlib

Question

У меня есть следующий фрейм данных

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({'var': ['bid', 'on', 'off', 'bid', 'on', 'off'],
                   'aud': ['H', 'H', 'H', 'L', 'L', 'L'],
                   'eff': [0.1, 0.2, 0.3, 0.01, 0.02, 0.03],
                   'spend': [10, 20, 30, 1, 2, 3],
                   'marg': [0.001, 0.002, 0.003, 0.0001, 0.0002, 0.0003]})

Моя конечная цель - создать пузырьковый график с определенными color с для каждого var, marg наy_axis, eff на x_axis графика и size пузырьков равны spend с соответствующими легендами для color и size

Я использую следующий код

x = df.loc[df.aud == 'H']['eff']
y = df.loc[df.aud == 'H']['marg']
z = df.loc[df.aud == 'H']['spend']
colours_dict = dict(zip(['bid', 'on', 'off'], ['#185177', '#FAA22C', '#8FC5E8']))

g, ax = plt.subplots()

scatter = ax.scatter(x, y, c=[ colours_dict[i] for i in df.loc[df.aud == 'H']['var'] ], s=z*10)

# produce a legend with the unique colors from the scatter
legend1 = ax.legend(*scatter.legend_elements(),
                    loc="center", title="var")
ax.add_artist(legend1)

# produce a legend with a cross section of sizes from the scatter
handles, labels = scatter.legend_elements(prop="sizes", alpha=0.6)
legend2 = ax.legend(handles, labels, loc="upper right", title="Spend")

plt.savefig('bubbles.png')

Проблема заключается в том, что легенда с цветами не отображает метки.

Любая помощь?

Answer 1

У вас есть два варианта:

Создание маркеров вручную

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({'var': ['bid', 'on', 'off', 'bid', 'on', 'off'],
                   'aud': ['H', 'H', 'H', 'L', 'L', 'L'],
                   'eff': [0.1, 0.2, 0.3, 0.01, 0.02, 0.03],
                   'spend': [10, 20, 30, 1, 2, 3],
                   'marg': [0.001, 0.002, 0.003, 0.0001, 0.0002, 0.0003]})

x = df.loc[df.aud == 'H']['eff']
y = df.loc[df.aud == 'H']['marg']
z = df.loc[df.aud == 'H']['spend']
labels = ['bid', 'on', 'off']
colors = ['#185177', '#FAA22C', '#8FC5E8']
colours_dict = dict(zip(labels, colors))

fig, ax = plt.subplots()

c = [colours_dict[i] for i in df.loc[df.aud == 'H']['var']]
scatter = ax.scatter(x, y, c=c, s=z*10)

# produce a legend with the unique colors from the scatter
handles = [plt.Line2D([],[], ls="", marker="o", color=c) for c in colors] 
legend1 = ax.legend(handles, labels, loc="lower right", title="var")

plt.show()

Использование цветового сопоставления

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, BoundaryNorm

df = pd.DataFrame({'var': ['bid', 'on', 'off', 'bid', 'on', 'off'],
                   'aud': ['H', 'H', 'H', 'L', 'L', 'L'],
                   'eff': [0.1, 0.2, 0.3, 0.01, 0.02, 0.03],
                   'spend': [10, 20, 30, 1, 2, 3],
                   'marg': [0.001, 0.002, 0.003, 0.0001, 0.0002, 0.0003]})

x = df.loc[df.aud == 'H']['eff']
y = df.loc[df.aud == 'H']['marg']
z = df.loc[df.aud == 'H']['spend']

labels = ['bid', 'on', 'off',]
colors = ['#185177', '#FAA22C', '#8FC5E8']
inv = [labels.index(i) for i  in df.loc[df.aud == 'H']['var']]

cmap = ListedColormap(colors)
norm = BoundaryNorm(np.arange(len(colors)+1)-0.5, len(colors))

fig, ax = plt.subplots()

scatter = ax.scatter(x, y, c=inv, s=z*10, cmap=cmap, norm=norm)

legend1 = ax.legend(scatter.legend_elements(num=len(labels))[0], labels, 
                    loc="lower right", title="var")

plt.show()