progress

src/strategies.py

@@ -4,7 +4,7 @@ import sys
 
 from typing import Dict, Optional
 
-import pandas as pd
+import pandas as pd  # type: ignore
 
 from tools.trading_pair import TradingPair
 from results import BacktestResult
@@ -22,7 +22,7 @@ class PairsTradingStrategy(ABC):
         "pair",
     ]
     @abstractmethod
-    def run_pair(self, pair: TradingPair, bt_result: BacktestResult) -> Optional[pd.DataFrame]:
+    def run_pair(self, config: Dict, pair: TradingPair, bt_result: BacktestResult) -> Optional[pd.DataFrame]:
         ...
 
 class StaticFitStrategy(PairsTradingStrategy):
@@ -49,7 +49,7 @@ class StaticFitStrategy(PairsTradingStrategy):
         return pair_trades
 
     def create_trading_signals(self, pair: TradingPair, config: Dict, result: BacktestResult) -> pd.DataFrame:
-        beta = pair.vecm_fit_.beta
+        beta = pair.vecm_fit_.beta  # type: ignore
         colname_a, colname_b = pair.colnames()
 
         predicted_df = pair.predicted_df_
@@ -229,7 +229,7 @@ class SlidingFitStrategy(PairsTradingStrategy):
                 testing_size=1
             )
 
-            if len(pair.training_df_) < training_minutes:
+            if len(pair.training_df_) < training_minutes:  # type: ignore
                 print(f"{pair}:  {self.curr_training_start_idx_} Not enough training data. Completing the job.")
                 if pair.user_data_["state"] == PairState.OPEN:
                     print(f"{pair}:  {self.curr_training_start_idx_} Position is not closed.")
@@ -251,7 +251,7 @@ class SlidingFitStrategy(PairsTradingStrategy):
             try:
                 is_cointegrated = pair.train_pair()
             except Exception as e:
-                raise Exception(f"{pair}: Training failed: {str(e)}") from e
+                raise RuntimeError(f"{pair}: Training failed: {str(e)}") from e
 
             if pair.user_data_["is_cointegrated"] != is_cointegrated:
                     pair.user_data_["is_cointegrated"] = is_cointegrated
@@ -271,7 +271,7 @@ class SlidingFitStrategy(PairsTradingStrategy):
             try:
                 pair.predict()
             except Exception as e:
-                raise Exception(f"{pair}: Prediction failed: {str(e)}") from e
+                raise RuntimeError(f"{pair}: Prediction failed: {str(e)}") from e
 
             if pair.user_data_["state"] == PairState.INITIAL:
 
@@ -295,8 +295,12 @@ class SlidingFitStrategy(PairsTradingStrategy):
         colname_a, colname_b = pair.colnames()
 
         predicted_df = pair.predicted_df_
+        
+        # Check if we have any data to work with
+        if len(predicted_df) == 0:
+            return None
 
-        open_row = predicted_df.loc[0]
+        open_row = predicted_df.iloc[0]
         open_tstamp = open_row["tstamp"]
         open_disequilibrium = open_row["disequilibrium"]
         open_scaled_disequilibrium = open_row["scaled_disequilibrium"]
@@ -359,7 +363,11 @@ class SlidingFitStrategy(PairsTradingStrategy):
     def _get_close_trades(self, pair: TradingPair, close_threshold: float) -> Optional[pd.DataFrame]:
         colname_a, colname_b = pair.colnames()
 
-        close_row = pair.predicted_df_.loc[0]
+        # Check if we have any data to work with
+        if len(pair.predicted_df_) == 0:
+            return None
+
+        close_row = pair.predicted_df_.iloc[0]
         close_tstamp = close_row["tstamp"]
         close_disequilibrium = close_row["disequilibrium"]
         close_scaled_disequilibrium = close_row["scaled_disequilibrium"]