wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-07 23:25:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

WIP: Lot of refactoring and bugfixing.

Browse Source
This commit is contained in:
Thomas Wiecki
2012-09-17 18:35:21 -04:00
parent 35a7da6ee7
commit 280f122353
7 changed files with 199 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
zipline/gens/cov.py
+8 -4
View File
@@ -1,5 +1,9 @@
from zipline.gens.transform import EventWindowBatch
from zipline.gens.transform import BatchWindow, batch_transform
class CovEventWindow(EventWindowBatch):
def get_value(self, prices, volumes):
return prices.cov()
class CovEventWindow(BatchWindow):
def get_value(self, data):
return data.cov()
@batch_transform
def cov(data):
return data.cov()
zipline/gens/tradegens.py
+6 -6
View File
@@ -181,11 +181,11 @@ class DataFrameSource(SpecificEquityTrades):
self.data = data
# Unpack config dictionary with default values.
self.count = kwargs.get('count', 500)
self.sids = kwargs.get('sids', [0])
self.start = kwargs.get('start', datetime(1957, 1, 1, 0, tzinfo = pytz.utc))
self.end = kwargs.get('end', datetime(2010, 1, 1, tzinfo=pytz.utc))
self.delta = kwargs.get('delta', timedelta(days = 1))
self.count = kwargs.get('count', len(data))
self.sids = kwargs.get('sids', data.columns)
self.start = kwargs.get('start', data.index[0])
self.end = kwargs.get('end', data.index[-1])
self.delta = kwargs.get('delta', data.index[1]-data.index[0])
# Default to None for event_list and filter.
self.filter = kwargs.get('filter')
@@ -207,7 +207,7 @@ class DataFrameSource(SpecificEquityTrades):
for sid, price in series.iterkv():
event = copy(event)
event['sid'] = 0
event['sid'] = sid
event['price'] = price
yield ndict(event)
zipline/gens/tradesimulation.py
+2 -2
View File
@@ -2,6 +2,7 @@ from logbook import Logger, Processor
from datetime import datetime
from itertools import groupby
from operator import attrgetter
from zipline import ndict
from zipline.utils.timeout import Heartbeat, Timeout
@@ -226,8 +227,7 @@ class AlgorithmSimulator(object):
# Group together events with the same dt field. This depends on the
# events already being sorted.
for date, snapshot in groupby(stream_in, lambda e: e.dt):
for date, snapshot in groupby(stream_in, attrgetter('dt')):
# Set the simulation date to be the first event we see.
# This should only occur once, at the start of the test.
if self.simulation_dt == None:
zipline/gens/transform.py
+55 -32
View File
@@ -217,7 +217,7 @@ class EventWindow(object):
self.assert_well_formed(event)
# Add new event and increment totals.
self.ticks.append(event)
self.ticks.append(deepcopy(event))
# Subclasses should override handle_add to define behavior for
# adding new ticks.
@@ -266,6 +266,7 @@ class EventWindow(object):
# day in the window.
if oldest.time() > newest.time():
trading_days_between -= 1
return trading_days_between >= self.days
def out_of_delta(self, oldest, newest):
@@ -284,62 +285,84 @@ class EventWindow(object):
class BatchWindow(EventWindow):
def __init__(self, func, refresh_period=None, wind_length=None, sids=None):
super(BatchWindow, self).__init__(True, days=wind_length, delta=None)
def __init__(self, func=None, refresh_period=None, days=None, sids=None):
super(BatchWindow, self).__init__(True, days=days, delta=None)
self.func = func
self.sids = sids
self.refresh_period = refresh_period
self.wind_length = wind_length
self.days = days
self.last_calc = False
self.full = False
self.last_refresh = None
self.updated = False
self.data = None
# def handle_data(self, data):
# """
# New method to handle a data frame as sent to the algorithm's handle_data
# method.
# """
# dts = [data[sid].datetime for sid in self.sids]
# prices = [data[sid].price for sid in self.sids]
# volumes = [data[sid].volume for sid in self.sids]
def handle_data(self, data):
"""
New method to handle a data frame as sent to the algorithm's handle_data
method.
"""
# extract dates
dts = [data[sid].datetime for sid in self.sids]
# we have to provide the event with a dt. This is only for
# checking if the event is outside the window or not so a
# couple of seconds shouldn't matter
data.dt = max(dts)
# price_df = pd.DataFrame(prices, columns=self.sids, index=dts)
# volume_df = pd.DataFrame(volumes, columns=self.sids, index=dts)
# append data frame to window
self.update(data)
# event = ndict({
# 'dt' : max(dts),
# 'prices': price_df,
# 'volumes': volume_df,
# })
# self.update(event)
# return newly computed or cached value
return self.compute()
def handle_add(self, event):
import pdb; pdb.set_trace()
if not self.last_calc:
self.last_calc = event.dt
if not self.last_refresh:
self.last_refresh = event.dt
return
age = event.dt - self.last_refresh
if age.days >= self.refresh_period:
self.prices = pd.concat(self.ticks.prices)
self.volumes = pd.concat(self.ticks.volumes)
# create Series price object
data_sids = {}
for sid in self.sids:
dts = [tick[sid].dt for tick in self.ticks]
prices = [tick[sid].price for tick in self.ticks]
data_sids[sid] = pd.Series(prices, index=dts)
# concatenate different sids into one df
self.data = pd.concat(data_sids, axis=1)
self.updated = True
self.last_refresh = event.dt
else:
self.updated = False
self.last_refresh = event.dt
def handle_remove(self, event):
# since an event is expiring, we know the window is full
self.full = True
def __call__(self, *args, **kwargs):
def get_value(self, *args, **kwargs):
raise NotImplementedError("Either overwrite get_value or provide a func argument.")
def compute(self, *args, **kwargs):
if self.data is None:
return False
if self.updated:
self.cached = self.get_value(self.prices, self.volumes, *args, **kwargs)
if self.func is not None:
# user supplied function
self.cached = self.func(self.data, *args, **kwargs)
else:
# assume inheritance
self.cached = self.get_value(self.data, *args, **kwargs)
return self.cached
# decorator for BatchWindow
def batch_transform(func):
def create_transform(*args, **kwargs):
return BatchWindow(*args, func=func, **kwargs)
return create_transform
zipline/optimize/algorithms.py
+111 -41
View File
@@ -48,6 +48,7 @@ class BuySellAlgorithm(object):
self.portfolio = portfolio
def handle_data(self, frame):
print frame.sid
order_size = self.buy_or_sell * (self.amount - (self.offset**2))
self.order(self.sid, order_size)
@@ -62,40 +63,114 @@ class BuySellAlgorithm(object):
def get_sid_filter(self):
return [self.sid]
# Algorithm base class, user algorithms inherit from this as they
# don't want to have to copy and know about set_order and
# set_portfolio
class TradingAlgorithm(object):
def _setup(self):
assert hasattr(self, 'source'), 'source not set.'
assert hasattr(self, 'sids'), "sids not set."
environment = create_trading_environment(start=self.data.index[0], end=self.data.index[-1])
class TradingAlgorithm(object):
"""
Base class for trading algorithms. Inherit and overload handle_data(data).
A new algorithm could look like this:
```
class MyAlgo(TradingAlgorithm):
def initialize(amount):
self.amount = amount
def handle_data(data):
sid = self.sids[0]
self.order(sid, amount)
```
To then run this algorithm:
>>> my_algo = MyAlgo(100)
>>> stats = my_algo.run(data)
"""
def __init__(self, sids, *args, **kwargs):
"""
Initialize sids and other state variables.
Calls user-defined initialize and forwarding *args and **kwargs.
"""
self.sids = sids
self.done = False
self.order = None
self.frame_count = 0
self.portfolio = None
self.registered_transforms = {}
# call to user-defined initialize method
self.initialize(*args, **kwargs)
def _create_simulator(self, source):
"""
Create trading environment, transforms and SimulatedTrading object.
Gets called by self.run(data).
"""
environment = create_trading_environment(start=source.data.index[0], end=source.data.index[-1])
# Create transforms by wrapping them into StatefulTransforms
transforms = []
if hasattr(self, 'registered_transforms'):
for namestring, trans_descr in self.registered_transforms.iteritems():
sf = StatefulTransform(
trans_descr['class'],
*trans_descr['args'],
**trans_descr['kwargs']
)
sf.namestring = namestring
for namestring, trans_descr in self.registered_transforms.iteritems():
sf = StatefulTransform(
trans_descr['class'],
*trans_descr['args'],
**trans_descr['kwargs']
)
sf.namestring = namestring
transforms.append(sf)
transforms.append(sf)
self.simulated_trading = SimulatedTrading(
[self.source],
# SimulatedTrading is the main class handling data streaming,
# application of transforms and calling of the user algo.
return SimulatedTrading(
[source],
transforms,
self,
environment,
FixedSlippage()
)
def run(self, data):
"""
Run the algorithm.
:Arguments:
data : pandas.DataFrame
* columns must consist of ints representing the different sids
* index must be TimeStamps
* array contents should be price
:Returns:
daily_stats : pandas.DataFrame
Daily performance metrics such as returns, alpha etc.
"""
assert isinstance(data, pd.DataFrame)
assert isinstance(data.index, pd.Timeseries)
source = DataFrameSource(data, sids=self.sids)
# create transforms and zipline
simulated_trading = self._create_simulator(source)
# loop through simulated_trading, each iteration returns a
# perf ndict
perfs = []
for perf in simulated_trading:
#from nose.tools import set_trace; set_trace()
perfs.append(perf)
#perfs = list(self.simulated_trading)
# convert perf ndict to pandas dataframe
daily_stats = self._create_daily_stats(perfs)
return daily_stats
def _create_daily_stats(self, perfs):
# create daily stats dataframe
# create daily and cumulative stats dataframe
daily_perfs = []
cum_perfs = []
for perf in perfs:
@@ -109,21 +184,23 @@ class TradingAlgorithm(object):
return daily_stats
def run(self, data, compute_risk_metrics=False):
self.source = DataFrameSource(data, sids=self.sids)
self.data = data
self._setup()
def add_transform(self, transform_class, tag, *args, **kwargs):
"""Add a single-sid, sequential transform to the model.
# drain simulated_trading
perfs = []
for perf in self.simulated_trading:
#from nose.tools import set_trace; set_trace()
perfs.append(perf)
:Arguments:
transform_class : class
Which transform to use. E.g. mavg.
tag : str
How to name the transform. Can later be access via:
data[sid].tag()
#perfs = list(self.simulated_trading)
Extra args and kwargs will be forwarded to the transform
instantiation.
daily_stats = self._create_daily_stats(perfs)
return daily_stats
"""
self.registered_transforms[tag] = {'class': transform_class,
'args': args,
'kwargs': kwargs}
def set_portfolio(self, portfolio):
self.portfolio = portfolio
@@ -137,19 +214,12 @@ class TradingAlgorithm(object):
def set_logger(self, logger):
self.logger = logger
def initialize(self):
def initialize(self, *args, **kwargs):
pass
def set_slippage_override(self, slippage_callable):
pass
def add_transform(self, transform_class, tag, *args, **kwargs):
if not hasattr(self, 'registered_transforms'):
self.registered_transforms = {}
self.registered_transforms[tag] = {'class': transform_class,
'args': args,
'kwargs': kwargs}
class BuySellAlgorithmNew(TradingAlgorithm):
zipline/optimize/example.py
+16 -12
View File
@@ -6,6 +6,7 @@ import numpy as np
import matplotlib.pyplot as plt
import cProfile
from zipline.gens.mavg import MovingAverage
from zipline.gens.cov import CovEventWindow, cov
from zipline.optimize.algorithms import TradingAlgorithm
from datetime import timedelta
@@ -15,15 +16,9 @@ from datetime import timedelta
class DMA(TradingAlgorithm):
"""Dual Moving Average algorithm.
"""
def __init__(self, sids, amount=100, short_window=20, long_window=40):
self.sids = sids
self.amount = amount
self.done = False
self.order = None
self.frame_count = 0
self.portfolio = None
def initialize(self, amount=100, short_window=20, long_window=40):
self.orders = []
self.amount = amount
self.prices = []
self.events = 0
@@ -33,15 +28,22 @@ class DMA(TradingAlgorithm):
self.add_transform(MovingAverage, 'short_mavg', ['price'],
market_aware=True,
days=short_window) #timedelta(days=int(short_window)))
days=short_window)
self.add_transform(MovingAverage, 'long_mavg', ['price'],
market_aware=True,
days=long_window) #timedelta(days=int(long_window)))
days=long_window)
self.cov = CovEventWindow(sids=self.sids, refresh_period=1, days=5)
self.cov2 = cov(sids=self.sids, refresh_period=1, days=5)
def handle_data(self, data):
self.events += 1
cov = self.cov.handle_data(data)
cov = self.cov2.handle_data(data)
print cov
for sid in self.sids:
# access transforms via their user-defined tag
if (data[sid].short_mavg['price'] > data[sid].long_mavg['price']) and not self.invested[sid]:
@@ -86,8 +88,8 @@ def load_close_px(indexes=None, stocks=None):
def run((short_window, long_window)):
#data = pd.DataFrame.from_csv('SP500.csv')
data = load_close_px()
myalgo = DMA([0], amount=100, short_window=short_window, long_window=long_window)
data = pd.DataFrame.from_csv('aapl.csv') #load_close_px()
myalgo = DMA([0, 1], amount=100, short_window=short_window, long_window=long_window)
stats = myalgo.run(data)
stats['sw'] = short_window
stats['lw'] = long_window
@@ -153,3 +155,5 @@ def plot_returns(port_returns, bmk_returns):
cum_bmk.plot(label='Benchmark')
plt.title('Portfolio performance')
plt.legend(loc='best')
print run((10, 20))
zipline/test_algorithms.py
+1 -1
View File
@@ -98,7 +98,7 @@ class TestAlgorithm():
def set_slippage_override(self, slippage_callable):
pass
#
class HeavyBuyAlgorithm():
"""
This algorithm will send a specified number of orders, to allow unit tests