progress

src/tools/trading_pair.py

@@ -4,6 +4,7 @@ import pandas as pd
 from statsmodels.tsa.vector_ar.vecm import VECM
 
 class TradingPair:
+    market_data_: pd.DataFrame
     symbol_a_: str
     symbol_b_: str
     price_column_: str
@@ -11,29 +12,59 @@ class TradingPair:
     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):
+    def __init__(self, market_data: pd.DataFrame, 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.market_data_ = self._transform_dataframe(market_data)[["tstamp"] + self.colnames()]
+        
+        
         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()
+    def _transform_dataframe(self, df: pd.DataFrame):
+        # Select only the columns we need
+        df_selected = df[["tstamp", "symbol", self.price_column_]]
+
+        # Start with unique timestamps
+        result_df: pd.DataFrame = pd.DataFrame(df_selected["tstamp"]).drop_duplicates().reset_index(drop=True)
+
+        # For each unique symbol, add a corresponding close price column
+        for symbol in df_selected["symbol"].unique():
+            # Filter rows for this symbol
+            df_symbol = df_selected[df_selected["symbol"] == symbol].reset_index(drop=True)
+
+            # Create column name like "close-COIN"
+            new_price_column = f"{self.price_column_}_{symbol}"
+
+            # Create temporary dataframe with timestamp and price
+            temp_df = pd.DataFrame({
+                "tstamp": df_symbol["tstamp"],
+                new_price_column: df_symbol[self.price_column_]
+            })
+
+            # Join with our result dataframe
+            result_df = pd.merge(result_df, temp_df, on="tstamp", how="left")
+            result_df = result_df.reset_index(drop=True) # do not dropna() since irrelevant symbol would affect dataset
+
+        return result_df
+    def get_datasets(self, training_minutes: int, training_start_index: int = 0, testing_size: Optional[int] = None) -> None:
+        self.training_df_ = self.market_data_.iloc[training_start_index:training_minutes - 1, :].copy()
         self.training_df_ = self.training_df_.dropna().reset_index(drop=True)
 
-        self.testing_df_ = market_data.iloc[training_minutes:, :].copy()
+        testing_start_index = training_start_index + training_minutes
+        if testing_size is None:
+            self.testing_df_ = self.market_data_.iloc[testing_start_index:, :].copy()
+        else:
+            self.testing_df_ = self.market_data_.iloc[testing_start_index:testing_start_index + testing_size, :].copy()
         self.testing_df_ = self.testing_df_.dropna().reset_index(drop=True)
 
     def colnames(self) -> List[str]:
@@ -70,7 +101,7 @@ class TradingPair:
             return False
             pass
 
-        print(f"*****\n**************** {self} IS COINTEGRATED ****************\n*****")
+        print('*' * 80 + '\n' + f"**************** {self} IS COINTEGRATED ****************\n" + '*' * 80)
         self.fit_VECM()
         diseq_series = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
         self.training_mu_ = diseq_series.mean().iloc[0]
@@ -84,6 +115,31 @@ class TradingPair:
 
         return True
 
+    def predict(self) -> None:
+        predicted_prices = self.vecm_fit_.predict(steps=len(self.testing_df_))
+
+        # Convert prediction to a DataFrame for readability
+        # predicted_df = 
+
+        self.predicted_df_ = pd.merge(
+            self.testing_df_.reset_index(drop=True),
+            pd.DataFrame(predicted_prices, columns=self.colnames()),
+            left_index=True,
+            right_index=True,
+            suffixes=("", "_pred"),
+        ).dropna()
+
+        self.predicted_df_["disequilibrium"] = self.predicted_df_[self.colnames()] @ self.vecm_fit_.beta
+
+        self.predicted_df_["scaled_disequilibrium"] = (
+            abs(self.predicted_df_["disequilibrium"] - self.training_mu_) / self.training_std_
+        )
+
+        # Reset index to ensure proper indexing
+        self.predicted_df_ = self.predicted_df_.reset_index()
+        return self.predicted_df_
+
+        
     def __repr__(self) ->str:
         return f"{self.symbol_a_} & {self.symbol_b_}"