progress

lib/pt_trading/trading_pair.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
 import pandas as pd  # type:ignore
-from statsmodels.tsa.vector_ar.vecm import VECM, VECMResults
+
 
 class PairState(Enum):
     INITIAL = 1
@@ -40,7 +41,7 @@ class CointegrationData:
         self.johansen_is_cointegrated_ = self.johansen_lr1_ > self.johansen_cvt_
 
         # Run Engle-Granger cointegration test
-        from statsmodels.tsa.stattools import coint #type: ignore
+        from statsmodels.tsa.stattools import coint  # type: ignore
 
         col1, col2 = pair.colnames()
         assert training_df is not None
@@ -68,7 +69,7 @@ class CointegrationData:
         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:
+class TradingPair(ABC):
     market_data_: pd.DataFrame
     symbol_a_: str
     symbol_b_: str
@@ -80,11 +81,9 @@ class TradingPair:
     training_df_: pd.DataFrame
     testing_df_: pd.DataFrame
 
-    vecm_fit_: VECMResults
-
     user_data_: Dict[str, Any]
 
-    predicted_df_: Optional[pd.DataFrame]
+    # predicted_df_: Optional[pd.DataFrame]
 
     def __init__(
         self, config: Dict[str, Any], market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
@@ -193,42 +192,6 @@ class TradingPair:
             f"{self.price_column_}_{self.symbol_b_}",
         ]
 
-    def fit_VECM(self) -> None:
-        assert self.training_df_ is not None
-        vecm_df = self.training_df_[self.colnames()].reset_index(drop=True)
-        vecm_model = VECM(vecm_df, coint_rank=1)
-        vecm_fit = vecm_model.fit()
-
-        assert vecm_fit is not None
-
-        # URGENT check beta and alpha
-
-        # 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
-        # print(f"{self}: beta={self.vecm_fit_.beta} alpha={self.vecm_fit_.alpha}"  )
-        # print(f"{self}: {self.vecm_fit_.summary()}")
-        pass
-
-    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
-        diseq_series = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
-        # print(diseq_series.shape)
-        self.training_mu_ = float(diseq_series[0].mean())
-        self.training_std_ = float(diseq_series[0].std())
-
-        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 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
@@ -267,45 +230,6 @@ class TradingPair:
     def get_trades(self) -> pd.DataFrame:
         return self.user_data_["trades"] if "trades" in self.user_data_ else pd.DataFrame()
 
-    def predict(self) -> pd.DataFrame:
-        assert self.testing_df_ is not None
-        assert self.vecm_fit_ is not None
-        predicted_prices = self.vecm_fit_.predict(steps=len(self.testing_df_))
-
-        # Convert prediction to a DataFrame for readability
-        predicted_df = pd.DataFrame(
-            predicted_prices, columns=pd.Index(self.colnames()), dtype=float
-        )
-
-
-        predicted_df = pd.merge(
-            self.testing_df_.reset_index(drop=True),
-            pd.DataFrame(
-                predicted_prices, columns=pd.Index(self.colnames()), dtype=float
-            ),
-            left_index=True,
-            right_index=True,
-            suffixes=("", "_pred"),
-        ).dropna()
-
-        predicted_df["disequilibrium"] = (
-            predicted_df[self.colnames()] @ self.vecm_fit_.beta
-        )
-
-        predicted_df["scaled_disequilibrium"] = (
-            abs(predicted_df["disequilibrium"] - self.training_mu_)
-            / self.training_std_
-        )
-        
-        predicted_df = predicted_df.reset_index(drop=True)
-        if self.predicted_df_ is None:
-            self.predicted_df_ = predicted_df
-        else:
-            self.predicted_df_ = pd.concat([self.predicted_df_, predicted_df], ignore_index=True)
-        # Reset index to ensure proper indexing 
-        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_
@@ -394,3 +318,10 @@ class TradingPair:
     def name(self) -> str:
         return f"{self.symbol_a_} & {self.symbol_b_}"
         # return f"{self.symbol_a_} & {self.symbol_b_}"
+
+    @abstractmethod
+    def predict(self) -> pd.DataFrame: ...
+
+    # @abstractmethod
+    # def predicted_df(self) -> Optional[pd.DataFrame]: ...
+