working on frame argument sent to algorithm.

zipline/finance/performance.py

@@ -12,11 +12,12 @@ import zipline.finance.risk as risk
 
 class PerformanceTracker():
     
-    def __init__(self, period_start, period_end, capital_base, trading_environment):
+    def __init__(self, trading_environment):
+        self.trading_environment    = trading_environment
         self.trading_day            = datetime.timedelta(hours=6, minutes=30)
         self.calendar_day           = datetime.timedelta(hours=24)
-        self.period_start           = period_start
-        self.period_end             = period_end
+        self.period_start           = self.trading_environment.period_start
+        self.period_end             = self.trading_environment.period_end
         self.market_open            = self.period_start 
         self.market_close           = self.market_open + self.trading_day
         self.progress               = 0.0
@@ -24,21 +25,20 @@ class PerformanceTracker():
         self.day_count              = 0
         self.cumulative_capital_used= 0.0
         self.max_capital_used       = 0.0
-        self.capital_base           = capital_base
-        self.trading_environment    = trading_environment
+        self.capital_base           = self.trading_environment.capital_base
         self.returns                = []
         self.txn_count              = 0 
         self.event_count            = 0
         self.cumulative_performance = PerformancePeriod(
             {}, 
-            capital_base, 
-            starting_cash = capital_base
+            self.capital_base, 
+            starting_cash = self.capital_base
         )
             
         self.todays_performance     = PerformancePeriod(
             {}, 
-            capital_base, 
-            starting_cash = capital_base
+            self.capital_base, 
+            starting_cash = self.capital_base
         )
         
     def to_dict(self):
@@ -121,16 +121,8 @@ class PerformanceTracker():
             'todays_perf'               : self.todays_perf.to_dict(),
             'cumulative_risk_metrics'   : self.cumulative_risk_metrics.to_dict()
         }
-    
-    def update(self, event_frame):
-        for dt, event_series in event_frame.iteritems():
-            data = {}
-            data.update(event_series)
-            event = zp.namedict(data)
-            self.process_event(event)
-        
+            
     def process_event(self, event):
-        qutil.LOGGER.debug("series is " + str(event))
         self.event_count += 1
         if(event.dt >= self.market_close):
             self.handle_market_close()

zipline/finance/risk.py

@@ -5,7 +5,6 @@ import numpy as np
 import numpy.linalg as la
 import zipline.util as qutil
 import zipline.protocol as zp
-from pymongo import ASCENDING, DESCENDING
 
 class DailyReturn():
     
@@ -137,7 +136,6 @@ class RiskMetrics():
         return period_returns, returns
         
     def calculate_volatility(self, daily_returns):
-        #qutil.LOGGER.debug("trading days {td}".format(td=self.trading_days))
         return np.std(daily_returns, ddof=1) * math.sqrt(self.trading_days)
         
     def calculate_sharpe(self):
@@ -326,11 +324,24 @@ def advance_by_months(dt, jump_in_months):
 
 class TradingEnvironment(object):
 
-    def __init__(self, benchmark_returns, treasury_curves):
+    def __init__(
+        self, 
+        benchmark_returns, 
+        treasury_curves, 
+        period_start=None, 
+        period_end=None, 
+        capital_base=None,
+        frame_index=None
+    ):
+    
         self.trading_days = []
         self.trading_day_map = {}
         self.treasury_curves = treasury_curves
         self.benchmark_returns = benchmark_returns
+        self.frame_index = frame_index
+        self.period_start = period_start
+        self.period_end = period_end
+        self.capital_base = capital_base
         for bm in benchmark_returns:
             self.trading_days.append(bm.date)
             self.trading_day_map[bm.date] = bm

zipline/finance/trading.py

@@ -8,19 +8,27 @@ from zmq.core.poll import select
 import zipline.messaging as qmsg
 import zipline.util as qutil
 import zipline.protocol as zp
+import zipline.finance.performance as perf
 
 class TradeSimulationClient(qmsg.Component):
     
-    def __init__(self, simulation_dt):
+    def __init__(self, trading_environment):
         qmsg.Component.__init__(self)
-        self.received_count     = 0
-        self.prev_dt            = None
-        self.event_queue        = None
-        self.event_callbacks    = []
-        self.txn_count          = 0
-        self.current_dt         = simulation_dt
-        self.last_iteration_duration = datetime.timedelta(seconds=0)
-        self.event_frame        = None
+        self.received_count         = 0
+        self.prev_dt                = None
+        self.event_queue            = None
+        self.event_callbacks        = []
+        self.txn_count              = 0
+        self.trading_environment    = trading_environment
+        self.current_dt             = trading_environment.period_start
+        self.last_iteration_dur     = datetime.timedelta(seconds=0)
+        
+        assert self.trading_environment.frame_index != None
+        self.event_frame = pandas.DataFrame(
+            index=self.trading_environment.frame_index
+        )
+        
+        self.perf = perf.PerformanceTracker(self.trading_environment)
     
     @property
     def get_id(self):
@@ -67,9 +75,9 @@ class TradeSimulationClient(qmsg.Component):
                     self.run_callbacks()
                 
                 #update time based on receipt of the order
-                self.last_iteration_duration = datetime.datetime.utcnow() - event_start
+                self.last_iteration_dur = datetime.datetime.utcnow() - event_start
                     
-                self.current_dt = self.current_dt + self.last_iteration_duration
+                self.current_dt = self.current_dt + self.last_iteration_dur
             
             #signal done to order source.
             self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))
@@ -95,15 +103,16 @@ class TradeSimulationClient(qmsg.Component):
         self.order_socket.send(str(zp.ORDER_PROTOCOL.DONE))
         
     def queue_event(self, event):
+        self.perf.process_event(event)
         if self.event_queue == None:
-            self.event_queue = {}
+            self.event_queue = []
         series = event.as_series()
-        self.event_queue[event.dt] = series
+        self.event_queue.append(series)
     
     def get_frame(self):
-        frame = pandas.DataFrame(self.event_queue)
-        self.event_queue = None
-        return frame
+        for event in self.event_queue:
+            self.event_frame[event['sid']] = event
+        return self.event_frame
         
 class OrderDataSource(qmsg.DataSource):
     """DataSource that relays orders from the client"""

zipline/test/test_finance.py

@@ -207,7 +207,11 @@ class FinanceTestCase(TestCase):
 
         set1 = SpecificEquityTrades("flat-133", trade_history)
         
-        trading_client = TradeSimulationClient(start_date)
+        self.trading_environment.period_start = trade_history[0].dt
+        self.trading_environment.period_end = trade_history[-1].dt
+        self.trading_environment.capital_base = 10000
+        
+        trading_client = TradeSimulationClient(self.trading_environment)
         #client will send 10 orders for 100 shares of 133
         test_algo = TestAlgorithm(133, 100, 10, trading_client)
 
@@ -280,25 +284,23 @@ class FinanceTestCase(TestCase):
             volume, 
             start_date, 
             trade_time_increment, 
-            self.trading_environment )
-
+            self.trading_environment 
+        )
+        
+        
+        self.trading_environment.period_start = trade_history[0].dt
+        self.trading_environment.period_end = trade_history[-1].dt
+        self.trading_environment.capital_base = 10000
+        
         set1 = SpecificEquityTrades("flat-133", trade_history)
 
         #client sill send 10 orders for 100 shares of 133
-        trading_client = TradeSimulationClient(start_date)
+        trading_client = TradeSimulationClient(self.trading_environment)
         test_algo = TestAlgorithm(133, 100, 10, trading_client)
 
         order_source = OrderDataSource()
         transaction_sim = TransactionSimulator()
-        perf_tracker = perf.PerformanceTracker(
-            trade_history[0]['dt'], 
-            trade_history[-1]['dt'], 
-            1000000.0, 
-            self.trading_environment)
         
-        #register perf_tracker to receive callbacks from the client.
-        trading_client.add_event_callback(perf_tracker.update)
-    
         sim.register_components([
             trading_client, 
             order_source, 
@@ -339,19 +341,19 @@ class FinanceTestCase(TestCase):
             
         self.assertEqual(
             transaction_sim.txn_count,
-            perf_tracker.txn_count,
+            trading_client.perf.txn_count,
             "The perf tracker should handle the same number of transactions \
             as the simulator emits."
         ) 
         
         self.assertEqual(
-            len(perf_tracker.cumulative_performance.positions), 
+            len(trading_client.perf.cumulative_performance.positions), 
             1, 
             "Portfolio should have one position."
         )
         
         self.assertEqual(
-            perf_tracker.cumulative_performance.positions[133].sid, 
+            trading_client.perf.cumulative_performance.positions[133].sid, 
             133, 
             "Portfolio should have one position in 133."
         )

zipline/test/test_perf_tracking.py

@@ -506,34 +506,46 @@ shares in position"
         
         trade_count = 100
         sid = 133
-        price = [10.1] * trade_count
+        price = 10.1 
+        price_list = [price] * trade_count
         volume = [100] * trade_count
         start_date = datetime.datetime.strptime("01/01/2011","%m/%d/%Y")
         start_date = start_date.replace(tzinfo=pytz.utc)
         trade_time_increment = datetime.timedelta(days=1)
         trade_history = factory.create_trade_history( 
             sid, 
-            price, 
+            price_list, 
             volume, 
             start_date, 
             trade_time_increment, 
             self.trading_environment 
         )
         
-        trade_client = TradeSimulationClient(start_date)
-        start = trade_history[0].dt
-        end = trade_history[-1].dt
-        tracker = perf.PerformanceTracker(
-            start, 
-            end, 
-            1000.0, 
-            self.trading_environment
+        sid2 = 134
+        price2 = 12.12
+        price2_list = [price2] * trade_count 
+        trade_history2 = factory.create_trade_history( 
+            sid2, 
+            price2_list, 
+            volume, 
+            start_date, 
+            trade_time_increment, 
+            self.trading_environment 
         )
         
+        trade_history.extend(trade_history2)
+        
+        self.trading_environment.period_start = trade_history[0].dt
+        self.trading_environment.period_end = trade_history[-1].dt
+        self.trading_environment.capital_base = 1000.0
+        self.trading_environment.frame_index = ['sid', 'volume', 'dt', \
+        'price', 'changed']
+        client = TradeSimulationClient(self.trading_environment)
+        
         for event in trade_history:
             #create a transaction for all but
-            #one trade, to simulate None transaction
-            if(event.dt != start):
+            #first trade in each sid, to simulate None transaction
+            if(event.dt != self.trading_environment.period_start):
                 txn = zp.namedict({
                     'sid'        : event.sid,
                     'amount'     : -25,
@@ -543,17 +555,19 @@ shares in position"
                 })
             else:
                 txn = None
-            event[zp.TRANSFORM_TYPE.TRANSACTION] = txn
-            trade_client.queue_event(event)
+            event[zp.TRANSFORM_TYPE.TRANSACTION] = txn  
+            client.queue_event(event)
             
-        df = trade_client.get_frame()    
-        tracker.update(df)
+        df = client.get_frame()
         
-        #we skip one trade, to test case of None transaction
-        txn_count = len(trade_history) - 1
-        self.assertEqual(tracker.txn_count, txn_count)
+        self.assertEqual(df[133]['price'], price)
+        self.assertEqual(df[134]['price'], price2)
         
-        cumulative_pos = tracker.cumulative_performance.positions[sid]
-        expected_size = txn_count * -25
+        #we skip two trades, to test case of None transaction
+        txn_count = len(trade_history) - 2
+        self.assertEqual(client.perf.txn_count, txn_count)
+        
+        cumulative_pos = client.perf.cumulative_performance.positions[sid]
+        expected_size = txn_count / 2 * -25
         self.assertEqual(cumulative_pos.amount, expected_size)