diff --git a/zipline/finance/movingaverage.py b/zipline/finance/movingaverage.py new file mode 100644 index 00000000..db495d9d --- /dev/null +++ b/zipline/finance/movingaverage.py @@ -0,0 +1,67 @@ +from datetime import timedelta +from collections import defaultdict + +from zipline.messaging import BaseTransform + +class MovingAverageTransform(BaseTransform): + + def init(self, daycount=3): + self.daycount = daycount + self.by_sid = defaultdict(MovingAverage) + + def transform(self, event): + cur = self.by_sid(event.sid) + cur.update(event) + self.state['value'] = cur.average + return self.state + + def create_vwap(self): + return DailyVWAP(self.daycount) + +class MovingAverage(object): + + def __init__(self, daycount): + self.window = EventWindow(daycount) + self.total = 0.0 + self.average = 0.0 + + def update(self, event): + self.window.update(event) + + self.total += event.price + + for dropped in self.window.dropped_ticks: + self.total -= dropped.price + + if len(self.window.ticks) > 0: + self.average = self.total / len(self.window.ticks) + else: + self.average = 0.0 + +class EventWindow(object): + """ + Tracks a window of the event history. Use an instance to track the events + inside your window to efficiently calculate rolling statistics. + """ + def __init__(self, daycount): + self.ticks = [] + self.dropped_ticks = [] + self.delta = timedelta(days=daycount) + + def update(self, event): + # add new event + self.ticks.append(event) + # determine which events are expired + last_date = event['dt'] + first_date = last_date - self.delta + + self.dropped_ticks = [] + for tick in self.ticks: + if tick['dt'] <= first_date: + self.dropped_ticks.append(tick) + + # remove the expired events + slice_index = len(self.dropped_ticks) + self.ticks = self.ticks[slice_index:] + + diff --git a/zipline/finance/returns.py b/zipline/finance/returns.py new file mode 100644 index 00000000..3c258bee --- /dev/null +++ b/zipline/finance/returns.py @@ -0,0 +1,44 @@ +import pandas +from datetime import timedelta +from collections import defaultdict + +from zipline.messaging import BaseTransform + +class WindowTransform(BaseTransform): + + def init(self, daycount=3): + self.daycount = daycount + self.by_sid = defaultdict(DailyReturns) + + def transform(self, event): + cur = self.by_sid(event.sid) + cur.update(event) + self.state['value'] = cur.vwap + return self.state + +class ReturnsFromPriorClose(object): + """ + Calculates a security's returns since the previous close, using the + current price. + """ + + def __init__(self): + self.last_close = None + self.last_event = None + self.returns = 0.0 + + def update(self, event): + next_close = None + if self.last_close: + change = event.price - self.last_close.price + self.returns = change / self.last_close.price + + if self.last_event: + if self.last_event.dt.day != event.dt.day: + # the current event is from the day after + # the last event. Therefore the last event was + # the last close + self.last_close = self.last_event + + # the current event is now the last_event + self.last_event = event \ No newline at end of file diff --git a/zipline/finance/transforms.py b/zipline/finance/vwap.py similarity index 53% rename from zipline/finance/transforms.py rename to zipline/finance/vwap.py index cdfcbdc8..f409ce4d 100644 --- a/zipline/finance/transforms.py +++ b/zipline/finance/vwap.py @@ -1,67 +1,61 @@ +import pandas from datetime import timedelta -from itertools import ifilter from collections import defaultdict from zipline.messaging import BaseTransform +from zipline.finance.movingaverage import EventWindow class VWAPTransform(BaseTransform): def init(self, daycount=3): self.daycount = daycount - self.by_sid = defaultdict(DailyVWAP) + self.by_sid = defaultdict(self.create_vwap) def transform(self, event): cur = self.by_sid(event.sid) cur.update(event) self.state['value'] = cur.vwap return self.state - + + def create_vwap(self): + return DailyVWAP(self.daycount) + class DailyVWAP: """A class that tracks the volume weighted average price based on tick updates.""" - def __init__(self, daycount=3): - self.ticks = [] - self.dropped_ticks = [] + def __init__(self, daycount): + self.window = EventWindow(daycount) self.flux = 0.0 self.volume = 0 - self.lastTick = None self.vwap = 0.0 self.delta = timedelta(days=daycount) - + def update(self, event): - - self.ticks.append(event) + + # update the event window + self.window.update(event) + + # add the current event's flux and volume to the tracker flux, volume = self.calculate_flux([event]) self.flux += flux self.volume += volume - - self.last_date = event['dt'] - self.first_date = self.last_date - self.delta - #use a list comprehension to filter the ticks to those within - #desired day range. The dt properties are full datetime objects - #and provide overloads for arithmetic operations. - self.dropped_ticks = [] - for tick in self.ticks: - if tick['dt'] < self.first_date: - self.dropped_ticks.append(tick) - - slice_index = len(self.dropped_ticks) - self.ticks = self.ticks[slice_index:] - dropped_flux, dropped_volume = self.calculate_flux(self.dropped_ticks) - + # subract the expired events flux and volume from the tracker + dropped = self.window.dropped_ticks + dropped_flux, dropped_volume = self.calculate_flux(dropped) + self.flux -= dropped_flux self.volume -= dropped_volume - + if(self.volume != 0): self.vwap = self.flux / self.volume else: self.vwap = None - + def calculate_flux(self, ticks): flux = 0.0 volume = 0 for tick in ticks: flux += tick['volume'] * tick['price'] volume += tick['volume'] - return flux, volume \ No newline at end of file + return flux, volume diff --git a/zipline/lines.py b/zipline/lines.py index 26d01f67..932d9e2f 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -143,6 +143,8 @@ class SimulatedTrading(object): sockets[7], logging = qutil.LOGGER ) + + self.con.cancel_socket = self.allocator.lease(1)[0] # TODO: Not freeform self.con.manage( diff --git a/zipline/test/test_transforms.py b/zipline/test/test_transforms.py new file mode 100644 index 00000000..6a2bf204 --- /dev/null +++ b/zipline/test/test_transforms.py @@ -0,0 +1,97 @@ +from datetime import timedelta +from collections import defaultdict +from unittest2 import TestCase + +import zipline.test.factory as factory +import zipline.util as qutil +from zipline.finance.vwap import DailyVWAP, VWAPTransform +from zipline.finance.returns import ReturnsFromPriorClose +from zipline.finance.movingaverage import MovingAverage +from zipline.lines import SimulatedTrading +from zipline.simulator import AddressAllocator, Simulator + + +allocator = AddressAllocator(1000) + +class ZiplineWithTransformsTestCase(TestCase): + leased_sockets = defaultdict(list) + + def setUp(self): + # skip ahead 100 spots + allocator.lease(100) + qutil.configure_logging() + self.trading_environment = factory.create_trading_environment() + self.zipline_test_config = { + 'allocator':allocator, + 'sid':133 + } + + def test_vwap_tnfm(self): + zipline = SimulatedTrading.create_test_zipline( + **self.zipline_test_config + ) + + vwap = VWAPTransform("vwap_10", daycount=10) + zipline.add_transform(vwap) + + zipline.simulate(blocking=True) + + self.assertTrue(zipline.sim.ready()) + self.assertFalse(zipline.sim.exception) + +class FinanceTransformsTestCase(TestCase): + def setUp(self): + self.trading_environment = factory.create_trading_environment() + + def test_vwap(self): + + trade_history = factory.create_trade_history( + 133, + [10.0, 10.0, 10.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + self.trading_environment + ) + + vwap = DailyVWAP(daycount=2) + for trade in trade_history: + vwap.update(trade) + + self.assertEqual(vwap.vwap, 10.75) + + + def test_returns(self): + trade_history = factory.create_trade_history( + 133, + [10.0, 10.0, 10.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + self.trading_environment + ) + + returns = ReturnsFromPriorClose() + for trade in trade_history: + returns.update(trade) + + + self.assertEqual(returns.returns, .1) + + + def test_moving_average(self): + trade_history = factory.create_trade_history( + 133, + [10.0, 10.0, 10.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + self.trading_environment + ) + + ma = MovingAverage(daycount=2) + for trade in trade_history: + ma.update(trade) + + + self.assertEqual(ma.average, 10.5) + + + \ No newline at end of file