fix division by zero bug when shorting

Lumiwealth · Dec 19, 2024 · 354e004 · 354e004
1 parent 1f73603
commit 354e004
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/lumibot/components/drift_rebalancer_logic.py b/lumibot/components/drift_rebalancer_logic.py
@@ -222,19 +222,23 @@ def _calculate_drift_row(self, row: pd.Series) -> Decimal:
             return Decimal(0)
 
         elif row["current_weight"] == Decimal(0) and row["target_weight"] == Decimal(0):
-            # Should nothing change?
+            # Do nothing
             return Decimal(0)
 
         elif row["current_quantity"] > Decimal(0) and row["target_weight"] == Decimal(0):
-            # Should we sell everything
+            # Sell everything
             return Decimal(-1)
 
+        elif row["current_quantity"] < Decimal(0) and row["target_weight"] == Decimal(0):
+            # Cover our short position
+            return Decimal(1)
+
         elif row["current_quantity"] == Decimal(0) and row["target_weight"] > Decimal(0):
             # We don't have any of this asset, but we want to buy some.
             return Decimal(1)
 
         elif row["current_quantity"] == Decimal(0) and row["target_weight"] == Decimal(-1):
-            # Should we short everything we have
+            # Short everything we have
             return Decimal(-1)
 
         elif row["current_quantity"] == Decimal(0) and row["target_weight"] < Decimal(0):