refactored code. before cleaning

lib/pt_strategy/models.py

@@ -27,10 +27,9 @@ class ZScoreOLSModel(PairsTradingModel):
 
         assert zscore_df is not None
         return Prediction(
-            tstamp_=pair.market_data_.index[-1],
-            disequilibrium_=self.training_df_["dis-equilibrium"].iloc[-1],
-            scaled_disequilibrium_=self.training_df_["scaled_dis-equilibrium"].iloc[-1],
-            pair_=pair,
+            tstamp=pair.market_data_.iloc[-1]["tstamp"],
+            disequilibrium=self.training_df_["dis-equilibrium"].iloc[-1],
+            scaled_disequilibrium=self.training_df_["scaled_dis-equilibrium"].iloc[-1],
         )
 
     def _fit_zscore(self, pair: TradingPair) -> pd.DataFrame:

lib/pt_strategy/pt_model.py

@@ -9,12 +9,15 @@ import pandas as pd
 
 from pt_strategy.trading_pair import TradingPair
 
-@dataclass
 class Prediction:
     tstamp_: pd.Timestamp
     disequilibrium_: float
     scaled_disequilibrium_: float
-    pair_: TradingPair
+
+    def __init__(self, tstamp: pd.Timestamp, disequilibrium: float, scaled_disequilibrium: float):
+        self.tstamp_ = tstamp
+        self.disequilibrium_ = disequilibrium
+        self.scaled_disequilibrium_ = scaled_disequilibrium
 
     def to_dict(self) -> Dict[str, Any]:
         return {
@@ -22,10 +25,10 @@ class Prediction:
             "disequilibrium": self.disequilibrium_,
             "signed_scaled_disequilibrium": self.scaled_disequilibrium_,
             "scaled_disequilibrium": abs(self.scaled_disequilibrium_),
-            "pair": self.pair_,
+            # "pair": self.pair_,
         }
-    def to_pd_series(self) -> pd.Series:
-        return pd.DataFrame([self.to_dict()]).iloc[0]
+    def to_df(self) -> pd.DataFrame:
+        return pd.DataFrame([self.to_dict()])
     
 class PairsTradingModel(ABC):
 

lib/pt_strategy/trading_strategy.py

@@ -26,6 +26,7 @@ class PtResearchStrategy:
     pt_mkt_data_: PtMarketData
 
     trades_: List[pd.DataFrame]
+    predictions_: pd.DataFrame
 
     def __init__(
         self,
@@ -41,6 +42,7 @@ class PtResearchStrategy:
         self.trades_ = []
         self.trading_pair_ = TradingPair(config=config, instruments=instruments)
         self.model_data_policy_ = ModelDataPolicy.create(config)
+        self.predictions_ = pd.DataFrame()
 
         import copy
 
@@ -83,6 +85,7 @@ class PtResearchStrategy:
             prediction = self.trading_pair_.run(
                 market_data_df, self.model_data_policy_.advance()
             )
+            self.predictions_ = pd.concat([self.predictions_, prediction.to_df()], ignore_index=True)
             assert prediction is not None
 
             trades = self._create_trades(