added window size optimization classes

lib/pt_strategy/model_data_policy.py

@@ -1,64 +1,235 @@
 from __future__ import annotations
 
+import copy
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from enum import Enum
-from typing import Any, Dict, Optional, cast, Generator, List
+from typing import Any, Dict, cast
+
+import numpy as np
+import pandas as pd
 
 
 @dataclass
-class DataParams:
+class DataWindowParams:
     training_size: int
     training_start_index: int
-        
+
+
 class ModelDataPolicy(ABC):
     config_: Dict[str, Any]
-    current_data_params_: DataParams
-    count_:int
+    current_data_params_: DataWindowParams
+    count_: int
 
-    
     def __init__(self, config: Dict[str, Any]):
         self.config_ = config
-        self.current_data_params_ = DataParams(
+        self.current_data_params_ = DataWindowParams(
             training_size=config.get("training_size", 120),
             training_start_index=0,
         )
         self.count_ = 0
-    
+
     @abstractmethod
-    def advance(self) -> DataParams:
+    def advance(self) -> DataWindowParams:
         self.count_ += 1
-        print(self.count_, end='\r')
-    
+        print(self.count_, end="\r")
+        return self.current_data_params_
+
     @staticmethod
-    def create(config: Dict[str, Any]) -> ModelDataPolicy:
+    def create(config: Dict[str, Any], *args: Any, **kwargs: Any) -> ModelDataPolicy:
         import importlib
-        
+
         model_data_policy_class_name = config.get("model_data_policy_class", None)
         assert model_data_policy_class_name is not None
         module_name, class_name = model_data_policy_class_name.rsplit(".", 1)
         module = importlib.import_module(module_name)
-        model_training_data_policy_object = getattr(module, class_name)(config=config)
+        model_training_data_policy_object = getattr(module, class_name)(
+            config=config, *args, **kwargs
+        )
         return cast(ModelDataPolicy, model_training_data_policy_object)
 
+
 class RollingWindowDataPolicy(ModelDataPolicy):
-    def __init__(self, config: Dict[str, Any]):
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
         super().__init__(config)
         self.count_ = 1
-        
-    def advance(self) -> DataParams:
+
+    def advance(self) -> DataWindowParams:
         super().advance()
         self.current_data_params_.training_start_index += 1
         return self.current_data_params_
-    
-    
+
+
 class ExpandingWindowDataPolicy(ModelDataPolicy):
-    def __init__(self, config: Dict[str, Any]):
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
         super().__init__(config)
-        
-    def advance(self) -> DataParams:
+
+    def advance(self) -> DataWindowParams:
         super().advance()
         self.current_data_params_.training_size += 1
         return self.current_data_params_
-    
 
+
+class OptimizedWndDataPolicy(ModelDataPolicy, ABC):
+    mkt_data_df_: pd.DataFrame
+    pair_: TradingPair  # type: ignore
+    min_training_size_: int
+    max_training_size_: int
+    end_index_: int
+    prices_a_: np.ndarray
+    prices_b_: np.ndarray
+
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
+        super().__init__(config)
+        assert (
+            kwargs.get("mkt_data") is not None and kwargs.get("pair") is not None
+        ), "mkt_data and/or pair must be provided"
+        assert (
+            "min_training_size" in config and "max_training_size" in config
+        ), "min_training_size and max_training_size must be provided"
+        self.min_training_size_ = cast(int, config.get("min_training_size"))
+        self.max_training_size_ = cast(int, config.get("max_training_size"))
+        assert self.min_training_size_ < self.max_training_size_
+
+        from pt_strategy.trading_pair import TradingPair
+
+        self.mkt_data_df_ = cast(pd.DataFrame, kwargs.get("mkt_data"))
+        self.pair_ = cast(TradingPair, kwargs.get("pair"))
+
+        self.end_index_ = (
+            self.current_data_params_.training_start_index + self.max_training_size_
+        )
+        col_a, col_b = self.pair_.colnames()
+        self.prices_a_ = np.array(self.mkt_data_df_[col_a])
+        self.prices_b_ = np.array(self.mkt_data_df_[col_b])
+        
+
+    def advance(self) -> DataWindowParams:
+        super().advance()
+        self.current_data_params_ = self.optimize_window_size()
+        self.end_index_ += 1
+        return self.current_data_params_
+
+    @abstractmethod
+    def optimize_window_size(self) -> DataWindowParams:
+        ...
+
+class EGOptimizedWndDataPolicy(OptimizedWndDataPolicy):
+    '''
+    # Engle-Granger cointegration test
+    *** VERY SLOW ***
+    '''
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
+        super().__init__(config, *args, **kwargs)
+
+    def optimize_window_size(self) -> DataWindowParams:
+        # Run Engle-Granger cointegration test
+        last_pvalue = 1.0
+        result = copy.copy(self.current_data_params_)
+        for trn_size in range(self.min_training_size_, self.max_training_size_):
+            from statsmodels.tsa.stattools import coint  # type: ignore
+
+            start_index = self.end_index_ - trn_size
+            series_a = self.prices_a_[start_index : self.end_index_]
+            series_b = self.prices_b_[start_index : self.end_index_]
+            eg_pvalue = float(coint(series_a, series_b)[1])
+            if eg_pvalue < last_pvalue:
+                last_pvalue = eg_pvalue
+                result.training_size = trn_size
+                result.training_start_index = start_index
+
+        # print(
+        #     f"*** DEBUG *** end_index={self.end_index_}, best_trn_size={self.current_data_params_.training_size}, {last_pvalue=}"
+        # )
+        return result
+
+class ADFOptimizedWndDataPolicy(OptimizedWndDataPolicy):
+    # Augmented Dickey-Fuller test
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
+        super().__init__(config, *args, **kwargs)
+
+    def optimize_window_size(self) -> DataWindowParams:
+        from statsmodels.regression.linear_model import OLS
+        from statsmodels.tools.tools import add_constant
+        from statsmodels.tsa.stattools import adfuller
+
+        last_pvalue = 1.0
+        result = copy.copy(self.current_data_params_)
+        for trn_size in range(self.min_training_size_, self.max_training_size_):
+            start_index = self.end_index_ - trn_size
+            y = self.prices_a_[start_index : self.end_index_]
+            x = self.prices_b_[start_index : self.end_index_]
+
+            # Add constant to x for intercept
+            x_with_const = add_constant(x)
+
+            # OLS regression: y = a + b*x + e
+            model = OLS(y, x_with_const).fit()
+            residuals = y - model.predict(x_with_const)
+
+            # ADF test on residuals
+            try:
+                adf_result = adfuller(residuals, maxlag=1, regression="c")
+                adf_pvalue = float(adf_result[1])
+            except Exception as e:
+                # Handle edge cases with exception (e.g., constant series, etc.)
+                adf_pvalue = 1.0
+
+            if adf_pvalue < last_pvalue:
+                last_pvalue = adf_pvalue
+                result.training_size = trn_size
+                result.training_start_index = start_index
+
+        # print(
+        #     f"*** DEBUG *** end_index={self.end_index_},"
+        #     f" best_trn_size={self.current_data_params_.training_size},"
+        #     f" {last_pvalue=}"
+        # )
+        return result
+    
+class JohansenOptdWndDataPolicy(OptimizedWndDataPolicy):
+    # Johansen test
+    def __init__(self, config: Dict[str, Any], *args: Any, **kwargs: Any):
+        super().__init__(config, *args, **kwargs)
+
+    def optimize_window_size(self) -> DataWindowParams:
+        from statsmodels.tsa.vector_ar.vecm import coint_johansen
+        import numpy as np
+
+        best_stat = -np.inf
+        best_trn_size = 0
+        best_start_index = -1
+
+        result = copy.copy(self.current_data_params_)
+        for trn_size in range(self.min_training_size_, self.max_training_size_):
+            start_index = self.end_index_ - trn_size
+            series_a = self.prices_a_[start_index:self.end_index_]
+            series_b = self.prices_b_[start_index:self.end_index_]
+
+            # Combine into 2D matrix for Johansen test
+            try:
+                data = np.column_stack([series_a, series_b])
+
+                # Johansen test: det_order=0 (no deterministic trend), k_ar_diff=1 (lag)
+                res = coint_johansen(data, det_order=0, k_ar_diff=1)
+
+                # Trace statistic for cointegration rank 1
+                trace_stat = res.lr1[0]  # test stat for rank=0 vs >=1
+                critical_value = res.cvt[0, 1]  # 5% critical value
+
+                if trace_stat > best_stat:
+                    best_stat = trace_stat
+                    best_trn_size = trn_size
+                    best_start_index = start_index
+            except Exception:
+                continue
+
+        if best_trn_size > 0:
+            result.training_size = best_trn_size
+            result.training_start_index = best_start_index
+        else:
+            print("*** WARNING: No valid cointegration window found.")
+
+        # print(
+        #     f"*** DEBUG *** end_index={self.end_index_}, best_trn_size={best_trn_size}, trace_stat={best_stat}"
+        # )
+        return result