matplotlib scatter plot using a 3rd data series to designate color

1
votes

I have 3 data series that share the same index values:

series a
A 0.6
B 0.4
C 0.7
D 0.5

series b
A 0.8
B 0.4
C 0.7
D 0.5

series c
A 10
B 23
C 50
D 100

series a and b are my x and y axis. I would like to use series c to designate the color of the dots (if value at c > 80 then colors = red elif value at c > 20 then colors = blue).

This is what my code looks like so far:

colors = 'black'  #default color
plt.scatter(a, b, s=np.pi*1, c=colors, alpha=0.5)
    #this is what I'm looking for
    #if value at c > 80 then colors = red elif value at c > 20 then colors = blue
plt.show()

this is what the finished graph would look like:

example

Thanks!

2
pandasmatplotlib

2 Answers

0
votes

You want np.select to define color:

colors = np.select((c>80, c>20), ('r','b'), default='g')
plt.scatter(a,b, c=colors)

Output:

enter image description here

0
votes

Another way, but not nearly as concise as @quang-hoang's.

Create a function with your criteria and then apply it to the df

def colors(row):
    if row.v > 80: return 'red'
    elif row.v > 20: return 'blue' 
    else: return 'green'

df['colors'] = df.apply(colors, axis=1)

Give you this:

     x    y      v colors
A  0.6  0.8   10.0  green
B  0.4  0.4   23.0   blue
C  0.7  0.7   50.0   blue
D  0.5  0.5  100.0    red

df.plot.scatter('x', 'y', c=df.colors)

enter image description here