How to put a multiclass legend in a scatter plot in matplotlib?

Question

I want to put the legend in my scatter plot, I used the c parameter to group the classes and assign the colors. My code is the following:

peral = data[['SiO2a', 'A/CNK','Type', 'Code']]**

Locality    SiO2a   A/CNK   Type    Code
0   CN  71.58   1.13    Rhyolite    7
1   CN  71.76   1.14    Rhyolite    7
2   CN  70.95   1.08    Rhyolite    7
3   CN  70.79   1.14    Rhyolite    7
4   CN  71.69   1.12    Rhyolite    7
...     ...     ...     ...     ...     ...
158     Pairique    63.74   1.09    mafic inclusion     11
159     Pairique    61.77   1.09    mafic inclusion     11
160     Huayra Huasi    65.43   1.10    mafic inclusion     11
161     Huayra Huasi    61.14   1.26    mafic inclusion     11
162     Huayra Huasi    62.53   1.21    mafic inclusion     11

plt.figure(figsize=(10,10))
plt.scatter(peral['SiO2a'], peral['A/CNK'], c=peral['Code'])
plt.legend()
plt.ylim(0.8,2.5)
plt.xlabel('SiO2 %wt')
plt.ylabel('A/CNK')

The column 'Code' is a result of applying LabelEncoder to the column 'Type', and I want to construct the legend with the original names of the column 'Type'.

Does anyone know how to do it?

Answer 1

You could loop over 'Type' values and apply each value as a filter:

for type_ in df['Type'].unique():
    ax.scatter(df[df['Type'] == type_]['SiO2a'], df[df['Type'] == type_]['A/CNK'], marker = 'o', label = type_)

Complete Code

import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv(r'data/data.csv')

fig, ax = plt.subplots(figsize = (10, 10))

for type_ in df['Type'].unique():
    ax.scatter(df[df['Type'] == type_]['SiO2a'], df[df['Type'] == type_]['A/CNK'], marker = 'o', label = type_)

ax.legend(frameon = True)
ax.set_ylim(0.8,2.5)
ax.set_xlabel('SiO2 %wt')
ax.set_ylabel('A/CNK')

plt.show()