automatic label and color in matplotlib scatter plot

Question

I have a dataframe like this: data =

sample  time (h)    Voltage
1   1   2.4
1   2   3
1   3   4
1   4   5.1
2   1   2.4
2   2   3.2
2   3   4.2
2   4   5.2
3   1   2.3
3   2   3.3
3   3   4.4
3   4   5.3

I want to make a scatter plot by x=time (h), and y=voltage, and I want to automatically make the label and color. So for example, the data with "sample 1" will be shown in blue, and will be labled 1, and then "sample 2" is red and is labeled 2, and so on... (by the way, I want to use a qualitative colormap like 'Set1'). This is something that I can de by adding "hue" in seaborne, but I want to do it in matplotlib.

This is the code I have:

fig, ax = plt.subplots()

color_labels = data['Sample'].unique()

rgb_values = sns.color_palette("Set2", 8)

color_map = dict(zip(color_labels, rgb_values))

ax.scatter(x='time (h)', y='Voltage', s=50, data=data, 
c=data['Sample'].map(color_map))

I appreciate your help

Could you please add your data as text, together with your best-effort code attempt? — JohanC, Oct 11 '22 at 13:08
does this help? https://stackoverflow.com/questions/30505407/create-legend-for-scatter-plot-using-the-label-of-the-samples-in-matplotlib?rq=1 — AlexWach, Oct 11 '22 at 13:09

score 0 · Answer 1 · answered Oct 13 '22 at 02:50

Here is an approach using pandas to filter the samples and plot them individually with their own label and color.

import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
import pandas as pd

data = pd.DataFrame({'Sample': [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
                     'Time (h)': [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4],
                     'Voltage': [2.4, 3, 4, 5.1, 2.4, 3.2, 4.2, 5.2, 2.3, 3.3, 4.4, 5.3]})
# Get colormap
cmap = get_cmap('Set2')
# Get samples info
sample_labels = data['Sample'].unique()

fig, ax = plt.subplots(layout='tight')

for idx_sample_value, sample_value in enumerate(sample_labels):
    # Plot each sample data
    sub_data = data.loc[data['Sample'] == sample_value]
    # Color tuple needs to be a list so that it applies only one color
    ax.scatter(x=sub_data['Time (h)'], y=sub_data['Voltage'], s=50, c=[cmap(idx_sample_value)],
               label=f"Sample {sample_value}")

ax.set_ylabel('Voltage')
ax.set_xlabel('Time (h)')
ax.set_ylim([0, 6])
ax.legend()
plt.show()

I hope this helps you out. Cheers!

automatic label and color in matplotlib scatter plot

1 Answers1