SPY - stock analysis engine with Quantopian zipline run_algorithm with portfolio and benchmark using matplotlib

run_daily_with_zipline.py

#!/usr/bin/env python

import pytz
import zipline
import datetime
import matplotlib.pyplot as plt
import pandas as pd
import analysis_engine.extract as extract
from collections import OrderedDict
from zipline.api import order
from zipline.api import record
from zipline.api import symbol
from zipline.api import set_benchmark

ticker = 'SPY'

extract_res = extract.extract(
    ticker=ticker)

df = extract_res[ticker]['daily']

data = OrderedDict()
use_columns = [
    'date',
    'high',
    'low',
    'open',
    'close',
    'volume'
]
use_columns = df.columns.values
index_col = 'date'
df[index_col] = pd.to_datetime(df[index_col], pytz.UTC)
df.set_index(index_col, drop=False, inplace=True)
data[ticker] = df[use_columns]
print(data)
print(data[ticker].head())
panel = pd.Panel(data)
panel.minor_axis = use_columns
panel.major_axis = panel.major_axis.tz_localize(pytz.utc)


def build_utc_date_from_value(
        date_val):
    return datetime.datetime(
        year=date_val.year,
        month=date_val.month,
        day=date_val.day,
        hour=date_val.hour,
        minute=date_val.minute,
        second=date_val.second,
        tzinfo=pytz.utc)


df_start_date = df['date'][0]
df_end_date = df['date'][-1]
start_date = build_utc_date_from_value(
    date_val=df_start_date)
end_date = build_utc_date_from_value(
    date_val=df_end_date)


def initialize(context):
    set_benchmark(symbol(ticker))


def handle_data(context, data):
    order(symbol(ticker), 10)
    record(ticker=data.current(symbol(ticker), 'price'))


print(start_date)
print(end_date)

perf = zipline.run_algorithm(
    start=start_date,
    end=end_date,
    initialize=initialize,
    capital_base=100000,
    handle_data=handle_data,
    data_frequency='daily',
    data=panel)


benchmark_title = (
    '{} - percent benchmark:\n{} to {}'.format(
        ticker,
        start_date.strftime('%Y-%m-%d %H:%M:%S'),
        end_date.strftime('%Y-%m-%d %H:%M:%S')))
perf.portfolio_value.pct_change().fillna(0).add(1).cumprod().sub(1).plot(
    label='portfolio')
perf.ticker.pct_change().fillna(0).add(1).cumprod().sub(1).plot(
    label=benchmark_title)

plt.legend(loc=2)

plt.show()