train_and_evaluate.py

"""
Benchmark pipeline with target-wise evaluation and normalized training.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import joblib
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
from src.preprocess import preprocess_data
from src.models import (
    get_polynomial_model,
    get_random_forest,
    get_xgboost,
    build_physics_consistent_ann
)
from src.config import ALL_TARGETS, PURITY_TARGETS, DUTY_TARGETS, MODEL_PATH, PLOT_PATH

def compute_detailed_metrics(y_true, y_pred):
    r2_vals = r2_score(y_true, y_pred, multioutput='raw_values')
    mae_vals = mean_absolute_error(y_true, y_pred, multioutput='raw_values')
    rmse_vals = np.sqrt(mean_squared_error(y_true, y_pred, multioutput='raw_values'))
    
    return {
        "R2_xD": r2_vals[0],
        "R2_xB": r2_vals[1],
        "R2_QC": r2_vals[2],
        "R2_QR": r2_vals[3],
        "Mean_R2": np.mean(r2_vals),
        "Purity_MAE": (mae_vals[0] + mae_vals[1]) / 2.0,
        "Duty_MAE_kW": (mae_vals[2] + mae_vals[3]) / 2.0,
    }

def main():
    X_train, X_test, y_train, y_test = preprocess_data()
    
    # Scale targets for uniform gradient optimization across all scales
    scaler_y = StandardScaler()
    y_train_scaled = scaler_y.fit_transform(y_train)
    y_test_scaled = scaler_y.transform(y_test)
    
    models = {
        "Polynomial (Deg 2)": get_polynomial_model(),
        "Random Forest": get_random_forest(),
        "XGBoost": get_xgboost()
    }
    
    results = {}
    predictions = {}
    
    # 1. Train Scikit-Learn / XGBoost Models
    for name, model in models.items():
        print(f"[*] Training {name}...")
        model.fit(X_train, y_train_scaled)
        y_pred_scaled = model.predict(X_test)
        y_pred = scaler_y.inverse_transform(y_pred_scaled)
        predictions[name] = y_pred
        results[name] = compute_detailed_metrics(y_test, y_pred)
        
    # 2. Train Physics-Consistent ANN
    print("[*] Training Normalized Multi-Head ANN...")
    ann = build_physics_consistent_ann(input_dim=X_train.shape[1])
    # Train directly on standard scaled target values
    ann.fit(X_train, y_train_scaled, validation_split=0.15, epochs=300, batch_size=32, verbose=0)
    y_pred_ann_scaled = ann.predict(X_test)
    y_pred_ann = scaler_y.inverse_transform(y_pred_ann_scaled)
    predictions["Multi-Head ANN"] = y_pred_ann
    results["Multi-Head ANN"] = compute_detailed_metrics(y_test, y_pred_ann)
    
    ann.save(MODEL_PATH)
    joblib.dump(scaler_y, "artifacts/scaler_y.pkl")
    
    # 3. Formatted Benchmark Table
    summary_df = pd.DataFrame(results).T
    pd.set_option('display.max_columns', None)
    pd.set_option('display.width', 1000)
    
    print("\n" + "="*80)
    print("DETAILED SURROGATE MODEL BENCHMARK RESULTS (PER-TARGET BREAKDOWN)")
    print("="*80)
    print(summary_df.round(4).to_string())
    print("="*80)
    
    # 4. Generate Parity Plots
    fig, axs = plt.subplots(2, 2, figsize=(11, 9))
    units = ["mol/mol", "mol/mol", "kW", "kW"]
    colors = ["#1E40AF", "#047857", "#D97706", "#DC2626"]
    
    for i, ax in enumerate(axs.flat):
        ax.scatter(y_test[:, i], y_pred_ann[:, i], alpha=0.6, edgecolors='k', s=25, color=colors[i])
        min_v = min(y_test[:, i].min(), y_pred_ann[:, i].min())
        max_v = max(y_test[:, i].max(), y_pred_ann[:, i].max())
        ax.plot([min_v, max_v], [min_v, max_v], 'r--', lw=1.5, label="Parity (y = x)")
        ax.set_title(f"{ALL_TARGETS[i]} ({units[i]})", fontweight="bold")
        ax.set_xlabel("DWSIM Simulation Value")
        ax.set_ylabel("ANN Predicted Value")
        ax.legend(loc="upper left")
        ax.grid(True, linestyle=":", alpha=0.6)
        
    plt.tight_layout()
    plt.savefig(PLOT_PATH, dpi=300)
    print(f"\n[+] Updated high-accuracy parity plots saved to: {PLOT_PATH}")

if __name__ == "__main__":
    main()