progress: added zscore fit

lib/pt_trading/rolling_window_fit.py

@@ -11,6 +11,14 @@ from statsmodels.tsa.vector_ar.vecm import VECM, VECMResults
 NanoPerMin = 1e9
     
 class RollingFit(PairsTradingFitMethod):
+    ''' 
+    N O T E:
+    =========
+    - This class remains to be abstract
+    - The following methods are to be implemented in the subclass:
+        - create_trading_pair()
+    =========
+    '''
     def __init__(self) -> None:
         super().__init__()
 

lib/pt_trading/vecm_rolling_fit.py

@@ -17,9 +17,8 @@ class VECMTradingPair(TradingPair):
         self.pair_predict_result_ = None
         
     def _train_pair(self) -> None:
-        # print('*' * 80 + '\n' + f"**************** {self} IS COINTEGRATED ****************\n" + '*' * 80)
         self._fit_VECM()
-        assert self.training_df_ is not None and self.vecm_fit_ is not None
+        assert self.vecm_fit_ is not None
         diseq_series = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
         # print(diseq_series.shape)
         self.training_mu_ = float(diseq_series[0].mean())

lib/pt_trading/z-score_rolling_fit.py

@@ -0,0 +1,78 @@
+from typing import Any, Dict, Optional, cast
+
+import pandas as pd
+from pt_trading.results import BacktestResult
+from pt_trading.rolling_window_fit import RollingFit
+from pt_trading.trading_pair import TradingPair
+import statsmodels.api as sm
+
+NanoPerMin = 1e9
+class ZScoreTradingPair(TradingPair):
+    zscore_model_: Optional[sm.regression.linear_model.RegressionResultsWrapper]
+    pair_predict_result_: Optional[pd.DataFrame]
+    zscore_df_: Optional[pd.DataFrame]
+    
+    def __init__(self, config: Dict[str, Any], market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str):
+        super().__init__(config, market_data, symbol_a, symbol_b, price_column)
+        self.zscore_model_ = None
+        self.pair_predict_result_ = None
+        self.zscore_df_ = None
+        
+    def _fit_zscore(self) -> None:
+        assert self.training_df_ is not None
+        a = self.training_df_[self.colnames()].iloc[:, 0]
+        b = self.training_df_[self.colnames()].iloc[:, 1]
+        
+        a,b = a.align(b, axis=0)
+        
+        
+        X = sm.add_constant(b)
+        self.zscore_model_ = sm.OLS(a, X).fit()
+        assert self.zscore_model_ is not None
+        hedge_ratio = self.zscore_model_.params.iloc[1]
+
+        # Calculate spread and Z-score
+        spread = a - hedge_ratio * b
+        self.zscore_df_  = (spread - spread.mean()) / spread.std()        
+
+    def predict(self) -> pd.DataFrame:
+        self._fit_zscore()
+        assert self.zscore_df_ is not None
+        self.training_df_["dis-equilibrium"] = self.zscore_df_
+        self.training_df_["scaled_dis-equilibrium"] = abs(self.zscore_df_)
+        
+        assert self.testing_df_ is not None
+        assert self.zscore_df_ is not None
+        predicted_df = self.testing_df_
+
+        predicted_df["disequilibrium"] = self.zscore_df_
+        predicted_df["scaled_disequilibrium"] = abs(self.zscore_df_)
+        
+        predicted_df = predicted_df.reset_index(drop=True)
+        if self.pair_predict_result_ is None:
+            self.pair_predict_result_ = predicted_df
+        else:
+            self.pair_predict_result_ = pd.concat([self.pair_predict_result_, predicted_df], ignore_index=True)
+        # Reset index to ensure proper indexing 
+        self.pair_predict_result_ = self.pair_predict_result_.reset_index(drop=True)
+        return self.pair_predict_result_
+
+    
+class ZScoreRollingFit(RollingFit):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def run_pair(
+        self, pair: TradingPair, bt_result: BacktestResult
+    ) -> Optional[pd.DataFrame]:
+        return super().run_pair(pair, bt_result)
+    def create_trading_pair(
+        self, config: Dict, market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
+    ) -> TradingPair:
+        return ZScoreTradingPair(
+            config=config,
+            market_data=market_data,
+            symbol_a=symbol_a,
+            symbol_b=symbol_b,
+            price_column=price_column
+        )