From bd268cc23c844bdcd9e6247e945f49ab2e3e6903 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Javier=20Rodr=C3=ADguez=20Chatruc?= <jrchatruc@gmail.com>
Date: Mon, 6 Jan 2020 12:35:12 -0300
Subject: [PATCH] Strategy now chooses to buy/sell as many contracts as initial
 capital allows

---
 backtester/backtester.py        |  9 +++++----
 backtester/strategy/strategy.py | 19 +++++++++++++------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/backtester/backtester.py b/backtester/backtester.py
index deab0c4..bc7b9b0 100644
--- a/backtester/backtester.py
+++ b/backtester/backtester.py
@@ -8,8 +8,7 @@ from .datahandler import HistoricalOptionsData
 
 class Backtest:
     """Processes signals from the Strategy object"""
-    def __init__(self, capital=1_000_000):
-        self.initial_capital = self.current_capital = capital
+    def __init__(self):
         self._strategy = None
         self._data = None
         self.inventory = pd.DataFrame()
@@ -23,6 +22,7 @@ class Backtest:
     def strategy(self, strat):
         assert isinstance(strat, Strategy)
         self._strategy = strat
+        self.current_capital = strat.initial_capital
 
     @property
     def data(self):
@@ -50,7 +50,7 @@ class Backtest:
         index = pd.MultiIndex.from_product(
             [[l.name for l in self._strategy.legs],
              ['contract', 'underlying', 'expiration', 'type', 'strike', 'cost', 'date', 'order']])
-        index_totals = pd.MultiIndex.from_product([['totals'], ['cost']])
+        index_totals = pd.MultiIndex.from_product([['totals'], ['cost', 'qty']])
         self.inventory = pd.DataFrame(columns=index.append(index_totals))
         self.trade_log = pd.DataFrame()
 
@@ -91,7 +91,8 @@ class Backtest:
         if not entry_signals.empty:
             # costs = entry_signals['totals']['cost']
             # return entry_signals.loc[costs.idxmin():costs.idxmin()], costs.min()
-            return entry_signals.iloc[0], entry_signals.iloc[0]['totals']['cost']
+            entry = entry_signals.iloc[0]
+            return entry, entry['totals']['cost'] * entry['totals']['qty']
         else:
             return entry_signals, 0
 
diff --git a/backtester/strategy/strategy.py b/backtester/strategy/strategy.py
index dfbb80c..991e4d7 100644
--- a/backtester/strategy/strategy.py
+++ b/backtester/strategy/strategy.py
@@ -18,11 +18,12 @@ class Strategy:
     Takes in a number of `StrategyLeg`'s (option contracts), and filters that determine
     entry and exit conditions.
     """
-    def __init__(self, schema, qty=1, shares_per_contract=100):
+    def __init__(self, schema, qty=1, shares_per_contract=100, initial_capital=1_000_000):
         assert isinstance(schema, Schema)
         self.schema = schema
         self.qty = qty
         self._shares_per_contract = shares_per_contract
+        self.initial_capital = initial_capital
         self.legs = []
         self.conditions = []
         self.exit_thresholds = (0.0, 0.0)
@@ -121,15 +122,17 @@ class Strategy:
             leg_candidates[i].columns = pd.MultiIndex.from_product([["leg_{}".format(i + 1)],
                                                                     leg_candidates[i].columns])
 
-        totals = pd.DataFrame.from_dict({"cost": total_costs})
+        qtys = inventory['totals']['qty']
+        totals = pd.DataFrame.from_dict({"cost": total_costs, "qty": qtys})
         totals.columns = pd.MultiIndex.from_product([["totals"], totals.columns])
         leg_candidates.append(totals)
         filter_mask = reduce(lambda x, y: x | y, filter_mask)
         exits_mask = threshold_exits | filter_mask
 
         exits = pd.concat([l[exits_mask] for l in leg_candidates], axis=1)
+        total_costs = total_costs[exits_mask] * exits['totals']['qty']
 
-        return (exits, exits_mask, total_costs[exits_mask])
+        return (exits, exits_mask, total_costs)
 
     def _filter_legs(self, options, signal):
         """Returns a hierarchically indexed `pd.DataFrame` containing signals for each
@@ -164,7 +167,7 @@ class Strategy:
             if leg.direction == Direction.SELL:
                 subset_df['cost'] = -subset_df['cost']
 
-            subset_df['cost'] *= self._shares_per_contract * self.qty
+            subset_df['cost'] *= self._shares_per_contract
 
             dfs.append(subset_df.reset_index(drop=True))
 
@@ -204,7 +207,11 @@ class Strategy:
             return pd.DataFrame()
 
         cost = sum(leg["cost"] for leg in dfs)
-        totals = pd.DataFrame.from_dict({"cost": cost})
+        # Put qty of contracts to buy/sell in ['totals']['qty']
+        qty = np.floor(self.initial_capital / cost)
+        qty = np.abs(qty)
+        # qty = qty.astype(int)
+        totals = pd.DataFrame.from_dict({"cost": cost, "qty": qty})
         totals.columns = pd.MultiIndex.from_product([["totals"], totals.columns])
 
         for i in range(len(dfs)):
@@ -242,7 +249,7 @@ class Strategy:
         if ~direction == Direction.SELL:
             candidates['cost'] = -candidates['cost']
 
-        candidates['cost'] *= self._shares_per_contract * self.qty
+        candidates['cost'] *= self._shares_per_contract
 
         return candidates