sliding fit fix

This commit is contained in:
Oleg Sheynin
2025-07-13 22:33:48 +00:00
parent 48f18f7b4f
commit b24285802a
6 changed files with 553 additions and 297 deletions
+35 -37
View File
@@ -1,4 +1,5 @@
from typing import Any, Dict, List, Optional
import pandas as pd # type:ignore
from statsmodels.tsa.vector_ar.vecm import VECM, VECMResults # type:ignore
@@ -19,6 +20,8 @@ class TradingPair:
user_data_: Dict[str, Any]
predicted_df_: Optional[pd.DataFrame]
def __init__(
self, market_data: pd.DataFrame, symbol_a: str, symbol_b: str, price_column: str
):
@@ -31,7 +34,7 @@ class TradingPair:
self.user_data_ = {}
self.predicted_df_ = pd.DataFrame()
self.predicted_df_ = None
def _transform_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
# Select only the columns we need
@@ -127,9 +130,9 @@ class TradingPair:
df = self.training_df_[self.colnames()].reset_index(drop=True)
result = coint_johansen(df, det_order=0, k_ar_diff=1)
print(
f"{self}: lr1={result.lr1[0]} > cvt={result.cvt[0, 1]}? {result.lr1[0] > result.cvt[0, 1]}"
)
# print(
# f"{self}: lr1={result.lr1[0]} > cvt={result.cvt[0, 1]}? {result.lr1[0] > result.cvt[0, 1]}"
# )
is_cointegrated: bool = bool(result.lr1[0] > result.cvt[0, 1])
return is_cointegrated
@@ -146,21 +149,22 @@ class TradingPair:
pvalue = coint(series1, series2)[1]
# Define cointegration if p-value < 0.05 (i.e., reject null of no cointegration)
is_cointegrated: bool = bool(pvalue < 0.05)
print(f"{self}: is_cointegrated={is_cointegrated} pvalue={pvalue}")
# print(f"{self}: is_cointegrated={is_cointegrated} pvalue={pvalue}")
return is_cointegrated
def train_pair(self) -> bool:
def check_cointegration(self) -> bool:
is_cointegrated_johansen = self.check_cointegration_johansen()
is_cointegrated_engle_granger = self.check_cointegration_engle_granger()
if not is_cointegrated_johansen and not is_cointegrated_engle_granger:
return False
pass
result = is_cointegrated_johansen or is_cointegrated_engle_granger
return result or True # TODO: remove this
def train_pair(self) -> bool:
result = self.check_cointegration()
# print('*' * 80 + '\n' + f"**************** {self} IS COINTEGRATED ****************\n" + '*' * 80)
self.fit_VECM()
assert self.training_df_ is not None and self.vecm_fit_ is not None
diseq_series = self.training_df_[self.colnames()] @ self.vecm_fit_.beta
print(diseq_series.shape)
# print(diseq_series.shape)
self.training_mu_ = float(diseq_series[0].mean())
self.training_std_ = float(diseq_series[0].std())
@@ -172,7 +176,16 @@ class TradingPair:
diseq_series - self.training_mu_
) / self.training_std_
return True
return result
def add_trades(self, trades: pd.DataFrame) -> None:
if self.user_data_["trades"] is None:
self.user_data_["trades"] = pd.DataFrame(trades)
else:
self.user_data_["trades"] = pd.concat([self.user_data_["trades"], pd.DataFrame(trades)], ignore_index=True)
def get_trades(self) -> pd.DataFrame:
return self.user_data_["trades"] if "trades" in self.user_data_ else pd.DataFrame()
def predict(self) -> pd.DataFrame:
assert self.testing_df_ is not None
@@ -184,24 +197,6 @@ class TradingPair:
predicted_prices, columns=pd.Index(self.colnames()), dtype=float
)
# self.predicted_df_ = pd.merge(
# self.testing_df_.reset_index(drop=True),
# pd.DataFrame(
# predicted_prices, columns=pd.Index(self.colnames()), dtype=float
# ),
# left_index=True,
# right_index=True,
# suffixes=("", "_pred"),
# ).dropna()
# self.predicted_df_["disequilibrium"] = (
# self.predicted_df_[self.colnames()] @ self.vecm_fit_.beta
# )
# self.predicted_df_["scaled_disequilibrium"] = (
# abs(self.predicted_df_["disequilibrium"] - self.training_mu_)
# / self.training_std_
# )
predicted_df = pd.merge(
self.testing_df_.reset_index(drop=True),
@@ -222,17 +217,20 @@ class TradingPair:
/ self.training_std_
)
print("*** PREDICTED DF")
print(predicted_df)
print("*" * 80)
print("*** SELF.PREDICTED_DF")
print(self.predicted_df_)
print("*" * 80)
# print("*** PREDICTED DF")
# print(predicted_df)
# print("*" * 80)
# print("*** SELF.PREDICTED_DF")
# print(self.predicted_df_)
# print("*" * 80)
predicted_df = predicted_df.reset_index(drop=True)
self.predicted_df_ = pd.concat([self.predicted_df_, predicted_df], ignore_index=True)
if self.predicted_df_ is None:
self.predicted_df_ = predicted_df
else:
self.predicted_df_ = pd.concat([self.predicted_df_, predicted_df], ignore_index=True)
# Reset index to ensure proper indexing
self.predicted_df_ = self.predicted_df_.reset_index()
self.predicted_df_ = self.predicted_df_.reset_index(drop=True)
return self.predicted_df_
def __repr__(self) -> str: