|
| 1 | +""" |
| 2 | +Illustrate the rec array utility funcitons by loading prices from a |
| 3 | +csv file, computing the daily returns, appending the results to the |
| 4 | +record arrays, joining on date |
| 5 | +""" |
| 6 | +import urllib |
| 7 | +import numpy as np |
| 8 | +import matplotlib.pyplot as plt |
| 9 | +import matplotlib.mlab as mlab |
| 10 | + |
| 11 | +# grab the price data off yahoo |
| 12 | +u1 = urllib.urlretrieve('http://ichart.finance.yahoo.com/table.csv?s=AAPL&d=9&e=14&f=2008&g=d&a=8&b=7&c=1984&ignore=.csv') |
| 13 | +u2 = urllib.urlretrieve('http://ichart.finance.yahoo.com/table.csv?s=GOOG&d=9&e=14&f=2008&g=d&a=8&b=7&c=1984&ignore=.csv') |
| 14 | + |
| 15 | +# load the CSV files into record arrays |
| 16 | +r1 = mlab.csv2rec(file(u1[0])) |
| 17 | +r2 = mlab.csv2rec(file(u2[0])) |
| 18 | + |
| 19 | +# compute the daily returns and add these columns to the arrays |
| 20 | +gains1 = np.zeros_like(r1.adj_close) |
| 21 | +gains2 = np.zeros_like(r2.adj_close) |
| 22 | +gains1[1:] = np.diff(r1.adj_close)/r1.adj_close[:-1] |
| 23 | +gains2[1:] = np.diff(r2.adj_close)/r2.adj_close[:-1] |
| 24 | +r1 = mlab.rec_append_fields(r1, 'gains', gains1) |
| 25 | +r2 = mlab.rec_append_fields(r2, 'gains', gains2) |
| 26 | + |
| 27 | +# now join them by date; the default postfixes are 1 and 2. The |
| 28 | +# default jointype is inner so it will do an intersection of dates and |
| 29 | +# drop the dates in AAPL which occurred before GOOG started trading in |
| 30 | +# 2004. r1 and r2 are reverse ordered by date since Yahoo returns |
| 31 | +# most recent first in the CSV files, but rec_join will sort by key so |
| 32 | +# r below will be properly sorted |
| 33 | +r = mlab.rec_join('date', r1, r2) |
| 34 | + |
| 35 | + |
| 36 | +# long appl, short goog |
| 37 | +g = r.gains1 - r.gains2 |
| 38 | +tr = (1 + g).cumprod() # the total return |
| 39 | + |
| 40 | +# plot the return |
| 41 | +fig, ax = plt.subplots() |
| 42 | +ax.plot(r.date, tr) |
| 43 | +ax.set_title('total return: long APPL, short GOOG') |
| 44 | +ax.grid() |
| 45 | +fig.autofmt_xdate() |
| 46 | +plt.show() |
0 commit comments