[Tut] How to Color a Scatter Plot by Category using Matplotlib in Python

[xSicKxBot]

How to Color a Scatter Plot by Category using Matplotlib in Python

Problem Formulation

Given three arrays:

• The first two arrays
  Code:

  x
  and
  Code:

  y
  of length
  Code:

  n
  contain the
  Code:

  (x_i, y_i)
  data of a 2D coordinate system.
  
  
• The third array
  Code:

  c
  provides categorical label information so we essentially get
  Code:

  n
  data bundles
  Code:

  (x_i, y_i, c_i)
  for an arbitrary number of categories
  Code:

  c_i
  .
  

Question: How to plot the data so that

Code:

(x_i, y_i)

and

Code:

(x_j, y_j)

with the same category

Code:

c_i == c_j

have the same color?

Solution: Use Pandas groupby() and Call plt.plot() Separately for Each Group

To plot data by category, you iterate over all groups separately by using the

Code:

data.groupby()

operation. For each group, you execute the

Code:

plt.plot()

operation to plot only the data in the group.

In particular, you perform the following steps:

1. Use the
   Code:

   data.groupby("Category")
   function assuming that data is a Pandas DataFrame containing the
   Code:

   x
   ,
   Code:

   y
   , and
   Code:

   category
   columns for n data points (rows).
   
   
2. Iterate over all
   Code:

   (name, group)
   tuples in the grouping operation result obtained from step one.
   
   
3. Use
   Code:

   plt.plot(group["X"], group["Y"], marker="o", linestyle="", label=name)
   to plot each group separately using the
   Code:

   x
   ,
   Code:

   y
   data and
   Code:

   name
   as a label.
   

Here’s what that looks like in code:

Code:

import pandas as pd import matplotlib.pyplot as plt # Generate the categorical data x = [1, 2, 3, 4, 5, 6] y = [42, 41, 40, 39, 38, 37] c = ['a', 'b', 'a', 'b', 'b', 'a'] data = pd.DataFrame({"X": x, "Y": y, "Category": c}) print(data) # Plot data by category groups = data.groupby("Category") for name, group in groups: plt.plot(group["X"], group["Y"], marker="o", linestyle="", label=name) plt.legend() plt.show()

Before I show you how the resulting plot looks, allow me to show you the data output from the

Code:

print()

function. Here’s the output of the categorical data:

Code:

X Y Category 0 1 42 a 1 2 41 b 2 3 40 a 3 4 39 b 4 5 38 b 5 6 37 a

Now, how does the colored category plot look like? Here’s how:



If you want to learn more about Matplotlib, feel free to check out our full blog tutorial series:

• Python Matplotlib Full Guide
  
  
• Matplotlib Cheat Sheets