mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 09:11:03 +08:00
scottsanderson
267 lines
8.7 KiB
Python
267 lines
8.7 KiB
Python
import pytz
|
|
|
|
from datetime import timedelta, datetime
|
|
from collections import defaultdict
|
|
from unittest2 import TestCase
|
|
|
|
from zipline import ndict
|
|
|
|
from zipline.lines import SimulatedTrading
|
|
|
|
from zipline.utils.test_utils import setup_logger, teardown_logger
|
|
from zipline.utils.date_utils import utcnow
|
|
|
|
from zipline.gens.tradegens import SpecificEquityTrades
|
|
from zipline.gens.transform import StatefulTransform, EventWindow
|
|
from zipline.gens.vwap import VWAP
|
|
from zipline.gens.mavg import MovingAverage
|
|
from zipline.gens.returns import Returns
|
|
|
|
import zipline.utils.factory as factory
|
|
|
|
def to_dt(msg):
|
|
return ndict({'dt': msg})
|
|
|
|
class NoopEventWindow(EventWindow):
|
|
"""
|
|
A no-op EventWindow subclass for testing the base EventWindow logic.
|
|
Keeps lists of all added and dropped events.
|
|
"""
|
|
def __init__(self, market_aware, days, delta):
|
|
EventWindow.__init__(self, market_aware, days, delta)
|
|
|
|
self.added = []
|
|
self.removed = []
|
|
|
|
def handle_add(self, event):
|
|
self.added.append(event)
|
|
|
|
def handle_remove(self, event):
|
|
self.removed.append(event)
|
|
|
|
class EventWindowTestCase(TestCase):
|
|
|
|
def setUp(self):
|
|
setup_logger(self)
|
|
|
|
# Constants calling before open, during the day, and after
|
|
# close on a valid trading day.
|
|
self.pre_open = datetime(2012, 8, 7, 13, tzinfo = pytz.utc)
|
|
self.mid_day = datetime(2012, 8, 7, 15, tzinfo = pytz.utc)
|
|
self.post_close = datetime(2012, 8, 7, 22, tzinfo = pytz.utc)
|
|
|
|
# Constants calling before open, during the day, and after
|
|
# close on a saturday.
|
|
self.pre_open_saturday = datetime(2012, 8, 11, 13, tzinfo = pytz.utc)
|
|
self.mid_day_saturday = datetime(2012, 8, 11, 15, tzinfo = pytz.utc)
|
|
self.post_close_saturday = datetime(2012, 8, 11, 22, tzinfo = pytz.utc)
|
|
|
|
# Constants calling before open, during the day, and after
|
|
# close on a holiday.
|
|
self.pre_open_holiday = datetime(2012, 12, 25, 13, tzinfo = pytz.utc)
|
|
self.mid_day_holiday = datetime(2012, 12, 25, tzinfo = pytz.utc)
|
|
self.post_close_holiday = datetime(2012, 12, 25, 22, tzinfo = pytz.utc)
|
|
|
|
def test_event_window_with_timedelta(self):
|
|
|
|
# Keep all events within a 5 minute window.
|
|
window = NoopEventWindow(
|
|
market_aware = False,
|
|
delta = timedelta(minutes = 5),
|
|
days = None
|
|
)
|
|
now = utcnow()
|
|
|
|
# 15 dates, increasing in 1 minute increments.
|
|
dates = [now + i * timedelta(minutes = 1)
|
|
for i in xrange(15)]
|
|
|
|
# Turn the dates into the format required by EventWindow.
|
|
dt_messages = [to_dt(date) for date in dates]
|
|
|
|
# Run all messages through the window and assert that we're adding
|
|
# and removing messages appropriately. We start the enumeration at 1
|
|
# for convenience.
|
|
for num, message in enumerate(dt_messages, 1):
|
|
window.update(message)
|
|
|
|
# Assert that we've added the correct number of events.
|
|
assert len(window.added) == num
|
|
|
|
# Assert that we removed only events that fall outside (or
|
|
# on the boundary of) the delta.
|
|
for dropped in window.removed:
|
|
assert message.dt - dropped.dt >= timedelta(minutes = 5)
|
|
|
|
def test_market_aware_window(self):
|
|
window = NoopEventWindow(
|
|
market_aware = True,
|
|
delta = None,
|
|
days = 1
|
|
)
|
|
dates = ([self.pre_open]*3)
|
|
dates += ([self.mid_day]*3)
|
|
dates += ([self.post_close]*3)
|
|
dates += [self.pre_open + timedelta(days = 1, seconds = 1)]
|
|
events = [to_dt(date) for date in dates]
|
|
|
|
# Run the events.
|
|
for event in events:
|
|
window.update(event)
|
|
|
|
# We should have removed the pre_open events on the first day.
|
|
# The rest should be intact.
|
|
|
|
assert window.added == events
|
|
assert window.removed == events[0:3]
|
|
assert list(window.ticks) == events[3:]
|
|
|
|
def test_market_aware_window_weekend(self):
|
|
window = NoopEventWindow(
|
|
market_aware = True,
|
|
delta = None,
|
|
days = 2
|
|
)
|
|
dates = [self.pre_open_saturday - timedelta(days = 1, seconds=1)]
|
|
dates += [self.mid_day_saturday - timedelta(days = 1, seconds=1)]
|
|
dates += [self.post_close_saturday - timedelta(days = 1, seconds=1)]
|
|
dates += [self.mid_day_saturday + timedelta(days = 1)]
|
|
|
|
events = [to_dt(date) for date in dates]
|
|
|
|
# Run the events.
|
|
for event in events:
|
|
window.update(event)
|
|
|
|
# We shouldn't remove any events.
|
|
assert window.added == events
|
|
assert window.removed == []
|
|
assert list(window.ticks) == events
|
|
|
|
extra = to_dt(self.mid_day_saturday + timedelta(days = 2))
|
|
window.update(extra)
|
|
|
|
# We should remove only the first event.
|
|
assert window.removed == [events[0]]
|
|
assert list(window.ticks) == events[1:] + [extra]
|
|
|
|
def tearDown(self):
|
|
setup_logger(self)
|
|
|
|
class FinanceTransformsTestCase(TestCase):
|
|
|
|
def setUp(self):
|
|
self.trading_environment = factory.create_trading_environment()
|
|
setup_logger(self)
|
|
|
|
trade_history = factory.create_trade_history(
|
|
133,
|
|
[10.0, 10.0, 11.0, 11.0],
|
|
[100, 100, 100, 300],
|
|
timedelta(days=1),
|
|
self.trading_environment
|
|
)
|
|
self.source = SpecificEquityTrades(event_list=trade_history)
|
|
|
|
def tearDown(self):
|
|
self.log_handler.pop_application()
|
|
|
|
def test_vwap(self):
|
|
|
|
vwap = StatefulTransform(
|
|
VWAP,
|
|
market_aware = False,
|
|
delta = timedelta(days = 2)
|
|
)
|
|
transformed = list(vwap.transform(self.source))
|
|
|
|
# Output values
|
|
tnfm_vals = [message.tnfm_value for message in transformed]
|
|
# "Hand calculated" values.
|
|
expected = [
|
|
(10.0 * 100) / 100.0,
|
|
((10.0 * 100) + (10.0 * 100)) / (200.0),
|
|
# We should drop the first event here.
|
|
((10.0 * 100) + (11.0 * 100)) / (200.0),
|
|
# We should drop the second event here.
|
|
((11.0 * 100) + (11.0 * 300)) / (400.0)
|
|
]
|
|
|
|
# Output should match the expected.
|
|
assert tnfm_vals == expected
|
|
|
|
def test_returns(self):
|
|
# Daily returns.
|
|
returns = StatefulTransform(Returns, 1)
|
|
|
|
transformed = list(returns.transform(self.source))
|
|
tnfm_vals = [message.tnfm_value for message in transformed]
|
|
|
|
# No returns for the first event because we don't have a
|
|
# previous close.
|
|
expected = [0.0, 0.0, 0.1, 0.0]
|
|
|
|
assert tnfm_vals == expected
|
|
|
|
# Two-day returns. An extra kink here is that the
|
|
# factory will automatically skip a weekend for the
|
|
# last event. Results shouldn't notice this blip.
|
|
|
|
trade_history = factory.create_trade_history(
|
|
133,
|
|
[10.0, 15.0, 13.0, 12.0, 13.0],
|
|
[100, 100, 100, 300, 100],
|
|
timedelta(days=1),
|
|
self.trading_environment
|
|
)
|
|
self.source = SpecificEquityTrades(event_list=trade_history)
|
|
|
|
returns = StatefulTransform(Returns, 2)
|
|
|
|
transformed = list(returns.transform(self.source))
|
|
tnfm_vals = [message.tnfm_value for message in transformed]
|
|
|
|
expected = [
|
|
0.0,
|
|
0.0,
|
|
(13.0 - 10.0) / 10.0,
|
|
(12.0 - 15.0) / 15.0,
|
|
(13.0 - 13.0) / 13.0
|
|
]
|
|
|
|
assert tnfm_vals == expected
|
|
|
|
def test_moving_average(self):
|
|
|
|
mavg = StatefulTransform(
|
|
MovingAverage,
|
|
market_aware = False,
|
|
fields = ['price', 'volume'],
|
|
delta = timedelta(days = 2),
|
|
)
|
|
transformed = list(mavg.transform(self.source))
|
|
# Output values.
|
|
tnfm_prices = [message.tnfm_value.price for message in transformed]
|
|
tnfm_volumes = [message.tnfm_value.volume for message in transformed]
|
|
|
|
# "Hand-calculated" values
|
|
expected_prices = [
|
|
((10.0) / 1.0),
|
|
((10.0 + 10.0) / 2.0),
|
|
# First event should get dropped here.
|
|
((10.0 + 11.0) / 2.0),
|
|
# Second event should get dropped here.
|
|
((11.0 + 11.0) / 2.0)
|
|
]
|
|
expected_volumes = [
|
|
((100.0) / 1.0),
|
|
((100.0 + 100.0) / 2.0),
|
|
# First event should get dropped here.
|
|
((100.0 + 100.0) / 2.0),
|
|
# Second event should get dropped here.
|
|
((100.0 + 300.0) / 2.0)
|
|
]
|
|
|
|
assert tnfm_prices == expected_prices
|
|
assert tnfm_volumes == expected_volumes
|