CODEWITHSUNDEEP
pythonmatplotlib
MT_
MT_

Reputation: 575

Change color of lineplot depending on data

Is there a way to change the color of a lineplot depending on data values using pyplot ? For example, a red line when data is negative and a black line when data is positive.

I tried to split the data into two sets and plot them separately but there may be a better way.

Upvotes: 15

Views: 25218

Answers (4)

Michael Szczesny
Michael Szczesny

Reputation: 5036

I couldn't find a clean solution on stackoverflow without vanishing line segments that cross the x-axis.

Following this approach to compute new x values at the crossing point and updating the point arrays

import numpy as np
import matplotlib.pyplot as plt
plt.style.use('ggplot')


x = [2, 5]
y = [5,-3]

x1, x2 = x
y1, y2 = y
xf = x1 + -y1 * (x2 - x1)/(y2 - y1)

xn = [2, xf, 5]
yn = [5, 0, -3]

we get a crossing line segment with two parts.

plt.figure(figsize=(7,6))

plt.plot(xn, yn, 'X:')
plt.show()

two part line segment

Vectorizing this approach can be done by finding the crossing line segments, computing the crossing points and updating both point arrays at the appropriate indices for the x-axis (needs to be sorted).

x = np.linspace(-np.pi*2, np.pi*2, 40)
y = np.cos(x)

x1, x2, y1, y2 = np.stack([x[:-1],  x[1:], y[:-1], y[1:]])[:,np.diff(y < 0)]
xf = x1 + -y1 * (x2 - x1) / (y2 - y1)

i = np.searchsorted(x, xf)
x0 = np.insert(x, i, xf)
y0 = np.insert(y, i, 0)

Drawing the updated arrays as a line graph with masked arrays for non positive and non negative y-coordinates.

plt.figure(figsize=(7,6))
plt.plot(np.ma.masked_array(x0, mask=y0 < 0), np.ma.masked_array(y0, mask=y0 < 0), 'o:')
plt.plot(np.ma.masked_array(x0, mask=y0 > 0), np.ma.masked_array(y0, mask=y0 > 0), 'o:')
plt.show()

example cos graph

To draw the original data with colored lines

plt.figure(figsize=(7,6))
plt.plot(np.ma.masked_array(x0, mask=y0 < 0), np.ma.masked_array(y0, mask=y0 < 0), 'g-')
plt.plot(np.ma.masked_array(x0, mask=y0 > 0), np.ma.masked_array(y0, mask=y0 > 0), 'r-')
plt.plot(np.ma.masked_array(x, mask=y < 0), np.ma.masked_array(y, mask=y < 0), 'g.')
plt.plot(np.ma.masked_array(x, mask=y > 0), np.ma.masked_array(y, mask=y > 0), 'r.')
plt.show()

old data multicolor

Upvotes: 2

Mike
Mike

Reputation: 7203

If you use a scatter plot you can give each point a different color:

x = range(1)
x = range(10)
y = [i - 5 for i in x]
c = [i < 0 for i in y]
plt.scatter(x, y, c=c, s=80)

enter image description here

Upvotes: 5

tillsten
tillsten

Reputation: 14878

I would just make two datasets and setting the right masks. By using that approach i wont have lines between different positive parts.

import matplotlib.pyplot as plt
import numpy as np

signal = 1.2*np.sin(np.linspace(0, 30, 2000))
pos_signal = signal.copy()
neg_signal = signal.copy()

pos_signal[pos_signal <= 0] = np.nan
neg_signal[neg_signal > 0] = np.nan

#plotting
plt.style.use('fivethirtyeight')
plt.plot(pos_signal, color='r')
plt.plot(neg_signal, color='b')
plt.savefig('pos_neg.png', dpi=200)
plt.show()

Example

Upvotes: 17

Adam
Adam

Reputation: 191

You can conditionally plot data in your axes object, using a where like syntax (if you're used to something like Pandas).

ax.plot(x[f(x)>=0], f(x)[f(x)>=0], 'g')
ax.plot(x[f(x)<0],  f(x)[f(x)<0],  'r')

Technically, it's splitting and plotting your data in two sets, but it's fairly compact and nice.

Upvotes: 8

Related Questions