From 44fbdff4ac8de52c21fcde378518d826e54b95f0 Mon Sep 17 00:00:00 2001 From: warren-oneill Date: Thu, 4 Jun 2015 14:23:46 +0200 Subject: [PATCH] added CLOSE_POSITION as source type, added pt.close_position_event(), added process_close_postion(), added close processing to tradesimulation, added unittest for close_position_event --- tests/test_perf_tracking.py | 34 ++++++++++++++++++- .../finance/performance/position_tracker.py | 14 ++++++++ zipline/finance/performance/tracker.py | 4 +++ zipline/gens/tradesimulation.py | 9 +++++ zipline/protocol.py | 3 +- 5 files changed, 62 insertions(+), 2 deletions(-) diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 9082b27a..1bb90802 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -46,7 +46,8 @@ from zipline.finance.commission import PerShare, PerTrade, PerDollar from zipline.finance import trading from zipline.utils.factory import create_random_simulation_parameters import zipline.protocol as zp -from zipline.protocol import Event +from zipline.protocol import Event, DATASOURCE_TYPE +from zipline.sources.data_frame_source import DataPanelSource logger = logging.getLogger('Test Perf Tracking') @@ -1937,6 +1938,37 @@ class TestPerformanceTracker(unittest.TestCase): check_perf_tracker_serialization(tracker) + def test_close_position_event(self): + pt = perf.PositionTracker() + dt = pd.Timestamp("1984/03/06 3:00PM") + pos1 = perf.Position(1, amount=np.float64(120.0), + last_sale_date=dt, last_sale_price=3.4) + pos2 = perf.Position(2, amount=np.float64(-100.0), + last_sale_date=dt, last_sale_price=3.4) + pt.update_positions({1: pos1, 2: pos2}) + + event_type = DATASOURCE_TYPE.CLOSE_POSITION + index = [dt + timedelta(days=1)] + pan = pd.Panel({1: pd.DataFrame({'price': 1, 'volume': 0, + 'type': event_type}, index=index), + 2: pd.DataFrame({'price': 1, 'volume': 0, + 'type': event_type}, index=index), + 3: pd.DataFrame({'price': 1, 'volume': 0, + 'type': event_type}, index=index)}) + + source = DataPanelSource(pan) + for i, event in enumerate(source): + txn = pt.create_close_position_transaction(event) + if event.sid == 1: + # Test owned long + self.assertEqual(-120, txn.amount) + elif event.sid == 2: + # Test owned short + self.assertEqual(100, txn.amount) + elif event.sid == 3: + # Test not-owned SID + self.assertIsNone(txn) + def test_serialization(self): start_dt = datetime(year=2008, month=10, diff --git a/zipline/finance/performance/position_tracker.py b/zipline/finance/performance/position_tracker.py index 9fc68b61..bd3f674f 100644 --- a/zipline/finance/performance/position_tracker.py +++ b/zipline/finance/performance/position_tracker.py @@ -13,6 +13,7 @@ except ImportError: from six import iteritems from six.moves import map, filter +from zipline.finance.slippage import Transaction from zipline.utils.serialization_utils import ( VERSION_LABEL ) @@ -217,6 +218,19 @@ class PositionTracker(object): net_cash_payment = payments['cash_amount'].fillna(0).sum() return net_cash_payment + def create_close_position_transaction(self, event): + if not self._position_amounts.get(event.sid): + return None + txn = Transaction( + sid=event.sid, + amount=(-1 * self._position_amounts[event.sid]), + dt=event.dt, + price=event.price, + commission=0, + order_id=0 + ) + return txn + def get_positions(self): positions = self._positions_store diff --git a/zipline/finance/performance/tracker.py b/zipline/finance/performance/tracker.py index 75988003..96410d7f 100644 --- a/zipline/finance/performance/tracker.py +++ b/zipline/finance/performance/tracker.py @@ -337,6 +337,10 @@ class PerformanceTracker(object): self.all_benchmark_returns[midnight] = event.returns + def process_close_position(self, event): + txn = self.position_tracker.create_close_position_transaction(event) + self.process_transaction(txn) + def check_upcoming_dividends(self, midnight_of_date_that_just_ended): """ Check if we currently own any stocks with dividends whose ex_date is diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index c23b6ca1..9f926ac4 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -204,6 +204,8 @@ class AlgorithmSimulator(object): perf_process_split = self.algo.perf_tracker.process_split perf_process_dividend = self.algo.perf_tracker.process_dividend perf_process_commission = self.algo.perf_tracker.process_commission + perf_process_close_position = \ + self.algo.perf_tracker.process_close_position blotter_process_trade = self.algo.blotter.process_trade blotter_process_benchmark = self.algo.blotter.process_benchmark @@ -219,6 +221,7 @@ class AlgorithmSimulator(object): # custom events. trades = [] customs = [] + closes = [] # splits and dividends are processed once a day. # @@ -247,6 +250,8 @@ class AlgorithmSimulator(object): if dividends is None: dividends = [] dividends.append(event) + elif event.type == DATASOURCE_TYPE.CLOSE_POSITION: + closes.append(event) else: raise log.warn("Unrecognized event=%s".format(event)) @@ -282,6 +287,10 @@ class AlgorithmSimulator(object): for custom in customs: self.update_universe(custom) + for close in closes: + self.update_universe(close) + perf_process_close_position(close) + if splits is not None: for split in splits: # process_split is not assigned to a variable since it is diff --git a/zipline/protocol.py b/zipline/protocol.py index 1b642e1a..7fcfef05 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -42,7 +42,8 @@ DATASOURCE_TYPE = Enum( 'DONE', 'CUSTOM', 'BENCHMARK', - 'COMMISSION' + 'COMMISSION', + 'CLOSE_POSITION' ) # Expected fields/index values for a dividend Series.