Color by category in matplotlib using np.where

user_13628384 0
Viewed 419

I'm trying to create a scatter plot with 100 data points and three variables: x value, y value, and category. This information is stored in an ndarray.

I can create the scatter plot, but I don't know how to use a different color for each category. I used the following code for the plot, which seems to work fine (although it's not finished):

def my_plot(data, color_map):
    f, ax = plt.subplots()
    ax.scatter(data.x, data.y, s = 150, edgecolors = "r")
    return f

In my function, color_map is a parameter which refers to a dictionary I created to color the different categories (there are four in total). This is the dictionary:

color_map = {"winter":(15, 28, 75), "spring":(92, 57, 32), "summer":(255, 253, 211), "fall":(174, 12, 12)}

What I would like to do is to somehow integrate this color_map in my function so that each dot in my plot receives a different color.

I think this could be done using np.where to create a mask, but I'm not sure how to proceed...

matplotlib python
1 Answers

The color values need to be divided by 255 because matplotlib likes them between 0 and 1. With this dict you can create an array of colors for the categories:

from matplotlib import pyplot as plt
from matplotlib.lines import Line2D
import pandas as pd
import numpy as np

color_map = {"winter": (15, 28, 75), "spring": (92, 57, 32), "summer": (255, 253, 211), "fall": (174, 12, 12)}
color_map = {key: (r / 255, g / 255, b / 255,) for key, (r, g, b) in color_map.items()}
N = 200
data = pd.DataFrame({'x': np.random.uniform(1, 9, N), 'y': np.random.uniform(1, 5, N),
                     'cat': np.random.choice([*color_map.keys()], N)})

fig, ax = plt.subplots()
ax.scatter(data.x, data.y, s=150, color=[color_map[c] for c in data.cat], ec='r')

handles = [Line2D([], [], marker='o', ls='', color=col, markeredgecolor='r', label=label)
           for label, col in color_map.items()]
plt.legend(handles=handles, bbox_to_anchor=[1.02, 1.02], loc='upper left')
plt.tight_layout()
plt.show()

example plot

PS: A similar plot can be generated with seaborn, which also automatically adds the corresponding legend. Note that the current version of matplotlib (3.3.1) has a problem with the hue parameter. Normally you would add it as hue='cat' but in this version a workaround via .to_list is needed.

import seaborn as sns
ax = sns.scatterplot(x='x', y='y', hue=data['cat'].to_list(), s=150, palette=color_map, edgecolor='r', data=data)
user_120464090
Related