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

Better unittest coverage. DataFrameSource is now filtering sids. Fixed outstanding issues.

Browse Source
This commit is contained in:
Thomas Wiecki
2012-09-21 11:59:46 -04:00
parent 2dbf905d43
commit 4324f95f36
10 changed files with 143 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_sources.py
+17 -8
View File
@@ -1,13 +1,22 @@
from unittest2 import TestCase
import zipline.utils.factory as factory
from zipline.gens.tradegens import DataFrameSource
def test_dataframe_source():
source, df = factory.create_test_df_source()
class TestDataFrameSource(TestCase):
def test_streaming_of_df(self):
source, df = factory.create_test_df_source()
for expected_dt, expected_price in df.iterrows():
sid0 = source.next()
sid1 = source.next()
for expected_dt, expected_price in df.iterrows():
sid0 = source.next()
sid1 = source.next()
assert expected_dt == sid0.dt == sid1.dt
assert expected_price[0] == sid0.price
assert expected_price[1] == sid1.price
assert expected_dt == sid0.dt == sid1.dt
assert expected_price[0] == sid0.price
assert expected_price[1] == sid1.price
def test_sid_filtering(self):
_, df = factory.create_test_df_source()
source = DataFrameSource(df, sids=[0])
assert 1 not in [event.sid for event in source], \
"DataFrameSource should only stream selected sid 0, not sid 1."
tests/test_transforms.py
+8 -49
View File
@@ -10,13 +10,14 @@ from zipline.utils.test_utils import setup_logger
from zipline.utils.date_utils import utcnow
from zipline.gens.tradegens import SpecificEquityTrades
from zipline.gens.transform import StatefulTransform, EventWindow, BatchTransform, batch_transform
from zipline.gens.transform import StatefulTransform, EventWindow
from zipline.gens.vwap import VWAP
from zipline.gens.mavg import MovingAverage
from zipline.gens.stddev import MovingStandardDev
from zipline.gens.returns import Returns
import zipline.utils.factory as factory
from zipline import TradingAlgorithm
from zipline.test_algorithms import BatchTransformAlgorithm
def to_dt(msg):
return ndict({'dt': msg})
@@ -288,36 +289,7 @@ class FinanceTransformsTestCase(TestCase):
############################################################
# Test BatchTransform
class NoopBatchTransform(BatchTransform):
def get_value(self, data):
return data.price
@batch_transform
def noop_batch_decorator(data):
return data.price
class BatchTransformAlgorithm(TradingAlgorithm):
def initialize(self, *args, **kwargs):
self.history_class = []
self.history_decorator = []
self.days = 3
self.noop_class = NoopBatchTransform(sids=[0, 1],
market_aware=False,
refresh_period=2,
delta=timedelta(days=self.days))
self.noop_decorator = noop_batch_decorator(sids=[0, 1],
market_aware=False,
refresh_period=2,
delta=timedelta(days=self.days))
def handle_data(self, data):
window_class = self.noop_class.handle_data(data)
window_decorator = self.noop_decorator.handle_data(data)
self.history_class.append(window_class)
self.history_decorator.append(window_decorator)
class BatchTransformTestCase():
class BatchTransformTestCase(TestCase):
def setUp(self):
setup_logger(self)
self.source, self.df = factory.create_test_df_source()
@@ -329,21 +301,8 @@ class BatchTransformTestCase():
assert algo.history_class[:2] == algo.history_decorator[:2] == [None, None], "First two iterations should return None"
# test overloaded class
# every 2nd event should be identical because of refresh_period=2
# not sure why actual length gets up to 4, bug in EventWindow?
assert np.all(algo.history_class[2][0].values == [2, 4, 6])
assert np.all(algo.history_class[2][1].values == [3, 5, 7])
assert np.all(algo.history_class[3][0].values == [2, 4, 6])
assert np.all(algo.history_class[3][1].values == [3, 5, 7])
assert np.all(algo.history_class[4][0].values == [4, 6, 8, 10])
assert np.all(algo.history_class[4][1].values == [5, 7, 9, 11])
# test decorator
assert np.all(algo.history_decorator[2][0].values == [2, 4, 6])
assert np.all(algo.history_decorator[2][1].values == [3, 5, 7])
assert np.all(algo.history_decorator[3][0].values == [2, 4, 6])
assert np.all(algo.history_decorator[3][1].values == [3, 5, 7])
assert np.all(algo.history_decorator[4][0].values == [4, 6, 8, 10])
assert np.all(algo.history_decorator[4][1].values == [5, 7, 9, 11])
for test_history in [algo.history_class, algo.history_decorator]:
self.assertTrue(np.all(test_history[2].values.flatten() == range(4, 10)))
self.assertTrue(np.all(test_history[3].values.flatten() == range(4, 10)))
self.assertTrue(np.all(test_history[4].values.flatten() == range(6, 14)))
zipline/__init__.py
+1 -3
View File
@@ -6,9 +6,7 @@ Zipline
# it is a place to expose the public interfaces.
from utils.protocol_utils import ndict
from algorithm import TradingAlgorithm
__all__ = [
ndict,
TradingAlgorithm
ndict
]
zipline/algorithm.py
+38 -17
View File
@@ -9,8 +9,8 @@ from zipline.finance.slippage import FixedSlippage
class TradingAlgorithm(object):
"""
Base class for trading algorithms. Inherit and overload handle_data(data).
"""Base class for trading algorithms. Inherit and overload
initialize() and handle_data(data).
A new algorithm could look like this:
```
@@ -22,7 +22,7 @@ class TradingAlgorithm(object):
sid = self.sids[0]
self.order(sid, amount)
```
To then run this algorithm:
To then to run this algorithm:
>>> my_algo = MyAlgo(100, sids=[0])
>>> stats = my_algo.run(data)
@@ -45,13 +45,13 @@ class TradingAlgorithm(object):
# call to user-defined initialize method
self.initialize(*args, **kwargs)
def _create_simulator(self, source):
def _create_simulator(self, start, end):
"""
Create trading environment, transforms and SimulatedTrading object.
Gets called by self.run().
"""
environment = create_trading_environment(start=source.data.index[0], end=source.data.index[-1])
environment = create_trading_environment(start=start, end=end)
# Create transforms by wrapping them into StatefulTransforms
transforms = []
@@ -68,39 +68,59 @@ class TradingAlgorithm(object):
# SimulatedTrading is the main class handling data streaming,
# application of transforms and calling of the user algo.
return SimulatedTrading(
[source],
self.sources,
transforms,
self,
environment,
FixedSlippage()
)
def run(self, source):
"""
Run the algorithm.
def run(self, source, start=None, end=None):
"""Run the algorithm.
:Arguments:
data : zipline source or pandas.DataFrame
pandas.DataFrame must have the following structure:
* column names must consist of ints representing the different sids
* index must be TimeStamps
* array contents should be price
source : can be either:
- pandas.DataFrame
- zipline source
- list of zipline sources
If pandas.DataFrame is provided, it must have the
following structure:
* column names must consist of ints representing the
different sids
* index must be DatetimeIndex
* array contents should be price info.
:Returns:
daily_stats : pandas.DataFrame
Daily performance metrics such as returns, alpha etc.
"""
if isinstance(source, pd.DataFrame):
if isinstance(source, (list, tuple)):
assert start is not None and end is not None, \
"When providing a list of sources, start and end date have to be specified."
elif isinstance(source, pd.DataFrame):
assert isinstance(source.index, pd.tseries.index.DatetimeIndex)
# if DataFrame provided, wrap in DataFrameSource
source = DataFrameSource(source, sids=self.sids)
# If values not set, try to extract from source.
if start is None:
start = source.start
if end is None:
end = source.end
if not isinstance(source, (list, tuple)):
self.sources = [source]
else:
self.sources = source
# create transforms and zipline
simulated_trading = self._create_simulator(source)
self.simulated_trading = self._create_simulator(start=start, end=end)
# loop through simulated_trading, each iteration returns a
# perf ndict
perfs = list(simulated_trading)
perfs = list(self.simulated_trading)
# convert perf ndict to pandas dataframe
daily_stats = self._create_daily_stats(perfs)
@@ -123,6 +143,7 @@ class TradingAlgorithm(object):
return daily_stats
def add_transform(self, transform_class, tag, *args, **kwargs):
"""Add a single-sid, sequential transform to the model.
zipline/gens/tradegens.py
+24 -11
View File
@@ -9,6 +9,7 @@ from itertools import chain, cycle, ifilter, izip, repeat
from datetime import datetime, timedelta
import pandas as pd
from copy import copy
import numpy as np
from zipline.protocol import DATASOURCE_TYPE
from zipline.utils import ndict
@@ -77,17 +78,31 @@ class SpecificEquityTrades(object):
# We shouldn't get any positional arguments.
assert len(args) == 0
# Unpack config dictionary with default values.
self.count = kwargs.get('count', 500)
self.sids = kwargs.get('sids', [1, 2])
self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc))
self.delta = kwargs.get('delta', timedelta(minutes = 1))
self.concurrent = kwargs.get('concurrent', False)
# Default to None for event_list and filter.
self.event_list = kwargs.get('event_list')
self.filter = kwargs.get('filter')
if self.event_list is not None:
# If event_list is provided, extract parameters from there
# This isn't really clean and ultimately I think this
# class should serve a single purpose (either take an
# event_list or autocreate events).
self.count = kwargs.get('count', len(self.event_list))
self.sids = kwargs.get('sids', np.unique([event.sid for event in self.event_list]).tolist())
self.start = kwargs.get('start', self.event_list[0].dt)
self.end = kwargs.get('start', self.event_list[-1].dt)
self.delta = kwargs.get('delta', self.event_list[1].dt - self.event_list[0].dt)
self.concurrent = kwargs.get('concurrent', False)
else:
# Unpack config dictionary with default values.
self.count = kwargs.get('count', 500)
self.sids = kwargs.get('sids', [1, 2])
self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc))
self.delta = kwargs.get('delta', timedelta(minutes = 1))
self.concurrent = kwargs.get('concurrent', False)
# Hash_value for downstream sorting.
self.arg_string = hash_args(*args, **kwargs)
@@ -188,9 +203,6 @@ class DataFrameSource(SpecificEquityTrades):
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')
# Hash_value for downstream sorting.
self.arg_string = hash_args(data, **kwargs)
@@ -214,4 +226,5 @@ class DataFrameSource(SpecificEquityTrades):
yield ndict(event)
# Return the filtered event stream.
return _generator()
drop_sids = lambda x: x.sid in self.sids
return ifilter(drop_sids, _generator())
zipline/gens/tradesimulation.py
+2 -4
View File
@@ -194,8 +194,6 @@ class AlgorithmSimulator(object):
'filled' : 0
})
log.debug(order)
# Tell the user if they try to buy 0 shares of something.
if order.amount == 0:
zero_message = "Requested to trade zero shares of {sid}".format(
@@ -296,8 +294,8 @@ class AlgorithmSimulator(object):
self.snapshot_dt = date
start_tic = datetime.now()
#with self.heartbeat_monitor:
self.algo.handle_data(self.universe)
with self.heartbeat_monitor:
self.algo.handle_data(self.universe)
stop_tic = datetime.now()
# How long did you take?
zipline/optimize/algorithms.py
+3 -64
View File
@@ -1,9 +1,9 @@
from logbook import Logger
from zipline import TradingAlgorithm
from zipline.algorithm import TradingAlgorithm
logger = Logger('Algo')
class BuySellAlgorithm(object):
class BuySellAlgorithm(TradingAlgorithm):
"""Algorithm that buys and sells alternatingly. The amount for
each order can be specified. In addition, an offset that will
quadratically reduce the amount that will be bought can be
@@ -15,69 +15,11 @@ class BuySellAlgorithm(object):
"""
def __init__(self, sid, amount, offset):
self.sid = sid
def initialize(self, amount=100, offset=0):
self.amount = amount
self.incr = 0
self.done = False
self.order = None
self.frame_count = 0
self.portfolio = None
self.buy_or_sell = -1
self.offset = offset
self.orders = []
self.prices = []
def initialize(self):
pass
def set_order(self, order_callable):
self.order = order_callable
def set_portfolio(self, portfolio):
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)
#sell next time around.
self.buy_or_sell *= -1
self.orders.append(order_size)
self.frame_count += 1
self.incr += 1
def get_sid_filter(self):
return [self.sid]
class BuySellAlgorithmNew(TradingAlgorithm):
"""Algorithm that buys and sells alternatingly. The amount for
each order can be specified. In addition, an offset that will
quadratically reduce the amount that will be bought can be
specified.
This algorithm is used to test the parameter optimization
framework. If combined with the UpDown trade source, an offset of
0 will produce maximum returns.
"""
def __init__(self, sids, amount, offset):
self.sids = sids
self.amount = amount
self.incr = 0
self.done = False
self.order = None
self.frame_count = 0
self.portfolio = None
self.buy_or_sell = -1
self.offset = offset
self.orders = []
self.prices = []
def handle_data(self, data):
order_size = self.buy_or_sell * (self.amount - (self.offset**2))
@@ -89,6 +31,3 @@ class BuySellAlgorithmNew(TradingAlgorithm):
self.orders.append(order_size)
self.frame_count += 1
self.incr += 1
zipline/optimize/factory.py
+2 -2
View File
@@ -9,7 +9,7 @@ import zipline.protocol as zp
from zipline.utils.factory import get_next_trading_dt, create_trading_environment
from zipline.gens.tradegens import SpecificEquityTrades
from zipline.optimize.algorithms import BuySellAlgorithmNew
from zipline.optimize.algorithms import BuySellAlgorithm
from zipline.finance.slippage import FixedSlippage
from copy import copy
@@ -120,7 +120,7 @@ def create_predictable_zipline(config, offset=0, simulate=True):
amplitude)
if 'algorithm' not in config:
algorithm = BuySellAlgorithmNew(sid, 100, offset)
algorithm = BuySellAlgorithm(sids=[sid], amount=100, offset=offset)
config['order_count'] = trade_count - 1
config['trade_count'] = trade_count
zipline/test_algorithms.py
+47
View File
@@ -52,6 +52,7 @@ The algorithm must expose methods:
"""
class TestAlgorithm():
"""
This algorithm will send a specified number of orders, to allow unit tests
@@ -382,3 +383,49 @@ class TestLoggingAlgorithm():
def set_slippage_override(self, slippage_callable):
pass
from datetime import timedelta
from zipline.algorithm import TradingAlgorithm
from zipline.gens.transform import BatchTransform, batch_transform
from zipline.gens.mavg import MovingAverage
class TestRegisterTransformAlgorithm(TradingAlgorithm):
def initialize(self):
self.add_transform(MovingAverage, 'mavg', ['price'],
market_aware=True,
days=2)
def handle_data(self, data):
pass
class NoopBatchTransform(BatchTransform):
def get_value(self, data):
return data.price
@batch_transform
def noop_batch_decorator(data):
return data.price
class BatchTransformAlgorithm(TradingAlgorithm):
def initialize(self, *args, **kwargs):
self.history_class = []
self.history_decorator = []
self.days = 3
self.noop_class = NoopBatchTransform(sids=[0, 1],
market_aware=False,
refresh_period=2,
delta=timedelta(days=self.days))
self.noop_decorator = noop_batch_decorator(sids=[0, 1],
market_aware=False,
refresh_period=2,
delta=timedelta(days=self.days))
def handle_data(self, data):
window_class = self.noop_class.handle_data(data)
window_decorator = self.noop_decorator.handle_data(data)
self.history_class.append(window_class)
self.history_decorator.append(window_decorator)
zipline/utils/factory.py
+1 -1
View File
@@ -240,7 +240,7 @@ def create_test_df_source():
start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc)
end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc)
index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day)
x = np.arange(0, 12).reshape((6, 2))
x = np.arange(2., 14.).reshape((6, 2))
df = pd.DataFrame(x, index=index, columns=[0, 1])
return DataFrameSource(df), df