shd_utils.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import auc
from scipy.stats import ttest_rel


def train_network(net, train_loader, epochs=50, lr=1e-2, momentum=0.9):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=lr)

    for epoch in range(epochs):
        net.train()
        running_loss = 0.0
        correct = 0
        total = 0

        for i, (inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device).float(), labels.to(device)
            labels = labels.long()

            inputs = inputs.squeeze()

            optimizer.zero_grad()
            outputs, _ = net(inputs, reset_states=True)
    

            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            if (i + 1) % 100 == 0:
                print(f"[Epoch {epoch+1}, Batch {i+1}] "
                      f"Loss: {running_loss/100:.3f}, "
                      f"Acc: {100*correct/total:.2f}%")
                running_loss = 0.0

        print(f"Epoch {epoch+1} completed | Training Accuracy: {100*correct/total:.2f}%")

    return net

def test_network(net, test_loader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net.to(device)
    net.eval()

    correct = 0
    total = 0
    all_preds, all_labels = [], []

    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device).float(), labels.to(device).long()
            outputs, _ = net(inputs, reset_states=True)

            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            all_preds.append(predicted.cpu())
            all_labels.append(labels.cpu())

    accuracy = 100 * correct / total
    print(f"Test Accuracy: {accuracy:.2f}%")

    return accuracy, torch.cat(all_preds), torch.cat(all_labels)

def get_correct_samples(dataloader, model, max_samples=1000):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    samples = []
    model.eval()
    for x, y in dataloader:
        x = x.to(device).float().squeeze()
        if y.ndim == 2: 
            y = y.argmax(dim=1)
        y = y.to(device)

        with torch.no_grad():
            z, _ = model(x, reset_states=True)
        preds = z.argmax(1)

        for i in torch.where(preds.eq(y))[0]:
            samples.append((x[i:i+1].cpu(), y[i].cpu()))
            if len(samples) >= max_samples:
                return samples

    return samples


def get_importance_for_label(cf_tracker, tsa_tracker, x, y_label,
                             spike_w, sil_w):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    cf_spike, cf_sil = cf_tracker.forward(x.to(device))


    S_spike = cf_spike[0, y_label].abs()
    S_sil   = cf_sil[0, y_label].abs()


    S_spike = S_spike / (S_spike.sum() + 1e-8)
    S_sil   = S_sil / (S_sil.sum() + 1e-8)

    S_cf = spike_w * S_spike + sil_w * S_sil


    tsa_out = tsa_tracker.forward(x.to(device))
    tsa_out = tsa_out[-1]
    
    S_tsa = tsa_out[0, y_label].abs()
    S_tsa = S_tsa / (S_tsa.sum() + 1e-8)

    S_random = torch.rand_like(S_cf)
    S_random = S_random / S_random.sum()

    return S_cf, S_tsa, S_random


def deletion_curve_logit_curve(model, x, y_label, importance, n_points=50):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    x = x.clone()
    order = np.argsort(-importance)

    # Base probability
    with torch.no_grad():
        base_logits,_ = model(x.to(device), reset_states=True)
        base_prob = base_logits.softmax(-1)[0, y_label].item()

    probs = [base_prob]

    for k in np.linspace(1, len(order), n_points, dtype=int):
        x_mod = x.clone()
        x_mod[:,:,order[:k]] = 0
        with torch.no_grad():
            logits,_ = model(x_mod.to(device), reset_states=True)
            p = logits.softmax(-1)[0, y_label].item()
        probs.append(p)

    return np.array(probs)

def compute_curve_auc(curve):
    N = len(curve)
    x = np.linspace(0, 1, N)
    return auc(x, curve)

def output_completeness_score(model, x, y_label, importance):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    zero_idx = np.where(np.abs(importance) < 1e-9)[0]

    x_mod = x.clone()
    if len(zero_idx) > 0:
        x_mod[:, :, zero_idx] = 0

    with torch.no_grad():
        logits, _ = model(x_mod.to(device), reset_states=True)
        prob = logits.softmax(-1)[0, y_label].item()
        pred = logits.argmax(-1).item()

    return prob, float(pred == y_label)


def compactness_metrics(importance, k_frac=0.1):
    imp = np.abs(importance)
    total = imp.sum() + 1e-12
    k = int(len(imp) * k_frac)

    topk_idx = np.argsort(-imp)[:k]
    topk_mass = imp[topk_idx].sum() / total

    nz_percent = np.mean(imp > 1e-9)

    return topk_mass, nz_percent
def append_dict(d, key, val):
    if key not in d:
        d[key] = []
    d[key].append(val)
    
def paired_t_tests(metric_dict):

    methods = list(metric_dict.keys())
    results = {}

    for i in range(len(methods)):
        for j in range(i+1, len(methods)):
            m1, m2 = methods[i], methods[j]
            a = np.array(metric_dict[m1])
            b = np.array(metric_dict[m2])
            t, p = ttest_rel(a, b)
            results[(m1, m2)] = {"t": t, "p": p}
            print(f"{m1} vs {m2}: t = {t:.4f}, p = {p:.4e}")

    return results

def plot_average_deletion_curves(curves_dict, title, ylabel, savefile):
    plt.figure(figsize=(7, 5))

    for method, curves in curves_dict.items():
        curves = np.array(curves)   # shape = [num_examples, num_points]

        mean_curve = curves.mean(axis=0)
        sem_curve  = curves.std(axis=0, ddof=1) / np.sqrt(len(curves))

        x = np.linspace(0, 1, len(mean_curve))

        if method == "CF":
            color = "blue"
            label = f"CF (Spike+Silence)"
        elif method == "TSA":
            color = "black"
            label = "TSA"
        else:
            color = "gray"
            label = "Random"

        plt.plot(x, mean_curve, color=color, linewidth=2, label=label)
        plt.fill_between(x,
                         mean_curve - 1.96*sem_curve,
                         mean_curve + 1.96*sem_curve,
                         color=color, alpha=0.15)

    plt.xlabel("Fraction of Features Removed", fontsize=13)
    plt.ylabel(ylabel, fontsize=13)
    plt.title(title, fontsize=14, fontweight='bold')
    plt.legend(fontsize=12)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(savefile, dpi=300, bbox_inches='tight')
    plt.show()