1

I'm plotting a scatterplot using matplotlib in python. I want to color the points based on some function, like so:

import matplotlib.pyplot as plt

def color(x, y):
    # based on some rules, return a color
    if(condition):
        return 'red'
    else:
        return 'blue'

plt.scatter(index, data) #c= something?

I'm aware of the matplotlib.from_levels_and_colors function, but the problem is that the mapping isn't based on levels of the values on the x or y axes. There's a third value associated with each data point that is calculated by the function, and that's what I want to color the dots based on.

Is there a way to do this?

bobroxsox
  • 902
  • 2
  • 10
  • 26

2 Answers2

4

Why don't you just make your c array an indicator function for a default colormap. For example:

import numpy as np
import matplotlib.pyplot as plt

x = np.arange(100)
y = np.arange(100)
# Colors whether or not x+y is a multiple of 5
c = (x + y)%5 == 0

# Use 'jet' colormap for red/blue.
plt.scatter(x, y, c=c, cmap='jet')

enter image description here

Of course you could use a colormap with different extremes that will get mapped to 0 and 1.

If your desired result has more than 2 colors, then it's totally fine to pass as c an array with many different values (and it doesn't need to be normalized). See here for an example.

wflynny
  • 18,065
  • 5
  • 46
  • 67
1

This is the way I ended up doing it. I created a list of colors based on the x and y values, and then passed that to the scatter function. Not as nice as wflynny's answer, but it does mean you can do as much computation as needed to come up with the array, rather than having to create a single function to do it.

import matplotlib.pyplot as plt

colors = calculate_colors(x, y)
plt.scatter(index, data, c=colors)
bobroxsox
  • 902
  • 2
  • 10
  • 26