CODEWITHSUNDEEP
pythonmatplotlib
jtitusj
jtitusj

Reputation: 3086

Changing line plot color by date

I wanted to plot a data which has datetime values for the x axis and another set of values as y. As an example, I will use the example from matplotlib where y in this case are stock prices. Here is the code for that.

import matplotlib.pyplot as plt
from matplotlib.finance import quotes_historical_yahoo_ochl
from matplotlib.dates import YearLocator, MonthLocator, DateFormatter
import datetime
date1 = datetime.date(1995, 1, 1)
date2 = datetime.date(2004, 4, 12)

years = YearLocator()   # every year
months = MonthLocator()  # every month
yearsFmt = DateFormatter('%Y')

quotes = quotes_historical_yahoo_ochl('INTC', date1, date2)
if len(quotes) == 0:
    raise SystemExit

dates = [q[0] for q in quotes]
opens = [q[1] for q in quotes]

fig, ax = plt.subplots()
ax.plot_date(dates, opens, '-')

# format the ticks
ax.xaxis.set_major_locator(years)
ax.xaxis.set_major_formatter(yearsFmt)
ax.xaxis.set_minor_locator(months)
ax.autoscale_view()

# format the coords message box
def price(x):
   return '$%1.2f' % x
ax.fmt_xdata = DateFormatter('%Y-%m-%d')
ax.fmt_ydata = price
ax.grid(True)

fig.autofmt_xdate()
plt.show()

Now, what I want to do is color each value in the graph based on some criterion. For simplicity's sake, let's say that the criterion in the case of the example is based on the year. That is, prices belonging to the same year will be colored the same. How would I do that? Thanks!

Upvotes: 1

Views: 3208

Answers (2)

Ed Smith
Ed Smith

Reputation: 13206

You can use numpy arrays with masks over the range you want (in this case a year). In order to use the inbuilt YearLocator function from your example, you need to plot the graph first and set the ticks, then remove and replace with the range per year, from your example,

import matplotlib.pyplot as plt
from matplotlib.finance import quotes_historical_yahoo_ochl
from matplotlib.dates import YearLocator, MonthLocator, DateFormatter
import datetime
import numpy 

date1 = datetime.date(1995, 1, 1)
date2 = datetime.date(2004, 4, 12)

years = YearLocator()   # every year
months = MonthLocator()  # every month
yearsFmt = DateFormatter('%Y')

quotes = quotes_historical_yahoo_ochl('INTC', date1, date2)
if len(quotes) == 0:
    raise SystemExit

dates = np.array([q[0] for q in quotes])
opens = np.array([q[1] for q in quotes])

fig, ax = plt.subplots()
l = ax.plot_date(dates, opens, '-')

# format the ticks
ax.xaxis.set_major_locator(years)
ax.xaxis.set_major_formatter(yearsFmt)
ax.xaxis.set_minor_locator(months)
ax.autoscale_view()

l[0].remove()
py = years()[0]
for year in years()[1:]:
    mask = (py < dates) & (dates < year)
    ax.plot_date(dates[mask], opens[mask], '-')
    py = year

# format the coords message box
def price(x):
   return '$%1.2f' % x
ax.fmt_xdata = DateFormatter('%Y-%m-%d')
ax.fmt_ydata = price
ax.grid(True)

fig.autofmt_xdate()
plt.show()

which gives,

enter image description here

Upvotes: 3

James
James

Reputation: 36623

The way I typically do this is by using a for loop to plot different sections of the data, coloring each section as I go. In your example, this section:

fig, ax = plt.subplots()
ax.plot_date(dates, opens, '-')

becomes:

# import the colormaps
from maplotlib import cm

fig, ax = plt.subplots()

for y in years:
    y_indices = [i for i in range(len(dates)) if dates[i].year==y]

    # subset the data, there are better ways to do this
    sub_dates = [dates[i] for i in y_indices]
    sub_opens = [opens[i] for i in y_indices]

    # plot each section of data, using a colormap to change the color for
    # each iteration.
    ax.plot_date(sub_dates, sub_opens, '-', linecolor=cm.spring((y-2000)/10.0)

Upvotes: 2

Related Questions