Refactoring Backtest class

backtester/backtester.py

@@ -2,66 +2,118 @@ import pandas as pd
 import numpy as np
 import pyprind
 
-from .strategy import Strategy, Order
-from .datahandler import HistoricalOptionsData
+from .strategy import Strategy
+from .enums import Order, Stock
+from .datahandler import HistoricalOptionsData, TiingoData
 
 
 class Backtest:
     """Processes signals from the Strategy object"""
-    def __init__(self):
-        self._strategy = None
-        self._data = None
+    def __init__(self, allocation, initial_capital=1_000_000):
+        assert isinstance(allocation, dict)
+
+        assets = ('stocks', 'options', 'cash')
+        total_allocation = sum(allocation.get(a, 0.0) for a in assets)
+
+        self.allocation = {}
+        for asset in assets:
+            self.allocation[asset] = allocation.get(asset, 0.0) / total_allocation
+
+        self.current_cash = self.initial_capital = initial_capital
         self.stop_if_broke = True
+        self._stocks = []
+        self._options_strategy = None
+        self._stock_data = None
+        self._options_data = None
+
+    def add_stock(self, stock):
+        """Adds stock to the backtest"""
+        assert isinstance(stock, Stock)
+        self.stocks.append(stock)
+        return self
+
+    def add_stocks(self, stocks):
+        """Adds stocks to the backtest"""
+        for stock in stocks:
+            self.add_stock(stock)
+        return self
 
     @property
     def strategy(self):
-        return self._strategy
+        return self._options_strategy
 
     @strategy.setter
     def strategy(self, strat):
         assert isinstance(strat, Strategy)
-        self._strategy = strat
+        self._options_strategy = strat
         self.current_cash = strat.initial_capital
 
     @property
-    def data(self):
-        return self._data
+    def stock_data(self):
+        return self._stock_data
 
-    @data.setter
-    def data(self, data):
+    @stock_data.setter
+    def stock_data(self, data):
+        assert isinstance(data, TiingoData)
+        self._stock_data = data
+
+    @property
+    def options_data(self):
+        return self._options_data
+
+    @options_data.setter
+    def options_data(self, data):
         assert isinstance(data, HistoricalOptionsData)
-        self._data = data
+        self._options_data = data
 
-    def run(self, monthly=False):
+    def run(self, rebalance_freq=0, monthly=False):
         """Runs the backtest and returns a `pd.DataFrame` of the orders executed (`self.trade_log`)
 
         Args:
-            monthly (bool, optional):   Iterates through data monthly rather than daily. Defaults to False.
+            rebalance_freq (int, optional): Determines the frequency of portfolio rebalances. Defaults to 0.
+            monthly (bool, optional):       Iterates through data monthly rather than daily. Defaults to False.
 
         Returns:
-            pd.DataFrame:               Log of the trades executed.
+            pd.DataFrame:                   Log of the trades executed.
         """
 
-        assert self._data is not None
-        assert self._strategy is not None
-        assert self._data.schema == self._strategy.schema
+        assert self._stock_data, 'Stock data not set'
+        assert self._options_data, 'Options data not set'
+        assert self._options_strategy, 'Options Strategy not set'
+        assert self._options_data.schema == self._options_strategy.schema
+
+        option_dates = self._options_data['date'].unique()
+        stock_dates = self._stock_data['date'].unique()
+        assert np.array_equal(stock_dates, option_dates), 'Stock and options dates do not match'
 
         columns = pd.MultiIndex.from_product(
-            [[l.name for l in self._strategy.legs],
+            [[l.name for l in self._options_strategy.legs],
              ['contract', 'underlying', 'expiration', 'type', 'strike', 'cost', 'order']])
         totals = pd.MultiIndex.from_product([['totals'], ['cost', 'qty', 'date']])
-        self.inventory = pd.DataFrame(columns=columns.append(totals))
+        self.options_inventory = pd.DataFrame(columns=columns.append(totals))
+
+        self.stock_inventory = pd.DataFrame(columns=['symbol', 'cost', 'qty'])
+
+        rebalancing_days = pd.date_range(
+            self.stock_data.first_date, self.stock_data.end_date, freq=str(rebalance_freq) +
+            'BMS') if rebalance_freq else []
+
         self.trade_log = pd.DataFrame()
         self.balance = pd.DataFrame({
             'capital': self.current_cash,
             'cash': self.current_cash
         },
-                                    index=[self.data.start_date - pd.Timedelta(1, unit='day')])
+                                    index=[self.stock_data.start_date - pd.Timedelta(1, unit='day')])
 
-        data_iterator = self._data.iter_months() if monthly else self._data.iter_dates()
+        data_iterator = self._data_iterator(monthly)
         bar = pyprind.ProgBar(data_iterator.ngroups, bar_char='█')
 
-        for date, options in data_iterator:
+        for date, stocks, options in data_iterator:
+            if date == first_day:
+                self._rebalance_portfolio(data, sma_days)
+            self._update_balance(date, data)
+            if date in rebalancing_days:
+                self._rebalance_portfolio(data, sma_days)
             entry_signals = self._strategy.filter_entries(options, self.inventory, date)
             exit_signals = self._strategy.filter_exits(options, self.inventory, date)
 
@@ -76,6 +128,21 @@ class Backtest:
 
         return self.trade_log
 
+    def _data_iterator(self, monthly):
+        """Returns combined iterator for stock and options data.
+        Each step, it produces a tuple like the following:
+            (date, stocks, options)
+
+        Returns:
+            generator: Daily/monthly iterator over `self.stock_data` and `self.options_data`
+        """
+        if monthly:
+            it = zip(self._stock_data.iter_months(), self._options_data.iter_months())
+        else:
+            it = zip(self._stock_data.iter_dates(), self._options_data.iter_dates())
+
+        return ((date, stocks, options) for (date, stocks), (_, options) in it)
+
     def _execute_entry(self, entry_signals):
         """Executes entry orders and updates `self.inventory` and `self.trade_log`"""
         entry, total_price = self._process_entry_signals(entry_signals)