import numpy as np import pandas as pd import matplotlib.pyplot as plt import seaborn as sns from sklearn.ensemble import ( RandomForestRegressor, GradientBoostingRegressor, ExtraTreesRegressor, ) from sklearn.linear_model import Ridge, Lasso, ElasticNet from sklearn.svm import SVR from sklearn.preprocessing import StandardScaler from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score import xgboost as xgb import lightgbm as lgb from statsmodels.tsa.arima.model import ARIMA from statsmodels.tsa.statespace.sarimax import SARIMAX from statsmodels.tsa.holtwinters import ExponentialSmoothing import warnings import json import os from datetime import datetime import pickle warnings.filterwarnings("ignore") def create_sequences(features, targets, seq_length=20): """Create sequences for time series prediction""" X, y = [], [] for i in range(len(features) - seq_length): X.append(features[i : i + seq_length].flatten()) # Flatten sequence y.append(targets[i + seq_length]) # Predict next value return np.array(X), np.array(y) def create_lagged_features(df, target_col, lags=[1, 2, 3, 5, 10, 20]): """Create lagged features for time series""" df_lagged = df.copy() for lag in lags: df_lagged[f"{target_col}_lag_{lag}"] = df_lagged[target_col].shift(lag) # Add rolling statistics for window in [5, 10, 20]: df_lagged[f"{target_col}_rolling_mean_{window}"] = ( df_lagged[target_col].rolling(window).mean() ) df_lagged[f"{target_col}_rolling_std_{window}"] = ( df_lagged[target_col].rolling(window).std() ) # Drop NaN values created by lagging df_lagged = df_lagged.dropna() return df_lagged class ModelTrainer: def __init__(self, model_name, model, save_dir="./checkpoints_classical"): self.model_name = model_name self.model = model self.save_dir = save_dir self.metrics = {} self.predictions = None def train(self, X_train, y_train): """Train the model""" print(f"\nTraining {self.model_name}...") self.model.fit(X_train, y_train) def predict(self, X): """Make predictions""" return self.model.predict(X) def evaluate(self, X_train, y_train, X_val, y_val): """Evaluate model on train and validation sets""" train_pred = self.predict(X_train) val_pred = self.predict(X_val) self.metrics = { "train_mse": mean_squared_error(y_train, train_pred), "train_rmse": np.sqrt(mean_squared_error(y_train, train_pred)), "train_mae": mean_absolute_error(y_train, train_pred), "train_r2": r2_score(y_train, train_pred), "val_mse": mean_squared_error(y_val, val_pred), "val_rmse": np.sqrt(mean_squared_error(y_val, val_pred)), "val_mae": mean_absolute_error(y_val, val_pred), "val_r2": r2_score(y_val, val_pred), } self.predictions = {"train": train_pred, "val": val_pred} return self.metrics def save_model(self, run_dir): """Save model to disk""" model_path = os.path.join(run_dir, f"{self.model_name}_model.pkl") with open(model_path, "wb") as f: pickle.dump(self.model, f) print(f"✓ Saved {self.model_name} model") class ARIMAModel: def __init__(self, order=(1, 1, 1)): self.order = order self.model = None self.model_fit = None def fit(self, X_train, y_train): """Fit ARIMA model - uses only target variable""" # ARIMA works on univariate time series self.model = ARIMA(y_train, order=self.order) self.model_fit = self.model.fit() def predict(self, X): """Make predictions""" n_periods = len(X) forecast = self.model_fit.forecast(steps=n_periods) return np.array(forecast) class SARIMAXModel: def __init__(self, order=(1, 1, 1), seasonal_order=(0, 0, 0, 0)): self.order = order self.seasonal_order = seasonal_order self.model = None self.model_fit = None def fit(self, X_train, y_train): """Fit SARIMAX model""" self.model = SARIMAX( y_train, order=self.order, seasonal_order=self.seasonal_order ) self.model_fit = self.model.fit(disp=False) def predict(self, X): """Make predictions""" n_periods = len(X) forecast = self.model_fit.forecast(steps=n_periods) return np.array(forecast) class ExponentialSmoothingModel: def __init__(self, seasonal_periods=None): self.seasonal_periods = seasonal_periods self.model = None self.model_fit = None def fit(self, X_train, y_train): """Fit Exponential Smoothing model""" self.model = ExponentialSmoothing( y_train, seasonal_periods=self.seasonal_periods, trend="add", seasonal="add" if self.seasonal_periods else None, ) self.model_fit = self.model.fit() def predict(self, X): """Make predictions""" n_periods = len(X) forecast = self.model_fit.forecast(steps=n_periods) return np.array(forecast) def get_ml_models(): """Get dictionary of classical ML models""" models = { # Linear Models "Ridge": Ridge(alpha=1.0), "Lasso": Lasso(alpha=0.1), "ElasticNet": ElasticNet(alpha=0.1, l1_ratio=0.5), "RandomForest": RandomForestRegressor( n_estimators=100, max_depth=10, min_samples_split=5, random_state=42, n_jobs=-1, ), "ExtraTrees": ExtraTreesRegressor( n_estimators=100, max_depth=10, min_samples_split=5, random_state=42, n_jobs=-1, ), "GradientBoosting": GradientBoostingRegressor( n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 ), "XGBoost": xgb.XGBRegressor( n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42, n_jobs=-1 ), "LightGBM": lgb.LGBMRegressor( n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42, n_jobs=-1, verbose=-1, ), "SVR": SVR(kernel="rbf", C=1.0, epsilon=0.1), } return models def get_time_series_models(): """Get dictionary of time series models""" models = { "ARIMA": ARIMAModel(order=(2, 1, 2)), "SARIMAX": SARIMAXModel(order=(1, 1, 1), seasonal_order=(1, 1, 1, 5)), "ExpSmoothing": ExponentialSmoothingModel(seasonal_periods=5), } return models def train_ml_models(X_train, y_train, X_val, y_val, save_dir): """Train all classical ML models""" models = get_ml_models() results = {} trained_models = {} print("\n" + "=" * 60) print("TRAINING CLASSICAL ML MODELS") print("=" * 60) for name, model in models.items(): try: trainer = ModelTrainer(name, model, save_dir) trainer.train(X_train, y_train) metrics = trainer.evaluate(X_train, y_train, X_val, y_val) trainer.save_model(save_dir) results[name] = metrics trained_models[name] = trainer print(f"\n{name}:") print( f" Train - RMSE: {metrics['train_rmse']:.6f}, MAE: {metrics['train_mae']:.6f}, R²: {metrics['train_r2']:.4f}" ) print( f" Val - RMSE: {metrics['val_rmse']:.6f}, MAE: {metrics['val_mae']:.6f}, R²: {metrics['val_r2']:.4f}" ) except Exception as e: print(f"\n{name}: FAILED - {str(e)}") results[name] = None return results, trained_models def train_time_series_models(y_train, y_val, save_dir): """Train time series models (univariate)""" models = get_time_series_models() results = {} trained_models = {} print("\n" + "=" * 60) print("TRAINING TIME SERIES MODELS") print("=" * 60) for name, model in models.items(): try: trainer = ModelTrainer(name, model, save_dir) # Time series models use only target variable trainer.train(None, y_train) # Make predictions train_pred = trainer.predict(np.arange(len(y_train))) val_pred = trainer.predict(np.arange(len(y_val))) # Calculate metrics metrics = { "train_mse": mean_squared_error(y_train, train_pred), "train_rmse": np.sqrt(mean_squared_error(y_train, train_pred)), "train_mae": mean_absolute_error(y_train, train_pred), "train_r2": r2_score(y_train, train_pred), "val_mse": mean_squared_error(y_val, val_pred), "val_rmse": np.sqrt(mean_squared_error(y_val, val_pred)), "val_mae": mean_absolute_error(y_val, val_pred), "val_r2": r2_score(y_val, val_pred), } trainer.metrics = metrics trainer.predictions = {"train": train_pred, "val": val_pred} trainer.save_model(save_dir) results[name] = metrics trained_models[name] = trainer print(f"\n{name}:") print( f" Train - RMSE: {metrics['train_rmse']:.6f}, MAE: {metrics['train_mae']:.6f}, R²: {metrics['train_r2']:.4f}" ) print( f" Val - RMSE: {metrics['val_rmse']:.6f}, MAE: {metrics['val_mae']:.6f}, R²: {metrics['val_r2']:.4f}" ) except Exception as e: print(f"\n{name}: FAILED - {str(e)}") results[name] = None return results, trained_models def plot_model_comparison(results, save_dir): """Plot comparison of all models""" # Filter out failed models results = {k: v for k, v in results.items() if v is not None} if not results: print("No successful models to plot") return models = list(results.keys()) # Extract metrics train_rmse = [results[m]["train_rmse"] for m in models] val_rmse = [results[m]["val_rmse"] for m in models] train_mae = [results[m]["train_mae"] for m in models] val_mae = [results[m]["val_mae"] for m in models] train_r2 = [results[m]["train_r2"] for m in models] val_r2 = [results[m]["val_r2"] for m in models] # Create comparison plots fig, axes = plt.subplots(2, 2, figsize=(16, 12)) # RMSE comparison ax = axes[0, 0] x = np.arange(len(models)) width = 0.35 ax.bar(x - width / 2, train_rmse, width, label="Train", alpha=0.8) ax.bar(x + width / 2, val_rmse, width, label="Validation", alpha=0.8) ax.set_xlabel("Model") ax.set_ylabel("RMSE") ax.set_title("Root Mean Squared Error Comparison") ax.set_xticks(x) ax.set_xticklabels(models, rotation=45, ha="right") ax.legend() ax.grid(True, alpha=0.3) # MAE comparison ax = axes[0, 1] ax.bar(x - width / 2, train_mae, width, label="Train", alpha=0.8) ax.bar(x + width / 2, val_mae, width, label="Validation", alpha=0.8) ax.set_xlabel("Model") ax.set_ylabel("MAE") ax.set_title("Mean Absolute Error Comparison") ax.set_xticks(x) ax.set_xticklabels(models, rotation=45, ha="right") ax.legend() ax.grid(True, alpha=0.3) # R² comparison ax = axes[1, 0] ax.bar(x - width / 2, train_r2, width, label="Train", alpha=0.8) ax.bar(x + width / 2, val_r2, width, label="Validation", alpha=0.8) ax.set_xlabel("Model") ax.set_ylabel("R² Score") ax.set_title("R² Score Comparison") ax.set_xticks(x) ax.set_xticklabels(models, rotation=45, ha="right") ax.legend() ax.grid(True, alpha=0.3) # Validation RMSE sorted ax = axes[1, 1] sorted_idx = np.argsort(val_rmse) sorted_models = [models[i] for i in sorted_idx] sorted_rmse = [val_rmse[i] for i in sorted_idx] colors = plt.cm.RdYlGn_r(np.linspace(0.3, 0.9, len(sorted_models))) ax.barh(range(len(sorted_models)), sorted_rmse, color=colors) ax.set_yticks(range(len(sorted_models))) ax.set_yticklabels(sorted_models) ax.set_xlabel("Validation RMSE") ax.set_title("Models Ranked by Validation RMSE") ax.grid(True, alpha=0.3, axis="x") plt.tight_layout() plt.savefig( os.path.join(save_dir, "model_comparison.png"), dpi=300, bbox_inches="tight" ) print(f"\n✓ Saved model comparison plot") plt.close() def plot_predictions_comparison(trained_models, y_val, save_dir, n_samples=200): """Plot predictions from top models""" # Get top 5 models by validation RMSE model_scores = [ (name, trainer.metrics["val_rmse"]) for name, trainer in trained_models.items() if trainer.metrics is not None ] model_scores.sort(key=lambda x: x[1]) top_models = model_scores[:5] fig, axes = plt.subplots(len(top_models), 1, figsize=(14, 4 * len(top_models))) if len(top_models) == 1: axes = [axes] plot_len = min(n_samples, len(y_val)) for i, (name, score) in enumerate(top_models): ax = axes[i] trainer = trained_models[name] val_pred = trainer.predictions["val"] ax.plot(y_val[:plot_len], label="Actual", alpha=0.7, linewidth=1.5) ax.plot(val_pred[:plot_len], label="Predicted", alpha=0.7, linewidth=1.5) ax.set_xlabel("Time Step") ax.set_ylabel("Value") ax.set_title(f"{name} Predictions (Val RMSE: {score:.6f})") ax.legend() ax.grid(True, alpha=0.3) plt.tight_layout() plt.savefig( os.path.join(save_dir, "top_model_predictions.png"), dpi=300, bbox_inches="tight", ) print(f"✓ Saved top model predictions plot") plt.close() def create_results_table(results, save_dir): """Create and save results table""" # Filter out failed models results = {k: v for k, v in results.items() if v is not None} df = pd.DataFrame(results).T df = df.sort_values("val_rmse") print("\n" + "=" * 80) print("MODEL COMPARISON RESULTS (sorted by validation RMSE)") print("=" * 80) print(df.to_string()) print("=" * 80) # Save to CSV df.to_csv(os.path.join(save_dir, "results_comparison.csv")) print(f"\n✓ Saved results table") return df # ========================= ABLATION STUDIES ========================= def run_ablation_study(X_train, y_train, X_val, y_val, save_dir): """Run ablation studies on feature importance and model configurations""" print("\n" + "=" * 60) print("ABLATION STUDY: Feature Importance") print("=" * 60) # Train a Random Forest to get feature importances rf_model = RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1) rf_model.fit(X_train, y_train) # Get feature importances importances = rf_model.feature_importances_ # Test with different number of features n_features_list = [10, 20, 50, 100, X_train.shape[1]] ablation_results = {} for n_features in n_features_list: if n_features > X_train.shape[1]: continue # Select top n features top_indices = np.argsort(importances)[-n_features:] X_train_subset = X_train[:, top_indices] X_val_subset = X_val[:, top_indices] # Train model with subset model = RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1) model.fit(X_train_subset, y_train) val_pred = model.predict(X_val_subset) rmse = np.sqrt(mean_squared_error(y_val, val_pred)) mae = mean_absolute_error(y_val, val_pred) r2 = r2_score(y_val, val_pred) ablation_results[f"Top_{n_features}_features"] = { "val_rmse": rmse, "val_mae": mae, "val_r2": r2, } print( f"\nTop {n_features} features: RMSE={rmse:.6f}, MAE={mae:.6f}, R²={r2:.4f}" ) # Save ablation results ablation_df = pd.DataFrame(ablation_results).T ablation_df.to_csv(os.path.join(save_dir, "ablation_feature_importance.csv")) # Plot ablation results fig, ax = plt.subplots(figsize=(10, 6)) x = range(len(ablation_results)) ax.plot( list(ablation_results.keys()), [v["val_rmse"] for v in ablation_results.values()], "o-", linewidth=2, markersize=8, ) ax.set_xlabel("Number of Features") ax.set_ylabel("Validation RMSE") ax.set_title("Ablation Study: Impact of Feature Count on Performance") ax.grid(True, alpha=0.3) plt.xticks(rotation=45, ha="right") plt.tight_layout() plt.savefig(os.path.join(save_dir, "ablation_feature_importance.png"), dpi=300) plt.close() print(f"\n✓ Saved ablation study results") return ablation_results # ========================= MAIN EXECUTION ========================= def main(): from data_prep.data_clean import clean_indicator from data_prep.data_load import prepare_data # Configuration config = { "data_path": "/home/aman/code/ml_fr/ml_stocks/data/NIFTY_5_years.csv", "seq_length": 20, "train_split": 0.8, "save_dir": "./checkpoints_classical", "target_col": "Daily_Return", } # Create save directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") save_dir = os.path.join(config["save_dir"], f"run_{timestamp}") os.makedirs(save_dir, exist_ok=True) print(f"\n{'='*60}") print(f"CLASSICAL ML & TIME SERIES MODEL TRAINING") print(f"{'='*60}") print(f"Save directory: {save_dir}") print(f"{'='*60}\n") # Load and prepare data print("Loading data...") load_df = prepare_data(config["data_path"]) df = clean_indicator(load_df) target_col = config["target_col"] feature_cols = [col for col in df.columns if col != target_col] # Split data train_size = int(len(df) * config["train_split"]) train_df = df[:train_size] val_df = df[train_size:] print(f"Train samples: {len(train_df)}") print(f"Validation samples: {len(val_df)}") print(f"Number of features: {len(feature_cols)}") # Prepare features for ML models (with sequences) scaler = StandardScaler() train_features = scaler.fit_transform(train_df[feature_cols].values) val_features = scaler.transform(val_df[feature_cols].values) train_targets = train_df[target_col].values val_targets = val_df[target_col].values # Create sequences X_train, y_train = create_sequences( train_features, train_targets, config["seq_length"] ) X_val, y_val = create_sequences(val_features, val_targets, config["seq_length"]) print(f"\nSequence shape: {X_train.shape}") print(f"Target shape: {y_train.shape}") # Save config with open(os.path.join(save_dir, "config.json"), "w") as f: json.dump(config, f, indent=4) # Train ML models ml_results, ml_models = train_ml_models(X_train, y_train, X_val, y_val, save_dir) # Train time series models (using non-sequenced data) ts_results, ts_models = train_time_series_models( train_targets[config["seq_length"] :], # Align with ML model targets val_targets[config["seq_length"] :], save_dir, ) # Combine results all_results = {**ml_results, **ts_results} all_models = {**ml_models, **ts_models} # Create visualizations print("\n" + "=" * 60) print("CREATING VISUALIZATIONS") print("=" * 60) plot_model_comparison(all_results, save_dir) plot_predictions_comparison(all_models, y_val, save_dir) results_df = create_results_table(all_results, save_dir) # Run ablation study ablation_results = run_ablation_study(X_train, y_train, X_val, y_val, save_dir) print(f"\n{'='*60}") print("TRAINING COMPLETE!") print(f"Results saved to: {save_dir}") print(f"{'='*60}\n") # Print best model best_model = results_df.index[0] best_rmse = results_df.loc[best_model, "val_rmse"] print(f"🏆 Best Model: {best_model}") print(f" Validation RMSE: {best_rmse:.6f}") print(f" Validation MAE: {results_df.loc[best_model, 'val_mae']:.6f}") print(f" Validation R²: {results_df.loc[best_model, 'val_r2']:.4f}") return all_results, all_models, save_dir if __name__ == "__main__": results, models, save_dir = main() print("\n" + "=" * 60) print("All models trained successfully!") print("Check the save directory for:") print(" - Model comparison plots") print(" - Results CSV") print(" - Saved model files (.pkl)") print(" - Ablation study results") print("=" * 60)