ivan
ivan

Reputation: 331

Legend in matplotlib

I'm wondering how can I do the following: I have a DataFrame with points and classes. I'd like to draw all points and use one color for each class. How can I specify how classes refer to colors in the legend?

fig = plt.figure(figsize=(18,10), dpi=1600) 
df = pd.DataFrame(dict(points1 = data_plot[:,0], points2 = data_plot[:,1], \
             target = target[0:2000]))
colors = {1: 'green', 2:'red', 3:'blue', 4:'yellow', 5:'orange', 6:'pink', \                             
              7:'brown', 8:'black', 9:'white'}
fig, ax = plt.subplots()
ax.scatter(df['points1'], df['points2'], c = df['target'].apply(lambda x: colors[x]))

Upvotes: 1

Views: 665

Answers (1)

Suever
Suever

Reputation: 65430

The easiest way to get your legend to have separate entries for each color (and therefore it's target value) is to create a separate plot object for each target value.

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

x = np.random.rand(100)
y = np.random.rand(100)
target = np.random.randint(1,9, size=100)

df = pd.DataFrame(dict(points1=x, points2=y, target=target))
colors = {1: 'green', 2:'red', 3:'blue', 4:'yellow', 5:'orange', 6:'pink', \
              7:'brown', 8:'black', 9:'white'}
fig, ax = plt.subplots()

for k,v in colors.items():
    series = df[df['target'] == k]
    scat = ax.scatter(series['points1'], series['points2'], c=v, label=k)

plt.legend()

enter image description here

Upvotes: 2

Related Questions