progress

src/tools/trading_pair.py

@@ -1,5 +1,5 @@
 
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 import pandas as pd
 from statsmodels.tsa.vector_ar.vecm import VECM
 
@@ -16,6 +16,8 @@ class TradingPair:
     testing_df_: Optional[pd.DataFrame]
 
     vecm_fit_: Optional[VECM]
+    
+    user_data_: Dict[str, Any]
 
     def __init__(self, market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str):
         self.symbol_a_ = symbol_a
@@ -23,12 +25,13 @@ class TradingPair:
         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.training_df_ = None
         self.testing_df_ = None
         self.vecm_fit_ = None
+        
+        self.user_data_ = {}
 
     def _transform_dataframe(self, df: pd.DataFrame):
         # Select only the columns we need
@@ -57,7 +60,9 @@ class TradingPair:
 
         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()
+        
+        testing_start_index = training_start_index + training_minutes
+        self.training_df_ = self.market_data_.iloc[training_start_index:testing_start_index, :].copy()
         self.training_df_ = self.training_df_.dropna().reset_index(drop=True)
 
         testing_start_index = training_start_index + training_minutes
@@ -101,7 +106,7 @@ class TradingPair:
             return False
             pass
 
-        print('*' * 80 + '\n' + f"**************** {self} IS COINTEGRATED ****************\n" + '*' * 80)
+        # 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]