progress

2025-07-25 20:20:23 +00:00
parent 21a473a4c2
commit c2f701e3a2
15 changed files with 5693 additions and 109 deletions
@@ -7,15 +7,23 @@ from pt_trading.trading_pair import TradingPair
 from statsmodels.tsa.vector_ar.vecm import VECM, VECMResults

 NanoPerMin = 1e9
+
+
 class VECMTradingPair(TradingPair):
    vecm_fit_: Optional[VECMResults]
    pair_predict_result_: Optional[pd.DataFrame]
-    
-    def __init__(self, config: Dict[str, Any], market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str):
-        super().__init__(config, market_data, symbol_a, symbol_b, price_column)
+
+    def __init__(
+        self,
+        config: Dict[str, Any],
+        market_data: pd.DataFrame,
+        symbol_a: str,
+        symbol_b: str,
+    ):
+        super().__init__(config, market_data, symbol_a, symbol_b)
        self.vecm_fit_ = None
        self.pair_predict_result_ = None
-        
+
    def _train_pair(self) -> None:
        self._fit_VECM()
        assert self.vecm_fit_ is not None
@@ -51,7 +59,7 @@ class VECMTradingPair(TradingPair):

    def predict(self) -> pd.DataFrame:
        self._train_pair()
-        
+
        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_))
@@ -79,31 +87,36 @@ class VECMTradingPair(TradingPair):
            predicted_df["disequilibrium"] - self.training_mu_
        ) / self.training_std_

-        predicted_df["scaled_disequilibrium"] = (
-            abs(predicted_df["signed_scaled_disequilibrium"])
+        predicted_df["scaled_disequilibrium"] = abs(
+            predicted_df["signed_scaled_disequilibrium"]
        )
-        
+
        predicted_df = predicted_df.reset_index(drop=True)
        if self.pair_predict_result_ is None:
            self.pair_predict_result_ = predicted_df
        else:
-            self.pair_predict_result_ = pd.concat([self.pair_predict_result_, predicted_df], ignore_index=True)
-        # Reset index to ensure proper indexing 
+            self.pair_predict_result_ = pd.concat(
+                [self.pair_predict_result_, predicted_df], ignore_index=True
+            )
+        # Reset index to ensure proper indexing
        self.pair_predict_result_ = self.pair_predict_result_.reset_index(drop=True)
        return self.pair_predict_result_

-    
+
 class VECMRollingFit(RollingFit):
    def __init__(self) -> None:
        super().__init__()

    def create_trading_pair(
-        self, config: Dict, market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
+        self,
+        config: Dict,
+        market_data: pd.DataFrame,
+        symbol_a: str,
+        symbol_b: str,
    ) -> TradingPair:
        return VECMTradingPair(
            config=config,
            market_data=market_data,
            symbol_a=symbol_a,
            symbol_b=symbol_b,
-            price_column=price_column
        )