initial cleaning after refactoring

lib/pt_strategy/model_data_policy.py

@@ -14,6 +14,8 @@ class DataParams:
 class ModelDataPolicy(ABC):
     config_: Dict[str, Any]
     current_data_params_: DataParams
+    count_:int
+
     
     def __init__(self, config: Dict[str, Any]):
         self.config_ = config
@@ -21,10 +23,12 @@ class ModelDataPolicy(ABC):
             training_size=config.get("training_size", 120),
             training_start_index=0,
         )
+        self.count_ = 0
     
     @abstractmethod
     def advance(self) -> DataParams:
-        ...
+        self.count_ += 1
+        print(self.count_, end='\r')
     
     @staticmethod
     def create(config: Dict[str, Any]) -> ModelDataPolicy:
@@ -43,9 +47,8 @@ class RollingWindowDataPolicy(ModelDataPolicy):
         self.count_ = 1
         
     def advance(self) -> DataParams:
+        super().advance()
         self.current_data_params_.training_start_index += 1
-        print(self.count_, end='\r')
-        self.count_ += 1
         return self.current_data_params_
     
     
@@ -54,6 +57,7 @@ class ExpandingWindowDataPolicy(ModelDataPolicy):
         super().__init__(config)
         
     def advance(self) -> DataParams:
+        super().advance()
         self.current_data_params_.training_size += 1
         return self.current_data_params_
     

lib/pt_strategy/models.py

@@ -10,7 +10,7 @@ from pt_strategy.pt_model import PairsTradingModel, Prediction
 from pt_strategy.trading_pair import TradingPair
 
 
-class ZScoreOLSModel(PairsTradingModel):
+class OLSModel(PairsTradingModel):
     zscore_model_: Optional[sm.regression.linear_model.RegressionResultsWrapper]
     pair_predict_result_: Optional[pd.DataFrame]
     zscore_df_: Optional[pd.DataFrame]
@@ -67,10 +67,9 @@ class VECMModel(PairsTradingModel):
         disequilibrium = (predicted_df[pair.colnames()] @ vecm_fit.beta)[0][0]
         scaled_disequilibrium = (disequilibrium - self.training_mu_) / self.training_std_
         return Prediction(
-            tstamp_=pair.market_data_.index[-1],
-            disequilibrium_=disequilibrium,
-            scaled_disequilibrium_=scaled_disequilibrium,
-            pair_=pair,
+            tstamp=pair.market_data_.iloc[-1]["tstamp"],
+            disequilibrium=disequilibrium,
+            scaled_disequilibrium=scaled_disequilibrium,
         )
         
     def _fit_VECM(self, pair: TradingPair) -> VECMResults: # type: ignore

lib/pt_strategy/pt_market_data.py

@@ -82,6 +82,8 @@ class ResearchMarketData(PtMarketData):
             )
             self.origin_mkt_data_df_ = pd.concat([self.origin_mkt_data_df_, md_df])
 
+        self.origin_mkt_data_df_ = self.origin_mkt_data_df_.sort_values(by="tstamp")
+        self.origin_mkt_data_df_ = self.origin_mkt_data_df_.dropna().reset_index(drop=True)
         self._set_market_data()
         
     def _set_market_data(self, ) -> None:

lib/pt_strategy/results.py

@@ -4,7 +4,7 @@ from datetime import date, datetime
 from typing import Any, Dict, List, Optional, Tuple
 
 import pandas as pd
-from pt_trading.trading_pair import TradingPair
+from pt_strategy.trading_pair import TradingPair
 
 
 # Recommended replacement adapters and converters for Python 3.12+