plot_all_model_corner_errors_scatter.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from xgboost import XGBClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.utils import shuffle
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Flatten, Dense, Dropout, BatchNormalization, Activation, Conv1D, MaxPooling1D, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping
import corner
import os
import re
import gc

# --- Configuration ---
FILL_GAS = "H2"
TRAIN_FILE = f"multirex_spectra_{FILL_GAS}_train.parquet"
TEST_FILE = f"multirex_spectra_{FILL_GAS}_test.parquet"
RESULTS_DIR = "final_results/plots"
PARAMS_TO_PLOT = ['p_radius', 'p_mass', 's temperature', 'atm temperature']
LABELS = ['Planet Radius (R_earth)', 'Planet Mass (M_earth)', 'Star Temp (K)', 'Atmosphere Temp (K)']

# --- Load and Prepare Data ---
print("--- Loading and Preparing Data ---")
df_train = pd.read_parquet(TRAIN_FILE)
df_test = pd.read_parquet(TEST_FILE)

df_train['label'] = df_train['biosignature'].apply(lambda x: 1 if x == 'yes' else 0)
df_test['label'] = df_test['biosignature'].apply(lambda x: 1 if x == 'yes' else 0)

float_pattern = re.compile(r"^-?\d+\.\d+$")
cols = [c for c in df_train.columns if isinstance(c, float) or (isinstance(c, str) and float_pattern.match(c))]

X_train_raw = df_train[cols].values
y_train = df_train['label'].values
X_test_raw = df_test[cols].values
y_test = df_test['label'].values

# --- Preprocessing (PCA) ---
scaler_raw = StandardScaler()
X_train_s = scaler_raw.fit_transform(X_train_raw)
X_test_s = scaler_raw.transform(X_test_raw)

pca = PCA(n_components=102, random_state=42)
X_train_p = pca.fit_transform(X_train_s)
X_test_p = pca.transform(X_test_s)

# Feature Sets (Unified 0-101 for all models)
X_train_final = X_train_p[:, 0:102]
X_test_final = X_test_p[:, 0:102]

# Apply second scaler (Whitening)
scaler_pca = StandardScaler()
X_train_final = scaler_pca.fit_transform(X_train_final)
X_test_final = scaler_pca.transform(X_test_final)

# Shuffle training data
X_train_final, y_train = shuffle(X_train_final, y_train, random_state=42)

# --- Model Training Functions ---
def train_and_predict_xgb():
    print("--- Training XGBoost ---")
    model = XGBClassifier(n_estimators=200, max_depth=5, learning_rate=0.2, subsample=0.8, random_state=42, n_jobs=-1, eval_metric='logloss')
    model.fit(X_train_final, y_train)
    return model.predict(X_test_final)

def train_and_predict_rf():
    print("--- Training Random Forest ---")
    model = RandomForestClassifier(n_estimators=300, max_depth=None, min_samples_leaf=2, min_samples_split=2, random_state=42, n_jobs=-1)
    model.fit(X_train_final, y_train)
    return model.predict(X_test_final)

def train_and_predict_mlp():
    print("--- Training MLP ---")
    model = Sequential([
        Input(shape=(102,)),
        Dense(256), BatchNormalization(), Activation('relu'), Dropout(0.3),
        Dense(128), BatchNormalization(), Activation('relu'), Dropout(0.3),
        Dense(64), BatchNormalization(), Activation('relu'), Dropout(0.3),
        Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss='binary_crossentropy')
    es = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
    model.fit(X_train_final, y_train, epochs=100, batch_size=128, validation_split=0.2, callbacks=[es], verbose=0)
    return (model.predict(X_test_final, verbose=0) > 0.5).astype(int).flatten()

def train_and_predict_cnn():
    print("--- Training CNN ---")
    model = Sequential([
        Input(shape=(102, 1)),
        Conv1D(filters=32, kernel_size=5, padding='same'), BatchNormalization(), Activation('relu'), MaxPooling1D(pool_size=2), Dropout(0.3),
        Conv1D(filters=64, kernel_size=5, padding='same'), BatchNormalization(), Activation('relu'), MaxPooling1D(pool_size=2), Dropout(0.3),
        Flatten(),
        Dense(100), BatchNormalization(), Activation('relu'), Dropout(0.5),
        Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer=Adam(learning_rate=0.0005), loss='binary_crossentropy')
    es = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
    model.fit(X_train_final.reshape(-1, 102, 1), y_train, epochs=100, batch_size=128, validation_split=0.2, callbacks=[es], verbose=0)
    return (model.predict(X_test_final.reshape(-1, 102, 1), verbose=0) > 0.5).astype(int).flatten()

# --- Main Loop ---
models_to_run = {
    "XGBoost": train_and_predict_xgb,
    "RandomForest": train_and_predict_rf,
    "MLP": train_and_predict_mlp,
    "CNN": train_and_predict_cnn
}

for name, train_func in models_to_run.items():
    gc.collect()
    y_pred = train_func()
    
    # Identify errors
    correct_mask = (y_pred == y_test)
    incorrect_mask = ~correct_mask
    params_correct = df_test[correct_mask][PARAMS_TO_PLOT].values
    params_incorrect = df_test[incorrect_mask][PARAMS_TO_PLOT].values

    # Generate Plot
    print(f"--- Generating Scatter Corner Plot for {name} ---")
    
    figure = plt.figure(figsize=(15, 15))

    # This is the corrected plotting logic
    # Plot Incorrect points first, so they are more visible
    corner.corner(params_incorrect, fig=figure, labels=LABELS,
                  plot_contours=False, plot_density=False, plot_datapoints=True,
                  data_kwargs={'color': 'crimson', 'alpha': 0.7})
    # Overlay Correct points
    corner.corner(params_correct, fig=figure, labels=LABELS,
                  plot_contours=False, plot_density=False, plot_datapoints=True,
                  data_kwargs={'color': 'navy', 'alpha': 0.3})

    from matplotlib.lines import Line2D
    legend_elements = [Line2D([0], [0], marker='o', color='w', label='Correct', markerfacecolor='navy', markersize=10, alpha=0.5),
                       Line2D([0], [0], marker='o', color='w', label='Incorrect', markerfacecolor='crimson', markersize=10, alpha=0.8)]
    plt.legend(handles=legend_elements, loc='upper right', fontsize=14, bbox_to_anchor=(0.95, 0.95), frameon=False)
    plt.suptitle(f"Scatter Corner Plot of Prediction Errors ({name})", fontsize=20)
    
    filename = os.path.join(RESULTS_DIR, f'corner_plot_errors_scatter_{name.lower()}.png')
    plt.savefig(filename, dpi=300)
    plt.close()
    print(f"Plot saved to: {filename}")

print("\n--- All Scatter Corner Plots Generated ---")