progress

lib/pt_strategy/pt_model.py

@@ -1,39 +1,15 @@
 from __future__ import annotations
 
 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 pandas as pd
+from pt_strategy.prediction import Prediction
 
-from pt_strategy.trading_pair import TradingPair
-
-class Prediction:
-    tstamp_: pd.Timestamp
-    disequilibrium_: float
-    scaled_disequilibrium_: float
-
-    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 {
-            "tstamp": self.tstamp_,
-            "disequilibrium": self.disequilibrium_,
-            "signed_scaled_disequilibrium": self.scaled_disequilibrium_,
-            "scaled_disequilibrium": abs(self.scaled_disequilibrium_),
-            # "pair": self.pair_,
-        }
-    def to_df(self) -> pd.DataFrame:
-        return pd.DataFrame([self.to_dict()])
     
 class PairsTradingModel(ABC):
 
     @abstractmethod
-    def predict(self, pair: TradingPair) -> Prediction:
+    def predict(self, pair: TradingPair) -> Prediction: # type: ignore[assignment]
         ...
     
     @staticmethod