Merge pull request #37 from quantopian/fix_transforms

Fix transforms
This commit is contained in:
Eddie Hebert
2012-12-06 14:00:35 -08:00
5 changed files with 109 additions and 176 deletions
+1 -3
View File
@@ -25,11 +25,9 @@ class TestDataFrameSource(TestCase):
for expected_dt, expected_price in df.iterrows():
sid0 = source.next()
sid1 = source.next()
assert expected_dt == sid0.dt == sid1.dt
assert expected_dt == sid0.dt
assert expected_price[0] == sid0.price
assert expected_price[1] == sid1.price
def test_sid_filtering(self):
_, df = factory.create_test_df_source()
+20 -61
View File
@@ -123,11 +123,7 @@ class TestEventWindow(TestCase):
# Record the length of the window after each event.
lengths.append(len(window.ticks))
# The window stretches out during the weekend because we wait
# to drop events until the weekend ends. The last window is
# briefly longer because it doesn't complete a full day. The
# window then shrinks once the day completes
assert lengths == [1, 2, 3, 3, 3, 4, 5, 5, 5, 3, 4, 3]
assert lengths == [1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]
assert window.added == events
assert window.removed == events[:-3]
@@ -146,7 +142,7 @@ class TestEventWindow(TestCase):
# Record the length of the window after each event.
lengths.append(len(window.ticks))
assert lengths == [1, 2, 3, 3, 2]
assert lengths == [1, 2, 2, 2, 2]
assert window.added == events
assert window.removed == events[:-2]
@@ -317,23 +313,16 @@ class TestBatchTransform(TestCase):
def test_event_window(self):
algo = BatchTransformAlgorithm()
algo.run(self.source)
self.assertEqual(algo.history_return_price_class[:2],
[None, None],
wl = algo.window_length
self.assertEqual(algo.history_return_price_class[:wl],
[None] * wl,
"First two iterations should return None")
self.assertEqual(algo.history_return_price_decorator[:2],
[None, None],
self.assertEqual(algo.history_return_price_decorator[:wl],
[None] * wl,
"First two iterations should return None")
self.assertEqual(algo.history_return_price_market_aware[:2],
[None, None],
"First two iterations should return None")
self.assertEqual(algo.history_return_more_days_than_refresh[:3],
[None, None, None],
"First five iterations should return None")
self.assertTrue(isinstance(
algo.history_return_more_days_than_refresh[4],
pd.DataFrame),
"Sixth iteration should not be None"
algo.history_return_price_class[wl + 1],
pd.DataFrame)
)
# Test whether arbitrary fields can be added to datapanel
@@ -344,27 +333,21 @@ class TestBatchTransform(TestCase):
)
self.assertTrue(all(
field['arbitrary'].values.flatten() == ['test'] * 8),
field['arbitrary'].values.flatten() ==
[123] * algo.window_length),
'arbitrary dataframe should contain only "test"'
)
# test overloaded class
for test_history in [algo.history_return_price_class,
algo.history_return_price_decorator]:
np.testing.assert_array_equal(
range(2, 8),
test_history[2].values.flatten()
)
np.testing.assert_array_equal(
range(2, 8),
test_history[3].values.flatten()
)
np.testing.assert_array_equal(
range(4, 12),
test_history[4].values.flatten()
)
# starting at window length, the window should contain
# consecutive (of window length) numbers up till the end.
for i in range(algo.window_length, len(test_history)):
np.testing.assert_array_equal(
range(i - algo.window_length + 1, i + 1),
test_history[i].values.flatten()
)
def test_passing_of_args(self):
algo = BatchTransformAlgorithm(1, kwarg='str')
@@ -375,29 +358,5 @@ class TestBatchTransform(TestCase):
expected_item = ((1, ), {'kwarg': 'str'})
self.assertEqual(
algo.history_return_args,
[None, None, expected_item, expected_item,
expected_item, expected_item])
class TestBatchTransformMarketAware(TestCase):
def setUp(self):
setup_logger(self)
start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc)
end = pd.datetime(1994, 1, 1, 0, 0, 0, 0, pytz.utc)
self.data = factory.load_from_yahoo(stocks=['AAPL'],
indexes={},
start=start, end=end)
def test_event_window(self):
days = 50
algo = BatchTransformAlgorithm(days=days, refresh_period=days)
algo.run(self.data)
self.assertEqual(algo.history_return_price_market_aware[:days],
[None] * days,
"First {days} iterations should return None"
.format(days=days))
self.assertFalse(algo.history_return_price_market_aware[days + 1]
is None,
"Window is contains too many Nones.")
[None, None, None, expected_item, expected_item,
expected_item])
+37 -26
View File
@@ -214,7 +214,6 @@ class TimeoutAlgorithm(TradingAlgorithm):
time.sleep(100)
pass
from datetime import timedelta
from zipline.algorithm import TradingAlgorithm
from zipline.transforms import BatchTransform, batch_transform
from zipline.transforms import MovingAverage
@@ -237,6 +236,7 @@ class TestRegisterTransformAlgorithm(TradingAlgorithm):
class ReturnPriceBatchTransform(BatchTransform):
def get_value(self, data):
assert data.shape[1] == self.window_length
return data.price
@@ -257,7 +257,7 @@ def return_data(data, *args, **kwargs):
class BatchTransformAlgorithm(TradingAlgorithm):
def initialize(self, *args, **kwargs):
self.refresh_period = kwargs.pop('refresh_period', 2)
self.refresh_period = kwargs.pop('refresh_period', 1)
self.window_length = kwargs.pop('window_length', 3)
self.args = args
@@ -266,46 +266,47 @@ class BatchTransformAlgorithm(TradingAlgorithm):
self.history_return_price_class = []
self.history_return_price_decorator = []
self.history_return_args = []
self.history_return_price_market_aware = []
self.history_return_more_days_than_refresh = []
self.history_return_arbitrary_fields = []
self.history_return_nan = []
self.return_price_class = ReturnPriceBatchTransform(
market_aware=False,
refresh_period=self.refresh_period,
delta=timedelta(days=self.window_length)
window_length=self.window_length,
fillna=False
)
self.return_price_decorator = return_price_batch_decorator(
market_aware=False,
refresh_period=self.refresh_period,
delta=timedelta(days=self.window_length)
window_length=self.window_length,
fillna=False
)
self.return_args_batch = return_args_batch_decorator(
market_aware=False,
refresh_period=self.refresh_period,
delta=timedelta(days=self.window_length)
window_length=self.window_length,
fillna=False
)
self.return_price_market_aware = ReturnPriceBatchTransform(
market_aware=True,
refresh_period=self.refresh_period,
window_length=self.window_length
)
self.return_price_more_days_than_refresh = ReturnPriceBatchTransform(
market_aware=True,
refresh_period=1,
window_length=3
window_length=self.window_length,
fillna=False
)
self.return_arbitrary_fields = return_data(
market_aware=True,
refresh_period=1,
window_length=3
refresh_period=self.refresh_period,
window_length=self.window_length,
fillna=False
)
self.return_nan = return_price_batch_decorator(
refresh_period=self.refresh_period,
window_length=self.window_length,
fillna=True
)
self.iter = 0
self.set_slippage(FixedSlippage())
def handle_data(self, data):
@@ -316,18 +317,28 @@ class BatchTransformAlgorithm(TradingAlgorithm):
self.history_return_args.append(
self.return_args_batch.handle_data(
data, *self.args, **self.kwargs))
self.history_return_price_market_aware.append(
self.return_price_market_aware.handle_data(data))
self.history_return_more_days_than_refresh.append(
self.return_price_more_days_than_refresh.handle_data(data))
new_data = deepcopy(data)
for sid in new_data:
new_data[sid]['arbitrary'] = 'test'
new_data[sid]['arbitrary'] = 123
self.history_return_arbitrary_fields.append(
self.return_arbitrary_fields.handle_data(new_data))
# nan every second event price
if self.iter % 2 == 0:
self.history_return_nan.append(
self.return_nan.handle_data(data))
else:
nan_data = deepcopy(data)
import numpy as np
for sid in nan_data.iterkeys():
nan_data[sid].price = np.nan
self.history_return_nan.append(
self.return_nan.handle_data(nan_data))
self.iter += 1
class SetPortfolioAlgorithm(TradingAlgorithm):
"""
+49 -84
View File
@@ -239,17 +239,9 @@ class EventWindow(object):
# adding new ticks.
self.handle_add(event)
if self.market_aware:
self.add_new_holidays(event.dt)
# Clear out any expired events. drop_condition changes depending
# on whether or not we are running in market_aware mode.
#
# oldest newest
# | |
# V V
while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt):
while self.drop_condition():
# popleft removes and returns the oldest tick in self.ticks
popped = self.ticks.popleft()
@@ -257,36 +249,17 @@ class EventWindow(object):
# behavior for removing ticks.
self.handle_remove(popped)
def add_new_holidays(self, newest):
# Add to our tracked window any untracked holidays that are
# older than our newest event. (newest should always be
# self.ticks[-1])
while len(self.all_holidays) > 0 and self.all_holidays[0] <= newest:
self.cur_holidays.append(self.all_holidays.popleft())
def out_of_market_window(self):
# Find number of unique days in window
# Note that this assumes that each day we received an
# event is a trading day.
unique_dts = set([event.dt.date() for event in self.ticks])
def drop_old_holidays(self, oldest):
# Drop from our tracked window any holidays that are older
# than our oldest tracked event. (oldest should always
# be self.ticks[0])
while len(self.cur_holidays) > 0 and self.cur_holidays[0] < oldest:
self.cur_holidays.popleft()
return len(unique_dts) > self.window_length
def out_of_market_window(self, oldest, newest):
self.drop_old_holidays(oldest)
calendar_dates_between = (newest.date() - oldest.date()).days
holidays_between = len(self.cur_holidays)
trading_days_between = calendar_dates_between - holidays_between
# "Put back" a day if oldest is earlier in its day than newest,
# reflecting the fact that we haven't yet completed the last
# day in the window.
if oldest.time() > newest.time():
trading_days_between -= 1
return trading_days_between >= self.window_length
def out_of_delta(self, oldest, newest):
return (newest - oldest) >= self.delta
def out_of_delta(self):
# newest - oldest
return (self.ticks[-1].dt - self.ticks[0].dt) >= self.delta
# All event windows expect to receive events with datetime fields
# that arrive in sorted order.
@@ -344,19 +317,19 @@ class BatchTransform(EventWindow):
def __init__(self,
func=None,
refresh_period=None,
market_aware=True,
delta=None,
window_length=None):
window_length=None,
fillna=True):
super(BatchTransform, self).__init__(market_aware,
window_length=window_length,
delta=delta)
super(BatchTransform, self).__init__(True,
window_length=window_length)
if func is not None:
self.compute_transform_value = func
else:
self.compute_transform_value = self.get_value
self.fillna = fillna
self.refresh_period = refresh_period
self.window_length = window_length
self.trading_days_since_update = 0
@@ -366,7 +339,7 @@ class BatchTransform(EventWindow):
self.last_dt = None
self.updated = False
self.data = None
self.cached = None
self.field_names = None
@@ -394,20 +367,22 @@ class BatchTransform(EventWindow):
# return newly computed or cached value
return self.get_transform_value(*args, **kwargs)
def handle_add(self, event):
if not self.last_dt:
self.last_dt = event.dt
return
def _extract_field_names(self, event):
# extract field names from sids (price, volume etc), make sure
# every sid has the same fields.
sid_keys = [set(sid.keys()) for sid in event.data.itervalues()]
assert sid_keys[0] == set.intersection(*sid_keys),\
"Each sid must have the same keys."
if self.field_names is None:
unwanted_fields = set(['portfolio', 'sid', 'dt', 'type',
'datetime'])
self.field_names = sid_keys[0] - unwanted_fields
unwanted_fields = set(['portfolio', 'sid', 'dt', 'type',
'datetime', 'source_id'])
return sid_keys[0] - unwanted_fields
def handle_add(self, event):
if not self.last_dt:
self.field_names = self._extract_field_names(event)
self.last_dt = event.dt
return
# update trading day counters
if self.last_dt.day != event.dt.day:
@@ -419,13 +394,11 @@ class BatchTransform(EventWindow):
self.trading_days_total >= self.window_length and
self.trading_days_since_update >= self.refresh_period
):
# Create datapanel of running event window.
self.data = self.get_data()
# Setting updated to True will cause get_transform_value()
# to call the user-defined batch-transform with the most
# recent datapanel
self.updated = True
self.full = True
self.trading_days_since_update = 0
else:
self.updated = False
@@ -442,36 +415,28 @@ class BatchTransform(EventWindow):
"""
# This Panel data structure ultimately gets passed to the
# user-overloaded get_value() method.
#
# self.ticks contains ndicts with data, dt keys.
# event parameter is an ndict with data, dt keys.
fields = {}
sids = set.union(*[set(tick.data.keys()) for tick in self.ticks])
dts = [tick.dt for tick in self.ticks]
for field_name in self.field_names:
sids = self.ticks[0].data.keys()
data = pd.Panel(items=self.field_names, major_axis=dts,
minor_axis=sids)
values_per_sid = {}
# Fill data panel
for tick in self.ticks:
dt = tick.dt
for sid, fields in tick.data.iteritems():
for field_name in self.field_names:
data[field_name][sid].ix[dt] = fields[field_name]
for sid in sids:
values_per_sid[sid] = pd.Series(
{tick.data[sid].dt: tick.data[sid][field_name]
for tick in self.ticks}
)
if self.fillna:
# Fills in gaps of missing data during transform
# of multiple stocks. E.g. we may be missing
# minute data because of illiquidity of one stock
data = data.fillna(method='ffill')
# concatenate different sids into one df
df = pd.DataFrame.from_dict(values_per_sid)
# Fills in gaps of missing data during transform of multiple
# stocks.
# e.g. we may be missing minute data because of illiquidity
# of one stock
df = df.fillna(method='ffill')
# Drop any empty rows after the fill.
# This will drop a leading row of N/A
df = df.dropna()
fields[field_name] = df
data = pd.Panel.from_dict(fields, orient='items')
# Drop any empty rows after the fill.
# This will drop a leading row of N/A
data = data.dropna(axis=1)
return data
@@ -491,11 +456,11 @@ class BatchTransform(EventWindow):
has actually been updated. Otherwise, the previously, cached
value will be returned.
"""
if self.data is None:
if not self.full:
return None
if self.updated:
self.cached = self.compute_transform_value(self.data,
self.cached = self.compute_transform_value(self.get_data(),
*args, **kwargs)
return self.cached
+2 -2
View File
@@ -272,9 +272,9 @@ 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(2., len(index) * 2 + 2).reshape((-1, 2))
x = np.arange(0, len(index))
df = pd.DataFrame(x, index=index, columns=[0, 1])
df = pd.DataFrame(x, index=index, columns=[0])
return DataFrameSource(df), df