diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index 78b1bea8..621c45ad 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -19,13 +19,35 @@ import numpy as np from zipline.utils.test_utils import setup_logger import zipline.utils.factory as factory -from zipline.test_algorithms import TestRegisterTransformAlgorithm +from zipline.test_algorithms import (TestRegisterTransformAlgorithm, + RecordAlgorithm) from zipline.sources import (SpecificEquityTrades, DataFrameSource, DataPanelSource) from zipline.transforms import MovingAverage +class TestRecordAlgorithm(TestCase): + def setUp(self): + self.trading_environment = factory.create_trading_environment() + 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) + self.df_source, self.df = \ + factory.create_test_df_source(self.trading_environment) + + def test_record_incr(self): + algo = RecordAlgorithm() + output = algo.run(self.source) + np.testing.assert_array_equal(output['incr'].values, + range(1, len(output) + 1)) + + class TestTransformAlgorithm(TestCase): def setUp(self): setup_logger(self) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 032af464..d950e8c2 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -87,6 +87,8 @@ class TradingAlgorithm(object): self.transforms = [] self.sources = [] + self._registered_vars = set() + self.logger = None # default components for transact @@ -222,8 +224,15 @@ class TradingAlgorithm(object): # create daily and cumulative stats dataframe daily_perfs = [] cum_perfs = [] + # TODO: the loop here could overwrite expected properties + # of daily_perf. Could potentially raise or log a + # warning. for perf in perfs: if 'daily_perf' in perf: + + perf['daily_perf'].update( + perf['daily_perf'].pop('recorded_vars') + ) daily_perfs.append(perf['daily_perf']) else: cum_perfs.append(perf) @@ -252,6 +261,37 @@ class TradingAlgorithm(object): 'args': args, 'kwargs': kwargs} + def record_variables(self, names): + """Track and record local variables (i.e. attributes) each + day. + + :Arguments: + names : str or list + List of variable names (strings) to record. + + :Notes: + You are responsible for making sure the attributes + exist. + + The corresponding variable name and its values will be + appended to the results returned by the .run() method. + + :Example: + + In initialize you would call + self.record_variables('mavg'). In handle_data you could + then set self.mavg to some value and it will be recorded. + + """ + if isinstance(names, basestring): + names = [names] + + self._registered_vars.update(set(names)) + + @property + def recorded_vars(self): + return {name: getattr(self, name) for name in self._registered_vars} + @property def portfolio(self): return self._portfolio diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 0b41b07a..df698e13 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -227,6 +227,10 @@ class AlgorithmSimulator(object): else: for event in snapshot: for perf_message in event.perf_messages: + # append current values of recorded vars + # to emitted message + perf_message['daily_perf']['recorded_vars'] =\ + self.algo.recorded_vars yield perf_message del event['perf_messages'] @@ -240,6 +244,8 @@ class AlgorithmSimulator(object): self.perf_tracker.handle_simulation_end() for message in perf_messages: + message['daily_perf']['recorded_vars'] =\ + self.algo.recorded_vars yield message yield risk_message diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index e86e1a21..e32b62db 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -215,6 +215,15 @@ class TimeoutAlgorithm(TradingAlgorithm): time.sleep(100) pass + +class RecordAlgorithm(TradingAlgorithm): + def initialize(self): + self.incr = 0 + self.record_variables(['incr']) + + def handle_data(self, data): + self.incr += 1 + from zipline.algorithm import TradingAlgorithm from zipline.transforms import BatchTransform, batch_transform from zipline.transforms import MovingAverage