progress

src/tools/trading_pair.py

@@ -51,28 +51,38 @@ class TradingPair:
             print(f"{self}: VECM model failed to converge properly")
 
         self.vecm_fit_ = vecm_fit
+        # print(f"{self}: beta={self.vecm_fit_.beta} alpha={self.vecm_fit_.alpha}"  )
+        # print(f"{self}: {self.vecm_fit_.summary()}")
+        pass
 
+    def check_cointegration(self):
+        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]}.")
+        is_cointegrated = result.lr1[0] > result.cvt[0, 1]
 
-    def train_pair(self):
+        return is_cointegrated
+
+    def train_pair(self) -> bool:
+        is_cointegrated = self.check_cointegration()
+        if not is_cointegrated:
+            return False
+            pass
+
+        print(f"*****\n**************** {self} IS COINTEGRATED ****************\n*****")
         self.fit_VECM()
         diseq_series = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
         self.training_mu_ = diseq_series.mean().iloc[0]
         self.training_std_ = diseq_series.std().iloc[0]
 
-        self.training_df_["disequilibrium"] = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
-        # Normalize the disequilibrium
-        self.training_df_["scaled_disequilibrium"] = (
+        self.training_df_["dis-equilibrium"] = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
+        # Normalize the dis-equilibrium
+        self.training_df_["scaled_dis-equilibrium"] = (
             diseq_series - self.training_mu_
         ) / self.training_std_
 
-
-    # def mu(self) -> float:
-    #     assert self.training_mu_ is not None
-    #     return self.training_mu_
-
-    # def std(self) -> float:
-    #     assert self.training_std_ is not None
-    #     return self.training_std_
+        return True
 
     def __repr__(self) ->str:
         return f"{self.symbol_a_} & {self.symbol_b_}"