fin/momentum.py at master · vsmolyakov/fin · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

from zipline.api import order_target, order, record, symbol, history, add_history
import matplotlib.pyplot as plt

from datetime import datetime
import pytz

from zipline.algorithm import TradingAlgorithm
from zipline.utils.factory import load_from_yahoo
from zipline.utils.factory import load_bars_from_yahoo

import numpy as np


def initialize(context):

    #register 2 histories to track daily prices
    add_history(100, '1d', 'price')
    add_history(300, '1d', 'price')

    context.i = 0
    context.invested = False

def handle_data(context, data):

    #trading algorithm (executed on every event)

    #skip first 300 days to get full windows
    context.i += 1
    if context.i < 300:
        return

    #compute short and long moving averages:
    short_mavg = history(100, '1d', 'price').mean()
    long_mavg = history(300, '1d', 'price').mean()

    buy = False
    sell = False

    #trading logic
    if (short_mavg[0] > long_mavg[0]) and not context.invested:
        buy = True
        context.invested = True
        order_target(symbol('AAPL'), 100)
    elif (short_mavg[0] < long_mavg[0]) and context.invested:
        sell = True
        context.invested = False
        order_target(symbol('AAPL'), -100)

    #save values for plotting
    record(AAPL = data[symbol('AAPL')].price,
           short_mavg = short_mavg[0],
           long_mavg = long_mavg[0],
           buy=buy,
           sell=sell)


def analyze(context=None, results=None, benchmark=None):

    hist_size = 300

    f, (ax1, ax2) = plt.subplots(2, sharex = True)
    ax1.plot(results.portfolio_value[hist_size:], linewidth = 2.0, label = 'porfolio')
    ax1.set_title('Dual Moving Average Strategy')
    ax1.set_ylabel('Portfolio value (USD)')
    ax1.legend(loc=0)
    ax1.grid(True)

    ax2.plot(results['AAPL'][hist_size:], linewidth = 2.0, label = 'AAPL')
    ax2.plot(results['short_mavg'][hist_size:], color = 'r', linestyle = '-', linewidth = 2.0, label = 'short mavg')
    ax2.plot(results['long_mavg'][hist_size:], color = 'g', linestyle = '-', linewidth = 2.0, label = 'long mavg')
    ax2.set_ylabel('AAPL price (USD)')
    ax2.legend(loc=0)
    ax2.grid(True)

    plt.show()


if __name__ == "__main__":

    plt.close('all')

    #load data
    #year, month, day, hour, minute, second, microsecond
    start = datetime(2012, 1, 1, 0, 0, 0, 0, pytz.utc)
    end = datetime(2015, 1, 1, 0, 0, 0, 0, pytz.utc)

    data = load_from_yahoo(stocks=['AAPL'], indexes={}, start=start, end=end)
    #data = load_bars_from_yahoo(stocks=['AAPL'], start=start, end=end)

    data.AAPL.tail(5)

    #run the algorithm
    algo = TradingAlgorithm(initialize=initialize, handle_data = handle_data)
    results = algo.run(data)

    #generate plots
    analyze(results = results)

    idx = np.where(results['buy'] == True)
    print "number of buys: %d" % np.size(idx[0])

    ##notes
    # see package requirements at:
    # etc/requirements.txt