Added support for shorting in the drift rebalancer; fixed a bug in limit price calculation.

README.md

@@ -188,7 +188,6 @@ git fetch origin
 git merge origin/dev
 git checkout my-feature
 git rebase dev
-git checkout my-feature
 git push --force-with-lease origin my-feature
 ```
 

lumibot/example_strategies/classic_60_40.py

@@ -17,13 +17,14 @@
     parameters = {
         "market": "NYSE",
         "sleeptime": "1D",
-        "absolute_drift_threshold": "0.15",
-        "acceptable_slippage": "0.0005",
+        "drift_threshold": "0.15",
+        "acceptable_slippage": "0.005",  # 50 BPS
         "fill_sleeptime": 15,
         "target_weights": {
             "SPY": "0.60",
             "TLT": "0.40"
-        }
+        },
+        "shorting": False
     }
 
     if is_live:

lumibot/strategies/drift_rebalancer.py

@@ -30,7 +30,7 @@ class DriftRebalancer(Strategy):
     """The DriftRebalancer strategy rebalances a portfolio based on drift from target weights.
 
     The strategy calculates the drift of each asset in the portfolio and triggers a rebalance if the drift exceeds
-    the absolute_drift_threshold. The strategy will sell assets that have drifted above the threshold and
+    the drift_threshold. The strategy will sell assets that have drifted above the threshold and
     buy assets that have drifted below the threshold.
 
     The current version of the DriftRebalancer strategy only supports limit orders and whole share quantities.
@@ -46,14 +46,14 @@ class DriftRebalancer(Strategy):
 
         ### DriftRebalancer parameters
 
-        # This is the absolute drift threshold that will trigger a rebalance. If the target_weight is 0.30 and the
-        # absolute_drift_threshold is 0.05, then the rebalance will be triggered when the assets current_weight
+        # This is the drift threshold that will trigger a rebalance. If the target_weight is 0.30 and the
+        # drift_threshold is 0.05, then the rebalance will be triggered when the assets current_weight
         # is less than 0.25 or greater than 0.35.
-        "absolute_drift_threshold": "0.05",
+        "drift_threshold": "0.05",
 
         # This is the acceptable slippage that will be used when calculating the number of shares to buy or sell.
-        # The default is 0.0005 (5 BPS)
-        "acceptable_slippage": "0.0005",
+        # The default is 0.005 (50 BPS)
+        "acceptable_slippage": "0.005",  # 50 BPS
 
          # The amount of time to sleep between the sells and buys to give enough time for the orders to fill
         "fill_sleeptime": 15,
@@ -64,35 +64,41 @@ class DriftRebalancer(Strategy):
             "TLT": "0.40",
             "USD": "0.00",
         }
+
+        # If you want to allow shorting, set this to True.
+        shorting: False
     }
     """
 
     # noinspection PyAttributeOutsideInit
     def initialize(self, parameters: Any = None) -> None:
         self.set_market(self.parameters.get("market", "NYSE"))
         self.sleeptime = self.parameters.get("sleeptime", "1D")
-        self.absolute_drift_threshold = Decimal(self.parameters.get("absolute_drift_threshold", "0.20"))
-        self.acceptable_slippage = Decimal(self.parameters.get("acceptable_slippage", "0.0005"))
+        self.drift_threshold = Decimal(self.parameters.get("drift_threshold", "0.20"))
+        self.acceptable_slippage = Decimal(self.parameters.get("acceptable_slippage", "0.005"))
         self.fill_sleeptime = self.parameters.get("fill_sleeptime", 15)
         self.target_weights = {k: Decimal(v) for k, v in self.parameters["target_weights"].items()}
+        self.shorting = self.parameters.get("shorting", False)
         self.drift_df = pd.DataFrame()
 
         # Sanity checks
-        if self.acceptable_slippage >= self.absolute_drift_threshold:
-            raise ValueError("acceptable_slippage must be less than absolute_drift_threshold")
-        if self.absolute_drift_threshold >= Decimal("1.0"):
-            raise ValueError("absolute_drift_threshold must be less than 1.0")
+        if self.acceptable_slippage >= self.drift_threshold:
+            raise ValueError("acceptable_slippage must be less than drift_threshold")
+        if self.drift_threshold >= Decimal("1.0"):
+            raise ValueError("drift_threshold must be less than 1.0")
         for key, target_weight in self.target_weights.items():
-            if self.absolute_drift_threshold >= target_weight:
+            if self.drift_threshold >= target_weight:
                 logger.warning(
-                    f"absolute_drift_threshold of {self.absolute_drift_threshold} is "
+                    f"drift_threshold of {self.drift_threshold} is "
                     f">= target_weight of {key}: {target_weight}. Drift in this asset will never trigger a rebalance."
                 )
 
     # noinspection PyAttributeOutsideInit
     def on_trading_iteration(self) -> None:
         dt = self.get_datetime()
-        logger.info(f"{dt} on_trading_iteration called")
+        msg = f"{dt} on_trading_iteration called"
+        logger.info(msg)
+        self.log_message(msg, broadcast=True)
         self.cancel_open_orders()
 
         if self.cash < 0:
@@ -123,14 +129,17 @@ def on_trading_iteration(self) -> None:
         # Check if the absolute value of any drift is greater than the threshold
         rebalance_needed = False
         for index, row in self.drift_df.iterrows():
-            if row["absolute_drift"] > self.absolute_drift_threshold:
+            msg = (
+                f"Symbol: {row['symbol']} current_weight: {row['current_weight']:.2%} "
+                f"target_weight: {row['target_weight']:.2%} drift: {row['drift']:.2%}"
+            )
+            if abs(row["drift"]) > self.drift_threshold:
                 rebalance_needed = True
-                msg = (
-                    f"Absolute drift for {row['symbol']} is {row['absolute_drift']:.2f} "
-                    f"and exceeds threshold of {self.absolute_drift_threshold:.2f}"
+                msg += (
+                    f" Absolute drift exceeds threshold of {self.drift_threshold:.2%}. Rebalance needed."
                 )
-                logger.info(msg)
-                self.log_message(msg, broadcast=True)
+            logger.info(msg)
+            self.log_message(msg, broadcast=True)
 
         if rebalance_needed:
             msg = f"Rebalancing portfolio."
@@ -140,7 +149,8 @@ def on_trading_iteration(self) -> None:
                 strategy=self,
                 df=self.drift_df,
                 fill_sleeptime=self.fill_sleeptime,
-                acceptable_slippage=self.acceptable_slippage
+                acceptable_slippage=self.acceptable_slippage,
+                shorting=self.shorting
             )
             rebalance_logic.rebalance()
 
@@ -167,7 +177,7 @@ def __init__(self, target_weights: Dict[str, Decimal]) -> None:
             "current_weight": Decimal(0),
             "target_weight": [Decimal(weight) for weight in target_weights.values()],
             "target_value": Decimal(0),
-            "absolute_drift": Decimal(0)
+            "drift": Decimal(0)
         })
 
     def add_position(self, *, symbol: str, is_quote_asset: bool, current_quantity: Decimal, current_value: Decimal) -> None:
@@ -184,7 +194,7 @@ def add_position(self, *, symbol: str, is_quote_asset: bool, current_quantity: D
                 "current_weight": Decimal(0),
                 "target_weight": Decimal(0),
                 "target_value": Decimal(0),
-                "absolute_drift": Decimal(0)
+                "drift": Decimal(0)
             }
             # Convert the dictionary to a DataFrame
             new_row_df = pd.DataFrame([new_row])
@@ -205,14 +215,20 @@ def calculate_drift_row(row: pd.Series) -> Decimal:
             if row["is_quote_asset"]:
                 # We can never buy or sell the quote asset
                 return Decimal(0)
+
+            # Check if we should sell everything
             elif row["current_quantity"] > Decimal(0) and row["target_weight"] == Decimal(0):
                 return Decimal(-1)
+
+            # Check if we need to buy for the first time
             elif row["current_quantity"] == Decimal(0) and row["target_weight"] > Decimal(0):
                 return Decimal(1)
+
+            # Otherwise we just need to adjust our holding
             else:
                 return row["target_weight"] - row["current_weight"]
 
-        self.df["absolute_drift"] = self.df.apply(calculate_drift_row, axis=1)
+        self.df["drift"] = self.df.apply(calculate_drift_row, axis=1)
         return self.df.copy()
 
 
@@ -223,41 +239,64 @@ def __init__(
             strategy: Strategy,
             df: pd.DataFrame,
             fill_sleeptime: int = 15,
-            acceptable_slippage: Decimal = Decimal("0.0005")
+            acceptable_slippage: Decimal = Decimal("0.005"),
+            shorting: bool = False
     ) -> None:
         self.strategy = strategy
         self.df = df
         self.fill_sleeptime = fill_sleeptime
         self.acceptable_slippage = acceptable_slippage
+        self.shorting = shorting
 
     def rebalance(self) -> None:
         # Execute sells first
+        sell_orders = []
         for index, row in self.df.iterrows():
-            if row["absolute_drift"] == -1:
+            if row["drift"] == -1:
                 # Sell everything
                 symbol = row["symbol"]
                 quantity = row["current_quantity"]
                 last_price = Decimal(self.strategy.get_last_price(symbol))
                 limit_price = self.calculate_limit_price(last_price=last_price, side="sell")
-                self.place_limit_order(symbol=symbol, quantity=quantity, limit_price=limit_price, side="sell")
-            elif row["absolute_drift"] < 0:
+                if quantity > 0 or (quantity == 0 and self.shorting):
+                    order = self.place_limit_order(
+                        symbol=symbol,
+                        quantity=quantity,
+                        limit_price=limit_price,
+                        side="sell"
+                    )
+                    sell_orders.append(order)
+
+            elif row["drift"] < 0:
                 symbol = row["symbol"]
                 last_price = Decimal(self.strategy.get_last_price(symbol))
                 limit_price = self.calculate_limit_price(last_price=last_price, side="sell")
                 quantity = ((row["current_value"] - row["target_value"]) / limit_price).quantize(Decimal('1'), rounding=ROUND_DOWN)
-                if quantity > 0:
-                    self.place_limit_order(symbol=symbol, quantity=quantity, limit_price=limit_price, side="sell")
+                if quantity > 0 and (quantity < row["current_quantity"] or self.shorting):
+                    order = self.place_limit_order(
+                        symbol=symbol,
+                        quantity=quantity,
+                        limit_price=limit_price,
+                        side="sell"
+                    )
+                    sell_orders.append(order)
+
+        for order in sell_orders:
+            logger.info(f"Submitted sell order: {order}")
 
         if not self.strategy.is_backtesting:
             # Sleep to allow sell orders to fill
             time.sleep(self.fill_sleeptime)
+            orders = self.strategy.broker._pull_all_orders(self.strategy.name, self.strategy)
+            for order in orders:
+                logger.info(f"Order at broker: {order}")
 
         # Get current cash position from the broker
         cash_position = self.get_current_cash_position()
 
         # Execute buys
         for index, row in self.df.iterrows():
-            if row["absolute_drift"] > 0:
+            if row["drift"] > 0:
                 symbol = row["symbol"]
                 last_price = Decimal(self.strategy.get_last_price(symbol))
                 limit_price = self.calculate_limit_price(last_price=last_price, side="buy")
@@ -271,19 +310,19 @@ def rebalance(self) -> None:
 
     def calculate_limit_price(self, *, last_price: Decimal, side: str) -> Decimal:
         if side == "sell":
-            return last_price * (1 - self.acceptable_slippage / Decimal(10000))
+            return last_price * (1 - self.acceptable_slippage)
         elif side == "buy":
-            return last_price * (1 + self.acceptable_slippage / Decimal(10000))
+            return last_price * (1 + self.acceptable_slippage)
 
     def get_current_cash_position(self) -> Decimal:
         self.strategy.update_broker_balances(force_update=True)
         return Decimal(self.strategy.cash)
 
-    def place_limit_order(self, *, symbol: str, quantity: Decimal, limit_price: Decimal, side: str) -> None:
+    def place_limit_order(self, *, symbol: str, quantity: Decimal, limit_price: Decimal, side: str) -> Any:
         limit_order = self.strategy.create_order(
             asset=symbol,
             quantity=quantity,
             side=side,
             limit_price=float(limit_price)
         )
-        self.strategy.submit_order(limit_order)
+        return self.strategy.submit_order(limit_order)