基于回归算法预测股市价格（Part 2）

# split the data into train, validation and test sets

total_len = len(train_data)
train_end = int(total_len * 0.8)  # first 80% for training

X_train, y_train = train_data.iloc[:train_end], train_data['Close'][:train_end]
X_test, y_test = train_data.iloc[train_end:], train_data['Close'][train_end:]

X_train.shape, y_train.shape, X_test.shape, y_test.shape    # Display the shapes of the train and test sets

((10684, 6), (10684,), (2672, 6), (2672,))

import logging
import warnings
from typing import Optional, Tuple

import numpy as np
import pandas as pd
from statsmodels.stats.outliers_influence import variance_inflation_factor

from balabs.utils.logger import setup_logging


class FeatureGenerator:
    def __init__(
        self,
        ma_windows: list = [5, 7, 14],
        ema_windows: list = [12, 26],
        rsi_window: int = 14,
        bb_window: int = 20,
        hv_window: int = 10,
        atr_window: int = 14,
        log_level=logging.INFO,
    ):

        self.ma_windows = ma_windows
        self.ema_windows = ema_windows
        self.rsi_window = rsi_window
        self.bb_window = bb_window
        self.hv_window = hv_window
        self.atr_window = atr_window

        self.logger = setup_logging(
            log_name=__name__, log_file=f"{__name__}.log", log_level=log_level
        )

    def fit(self, X: pd.DataFrame, y: pd.Series) -> "FeatureGenerator":
        self.columns = X.columns
        self.logger.info(f"Fitted FeatureGenerator with {len(self.columns)} columns.")
        return self

    def transform(
        self, X: pd.DataFrame, y: Optional[pd.Series] = None
    ) -> Tuple[pd.DataFrame, Optional[pd.Series]]:
        X_transformed = X.copy()

        self.logger.info("Starting feature transformation.")

        X_transformed = self._add_time_features(X_transformed)
        X_transformed = self._add_price_features(X_transformed)
        X_transformed = self._add_volume_features(X_transformed)
        X_transformed = self._add_technical_indicators(X_transformed)
        X_transformed = self._add_feature_interactions(X_transformed)

        if y is not None:
            X_transformed, y_aligned = self._handle_missing_values(X_transformed, y)
            if len(X_transformed) != len(y_aligned):
                self.logger.error(
                    "Data alignment failed; X and y have different lengths."
                )
                raise ValueError(
                    "Data alignment failed; X and y have different lengths."
                )
            self.logger.info(
                f"Feature transformation completed with target alignment: The shape of X is {X_transformed.shape}, and the shape of y is {y_aligned.shape}."
            )
            return X_transformed, y_aligned
        return X_transformed, None

    def _add_time_features(self, df: pd.DataFrame) -> pd.DataFrame:
        original_columns = df.columns

        df["DoW"] = df.index.dayofweek
        df["Month"] = df.index.month
        df["Quarter"] = df.index.quarter
        df["MonthEnd"] = df.index.is_month_end.astype(int)
        df["QuarterEnd"] = df.index.is_quarter_end.astype(int)

        self.logger.debug(
            f"Added time features: {list(df.columns.difference(original_columns))}."
        )
        return df

    def _add_price_features(self, df: pd.DataFrame) -> pd.DataFrame:
        original_columns = df.columns

        df["Close_t-1"] = df["Close"].shift(1)
        denominator = df["Close"].shift(1)
        df["Return_t-1"] = np.where(
            denominator == 0, 0, (df["Close"] - denominator) / denominator
        )
        for window in self.ma_windows:
            df[f"MA_{window}"] = df["Close"].rolling(window=window).mean()
        for window in self.ema_windows:
            df[f"EMA_{window}"] = df["Close"].ewm(span=window, adjust=False).mean()
        df["Gap"] = df["Open"] - df["Close"].shift(1)
        df[f"HV_{self.hv_window}"] = (
            df["Return_t-1"].rolling(window=self.hv_window).std()
        )
        df["Range_t-1"] = (df["High"] - df["Low"]).shift(1)

        self.logger.debug(
            f"Added time features: {list(df.columns.difference(original_columns))}."
        )
        return df

    def _add_volume_features(self, df: pd.DataFrame) -> pd.DataFrame:
        original_columns = df.columns

        for window in self.ma_windows:
            df[f"Vol_MA_{window}"] = df["Volume"].rolling(window=window).mean()
        denominator = df["Volume"].shift(1)
        df["Vol_Chg"] = np.where(
            denominator == 0, 0, (df["Volume"] - denominator) / denominator
        )
        df["VPT"] = df["Volume"].shift(1) * df["Return_t-1"]

        self.logger.debug(
            f"Added time features: {list(df.columns.difference(original_columns))}."
        )
        return df

    def _add_technical_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        original_columns = df.columns

        delta = df["Close"].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=self.rsi_window).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=self.rsi_window).mean()
        rs = gain / loss
        df["RSI"] = 100 - (100 / (1 + rs))

        ema_short = df["Close"].ewm(span=12, adjust=False).mean()
        ema_long = df["Close"].ewm(span=26, adjust=False).mean()
        df["MACD"] = ema_short - ema_long

        rolling_mean = df["Close"].rolling(window=self.bb_window).mean()
        rolling_std = df["Close"].rolling(window=self.bb_window).std()
        df["BB_upper"] = rolling_mean + (2 * rolling_std)
        df["BB_lower"] = rolling_mean - (2 * rolling_std)

        high_low = df["High"] - df["Low"]
        high_close = np.abs(df["High"] - df["Close"].shift(1))
        low_close = np.abs(df["Low"] - df["Close"].shift(1))
        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        df[f"ATR_{self.atr_window}"] = tr.rolling(window=self.atr_window).mean()

        self.logger.debug(
            f"Added time features: {list(df.columns.difference(original_columns))}."
        )
        return df

    def _add_feature_interactions(self, df: pd.DataFrame) -> pd.DataFrame:
        original_columns = df.columns

        df["PVT"] = df["MA_5"] * df["Vol_MA_5"]
        hv = df[f"HV_{self.hv_window}"]
        df["VAR"] = np.where(hv == 0, 0, df["Return_t-1"] / hv)
        df["GapDir"] = np.sign(df["Gap"])
        df["GapAbs"] = np.abs(df["Gap"])

        self.logger.debug(
            f"Added time features: {list(df.columns.difference(original_columns))}."
        )
        return df

    def _handle_missing_values(
        self, X: pd.DataFrame, y: pd.Series
    ) -> Tuple[pd.DataFrame, pd.Series]:
        mask = X.notna().all(axis=1)
        X_cleaned = X[mask]
        y_aligned = y[mask]
        removed_rows = len(X) - len(X_cleaned)
        if removed_rows > 0:
            self.logger.warning(f"Removed {removed_rows} rows due to NaN values.")
        else:
            self.logger.info("No rows removed; data is clean.")
        return X_cleaned, y_aligned

    def get_feature_correlation(
        self,
        X: pd.DataFrame,
        y: Optional[pd.Series] = None,
        threshold: float = 0.1,
        keep_all: bool = False,
    ) -> pd.DataFrame:

        if y is not None:
            corr_with_target = X.corrwith(y, method="pearson")
            corr_df = pd.DataFrame(corr_with_target, columns=["Correlation"])
            corr_df.index.name = "Feature"
            if not keep_all:
                corr_df = corr_df[abs(corr_df["Correlation"]) > threshold]
            result = corr_df.sort_values(by="Correlation", key=abs, ascending=False)
            self.logger.info(
                f"Computed correlation with target; {len(result)} features above threshold {threshold}."
            )
            self.logger.info(
                f"Feature Correlation Matrix (|r| > {threshold}):\n{result.to_string()}"
            )
            return result
        else:
            corr_matrix = X.corr(method="pearson")
            if not keep_all:
                result = corr_matrix[abs(corr_matrix) > threshold]
            else:
                result = corr_matrix
            self.logger.info("Computed feature-to-feature correlation matrix.")
            self.logger.info(
                f"Feature Correlation Matrix (|r| > {threshold}):\n{result}"
            )
            return result

    def get_vif(
        self, X: pd.DataFrame, threshold: float = 5.0, keep_all: bool = False
    ) -> pd.DataFrame:
        X_numeric = X.select_dtypes(include=[np.number])

        nan_count = X_numeric.isna().sum().sum()
        inf_count = np.isinf(X_numeric).sum().sum()
        self.logger.info(f"NaN count in data: {nan_count}")
        self.logger.info(f"Inf count in data: {inf_count}")

        if nan_count > 0 or inf_count > 0:
            self.logger.warning("Cleaning NaN and inf values...")
            X_cleaned = X_numeric.replace([np.inf, -np.inf], np.nan)
            X_cleaned = X_cleaned.dropna()
            self.logger.warning(
                f"Rows removed due to NaN/inf: {len(X_numeric) - len(X_cleaned)}"
            )
        else:
            self.logger.info("No NaN or inf values found, using original data.")
            X_cleaned = X_numeric.copy()

        if X_cleaned.shape[0] < 2 or X_cleaned.shape[1] < 1:
            self.logger.error("Insufficient data after cleaning for VIF calculation.")
            raise ValueError("Insufficient data after cleaning for VIF calculation")

        variances = X_cleaned.var()
        constant_cols = variances[variances == 0].index
        if len(constant_cols) > 0:
            self.logger.warning(f"Removing constant columns: {list(constant_cols)}")
            X_cleaned = X_cleaned.drop(columns=constant_cols)

        vif_data = pd.DataFrame()
        vif_data["Feature"] = X_cleaned.columns
        vif_values = []

        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=RuntimeWarning)
            for i in range(X_cleaned.shape[1]):
                vif = variance_inflation_factor(X_cleaned.values, i)
                if np.isinf(vif):
                    self.logger.warning(
                        f"Feature '{X_cleaned.columns[i]}' has perfect multicollinearity (VIF = inf)"
                    )
                vif_values.append(vif)

        vif_data["VIF"] = vif_values
        vif_data = vif_data.sort_values(by="VIF", ascending=True)

        if not keep_all:
            vif_data = vif_data[vif_data["VIF"] <= threshold]

        removed_features = vif_data[vif_data["VIF"] > threshold]["Feature"].tolist()
        if removed_features:
            self.logger.info(
                f"Features with VIF > {threshold} (removed): {removed_features}"
            )
        self.logger.info(
            f"Computed VIF; {len(vif_data)} features with VIF <= {threshold}."
        )

        self.logger.info(
            f"VIF Data (VIF <= {threshold}):\n{vif_data.to_string(index=False)}"
        )

        return vif_data

    def get_high_vif_high_corr_features(
        self,
        X: pd.DataFrame,
        y: pd.Series,
        vif_threshold: float = 10.0,
        corr_threshold: float = 0.1,
    ) -> pd.DataFrame:

        vif_df = self.get_vif(X, threshold=vif_threshold, keep_all=True)
        self.logger.info("Retrieved VIF data for all features.")

        corr_df = self.get_feature_correlation(
            X, y, threshold=corr_threshold, keep_all=True
        )
        self.logger.info("Retrieved correlation data with target for all features.")

        combined_df = vif_df.merge(corr_df, on="Feature", how="inner")
        self.logger.debug(
            f"Combined VIF and correlation data:\n{combined_df.to_string()}"
        )

        high_vif_high_corr = combined_df[
            (combined_df["VIF"] > vif_threshold)
            & (combined_df["Correlation"].abs() > corr_threshold)
        ]

        if high_vif_high_corr.empty:
            self.logger.warning("No features found with high VIF and high correlation.")
        else:
            self.logger.info(
                f"Found {len(high_vif_high_corr)} features with VIF > {vif_threshold} "
                f"and |Correlation| > {corr_threshold}:\n{high_vif_high_corr.to_string()}"
            )

        return high_vif_high_corr

from balabs.regression.processing.feature_engineering import FeatureGenerator
import logging

feature_gen = FeatureGenerator(log_level=logging.DEBUG)
feature_gen.fit(X_train, y_train)
X_train_transformed, y_train_aligned = feature_gen.transform(X_train, y_train)
X_test_transformed, y_test_aligned = feature_gen.transform(X_test, y_test)

Feature Correlation Matrix (|r| > 0.1):
                Correlation
Feature                    
Close              1.000000
High               0.999879
Low                0.999874
Open               0.999763
Close_t-1          0.999646
MA_5               0.999597
MA_7               0.999385
EMA_12             0.999190
MA_14              0.998674
EMA_26             0.998149
BB_upper           0.996720
BB_lower           0.996684
Adjusted Close     0.996421
PVT                0.853159
ATR_14             0.796893
Range_t-1          0.711059
GapAbs             0.461156
Vol_MA_14          0.306307
Vol_MA_7           0.290633
Vol_MA_5           0.282938
Volume             0.246980
MACD               0.214535
RSI                0.115112

VIF Data (VIF <= 5.0):
   Feature      VIF
    GapDir 1.486466
   Vol_Chg 1.544253
QuarterEnd 1.557439
  MonthEnd 1.571519
    GapAbs 2.080755
       DoW 3.000157
       VPT 3.339900