How to Plot Different Colors for Different Categorical Levels in Python Using Only Matplotlib

Introduction

This article addresses how to create a scatter plot in Python using matplotlib, where each color represents a different categorical level. This approach avoids using auxiliary plotting packages like seaborn and ggplot for Python.

With Matplotlib

Matplotlib provides the c argument in plt.scatter, which allows color customization. Here's an example:

import matplotlib.pyplot as plt
import pandas as pd

# Sample DataFrame
df = pd.DataFrame({'carat': [0.23, 0.21, 0.23],
                    'price': [326, 326, 327],
                    'color': ['E', 'E', 'E']})

# Color mapping
colors = {'D': 'tab:blue', 'E': 'tab:orange', 'F': 'tab:green', 'G': 'tab:red', 'H': 'tab:purple', 'I': 'tab:brown', 'J': 'tab:pink'}

# Scatter plot with colors
plt.scatter(df['carat'], df['price'], c=df['color'].map(colors))
plt.show()

The map(colors) function maps the "diamond" colors to the "plotting" colors.

With seaborn

Although this article focuses on matplotlib, it's worth mentioning that seaborn also offers a convenient solution:

import seaborn as sns

# Scatter plot with colors
sns.lmplot(x='carat', y='price', data=df, hue='color', fit_reg=False)

With pandas.DataFrame.groupby & pandas.DataFrame.plot

For a manual approach, you can use pandas to group by color and plot each group separately:

import matplotlib.pyplot as plt
import pandas as pd

# Sample DataFrame
df = pd.DataFrame({'carat': [0.23, 0.21, 0.23],
                    'price': [326, 326, 327],
                    'color': ['E', 'E', 'E']})

# Color mapping
colors = {'D': 'tab:blue', 'E': 'tab:orange', 'F': 'tab:green', 'G': 'tab:red', 'H': 'tab:purple', 'I': 'tab:brown', 'J': 'tab:pink'}

# Group by color and plot
grouped = df.groupby('color')
for key, group in grouped:
    group.plot(ax=plt.gca(), kind='scatter', x='carat', y='price', label=key, color=colors[key])

plt.show()

This assumes the same DataFrame as before and manually assigns colors during the plotting process.

Conclusion

This article has demonstrated how to plot different colors for different categorical levels in Python using matplotlib, along with additional options using seaborn and a manual approach with pandas.