progress

src/tools/trading_pair.py

@@ -1,36 +1,78 @@
 
 from typing import List, Optional
-
+import pandas as pd
+from statsmodels.tsa.vector_ar.vecm import VECM
 
 class TradingPair:
     symbol_a_: str
     symbol_b_: str
     price_column_: str
 
-    disequilibrium_mu_: Optional[float]
-    disequilibrium_std_: Optional[float]
+    training_mu_: Optional[float]
+    training_std_: Optional[float]
+
+    original_df_: Optional[pd.DataFrame]
+    training_df_: Optional[pd.DataFrame]
+    testing_df_: Optional[pd.DataFrame]
+
+    vecm_fit_: Optional[VECM]
 
     def __init__(self, symbol_a: str, symbol_b: str, price_column: str):
         self.symbol_a_ = symbol_a
         self.symbol_b_ = symbol_b
         self.price_column_ = price_column
-        self.disequilibrium_mu_ = None
-        self.disequilibrium_std_ = None
+        self.training_mu_ = None
+        self.training_std_ = None
+        self.original_df_ = None
+        self.training_df_ = None
+        self.testing_df_ = None
+        self.vecm_fit_ = None
+
+    def get_datasets(self, market_data: pd.DataFrame, training_minutes: int) -> None:
+        self.original_df_ = market_data[["tstamp"] + self.colnames()]
+        self.training_df_ = market_data.iloc[:training_minutes - 1, :].copy()
+        self.training_df_ = self.training_df_.dropna().reset_index(drop=True)
+
+        self.testing_df_ = market_data.iloc[training_minutes:, :].copy()
+        self.testing_df_ = self.testing_df_.dropna().reset_index(drop=True)
 
     def colnames(self) -> List[str]:
         return [f"{self.price_column_}_{self.symbol_a_}", f"{self.price_column_}_{self.symbol_b_}"]
 
-    def set_training_disequilibrium(self, disequilibrium_mu: float, disequilibrium_std: float):
-        self.disequilibrium_mu_ = disequilibrium_mu
-        self.disequilibrium_std_ = disequilibrium_std
+    def fit_VECM(self):
+        vecm_df = self.training_df_[self.colnames()].reset_index(drop=True)
+        vecm_model = VECM(vecm_df, coint_rank=1)
+        vecm_fit = vecm_model.fit()
 
-    def mu(self) -> float:
-        assert self.disequilibrium_mu_ is not None
-        return self.disequilibrium_mu_
+        # URGENT check beta and alpha
 
-    def std(self) -> float:
-        assert self.disequilibrium_std_ is not None
-        return self.disequilibrium_std_
+        # Check if the model converged properly
+        if not hasattr(vecm_fit, "beta") or vecm_fit.beta is None:
+            print(f"{self}: VECM model failed to converge properly")
+
+        self.vecm_fit_ = vecm_fit
+
+
+    def train_pair(self):
+        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"] = (
+            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_
 
     def __repr__(self) ->str:
         return f"{self.symbol_a_} & {self.symbol_b_}"