Cluster_analysis.py

import os

os.environ["OMP_NUM_THREADS"] = "1"

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.mixture import GaussianMixture
from sklearn.preprocessing import StandardScaler


def load_ff5_csv(file_path: str) -> pd.DataFrame:
    with open(file_path, "r", encoding="utf-8", errors="ignore") as file_obj:
        lines = file_obj.readlines()

    header_idx = None
    for idx, line in enumerate(lines):
        if "Mkt-RF" in line and "SMB" in line and "HML" in line:
            header_idx = idx
            break

    if header_idx is None:
        raise ValueError("Unable to find the FF5 header row.")

    df = pd.read_csv(file_path, skiprows=header_idx)
    df = df.rename(columns={df.columns[0]: "Date"}).dropna(how="all")
    df["Date"] = df["Date"].astype(str).str.strip()
    df = df[df["Date"].str.match(r"^\d{6}$", na=False)].copy()

    factor_cols = ["Mkt-RF", "SMB", "HML", "RMW", "CMA", "RF"]
    for column in factor_cols:
        df[column] = pd.to_numeric(df[column], errors="coerce")

    df["Date"] = pd.to_datetime(df["Date"], format="%Y%m")
    df[factor_cols] = df[factor_cols] / 100.0
    df = df.set_index("Date").dropna(subset=factor_cols)

    return df


class MarketRegimeClustering:
    def __init__(self, ff5_df: pd.DataFrame):
        self.df = ff5_df.copy()
        self.factor_cols = ["Mkt-RF", "SMB", "HML", "RMW", "CMA"]
        self.label_col = "Regime"
        self.factor_data = None
        self.scaler = None
        self.gmm_model = None
        self.cluster_summary = None

    def prepare_factor_data(self) -> pd.DataFrame:
        factor_data = self.df[self.factor_cols].copy()
        self.factor_data = factor_data
        self.scaler = StandardScaler().fit(factor_data)
        return factor_data

    def fit_clusters(self, n_clusters: int = 3) -> pd.DataFrame:
        if self.factor_data is None or self.scaler is None:
            self.prepare_factor_data()

        scaled = self.scaler.transform(self.factor_data)
        self.gmm_model = GaussianMixture(
            n_components=n_clusters,
            covariance_type="tied",
            random_state=42,
        )
        labels = self.gmm_model.fit_predict(scaled)

        self.factor_data = self.factor_data.copy()
        self.factor_data[self.label_col] = labels
        self.cluster_summary = self.factor_data.groupby(self.label_col)[self.factor_cols].mean()
        return self.factor_data

    def predict_regime(self, factor_row: pd.DataFrame, history_df: pd.DataFrame | None = None) -> int:
        if self.gmm_model is None or self.scaler is None:
            raise ValueError("Call fit_clusters() before predict_regime().")

        scaled = self.scaler.transform(factor_row[self.factor_cols])
        return int(self.gmm_model.predict(scaled)[0])

    def export_to_excel(self, output_file: str) -> pd.DataFrame:
        if self.factor_data is None or self.label_col not in self.factor_data.columns:
            raise ValueError("Call fit_clusters() before export_to_excel().")

        df_export = self.df.copy()
        df_export[self.label_col] = self.factor_data[self.label_col]
        df_export.to_excel(output_file)
        return df_export

    def plot_cluster_scatter(self, x_col: str = "Mkt-RF", y_col: str = "HML"):
        if self.factor_data is None or self.label_col not in self.factor_data.columns:
            raise ValueError("Call fit_clusters() before plot_cluster_scatter().")

        plt.figure(figsize=(8, 6))
        scatter = plt.scatter(
            self.factor_data[x_col],
            self.factor_data[y_col],
            c=self.factor_data[self.label_col],
            cmap="viridis",
            alpha=0.75,
        )
        plt.xlabel(x_col)
        plt.ylabel(y_col)
        plt.title(f"GMM Clustering: {x_col} vs {y_col}")
        plt.colorbar(scatter, label="Cluster")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    def plot_regime_time_series(self):
        if self.factor_data is None or self.label_col not in self.factor_data.columns:
            raise ValueError("Call fit_clusters() before plot_regime_time_series().")

        plt.figure(figsize=(12, 4))
        plt.plot(
            self.factor_data.index,
            self.factor_data[self.label_col],
            marker="o",
            linestyle="-",
        )
        plt.xlabel("Date")
        plt.ylabel("Regime Cluster")
        plt.title("Market Regime Over Time")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    def plot_cluster_3d(
        self,
        x_col: str = "Mkt-RF",
        y_col: str = "SMB",
        z_col: str = "HML",
    ):
        if self.factor_data is None or self.label_col not in self.factor_data.columns:
            raise ValueError("Call fit_clusters() before plot_cluster_3d().")

        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111, projection="3d")
        scatter = ax.scatter(
            self.factor_data[x_col],
            self.factor_data[y_col],
            self.factor_data[z_col],
            c=self.factor_data[self.label_col],
            cmap="viridis",
            alpha=0.8,
        )
        ax.set_xlabel(x_col)
        ax.set_ylabel(y_col)
        ax.set_zlabel(z_col)
        ax.set_title(f"GMM 3D Clustering: {x_col}, {y_col}, {z_col}")
        cbar = plt.colorbar(scatter, ax=ax, pad=0.1)
        cbar.set_label("Cluster")
        plt.tight_layout()
        plt.show()


def fit_regime_model(train_feature_df: pd.DataFrame, config):
    method = config.cluster_method.lower()
    if method == "gmm":
        model = MarketRegimeClustering(train_feature_df)
        model.fit_clusters(n_clusters=config.n_clusters)
        return model

    if method == "hmm":
        from HMM_regime import MarketRegimeHMM

        model = MarketRegimeHMM(train_feature_df)
        model.fit_clusters(
            n_clusters=config.n_clusters,
            n_iter=config.hmm_n_iter,
            tol=config.hmm_tol,
            stickiness=config.hmm_stickiness,
            covariance_floor=config.hmm_covariance_floor,
            random_state=config.random_state,
        )
        return model

    raise ValueError(f"Unsupported cluster_method: {config.cluster_method}")


def generate_historical_regime_labels(regime_model) -> pd.Series:
    if regime_model.factor_data is None or regime_model.label_col not in regime_model.factor_data.columns:
        raise ValueError("Regime model has not generated historical labels yet.")

    labels = regime_model.factor_data[regime_model.label_col].copy()
    labels.name = "Regime"
    return labels