train.py

# Copyright 2023 Andreas Koehler, Erik Myklebust, MIT license

import numpy as np
from sklearn.model_selection import GroupKFold, KFold
from sklearn.metrics import mean_absolute_error, accuracy_score, balanced_accuracy_score
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.utils.class_weight import compute_sample_weight
from scipy.optimize import minimize
import matplotlib.pyplot as plt

# code for training ArrayNet

REPEATS = 1
FOLDS = 5
BATCH_SIZE = 128

from sklearn.metrics.pairwise import paired_cosine_distances

import numpy as np

def _smallestSignedAngleBetween(x, y):
    """
    Helper function.
    Returns angle between two angles x and y in radians.
    """
    tau = 2 * np.pi
    a = (x - y) % tau
    b = (y - x) % tau
    return np.where(a < b, -a, b)

def circular_r2_score(true,pred):
    """
    Calculates R2 score between angles (in rad).
    Angles are shifted into (0,2pi) where angle differences are calculated.
    Thereafter, same calulation as for normal R2 score.
    params:
        true : np.array
        pred : np.array

    return:
        float : r2 score between the two sets of angles.
    """
    res = sum(_smallestSignedAngleBetween(true,pred)**2)
    tot = sum(_smallestSignedAngleBetween(true,true.mean())**2)
    return 1 - (res/tot)

def find_nearest_idx(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return idx
"""
def distribution_sample_weights(hist, bins, angle):
    idx = find_nearest_idx(bins, angle)
    idx %= len(hist)
    return np.clip(1/hist[idx], 1e-2, 1e9)
"""

from scipy.interpolate import interp1d

def distribution_sample_weights(f, angle):
    return 1/f(angle)

def load_data(X,y,y_class, weight_class=False, weight_angle=False):
    where_noise = np.where(y_class == 'N')[0]

    y[np.isnan(y)] = 0
    X[np.isnan(X)] = 0

    le = OneHotEncoder()
    if weight_class:
        class_weights = compute_sample_weight('balanced',y_class)
    else:
        class_weights = np.ones(y_class.shape[0])

    y_class = np.asarray(le.fit_transform(y_class.reshape((-1,1))).todense())

    angles = np.angle(y.view(complex))
    if weight_angle:
        num_bins = 36
        hist, bins = np.histogram(angles, bins=num_bins, density=True)
        x = np.linspace(-np.pi, np.pi, num=num_bins, endpoint=True)
        f = interp1d(x,hist,kind='cubic')
        sample_weights = 1/f(angles)
        sample_weights = (sample_weights - sample_weights.min()) / (sample_weights.max() - sample_weights.min())
        sample_weights = np.clip(sample_weights, 0.1, 1)
    else:
        sample_weights = np.ones(y.shape[0])

    sample_weights[where_noise] = 0
    sample_weights[(y == 0).all(axis=-1)] = 0

    return X, (y, y_class), (sample_weights, class_weights)

def load_data_sub(X, y, y_class, weight_class=False, weight_angle=False):

    where_p = np.where(['P' in s for s in y_class])[0]
    where_s = np.where(['S' in s for s in y_class])[0]
    where_only_p = np.where(['P' == s for s in y_class])[0]
    where_only_s = np.where(['S' == s for s in y_class])[0]

    p_class = y_class.copy()
    idx = np.union1d(np.setdiff1d(np.arange(len(p_class)), where_p), where_only_p)
    if weight_class:
        p_weights = compute_sample_weight('balanced', p_class)
    else:
        p_weights = np.ones(p_class.shape[0])
    p_class[idx] = 'PN'  # will be weighted by zero
    p_weights[idx] = 0

    s_class = y_class.copy()
    idx = np.union1d(np.setdiff1d(np.arange(len(s_class)), where_s), where_only_s)
    if weight_class:
        s_weights = compute_sample_weight('balanced', s_class)
    else:
        s_weights = np.ones(s_class.shape[0])
    s_class[idx] = 'SN'  # will be weighted by zero
    s_weights[idx] = 0

    y_class = np.array([c[0] for c in y_class]) #removes N/G from PN/PG etc.

    where_noise = np.setdiff1d(np.arange(len(y_class)),
                               np.union1d(np.where(y_class == 'P')[0], np.where(y_class == 'S')[0]))

    y[np.isnan(y)] = 0
    X[np.isnan(X)] = 0

    le = OneHotEncoder()
    lep = OneHotEncoder()
    les = OneHotEncoder()

    if weight_class:
        class_weights = compute_sample_weight('balanced', y_class)
    else:
        class_weights = np.ones(y_class.shape[0])

    y_class = np.asarray(le.fit_transform(y_class.reshape((-1, 1))).todense())
    p_class = np.asarray(lep.fit_transform(p_class.reshape((-1, 1))).todense())
    s_class = np.asarray(les.fit_transform(s_class.reshape((-1, 1))).todense())

    angles = np.angle(y.view(complex))
    if weight_angle:
        num_bins = 36
        hist, bins = np.histogram(angles, bins=num_bins, density=True)
        x = np.linspace(-np.pi, np.pi, num=num_bins, endpoint=True)
        f = interp1d(x,hist,kind='cubic')
        sample_weights = 1/f(angles)
        sample_weights = (sample_weights - sample_weights.min()) / (sample_weights.max() - sample_weights.min())
        sample_weights = np.clip(sample_weights, 0.1, 1)
    else:
        sample_weights = np.ones(y.shape[0])

    sample_weights[where_noise] = 0
    sample_weights[(y == 0).all(axis=-1)] = 0

    return X, (y, y_class, p_class, s_class), (sample_weights, class_weights, p_weights, s_weights)


def angle_diff_tf(true,pred,sample_weight=None):
    true = tf.math.angle(tf.complex(true[:, 0], true[:, 1]))
    pred = tf.math.angle(tf.complex(pred[:, 0], pred[:, 1]))

    diff = tf.math.atan2(tf.math.sin(true-pred), tf.math.cos(true-pred))
    if sample_weight:
        return sample_weight * diff
    return diff

def angle_mae_tf(true,pred,sample_weight=None):
    return tf.reduce_mean(tf.math.abs(angle_diff_tf(true,pred,sample_weight)))

def angle_mse_tf(true,pred,sample_weight=None):
    return tf.reduce_mean(tf.math.square(angle_diff_tf(true,pred,sample_weight)))

from tensorflow.keras import backend as K

def denseblock(x, units=128, k=3, act='relu'):

    features = [x]

    for _ in range(k):
        y = tf.keras.layers.Concatenate()(features) if len(features)>1 else features[0]
        y = tf.keras.layers.Dense(units)(y)
        y = tf.keras.layers.BatchNormalization()(y)
        y = tf.keras.layers.Activation(act)(y)
        features.append(y)

    return features[-1]

def superblock(x, num_blocks, name='superblock'):
    for i in range(num_blocks):
        x = denseblock(x, 512, 3, 'relu')
        x = tf.keras.layers.Dropout(0.2)(x)
    return x

import tensorflow as tf

def create_model(params):

    inp = tf.keras.layers.Input(params['input_shape'])

    x = inp

    x = superblock(x, num_blocks=2, name='common')
    cl_output = superblock(x, num_blocks=2, name='cl')
    reg_output = superblock(x, num_blocks=2, name='reg')

    cl_output = tf.keras.layers.Dense(params['num_classes'], activation='softmax', name='cl_output')(cl_output)
    reg_output = tf.keras.layers.Dense(params['num_outputs'], activation='linear',
                                       kernel_initializer='zeros',
                                       )(reg_output)
    
    model = tf.keras.Model(inputs=inp, outputs=[reg_output, cl_output])

    alpha = 0.5

    model.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                  loss=['mse',
                        tf.keras.losses.CategoricalCrossentropy()],
                  loss_weights=[alpha, 1-alpha],
                  weighted_metrics=([angle_mae_tf], ['accuracy'])
                  )
    return model

def create_model_sub(params):

    inp = tf.keras.layers.Input(params['input_shape'])

    act = 'relu'

    x = inp

    x = superblock(x, num_blocks=2, name='common')
    cl_output = superblock(x, num_blocks=2, name='cl')
    reg_output = superblock(x, num_blocks=2, name='reg')
    p_output = superblock(x, num_blocks=2, name='p')
    s_output = superblock(x, num_blocks=2, name='s')

    cl_output = tf.keras.layers.Dense(params['num_classes'], activation='softmax', name='cl_output')(cl_output)
    p_output = tf.keras.layers.Dense(params['num_p_classes'], activation='softmax', name='p_output')(p_output)
    s_output = tf.keras.layers.Dense(params['num_s_classes'], activation='softmax', name='s_output')(s_output)

    reg_output = tf.keras.layers.Dense(params['num_outputs'], activation='linear',
                                       kernel_initializer='zeros',
                                       )(reg_output)

    model = tf.keras.Model(inputs=inp, outputs=[reg_output, cl_output, p_output, s_output])

    alpha = 0.5

    model.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                  loss=['mse',
                        tf.keras.losses.CategoricalCrossentropy(),
                        tf.keras.losses.CategoricalCrossentropy(),
                        tf.keras.losses.CategoricalCrossentropy(),
                        ],
                  loss_weights=[alpha, 1-alpha, 1-alpha, 1-alpha],
                  weighted_metrics=([angle_mae_tf], ['accuracy'], ['accuracy'], ['accuracy'])
                  )
    return model


# ============================================================================
# Uncertainty Estimation Functions
# ============================================================================

# Neural Network Calibration for Phase Classification
def temperature_scale(logits, temperature):
    """Apply temperature scaling to logits."""
    return logits / temperature


def calculate_calibration_curve(model, X_val, y_val, temperature, num_bins=10):
    """
    Calculate calibration curve for a given temperature.
    
    Args:
        model: Keras model that outputs logits (before softmax)
        X_val: Validation input data
        y_val: Validation true labels (one-hot encoded)
        temperature: Temperature scaling parameter
        num_bins: Number of bins for calibration curve
        
    Returns:
        avg_confidences: Average predicted confidence per bin
        accuracies: Observed accuracy per bin
    """
    # Get the logits (before softmax)
    logits = model(X_val, training=False)
    
    # Apply temperature scaling
    scaled_logits = temperature_scale(logits, temperature)
    
    # Get predicted probabilities (using softmax)
    predicted_probs = tf.nn.softmax(scaled_logits)
    
    # Predicted confidences (highest predicted probability per sample)
    predicted_confidences = tf.reduce_max(predicted_probs, axis=1).numpy()
    
    # Get true class confidence (1 if correct class, 0 otherwise)
    # Ensure y_val is a numpy array
    if hasattr(y_val, 'numpy'):
        y_val = y_val.numpy()
    elif not isinstance(y_val, np.ndarray):
        y_val = np.asarray(y_val)
    
    # Check if predicted class matches true class
    predicted_probs_np = predicted_probs.numpy()
    predicted_class_indices = np.argmax(predicted_probs_np, axis=-1)
    true_class_indices = np.argmax(y_val, axis=-1)
    true_class_confidences = (predicted_class_indices == true_class_indices).astype(float)
    
    # Bin the predicted confidences
    bins = np.linspace(0, 1, num_bins + 1)
    
    # Store accuracy and average confidence for each bin
    accuracies = []
    avg_confidences = []
    
    for i in range(num_bins):
        bin_mask = (predicted_confidences >= bins[i]) & (predicted_confidences < bins[i + 1])
        bin_accuracies = np.mean(true_class_confidences[bin_mask])
        avg_confidence = np.mean(predicted_confidences[bin_mask])
        if not np.isnan(bin_accuracies) and not np.isnan(avg_confidence):
            accuracies.append(bin_accuracies)
            avg_confidences.append(avg_confidence)
    
    return np.array(avg_confidences), np.array(accuracies)


def loss_function_t_scaling(temperature, model, X_val, y_val):
    """
    Loss function to optimize temperature (minimize deviation from diagonal).
    """
    # Apply bounds manually to the temperature
    temperature = np.clip(temperature, 0.01, 4.0)
    # Calculate the calibration curve
    avg_confidences, accuracies = calculate_calibration_curve(model, X_val, y_val, temperature)
    # Calculate MSE between calibration curve and diagonal line (ideal calibration)
    mse = np.mean((accuracies - avg_confidences) ** 2)
    return mse


def plot_calibration_curve(res, res_org, optimal_temperature, output):
    """
    Plot the calibration curve for the optimal temperature.
    """
    avg_confidences, accuracies = res
    avg_confidences_org, accuracies_org = res_org
    optimal_temperature = round(optimal_temperature, 1)
    plt.close()
    plt.plot(avg_confidences_org, accuracies_org, marker='o', label='Original')
    plt.plot(avg_confidences, accuracies, marker='o', label=f'Temperature = {optimal_temperature}')
    plt.plot([0, 1], [0, 1], linestyle='--', label='Perfect Calibration')
    plt.xlabel('Average Predicted Confidence')
    plt.ylabel('Observed Accuracy')
    plt.legend()
    plt.title('Calibration Curve with Optimized Temperature')
    plt.savefig(output)
    plt.close()


def predict_with_calibration(model, X, ytest_cl, outputplot='calibration_curve.png', num_runs=1, 
                              optimal_temperatures=False, sub=False, clweights=False):
    """
    Predict phase classification with calibrated probabilities using temperature scaling.
    
    Args:
        model: Trained Keras model
        X: Input data
        ytest_cl: True classification labels (one-hot encoded). For sub-models, should be a list [cl, p, s]
        outputplot: Path to save calibration curve plot
        num_runs: If > 1, also use MC dropout for regression
        optimal_temperatures: If False, optimize temperature. If a value/array, use pre-computed 
                            temperatures. If a string, load temperatures from .npy file.
                            For sub-models, should be [temp_cl, temp_p, temp_s] or array of shape (3,)
        sub: If True, model has sub-classifications (P and S outputs)
        clweights: For sub-models, sample weights [sw, cw, p_weights, s_weights]
        
    Returns:
        p_cl: Calibrated classification probabilities
        optimal_temperature: Optimal temperature(s) found or used
        p_reg: Regression predictions (if num_runs > 1, uses MC dropout)
        p_reg_std: Regression uncertainty (if num_runs > 1)
        pp, ss: (if sub=True) Calibrated P and S sub-classification probabilities
    """
    # Load temperatures from file if string path provided
    if isinstance(optimal_temperatures, str):
        optimal_temperatures = np.load(optimal_temperatures)
        # If array has multiple values (from multiple folds), use mean
        if optimal_temperatures.ndim > 0:
            if sub and optimal_temperatures.ndim == 2 and optimal_temperatures.shape[1] == 3:
                # Array of shape (folds, 3) - take mean across folds
                optimal_temperatures = np.mean(optimal_temperatures, axis=0)
            elif optimal_temperatures.ndim > 0:
                optimal_temperatures = np.mean(optimal_temperatures)
    
    # Create a copy of the model without softmax activation to get logits
    last_layer = model.get_layer('cl_output')
    
    # Get the input to cl_output - use the layer's input property directly
    # Handle both single tensor and list/tuple cases
    if isinstance(last_layer.input, (list, tuple)) and len(last_layer.input) > 0:
        cl_input = last_layer.input[0]
    else:
        cl_input = last_layer.input
    
    # Create a new Dense layer (same weights as the original layer but no softmax activation)
    # When using pre-computed temperatures, need to use a different approach for building the model
    if optimal_temperatures is False:
        logits_layer = tf.keras.layers.Dense(last_layer.units, activation=None,
                        weights=last_layer.get_weights(), name='cl_output_log')(cl_input)
    else:
        # When using previously trained model with saved temperatures
        new_dense = tf.keras.layers.Dense(units=last_layer.units, activation=None, name='cl_output_log')
        logits_layer = new_dense(cl_input)
        new_dense.set_weights(last_layer.get_weights())
    
    # Build the new model
    log_cl_model = tf.keras.Model(inputs=model.input, outputs=logits_layer)
    
    # Optimize temperature or use pre-computed
    if optimal_temperatures is False:
        # Extract main classification labels (first element if sub-model, otherwise ytest_cl itself)
        if sub:
            ytest_cl_main = ytest_cl[0]
        else:
            ytest_cl_main = ytest_cl
        
        # Optimize temperature
        options = {
            'initial_simplex': np.array([[0.1], [4.0]]),
            'maxiter': 500,
            'xatol': 1e-6,
            'fatol': 1e-6
        }
        result = minimize(loss_function_t_scaling, x0=np.array([1.0]), 
                         args=(log_cl_model, X, ytest_cl_main),
                         bounds=[(0.01, 10.0)], method='Nelder-Mead', options=options)
        
        # The optimal temperature found
        optimal_temperature_cl = result.x[0]
        print(f"Optimal Temperature: {optimal_temperature_cl}")
        
        # Calculate and plot calibration curve
        res_org = calculate_calibration_curve(log_cl_model, X, ytest_cl_main, 1.0)
        res = calculate_calibration_curve(log_cl_model, X, ytest_cl_main, optimal_temperature_cl)
        plot_calibration_curve(res, res_org, optimal_temperature_cl, outputplot)
    else:
        # Use pre-computed temperature
        if isinstance(optimal_temperatures, (list, np.ndarray)) and len(optimal_temperatures) >= 1:
            optimal_temperature_cl = float(optimal_temperatures[0])
        else:
            optimal_temperature_cl = float(optimal_temperatures)
        print(f"Using pre-computed Temperature: {optimal_temperature_cl}")
    
    # Get calibrated predictions
    logits = log_cl_model(X, training=False)
    scaled_logits = temperature_scale(logits, optimal_temperature_cl)
    p_cl = tf.nn.softmax(scaled_logits)
    
    # Handle sub-model P and S calibration
    pp = None
    ss = None
    if sub:
        # Get P and S output layers
        p_layer = model.get_layer('p_output')
        s_layer = model.get_layer('s_output')
        
        # Get inputs to P and S outputs - use the layer's input property directly
        # Handle both single tensor and list/tuple cases
        if isinstance(p_layer.input, (list, tuple)) and len(p_layer.input) > 0:
            p_input = p_layer.input[0]
        else:
            p_input = p_layer.input
        
        if isinstance(s_layer.input, (list, tuple)) and len(s_layer.input) > 0:
            s_input = s_layer.input[0]
        else:
            s_input = s_layer.input
        
        # Create logit models for P and S
        if optimal_temperatures is False:
            logits_layer_p = tf.keras.layers.Dense(p_layer.units, activation=None,
                            weights=p_layer.get_weights(), name='p_output_log')(p_input)
            logits_layer_s = tf.keras.layers.Dense(s_layer.units, activation=None,
                            weights=s_layer.get_weights(), name='s_output_log')(s_input)
        else:
            new_dense_p = tf.keras.layers.Dense(units=p_layer.units, activation=None, name='p_output_log')
            logits_layer_p = new_dense_p(p_input)
            new_dense_p.set_weights(p_layer.get_weights())
            new_dense_s = tf.keras.layers.Dense(units=s_layer.units, activation=None, name='s_output_log')
            logits_layer_s = new_dense_s(s_input)
            new_dense_s.set_weights(s_layer.get_weights())
        
        log_cl_model_p = tf.keras.Model(inputs=model.input, outputs=logits_layer_p)
        log_cl_model_s = tf.keras.Model(inputs=model.input, outputs=logits_layer_s)
        
        # Get P and S labels and weights
        if optimal_temperatures is False:
            ytest_cl_p = ytest_cl[1]
            ytest_cl_s = ytest_cl[2]
            idx_p_true = np.where(clweights[2] != 0)[0]
            idx_s_true = np.where(clweights[3] != 0)[0]
            
            # Optimize temperatures for P and S
            options = {
                'initial_simplex': np.array([[0.1], [4.0]]),
                'maxiter': 500,
                'xatol': 1e-6,
                'fatol': 1e-6
            }
            result_p = minimize(loss_function_t_scaling, x0=np.array([1.0]), 
                             args=(log_cl_model_p, X[idx_p_true], ytest_cl_p[idx_p_true]),
                             bounds=[(0.01, 10.0)], method='Nelder-Mead', options=options)
            optimal_temperature_p = result_p.x[0]
            print(f"Optimal Temperature for P: {optimal_temperature_p}")
            
            result_s = minimize(loss_function_t_scaling, x0=np.array([1.0]), 
                             args=(log_cl_model_s, X[idx_s_true], ytest_cl_s[idx_s_true]),
                             bounds=[(0.01, 10.0)], method='Nelder-Mead', options=options)
            optimal_temperature_s = result_s.x[0]
            print(f"Optimal Temperature for S: {optimal_temperature_s}")
            
            # Plot calibration curves for P and S
            res_org_p = calculate_calibration_curve(log_cl_model_p, X[idx_p_true], ytest_cl_p[idx_p_true], 1.0)
            res_p = calculate_calibration_curve(log_cl_model_p, X[idx_p_true], ytest_cl_p[idx_p_true], optimal_temperature_p)
            plot_calibration_curve(res_p, res_org_p, optimal_temperature_p, f'{outputplot.split(".")[0]}_p.png')
            
            res_org_s = calculate_calibration_curve(log_cl_model_s, X[idx_s_true], ytest_cl_s[idx_s_true], 1.0)
            res_s = calculate_calibration_curve(log_cl_model_s, X[idx_s_true], ytest_cl_s[idx_s_true], optimal_temperature_s)
            plot_calibration_curve(res_s, res_org_s, optimal_temperature_s, f'{outputplot.split(".")[0]}_s.png')
            
            optimal_temperature = [optimal_temperature_cl, optimal_temperature_p, optimal_temperature_s]
        else:
            # Use pre-computed temperatures
            if isinstance(optimal_temperatures, (list, np.ndarray)) and len(optimal_temperatures) >= 3:
                optimal_temperature_p = float(optimal_temperatures[1])
                optimal_temperature_s = float(optimal_temperatures[2])
            else:
                # Fallback: use same temperature for all
                temp_val = float(optimal_temperatures) if not isinstance(optimal_temperatures, (list, np.ndarray)) else float(optimal_temperatures[0])
                optimal_temperature_p = temp_val
                optimal_temperature_s = temp_val
            print(f"Using pre-computed Temperatures - P: {optimal_temperature_p}, S: {optimal_temperature_s}")
            optimal_temperature = [optimal_temperature_cl, optimal_temperature_p, optimal_temperature_s]
        
        # Get calibrated P and S predictions
        logits_p = log_cl_model_p(X, training=False)
        scaled_logits_p = temperature_scale(logits_p, optimal_temperature_p)
        pp = tf.nn.softmax(scaled_logits_p)
        
        logits_s = log_cl_model_s(X, training=False)
        scaled_logits_s = temperature_scale(logits_s, optimal_temperature_s)
        ss = tf.nn.softmax(scaled_logits_s)
    
    # If num_runs > 1, also use MC dropout for regression
    p_reg = None
    p_reg_std = None
    if num_runs > 1:
        p_reg = []
        for _ in range(num_runs):
            outputs = model(X, training=True)
            tm_reg = outputs[0]  # Regression is always first output
            p_reg.append(tm_reg)
        p_reg = np.stack(p_reg)
        reg_mean = np.mean(p_reg, axis=0)
        
        # Calculate back-azimuth uncertainty
        baz_samples = np.stack([np.degrees(np.arctan2(sample.transpose()[0], 
                                                       sample.transpose()[1])) 
                                for sample in p_reg])
        baz_mean = np.degrees(np.arctan2(reg_mean.transpose()[0], 
                                          reg_mean.transpose()[1]))
        
        # Calculate circular standard deviation using mean resultant length
        R = np.sqrt(np.cos(np.radians(baz_samples)).mean(axis=0)**2 + 
                    np.sin(np.radians(baz_samples)).mean(axis=0)**2)
        p_reg_std = np.degrees(np.sqrt(-2 * np.log(R)))
        
        p_reg = reg_mean
    else:
        outputs = model.predict(X)
        p_reg = outputs[0]  # Regression is always first output
    
    if sub:
        return p_cl, optimal_temperature, p_reg, p_reg_std, pp, ss
    else:
        return p_cl, optimal_temperature_cl, p_reg, p_reg_std


from sklearn.model_selection import train_test_split

def run_experiment(NAME, X, y_reg, y_cl, weight_class=False, weight_angle=False, 
                   compute_uncertainties=False, ensemble_runs=20):
    """
    Run experiment with optional uncertainty estimation.
    
    Args:
        NAME: Experiment name
        X: Input features
        y_reg: Regression targets (back-azimuth)
        y_cl: Classification targets (phase labels)
        weight_class: Whether to use class weights
        weight_angle: Whether to use angle-based sample weights
        compute_uncertainties: If True, use calibration for phase classification and 
                              MC dropout for back-azimuth regression
        ensemble_runs: Number of MC dropout runs (if compute_uncertainties is True)
    """
    X, y, sw = load_data(X,y_reg,y_cl,weight_class,weight_angle)

    train_idx, test_idx = train_test_split(np.arange(len(X)),test_size=0.2)
    tmp = X[train_idx], (y[0][train_idx], y[1][train_idx]), (sw[0][train_idx], sw[1][train_idx])
    Xtest, ytest, swtest = X[test_idx], (y[0][test_idx], y[1][test_idx]), (sw[0][test_idx], sw[1][test_idx])
    X, y, sw = tmp

    params = {'input_shape':X.shape[1:],
              'num_outputs':y[0].shape[-1],
              'num_classes':y[1].shape[-1]}

    cmf = lambda : create_model(params)

    oof_reg = np.zeros_like(y[0])
    oof_cl = np.zeros_like(y[1])
    test_reg = []
    test_cl = []
    test_reg_std = None
    test_cl_confidence = None
    temperatures = []

    for i, (tr_idx, te_idx) in enumerate(KFold(n_splits=FOLDS).split(X,y[0])):
        y_train_reg, y_train_cl = y[0][tr_idx], y[1][tr_idx]
        y_test_reg, y_test_cl = y[0][te_idx], y[1][te_idx]

        X_train = X[tr_idx]
        X_test = X[te_idx]

        model = cmf()

        model.fit(X_train, (y_train_reg, y_train_cl),
                  validation_data=(X_test, (y_test_reg, y_test_cl), (sw[0][te_idx], sw[1][te_idx])),
                  callbacks=[tf.keras.callbacks.EarlyStopping('val_loss', patience=15),
                             tf.keras.callbacks.ReduceLROnPlateau('val_loss', patience=7, factor=0.5, verbose=1)],
                  sample_weight=(sw[0][tr_idx], sw[1][tr_idx]),
                  epochs=200,
                  verbose=1,
                  batch_size=BATCH_SIZE)

        p_reg, p_cl = model.predict(X_test)
        oof_reg[te_idx] += p_reg
        oof_cl[te_idx] += p_cl

        # Test predictions with uncertainty estimation
        if compute_uncertainties:
            # Use calibration for classification and MC dropout for regression
            # predict_with_calibration handles both when num_runs > 1
            result = predict_with_calibration(
                model, Xtest, ytest[1], 
                outputplot=f'output/{NAME}_calibration_fold_{i}.png',
                num_runs=ensemble_runs)
            p_cl_cal, temp, p_reg_pred, p_reg_std_fold = result
            temperatures.append(temp)
            test_cl.append(p_cl_cal.numpy())
            test_reg.append(p_reg_pred)
            if p_reg_std_fold is not None:
                if test_reg_std is None:
                    test_reg_std = np.zeros_like(p_reg_std_fold)
                test_reg_std += p_reg_std_fold / FOLDS
        else:
            # Standard prediction without uncertainty
            p_reg, p_cl = model.predict(Xtest)
            test_reg.append(p_reg)
            test_cl.append(p_cl)

        model.save_weights(f'output/models/{NAME}_model_fold_{i}.h5', save_format='h5')
        tf.keras.backend.clear_session()

    # Average predictions across folds
    test_reg = np.mean(np.stack(test_reg), axis=0)
    test_cl = np.mean(np.stack(test_cl), axis=0)
    
    # Ensure test_reg matches ytest[0] dtype and is contiguous for complex view
    test_reg = np.ascontiguousarray(test_reg, dtype=ytest[0].dtype)
    
    # Calculate uncertainties
    if compute_uncertainties:
        # Classification uncertainty is 1 - max probability (confidence)
        test_cl_confidence = 1.0 - np.max(test_cl, axis=-1)
        # test_reg_std already computed and averaged across folds

    np.save(f'output/{NAME}_oof_reg_predictions.npy', oof_reg)
    np.save(f'output/{NAME}_oof_class_predictions.npy', np.argmax(oof_cl, axis=-1))

    np.save(f'output/{NAME}_test_reg_predictions.npy', test_reg)
    np.save(f'output/{NAME}_test_class_predictions.npy', np.argmax(test_cl, axis=-1))
    
    # Save uncertainties if computed
    if test_reg_std is not None:
        np.save(f'output/{NAME}_test_reg_predictions_error.npy', test_reg_std)
    if test_cl_confidence is not None:
        np.save(f'output/{NAME}_test_class_predictions_error.npy', test_cl_confidence)
    if temperatures:
        np.save(f'output/{NAME}_calib_temperatures.npy', np.array(temperatures))

    true = np.angle(y[0].view(np.complex))
    pred = np.angle(oof_reg.view(np.complex))

    true_test = np.angle(ytest[0].view(np.complex))
    pred_test = np.angle(test_reg.view(np.complex))

    return {
        'OOF R2': circular_r2_score(true[np.where(sw[0] != 0)[0]], pred[np.where(sw[0] != 0)[0]]),
        'OOF (Balanced) Accuracy': accuracy_score(y[1].argmax(axis=-1), np.argmax(oof_cl, axis=-1)),
        'Test R2': circular_r2_score(true_test[np.where(swtest[0] != 0)[0]], pred_test[np.where(swtest[0] != 0)[0]]),
        'Test (Balanced) Accuracy': accuracy_score(ytest[1].argmax(axis=-1), np.argmax(test_cl, axis=-1))
    }


def run_experiment_sub(NAME, X, y_reg, y_cl, weight_class=False, weight_angle=False,
                       compute_uncertainties=False, ensemble_runs=20):
    """
    Run experiment with sub-model and optional uncertainty estimation.
    
    Args:
        NAME: Experiment name
        X: Input features
        y_reg: Regression targets (back-azimuth)
        y_cl: Classification targets (phase labels)
        weight_class: Whether to use class weights
        weight_angle: Whether to use angle-based sample weights
        compute_uncertainties: If True, use calibration for phase classification and 
                              MC dropout for back-azimuth regression
        ensemble_runs: Number of MC dropout runs (if compute_uncertainties is True)
    """
    X, y, sw = load_data_sub(X,y_reg,y_cl,weight_class,weight_angle)

    train_idx, test_idx = train_test_split(np.arange(len(X)),test_size=0.2)
    tmp = X[train_idx], [a[train_idx] for a in y], [s[train_idx] for s in sw]
    Xtest, ytest, swtest = X[test_idx], [a[test_idx] for a in y], [s[test_idx] for s in sw]
    X, y, sw = tmp

    params = {'input_shape':X.shape[1:],
              'num_outputs':y[0].shape[-1],
              'num_classes':y[1].shape[-1],
              'num_p_classes':y[2].shape[-1],
              'num_s_classes':y[3].shape[-1]}

    cmf = lambda : create_model_sub(params)

    oof_reg = np.zeros_like(y[0])
    oof_cl = np.zeros_like(y[1])
    oof_p = np.zeros_like(y[2])
    oof_s = np.zeros_like(y[3])

    test_reg = []
    test_cl = []
    test_p = []
    test_s = []
    test_reg_std = None
    test_cl_confidence = None
    temperatures = []

    for i, (tr_idx, te_idx) in enumerate(KFold(n_splits=FOLDS).split(X,y[0])):
        X_train = X[tr_idx]
        X_test = X[te_idx]

        model = cmf()

        model.fit(X_train, [a[tr_idx] for a in y],
                  validation_data=(X_test, [a[te_idx] for a in y], [a[te_idx] for a in sw]),
                  callbacks=[tf.keras.callbacks.EarlyStopping('val_loss', patience=15),
                             tf.keras.callbacks.ReduceLROnPlateau('val_loss', patience=7, factor=0.5, verbose=1)],
                  sample_weight=[a[tr_idx] for a in sw],
                  epochs=200,
                  verbose=1,
                  batch_size=BATCH_SIZE)

        p_reg, p_cl, pp, ss = model.predict(X_test)
        oof_reg[te_idx] += p_reg
        oof_cl[te_idx] += p_cl
        oof_p[te_idx] += pp
        oof_s[te_idx] += ss

        # Test predictions with uncertainty estimation
        if compute_uncertainties:
            # Use calibration for main classification, P, and S, plus MC dropout for regression
            result = predict_with_calibration(
                model, Xtest, [ytest[1], ytest[2], ytest[3]], 
                outputplot=f'output/{NAME}_calibration_fold_{i}.png',
                num_runs=ensemble_runs,
                sub=True,
                clweights=swtest)
            p_cl_cal, temp, p_reg_pred, p_reg_std_fold, pp_cal, ss_cal = result
            temperatures.append(temp)
            test_cl.append(p_cl_cal.numpy())
            test_reg.append(p_reg_pred)
            if p_reg_std_fold is not None:
                if test_reg_std is None:
                    test_reg_std = np.zeros_like(p_reg_std_fold)
                test_reg_std += p_reg_std_fold / FOLDS
            test_p.append(pp_cal.numpy())
            test_s.append(ss_cal.numpy())
        else:
            # Standard prediction without uncertainty
            p_reg, p_cl, pp, ss = model.predict(Xtest)
            test_reg.append(p_reg)
            test_cl.append(p_cl)
            test_p.append(pp)
            test_s.append(ss)

        model.save_weights(f'output/models/{NAME}_model_fold_{i}.h5', save_format='h5')
        tf.keras.backend.clear_session()

    # Average predictions across folds
    test_reg = np.mean(np.stack(test_reg), axis=0)
    test_cl = np.mean(np.stack(test_cl), axis=0)
    test_p = np.mean(np.stack(test_p), axis=0)
    test_s = np.mean(np.stack(test_s), axis=0)
    
    # Ensure test_reg matches ytest[0] dtype and is contiguous for complex view
    test_reg = np.ascontiguousarray(test_reg, dtype=ytest[0].dtype)
    
    # Calculate uncertainties
    if compute_uncertainties:
        # Note: Full calibration for sub-model requires separate temperatures for P and S
        # This is a simplified version - classification uncertainty is 1 - max probability
        test_cl_confidence = 1.0 - np.max(test_cl, axis=-1)

    np.save(f'output/{NAME}_oof_reg_predictions.npy', oof_reg)
    np.save(f'output/{NAME}_oof_class_predictions.npy', np.argmax(oof_cl, axis=-1))
    np.save(f'output/{NAME}_oof_p_class_predictions.npy', np.argmax(oof_p, axis=-1))
    np.save(f'output/{NAME}_oof_s_class_predictions.npy', np.argmax(oof_s, axis=-1))

    np.save(f'output/{NAME}_test_reg_predictions.npy', test_reg)
    np.save(f'output/{NAME}_test_class_predictions.npy', np.argmax(test_cl, axis=-1))
    np.save(f'output/{NAME}_test_p_class_predictions.npy', np.argmax(test_p, axis=-1))
    np.save(f'output/{NAME}_test_s_class_predictions.npy', np.argmax(test_s, axis=-1))
    
    # Save uncertainties if computed
    if test_reg_std is not None:
        np.save(f'output/{NAME}_test_reg_predictions_error.npy', test_reg_std)
    if test_cl_confidence is not None:
        np.save(f'output/{NAME}_test_class_predictions_error.npy', test_cl_confidence)
    if temperatures:
        np.save(f'output/{NAME}_calib_temperatures.npy', np.array(temperatures))

    true = np.angle(y[0].view(np.complex))
    pred = np.angle(oof_reg.view(np.complex))

    true_test = np.angle(ytest[0].view(np.complex))
    pred_test = np.angle(test_reg.view(np.complex))

    return {
        'OOF R2': circular_r2_score(true[np.where(sw[0] != 0)[0]], pred[np.where(sw[0] != 0)[0]]),
        'OOF (Balanced) Accuracy': accuracy_score(y[1].argmax(axis=-1), np.argmax(oof_cl, axis=-1)),
        'Test R2': circular_r2_score(true_test[np.where(swtest[0] != 0)[0]], pred_test[np.where(swtest[0] != 0)[0]]),
        'Test (Balanced) Accuracy': accuracy_score(ytest[1].argmax(axis=-1), np.argmax(test_cl, axis=-1)),
        'Test P (Balanced) Accuracy': accuracy_score(ytest[2].argmax(axis=-1)[np.where(swtest[2] != 0)[0]],
                                                     np.argmax(test_p[np.where(swtest[2] != 0)[0]], axis=-1)),
        'Test S (Balanced) Accuracy': accuracy_score(ytest[3].argmax(axis=-1)[np.where(swtest[3] != 0)[0]],
                                                     np.argmax(test_s[np.where(swtest[3] != 0)[0]], axis=-1))
    }