add account object to context

docs/release-notes/zipline-0.7.1.md

@@ -4,3 +4,15 @@
 
 * Fix a bug where the reported returns could sharply dip for random periods of time.
   [PR378](https://github.com/quantopian/zipline/pull/378)
+
+## Enhancements (ENH)
+
+* Account object: Adds an account object to conext to track information about the trading account. [PR396](https://github.com/quantopian/zipline/pull/396)
+
+  > Example:
+
+  > ```
+  > context.account.settled_cash
+  > ```
+
+  > Returns the settled cash value that is stored on the account object. This value is updated accordingly as the algorithm is run.

tests/test_algorithm.py

@@ -34,6 +34,7 @@ from zipline.errors import (
     TradingControlViolation,
 )
 from zipline.test_algorithms import (
+    access_account_in_init,
     access_portfolio_in_init,
     AmbitiousStopLimitAlgorithm,
     EmptyPositionsAlgorithm,
@@ -687,6 +688,24 @@ def handle_data(context, data):
 
         output, _ = drain_zipline(self, zipline)
 
+    def test_account_in_init(self):
+        """
+        Test that accessing account in init doesn't break.
+        """
+        test_algo = TradingAlgorithm(
+            script=access_account_in_init,
+            sim_params=self.sim_params,
+        )
+        set_algo_instance(test_algo)
+
+        self.zipline_test_config['algorithm'] = test_algo
+        self.zipline_test_config['trade_count'] = 1
+
+        zipline = simfactory.create_test_zipline(
+            **self.zipline_test_config)
+
+        output, _ = drain_zipline(self, zipline)
+
 
 class TestHistory(TestCase):
     def test_history(self):

tests/test_perf_tracking.py

@@ -29,6 +29,7 @@ import pytz
 import itertools
 
 import pandas as pd
+import numpy as np
 from six.moves import range, zip
 
 import zipline.utils.factory as factory
@@ -107,6 +108,7 @@ def calculate_results(host,
     splits = splits or []
 
     perf_tracker = perf.PerformanceTracker(host.sim_params)
+
     if dividend_events is not None:
         dividend_frame = pd.DataFrame(
             [
@@ -149,6 +151,7 @@ def calculate_results(host,
 
         if bm_updated:
             msg = perf_tracker.handle_market_close_daily()
+            msg['account'] = perf_tracker.get_account(True)
             results.append(msg)
             bm_updated = False
     return results
@@ -216,6 +219,28 @@ class TestSplitPerformance(unittest.TestCase):
             zp_math.tolerant_equals(8020,
                                     daily_perf['ending_cash'], 1))
 
+        # Validate that the account attributes were updated.
+        account = results[1]['account']
+        self.assertEqual(float('inf'), account['day_trades_remaining'])
+        np.testing.assert_allclose(0.198, account['leverage'], rtol=1e-3)
+        np.testing.assert_allclose(8020, account['regt_equity'], rtol=1e-3)
+        self.assertEqual(float('inf'), account['regt_margin'])
+        np.testing.assert_allclose(8020, account['available_funds'], rtol=1e-3)
+        self.assertEqual(0, account['maintenance_margin_requirement'])
+        np.testing.assert_allclose(10000,
+                                   account['equity_with_loan'], rtol=1e-3)
+        self.assertEqual(float('inf'), account['buying_power'])
+        self.assertEqual(0, account['initial_margin_requirement'])
+        np.testing.assert_allclose(8020, account['excess_liquidity'],
+                                   rtol=1e-3)
+        np.testing.assert_allclose(8020, account['settled_cash'], rtol=1e-3)
+        np.testing.assert_allclose(10000, account['net_liquidation'],
+                                   rtol=1e-3)
+        np.testing.assert_allclose(0.802, account['cushion'], rtol=1e-3)
+        np.testing.assert_allclose(1980, account['total_positions_value'],
+                                   rtol=1e-3)
+        self.assertEqual(0, account['accrued_interest'])
+
         for i, result in enumerate(results):
             for perf_kind in ('daily_perf', 'cumulative_perf'):
                 perf_result = result[perf_kind]
@@ -291,6 +316,30 @@ class TestCommissionEvents(unittest.TestCase):
             # Validate that the cost basis of our position changed.
             self.assertEqual(results[-1]['daily_perf']['positions']
                              [0]['cost_basis'], 320.0)
+            # Validate that the account attributes were updated.
+            account = results[1]['account']
+            self.assertEqual(float('inf'), account['day_trades_remaining'])
+            np.testing.assert_allclose(0.001, account['leverage'], rtol=1e-3,
+                                       atol=1e-4)
+            np.testing.assert_allclose(9680, account['regt_equity'], rtol=1e-3)
+            self.assertEqual(float('inf'), account['regt_margin'])
+            np.testing.assert_allclose(9680, account['available_funds'],
+                                       rtol=1e-3)
+            self.assertEqual(0, account['maintenance_margin_requirement'])
+            np.testing.assert_allclose(9690,
+                                       account['equity_with_loan'], rtol=1e-3)
+            self.assertEqual(float('inf'), account['buying_power'])
+            self.assertEqual(0, account['initial_margin_requirement'])
+            np.testing.assert_allclose(9680, account['excess_liquidity'],
+                                       rtol=1e-3)
+            np.testing.assert_allclose(9680, account['settled_cash'],
+                                       rtol=1e-3)
+            np.testing.assert_allclose(9690, account['net_liquidation'],
+                                       rtol=1e-3)
+            np.testing.assert_allclose(0.999, account['cushion'], rtol=1e-3)
+            np.testing.assert_allclose(10, account['total_positions_value'],
+                                       rtol=1e-3)
+            self.assertEqual(0, account['accrued_interest'])
 
     def test_commission_zero_position(self):
         """

zipline/algorithm.py

@@ -176,7 +176,10 @@ class TradingAlgorithm(object):
             self.blotter = Blotter()
 
         self.portfolio_needs_update = True
+        self.account_needs_update = True
+        self.performance_needs_update = True
         self._portfolio = None
+        self._account = None
 
         self.history_container = None
         self.history_specs = {}
@@ -356,6 +359,8 @@ class TradingAlgorithm(object):
             self.perf_tracker = PerformanceTracker(sim_params)
 
         self.portfolio_needs_update = True
+        self.account_needs_update = True
+        self.performance_needs_update = True
 
         self.data_gen = self._create_data_generator(source_filter, sim_params)
 
@@ -702,10 +707,24 @@ class TradingAlgorithm(object):
 
     def updated_portfolio(self):
         if self.portfolio_needs_update:
-            self._portfolio = self.perf_tracker.get_portfolio()
+            self._portfolio = \
+                self.perf_tracker.get_portfolio(self.performance_needs_update)
             self.portfolio_needs_update = False
+            self.performance_needs_update = False
         return self._portfolio
 
+    @property
+    def account(self):
+        return self.updated_account()
+
+    def updated_account(self):
+        if self.account_needs_update:
+            self._account = \
+                self.perf_tracker.get_account(self.performance_needs_update)
+            self.account_needs_update = False
+            self.performance_needs_update = False
+        return self._account
+
     def set_logger(self, logger):
         self.logger = logger
 

zipline/finance/performance/period.py

@@ -104,6 +104,7 @@ class PerformancePeriod(object):
         self.ending_value = 0.0
         self.period_cash_flow = 0.0
         self.pnl = 0.0
+
         # sid => position object
         self.positions = positiondict()
         self.ending_cash = starting_cash
@@ -122,6 +123,7 @@ class PerformancePeriod(object):
         # when returning portfolio information.
         # So as not to avoid creating a new object for each event
         self._portfolio_store = zp.Portfolio()
+        self._account_store = zp.Account()
         self._positions_store = zp.Positions()
         self.serialize_positions = serialize_positions
 
@@ -251,6 +253,9 @@ class PerformancePeriod(object):
     def adjust_cash(self, amount):
         self.period_cash_flow += amount
 
+    def adjust_field(self, field, value):
+        setattr(self, field, value)
+
     def calculate_performance(self):
         self.ending_value = self.calculate_positions_value()
 
@@ -407,6 +412,50 @@ class PerformancePeriod(object):
         portfolio.positions_value = self.ending_value
         return portfolio
 
+    def as_account(self):
+        account = self._account_store
+
+        # If no attribute is found on the PerformancePeriod resort to the
+        # following default values. If an attribute is found use the existing
+        # value. For instance, a broker may provide updates to these
+        # attributes. In this case we do not want to over write the broker
+        # values with the default values.
+        account.settled_cash = \
+            getattr(self, 'settled_cash', self.ending_cash)
+        account.accrued_interest = \
+            getattr(self, 'accrued_interest', 0.0)
+        account.buying_power = \
+            getattr(self, 'buying_power', float('inf'))
+        account.equity_with_loan = \
+            getattr(self, 'equity_with_loan',
+                    self.ending_cash + self.ending_value)
+        account.total_positions_value = \
+            getattr(self, 'total_positions_value', self.ending_value)
+        account.regt_equity = \
+            getattr(self, 'regt_equity', self.ending_cash)
+        account.regt_margin = \
+            getattr(self, 'regt_margin', float('inf'))
+        account.initial_margin_requirement = \
+            getattr(self, 'initial_margin_requirement', 0.0)
+        account.maintenance_margin_requirement = \
+            getattr(self, 'maintenance_margin_requirement', 0.0)
+        account.available_funds = \
+            getattr(self, 'available_funds', self.ending_cash)
+        account.excess_liquidity = \
+            getattr(self, 'excess_liquidity', self.ending_cash)
+        account.cushion = \
+            getattr(self, 'cushion',
+                    self.ending_cash / (self.ending_cash + self.ending_value))
+        account.day_trades_remaining = \
+            getattr(self, 'day_trades_remaining', float('inf'))
+        account.leverage = \
+            getattr(self, 'leverage',
+                    self.ending_value / (self.ending_value + self.ending_cash))
+        account.net_liquidation = \
+            getattr(self, 'net_liquidation',
+                    self.ending_cash + self.ending_value)
+        return account
+
     def get_positions(self):
 
         positions = self._positions_store

zipline/finance/performance/tracker.py

@@ -229,10 +229,16 @@ class PerformanceTracker(object):
         for perf_period in self.perf_periods:
             perf_period.calculate_performance()
 
-    def get_portfolio(self):
-        self.update_performance()
+    def get_portfolio(self, performance_needs_update):
+        if performance_needs_update:
+            self.update_performance()
         return self.cumulative_performance.as_portfolio()
 
+    def get_account(self, performance_needs_update):
+        if performance_needs_update:
+            self.update_performance()
+        return self.cumulative_performance.as_account()
+
     def to_dict(self, emission_type=None):
         """
         Creates a dictionary representing the state of this tracker.

zipline/gens/tradesimulation.py

@@ -177,6 +177,8 @@ class AlgorithmSimulator(object):
                             self._call_before_trading_start(next_day)
 
                     self.algo.portfolio_needs_update = True
+                    self.algo.account_needs_update = True
+                    self.algo.performance_needs_update = True
 
             risk_message = self.algo.perf_tracker.handle_simulation_end()
             yield risk_message
@@ -277,6 +279,7 @@ class AlgorithmSimulator(object):
         # dt before we emit a perf message.  This is a no-op if
         # updated_portfolio has already been called this dt.
         self.algo.updated_portfolio()
+        self.algo.updated_account()
 
         rvars = self.algo.recorded_vars
         if self.algo.perf_tracker.emission_rate == 'daily':

zipline/protocol.py

@@ -134,6 +134,38 @@ class Portfolio(object):
         return "Portfolio({0})".format(self.__dict__)
 
 
+class Account(object):
+    '''
+    The account object tracks information about the trading account. The
+    values are updated as the algorithm runs and its keys remain unchanged.
+    If connected to a broker, one can update these values with the trading
+    account values as reported by the broker.
+    '''
+
+    def __init__(self):
+        self.settled_cash = 0.0
+        self.accrued_interest = 0.0
+        self.buying_power = float('inf')
+        self.equity_with_loan = 0.0
+        self.total_positions_value = 0.0
+        self.regt_equity = 0.0
+        self.regt_margin = float('inf')
+        self.initial_margin_requirement = 0.0
+        self.maintenance_margin_requirement = 0.0
+        self.available_funds = 0.0
+        self.excess_liquidity = 0.0
+        self.cushion = 0.0
+        self.day_trades_remaining = float('inf')
+        self.leverage = 0.0
+        self.net_liquidation = 0.0
+
+    def __getitem__(self, key):
+        return self.__dict__[key]
+
+    def __repr__(self):
+        return "Account({0})".format(self.__dict__)
+
+
 class Position(object):
 
     def __init__(self, sid):

zipline/test_algorithms.py

@@ -980,6 +980,15 @@ def handle_data(context, data):
     pass
 """
 
+access_account_in_init = """
+def initialize(context):
+    var = context.account.settled_cash
+    pass
+
+def handle_data(context, data):
+    pass
+"""
+
 call_all_order_methods = """
 from zipline.api import (order,
                          order_value,