cointegration test initial

lib/pt_trading/trading_pair.py

@@ -1,8 +1,63 @@
+from __future__ import annotations
+
 from typing import Any, Dict, List, Optional
 
 import pandas as pd  # type:ignore
 from statsmodels.tsa.vector_ar.vecm import VECM, VECMResults  # type:ignore
 
+class CointegrationData:
+    EG_PVALUE_THRESHOLD = 0.05
+    
+    tstamp_: pd.Timestamp
+    pair_: str
+    eg_pvalue_: float
+    johansen_lr1_: float
+    johansen_cvt_: float
+    eg_is_cointegrated_: bool
+    johansen_is_cointegrated_: bool
+
+    def __init__(self, pair: TradingPair):
+        training_df = pair.training_df_
+
+        assert training_df is not None
+        from statsmodels.tsa.vector_ar.vecm import coint_johansen
+
+        df = training_df[pair.colnames()].reset_index(drop=True)
+
+        # Run Johansen cointegration test
+        result = coint_johansen(df, det_order=0, k_ar_diff=1)
+        self.johansen_lr1_ = result.lr1[0]
+        self.johansen_cvt_ = result.cvt[0, 1]
+        self.johansen_is_cointegrated_ = self.johansen_lr1_ > self.johansen_cvt_
+
+        # Run Engle-Granger cointegration test
+        from statsmodels.tsa.stattools import coint #type: ignore
+
+        col1, col2 = pair.colnames()
+        assert training_df is not None
+        series1 = training_df[col1].reset_index(drop=True)
+        series2 = training_df[col2].reset_index(drop=True)
+
+        self.eg_pvalue_ = float(coint(series1, series2)[1])
+        self.eg_is_cointegrated_ = bool(self.eg_pvalue_ < self.EG_PVALUE_THRESHOLD)
+
+        self.tstamp_ = training_df.index[-1]
+        self.pair_ = pair.name()
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "tstamp": self.tstamp_,
+            "pair": self.pair_,
+            "eg_pvalue": self.eg_pvalue_,
+            "johansen_lr1": self.johansen_lr1_,
+            "johansen_cvt": self.johansen_cvt_,
+            "eg_is_cointegrated": self.eg_is_cointegrated_,
+            "johansen_is_cointegrated": self.johansen_is_cointegrated_,
+        }   
+
+    def __repr__(self) -> str:
+        return f"CointegrationData(tstamp={self.tstamp_}, pair={self.pair_}, eg_pvalue={self.eg_pvalue_}, johansen_lr1={self.johansen_lr1_}, johansen_cvt={self.johansen_cvt_}, eg_is_cointegrated={self.eg_is_cointegrated_}, johansen_is_cointegrated={self.johansen_is_cointegrated_})"
+
 
 class TradingPair:
     market_data_: pd.DataFrame
@@ -148,42 +203,7 @@ class TradingPair:
         # print(f"{self}: {self.vecm_fit_.summary()}")
         pass
 
-    def check_cointegration_johansen(self) -> bool:
-        assert self.training_df_ is not None
-        from statsmodels.tsa.vector_ar.vecm import coint_johansen
-
-        df = self.training_df_[self.colnames()].reset_index(drop=True)
-        result = coint_johansen(df, det_order=0, k_ar_diff=1)
-        # print(
-        #     f"{self}: lr1={result.lr1[0]} > cvt={result.cvt[0, 1]}? {result.lr1[0] > result.cvt[0, 1]}"
-        # )
-        is_cointegrated: bool = bool(result.lr1[0] > result.cvt[0, 1])
-
-        return is_cointegrated
-
-    def check_cointegration_engle_granger(self) -> bool:
-        from statsmodels.tsa.stattools import coint
-
-        col1, col2 = self.colnames()
-        assert self.training_df_ is not None
-        series1 = self.training_df_[col1].reset_index(drop=True)
-        series2 = self.training_df_[col2].reset_index(drop=True)
-
-        # Run Engle-Granger cointegration test
-        pvalue = coint(series1, series2)[1]
-        # Define cointegration if p-value < 0.05 (i.e., reject null of no cointegration)
-        is_cointegrated: bool = bool(pvalue < 0.05)
-        # print(f"{self}: is_cointegrated={is_cointegrated} pvalue={pvalue}")
-        return is_cointegrated
-
-    def check_cointegration(self) -> bool:
-        is_cointegrated_johansen = self.check_cointegration_johansen()
-        is_cointegrated_engle_granger = self.check_cointegration_engle_granger()
-        result = is_cointegrated_johansen or is_cointegrated_engle_granger
-        return result or True # TODO: remove this
-
-    def train_pair(self) -> bool:
-        result = self.check_cointegration()
+    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
@@ -200,8 +220,6 @@ class TradingPair:
             diseq_series - self.training_mu_
         ) / self.training_std_
 
-        return result
-
     def add_trades(self, trades: pd.DataFrame) -> None:
         if self.user_data_["trades"] is None or len(self.user_data_["trades"]) == 0:
             # If trades is empty or None, just assign the new trades directly
@@ -286,6 +304,45 @@ class TradingPair:
         self.predicted_df_ = self.predicted_df_.reset_index(drop=True)
         return self.predicted_df_
 
+    def cointegration_check(self) -> Optional[pd.DataFrame]:
+        print(f"***{self}*** STARTING....")
+        config = self.config_
+
+        curr_training_start_idx = 0
+
+        COINTEGRATION_DATA_COLUMNS = {
+            "tstamp" : "datetime64[ns]",
+            "pair" : "string",
+            "eg_pvalue" : "float64",
+            "johansen_lr1" : "float64",
+            "johansen_cvt" : "float64",
+            "eg_is_cointegrated" : "bool",
+            "johansen_is_cointegrated" : "bool",
+        }
+        # Initialize trades DataFrame with proper dtypes to avoid concatenation warnings
+        result: pd.DataFrame = pd.DataFrame(columns=[col for col in COINTEGRATION_DATA_COLUMNS.keys()]) #.astype(COINTEGRATION_DATA_COLUMNS)
+        
+        training_minutes = config["training_minutes"]
+        while True:
+            print(curr_training_start_idx, end="\r")
+            self.get_datasets(
+                training_minutes=training_minutes,
+                training_start_index=curr_training_start_idx,
+                testing_size=1,
+            )
+
+            if len(self.training_df_) < training_minutes:
+                print(
+                    f"{self}: current offset={curr_training_start_idx}"
+                    f" * Training data length={len(self.training_df_)} < {training_minutes}"
+                    " * Not enough training data. Completing the job."
+                )
+                break
+            new_row = pd.Series(CointegrationData(self).to_dict())
+            result.loc[len(result)] = new_row
+            curr_training_start_idx += 1
+        return result
+
     def __repr__(self) -> str:
         return self.name()