cup_model.py

import numpy as np
import matplotlib.pyplot as plt

def fit_standardize_X(X_trn, eps=1e-8):
    """
    Computes the mean and standard deviation of the training input data (X_trn).
    These statistics are used to normalize the data to have zero mean and unit variance.
    Adds a small epsilon to the standard deviation to prevent division by zero.
    Returns: (mean, std_dev) tuple.
    """
    
    mu = X_trn.mean(axis=0, keepdims=True)
    sd = X_trn.std(axis=0, keepdims=True) + eps
    return mu, sd

def apply_standardize_X(X, mu, sd):
    """
    Applies Z-score normalization to input data X using the provided mean (mu) 
    and standard deviation (sd). 
    Formula: (X - mu) / sd
    """
    
    return (X - mu) / sd

def fit_standardize_Y(Y_trn, eps=1e-8):
    """
    Computes the mean and standard deviation of the training target data (Y_trn).
    Used to normalize regression targets, which helps gradient descent converge faster.
    Returns: (mean, std_dev) tuple.
    """
    
    mu = Y_trn.mean(axis=0, keepdims=True)
    sd = Y_trn.std(axis=0, keepdims=True) + eps
    return mu, sd

def apply_standardize_Y(Y, mu, sd):
    """
    Applies Z-score normalization to target data Y using the provided mean (mu) 
    and standard deviation (sd).
    """
    
    return (Y - mu) / sd

def relu(z):
    return np.maximum(0.0, z)

def init_mlp_3hidden(d, h1, h2, h3, k, seed=0):
    
    """
    Initializes all weights (W1-W4) and biases (b1-b4) for a 3-hidden-layer MLP.
    Uses the 'He' initialization strategy by default (via the internal 'xavier' helper, 
    which appears to be implemented as He/Xavier uniform in the script logic).
    Returns: All weight and bias matrices.
    """
    
    rng = np.random.default_rng(seed)

    def xavier(n_in, n_out):
        limit = np.sqrt(6.0 / (n_in + n_out))
        return rng.uniform(-limit, limit, size=(n_in, n_out))

    W1 = xavier(d, h1);  b1 = np.zeros((1, h1))
    W2 = xavier(h1, h2); b2 = np.zeros((1, h2))
    W3 = xavier(h2, h3); b3 = np.zeros((1, h3))
    W4 = xavier(h3, k);  b4 = np.zeros((1, k))
    return W1, b1, W2, b2, W3, b3, W4, b4

def predict_mlp3(model, stats, X, activation='relu'):
    
    """
    Performs a forward pass to generate predictions for input X using the trained model.
    1. Normalizes input X using stored stats.
    2. Propagates data through 3 hidden layers + 1 output layer.
    3. Denormalizes the output back to the original target scale.
    Returns: Predictions in the original scale.
    """
    
    W1, b1, W2, b2, W3, b3, W4, b4 = model
    muX, sdX, muY, sdY = stats
    
    X_n = (X - muX) / sdX
    Z1 = X_n @ W1 + b1; A1 = relu(Z1)
    Z2 = A1 @ W2 + b2;  A2 = relu(Z2)
    Z3 = A2 @ W3 + b3;  A3 = relu(Z3)
    Y_hat_n = A3 @ W4 + b4
    
    return Y_hat_n * sdY + muY

def train_mlp3_flexible(
    
    X_trn, Y_trn, X_val, Y_val, 
    h1=128, h2=64, h3=32, 
    lr=1e-3, momentum=0.8, l2=0.0, 
    activation='relu', init_type='xavier', dropout=0.0,
    epochs=4000, patience=200, min_delta=1e-4, 
    clip_value=5.0, seed=0, verbose_n=200, plot=False
):
    """
    Main training loop for the ML-CUP regression task.
    1. Normalizes Inputs and Targets.
    2. Initializes the network.
    3. Iterates for 'epochs' times:
       - Performs Forward Pass (computes predictions).
       - Performs Backward Pass (computes gradients of MSE).
       - Updates weights using SGD with Momentum.
       - Evaluates Train and Validation MEE (Mean Euclidean Error) *after* the update.
    4. Implements Early Stopping based on Validation MEE.
    Returns: Best model parameters, Normalization stats, Best Val MEE, and History dict.
    """
    
    rng = np.random.default_rng(seed)
    
    # 1. Normalization
    muX, sdX = fit_standardize_X(X_trn)
    X_trn_n = apply_standardize_X(X_trn, muX, sdX)
    X_val_n = apply_standardize_X(X_val, muX, sdX)
    muY, sdY = fit_standardize_Y(Y_trn)
    Y_trn_n = apply_standardize_Y(Y_trn, muY, sdY)
    Y_val_n = apply_standardize_Y(Y_val, muY, sdY)

    d = X_trn.shape[1]
    k_out = Y_trn.shape[1]
    
    # 2. Initialization
    W1, b1, W2, b2, W3, b3, W4, b4 = init_mlp_3hidden(d, h1, h2, h3, k_out, seed=seed)
    
    # Momentum Buffers
    vW1 = np.zeros_like(W1); vb1 = np.zeros_like(b1)
    vW2 = np.zeros_like(W2); vb2 = np.zeros_like(b2)
    vW3 = np.zeros_like(W3); vb3 = np.zeros_like(b3)
    vW4 = np.zeros_like(W4); vb4 = np.zeros_like(b4)

    best_val_mee = np.inf; best_epoch = 0; wait = 0
    best_params = (W1.copy(), b1.copy(), W2.copy(), b2.copy(), W3.copy(), b3.copy(), W4.copy(), b4.copy())
    history = {'train_mee': [], 'val_mee': []}
    N = X_trn_n.shape[0]

    def d_relu(A):
        return (A > 0).astype(float)

    for epoch in range(epochs):
        # --- FORWARD TRAIN ---
        Z1 = X_trn_n @ W1 + b1; A1 = relu(Z1)
        Z2 = A1 @ W2 + b2;      A2 = relu(Z2)
        Z3 = A2 @ W3 + b3;      A3 = relu(Z3)
        Y_hat_trn_n = A3 @ W4 + b4

        # --- BACKWARD ---
        dZ4 = (Y_hat_trn_n - Y_trn_n) * (2.0 / N)
        dW4 = A3.T @ dZ4 + 2.0 * l2 * W4; db4 = np.sum(dZ4, axis=0, keepdims=True)
        dA3 = dZ4 @ W4.T; dZ3 = dA3 * d_relu(A3)
        dW3 = A2.T @ dZ3 + 2.0 * l2 * W3; db3 = np.sum(dZ3, axis=0, keepdims=True)
        dA2 = dZ3 @ W3.T; dZ2 = dA2 * d_relu(A2)
        dW2 = A1.T @ dZ2 + 2.0 * l2 * W2; db2 = np.sum(dZ2, axis=0, keepdims=True)
        dA1 = dZ2 @ W2.T; dZ1 = dA1 * d_relu(A1)
        dW1 = X_trn_n.T @ dZ1 + 2.0 * l2 * W1; db1 = np.sum(dZ1, axis=0, keepdims=True)

        if clip_value:
            for g in (dW1, db1, dW2, db2, dW3, db3, dW4, db4): np.clip(g, -clip_value, clip_value, out=g)

        # --- UPDATE ---
        vW4 = momentum * vW4 - lr * dW4; vb4 = momentum * vb4 - lr * db4
        vW3 = momentum * vW3 - lr * dW3; vb3 = momentum * vb3 - lr * db3
        vW2 = momentum * vW2 - lr * dW2; vb2 = momentum * vb2 - lr * db2
        vW1 = momentum * vW1 - lr * dW1; vb1 = momentum * vb1 - lr * db1
        W4 += vW4; b4 += vb4; W3 += vW3; b3 += vb3; W2 += vW2; b2 += vb2; W1 += vW1; b1 += vb1

        # --- EVALUATION (Corrected placement: AFTER update) ---
        # 1. Train Metrics
        Z1_t = X_trn_n @ W1 + b1; A1_t = relu(Z1_t)
        Z2_t = A1_t @ W2 + b2;    A2_t = relu(Z2_t)
        Z3_t = A2_t @ W3 + b3;    A3_t = relu(Z3_t)
        Y_hat_t = (A3_t @ W4 + b4) * sdY + muY
        train_mee = float(np.mean(np.linalg.norm(Y_trn - Y_hat_t, axis=1)))
        history['train_mee'].append(train_mee)

        # 2. Validation Metrics
        Z1_v = X_val_n @ W1 + b1; A1_v = relu(Z1_v)
        Z2_v = A1_v @ W2 + b2;    A2_v = relu(Z2_v)
        Z3_v = A2_v @ W3 + b3;    A3_v = relu(Z3_v)
        Y_hat_v = (A3_v @ W4 + b4) * sdY + muY
        val_mee = float(np.mean(np.linalg.norm(Y_val - Y_hat_v, axis=1)))
        history['val_mee'].append(val_mee)

        # --- EARLY STOPPING ---
        if val_mee < best_val_mee - min_delta:
            best_val_mee = val_mee; best_epoch = epoch + 1; wait = 0
            best_params = (W1.copy(), b1.copy(), W2.copy(), b2.copy(), W3.copy(), b3.copy(), W4.copy(), b4.copy())
        else:
            wait += 1
            if wait >= patience: break
            
        if verbose_n and (epoch + 1) % verbose_n == 0:
            print(f"Epoch {epoch+1} | Train MEE: {train_mee:.4f} | Val MEE: {val_mee:.4f}")

    return best_params, (muX, sdX, muY, sdY), best_val_mee, best_epoch, history

def save_cup_model_npz(filepath, model, stats, config):
    """
    Saves the trained model weights, biases, and normalization statistics to a .npz file.
    Also saves the configuration dictionary as a string for reference.
    """
    
    W1, b1, W2, b2, W3, b3, W4, b4 = model
    muX, sdX, muY, sdY = stats
    np.savez(filepath, 
             W1=W1, b1=b1, W2=W2, b2=b2, W3=W3, b3=b3, W4=W4, b4=b4, 
             muX=muX, sdX=sdX, muY=muY, sdY=sdY,
             config=str(config))
    print(f"Model saved to {filepath}")