pyCthon
pyCthon

Reputation: 12361

multi colored plots in matplotlib plt based on certain properties of data

I want to color the line in a plot based on the following of a data set on the y axis.

if data > 0:
  color = 'r'
if data = 0:
  color = 'g'
if data < 0:
  color = 'b'

Unfortunately I only know how to color the entire data set one color. I also couldn't find anything on the web. I'm assuming there is a way to do this without breaking up the dataset for every time the color changes.

Below is an example of plotting the data with just one color.

import matplotlib.pyplot as plt
import numpy as np

# Simple data 
x = np.linspace(0, 2 * np.pi, 400)
data = np.sin(x ** 2)

#plot
f, ax = plt.subplots()
ax.plot(x, data, color='r')

plt.show()

Upvotes: 1

Views: 1754

Answers (1)

user3030010
user3030010

Reputation: 1867

The color parameter actually can take a list as an argument. For example, here's a simple bit of code that sets up a list of colors based on whether the data is positive or negative:

colors = []
for item in data:
    if item < 0:
        colors.append('r')
    else:
        colors.append('g')

then simply:

ax.bar(x, data, color=colors)

Edit: So I tested it, and it appears that my answer is only applicable for bar graphs. I couldn't find anything in the matplotlib documentation that seemed to indicate that coloring a line plot with multiple colors was possible. I did, however find this site, which I believe has the information you want. The guy there defines his own function to achieve it.

Using the file at my link, here is an equivalent version for a line graph:

cmap = ListedColormap(['r', 'g']) # use the colors red and green
norm = BoundaryNorm([-1000,0,1000], cmap.N) # map red to negative and green to positive
                                            # this may work with just 0 in the list
fig, axes = plt.subplots()
colorline(x, data, data, cmap=cmap, norm=norm)

plt.xlim(x.min(), x.max())
plt.ylim(data.min(), data.max())

plt.show()

The last three arguments of colorline here tell it the color data and how to map it.

Upvotes: 2

Related Questions