Bitcoin - MA cross over strategy - daily - mtgoxBTCUSD

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pytz
from collections import OrderedDict
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from zipline.algorithm import TradingAlgorithm
from zipline.transforms import MovingAverage, batch_transform
from zipline.utils.factory import load_from_yahoo

class DualMovingAverage(TradingAlgorithm):
    """Dual Moving Average Crossover algorithm.

    This algorithm buys btc once its short moving average crosses
    its long moving average (indicating upwards momentum) and sells
    its shares once the averages cross again (indicating downwards
    momentum).

    """
    def initialize(self):
        # Add 2 mavg transforms, one with a long window, one
        # with a short window.
        self.add_transform(MovingAverage, 'short_mavg', ['price'],
                           window_length=10)

        self.add_transform(MovingAverage, 'long_mavg', ['price'],
                           window_length=20)

        # To keep track of whether we invested in the stock or not
        self.invested = False

        self.short_mavg = []
        self.long_mavg = []
        self.buy_orders = []
        self.sell_orders = []

    def handle_data(self, data):
        if (data['BTC'].short_mavg['price'] > data['BTC'].long_mavg['price']) and not self.invested:
            self.order('BTC', 200)
            self.invested = True
            self.buy_orders.append(data['BTC'].datetime)
            print "{dt}: Buying 100 BTC shares.".format(dt=data['BTC'].datetime)
        elif (data['BTC'].short_mavg['price'] < data['BTC'].long_mavg['price']) and self.invested:
            self.order('BTC', -200)
            self.invested = False
            self.sell_orders.append(data['BTC'].datetime)
            print "{dt}: Selling 100 BTC shares.".format(dt=data['BTC'].datetime)

        # Save mavgs for later analysis.
        self.short_mavg.append(data['BTC'].short_mavg['price'])
        self.long_mavg.append(data['BTC'].long_mavg['price'])


def rename_col(df):
    df = df.rename(columns={'Weighted Price': 'price',
                            'High':'high',
                            'Low':'low',
                            'Open':'open',
                            'Close':'close',
                            'Volume (BTC)':'volume',
                            'Volume (Currency)':'volume_usd'
                           })
    df = df.fillna(method='ffill')
    #df = df[['price', 'sid']]
    return df


def display_dataframe(df):
    print(df.head())
    print(df)
    print(df.tail())

if __name__ == '__main__':
    data = OrderedDict()
    #data['BTC'] = pd.DataFrame.from_csv('http://www.quandl.com/api/v1/datasets/BITCOIN/MTGOXUSD.csv?trim_start=2012-01-01&sort_order=desc')
    #data['BTC'] = pd.DataFrame.from_csv('http://www.quandl.com/api/v1/datasets/BITCOIN/MTGOXUSD.csv?trim_start=2012-01-01&sort_order=desc')
    data['BTC'] = pd.DataFrame.from_csv('MTGOXUSD.csv')
    data['BTC'] = data['BTC'].tz_localize(pytz.utc)

    data['BTC'] = rename_col(data['BTC'])

    data = pd.DataFrame({key: d['price'] for key, d in data.iteritems()})

    print("="*10 + " Display price " + "="*10)
    data.plot()
    plt.show()

    #print(data)
    #display_dataframe(data['BTC'])

    print("="*10 + " Running strategy " + "="*10)
    dma = DualMovingAverage()
    results = dma.run(data)

    print("="*10 + " Display results " + "="*10)

    dma.short_mavg.append(0)
    dma.short_mavg.append(0)
    dma.long_mavg.append(0)
    dma.long_mavg.append(0)
    print len(dma.long_mavg)
    print len(dma.short_mavg)
    print len(data['BTC'])

    ax1 = plt.subplot(211)
    data['short'] = dma.short_mavg
    data['long'] = dma.long_mavg
    data[['BTC', 'short', 'long']].plot(ax=ax1)
    plt.plot(dma.buy_orders, data['short'].ix[dma.buy_orders], '^', c='m', markersize=10, label='buy')
    plt.plot(dma.sell_orders, data['short'].ix[dma.sell_orders], 'v', c='k', markersize=10, label='sell')
    plt.legend(loc=0)

    ax2 = plt.subplot(212)
    results.portfolio_value.plot(ax=ax2)

    plt.show()