explain_model.py

import numpy as np
import torch
import wandb
import random
import sys
import os
from pathlib import Path
from tqdm import tqdm
import matplotlib.pyplot as plt
import pandas as pd

import hydra
from omegaconf import DictConfig, OmegaConf

from sklearn.model_selection import train_test_split
from einops import rearrange
from dataset.utils import process_classification_dataset
from dataset.infection_utils import ground_truth
from utils.utils import (
    get_K,
    change_basis,
    get_weights_from_SINDy,
    get_weights_from_DMD,
    get_weights_from_TT_DMD,
    get_weights_from_PCA,
    run_saliency,
)
from utils.metrics import (
    threshold_based_detection,
    windowing_analysis,
    F1_baseline_saliency,
    auc_analysis_edges,
    auc_analysis_nodes,
    mann_whitney_test,
    mann_whitney_test_dataset,
    autocorrelation_distance,
)
from utils.config_utils import (
    load_saved_config,
    get_saved_paths,
    get_cli_overrides,
    is_sweep,
)


# Add the project root to the Python path
project_root = Path(__file__).parent
sys.path.append(str(project_root))

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


def run_explanation(cfg: DictConfig, run_id: str = None) -> None:
    """Core explanation logic. Can be called directly or via Hydra.

    Args:
        cfg: Configuration dictionary
        run_id: Optional wandb run ID to resume. If None, creates new run.
    """
    encoder_type = cfg.model.encoder_type
    dataset_name = cfg.dataset.name

    # Get model file path
    model_file, _ = get_saved_paths(encoder_type, dataset_name, run_id)

    # Load model weights
    if not os.path.exists(model_file):
        raise FileNotFoundError(
            f"Model file '{model_file}' not found. Please train the model first."
        )

    checkpoint = torch.load(model_file, map_location="cpu", weights_only=False)
    if "model_state_dict" in checkpoint:
        model_state_dict = checkpoint["model_state_dict"]
    else:
        raise ValueError(f"Model state dict not found in {model_file}")

    # Init wandb - skip if already running (e.g., called from run_experiment.py)
    if wandb.run is None:
        if run_id:
            wandb.init(project=cfg.wandb.project, id=run_id, resume="allow")
            wandb.config.update(
                OmegaConf.to_container(cfg, resolve=True), allow_val_change=True
            )
        else:
            wandb.init(
                project=cfg.wandb.project,
                config=OmegaConf.to_container(cfg, resolve=True),
            )

    # Log used configs
    if cfg.verbose:
        print("Configuration:")
        print(OmegaConf.to_yaml(cfg))

    seed = cfg.seed
    random.seed(seed)
    torch.manual_seed(seed)
    np.random.seed(seed)

    # Create the directory to save the plots
    if cfg.explainer.plot:
        if not os.path.exists("plots"):
            os.makedirs("plots")
        if not os.path.exists(f"plots/{cfg.dataset.name}"):
            os.makedirs(f"plots/{cfg.dataset.name}")
        if not os.path.exists(f"plots/{cfg.dataset.name}/time_gt"):
            os.makedirs(f"plots/{cfg.dataset.name}/time_gt")
        if not os.path.exists(f"plots/{cfg.dataset.name}/node_gt/global_dmd"):
            os.makedirs(f"plots/{cfg.dataset.name}/node_gt/global_dmd")
        if not os.path.exists(f"plots/{cfg.dataset.name}/node_gt/local_dmd"):
            os.makedirs(f"plots/{cfg.dataset.name}/node_gt/local_dmd")
        if not os.path.exists(f"plots/{cfg.dataset.name}/edge_gt/deg2"):
            os.makedirs(f"plots/{cfg.dataset.name}/edge_gt/deg2")
        if not os.path.exists(f"plots/{cfg.dataset.name}/edge_gt/deg3"):
            os.makedirs(f"plots/{cfg.dataset.name}/edge_gt/deg3")

    # Load the dataset
    states, node_states, node_labels, edge_indexes, _, graph_labels, _ = (
        process_classification_dataset(
            cfg,
            model_state_dict,
            device,
            verbose=cfg.get("verbose", False),
        )
    )
    indices = list(range(len(node_labels)))
    train_X, val_X, _, val_y, _, val_nodes, _, val_idx = train_test_split(
        states,
        graph_labels,
        node_states,
        indices,
        test_size=0.2,
        random_state=seed,
    )

    # Load ground-truth labels
    nodes_gt, node_sums_gt, times_gt, edges_gt = ground_truth(cfg.dataset.name)
    (
        _,
        val_nodes_gt,
        _,
        val_times_gt,
        _,
        val_edge_indexes,
    ) = train_test_split(
        nodes_gt,
        torch.stack(times_gt).numpy(),
        edge_indexes,
        test_size=0.2,
        random_state=seed,
    )

    ### Explanation on validation dataset

    ## Compute the DMD modes

    # Compute Koopman operator
    dmd, K = get_K(cfg, train_X)

    # Compute the projection on DMD eigenvectors
    mode_idx = cfg.explainer.mode_idx
    modes = dmd.compute_weights(mode_idx=mode_idx, Z=states)
    val_modes = dmd.compute_weights(mode_idx=mode_idx, Z=val_X)

    # Compute saliency weights as baseline
    sal_attr = run_saliency(
        edge_indexes,
        node_labels,
        graph_labels,
        cfg,
        model_state_dict,
        device,
        model_file,
        verbose=False,
    )

    # Compute the PCA weights as baseline
    v_id = np.eye(K.shape[0])
    pca_modes = change_basis(states, v_id, dmd.emb_engine)
    val_pca_modes = change_basis(val_X, v_id, dmd.emb_engine)

    # Compute the cosine similarity of the first mode and the first PC
    v = dmd.modes  # DMD modes #FIXME: check what it returns for TT
    v_normed = v / np.linalg.norm(v, axis=0)

    # Compute how similar is the basis provided by DMD
    # to the basis provided by PCA
    cosine_similarity_matr = np.abs(np.dot(v_normed.T, v_id))
    cosine_similarity_error = np.linalg.norm(
        cosine_similarity_matr - np.eye(v.shape[1]), "fro"
    )
    cosine_similarity = cosine_similarity_matr[mode_idx, mode_idx]

    fig, ax = plt.subplots()
    co_ax = ax.imshow(cosine_similarity_matr, cmap="viridis")
    ax.set_xlabel("PCA basis")
    ax.set_ylabel("DMD basis")
    plt.colorbar(co_ax, ax=ax)
    wandb.log({"cosine_sim_matrix_img": wandb.Image(fig)})
    plt.close(fig)

    wandb.log({"cosine_sim_error": cosine_similarity_error})  # range (0, dim_red)
    wandb.log({"cosine_sim": cosine_similarity})

    # Compute temporal offset for encoders that shorten the sequence (e.g. GraphWaveNet)
    # The receptive field consumes the first (receptive_field - 1) time steps,
    # so the ground truth must be trimmed from the beginning to stay aligned.
    gt_time_offset = 0
    if encoder_type == "graphwavenet":
        _k = cfg.model.get("temporal_kernel_size", 2)
        _d = cfg.model.get("dilation", 2)
        _d_mod = 2  # dilation_mod default in GraphWaveNet
        _n = cfg.model.rnn_layers
        gt_time_offset = sum(_d ** (i % _d_mod) * (_k - 1) for i in range(_n))

    gs = []
    r_thr_prec, r_thr_rec, r_thr_f1, r_thr_base = [], [], [], []
    r_mad_prec, r_mad_rec, r_mad_f1, r_mad_base = [], [], [], []
    r_win_prec, r_win_rec, r_win_f1, r_win_base = [], [], [], []
    r_sal_prec, r_sal_rec, r_sal_f1 = [], [], []
    r_thr_pca_base, r_win_pca_base = [], []
    r_mann = []
    r_acr_distance = []
    aucs_nodes, aucs_nodes_pca_val = [], []

    for g in tqdm(range(len(val_modes)), desc="Validation dataset", leave=False):

        # Trim ground truth to match the model's output length
        times_gt_g = val_times_gt[g][gt_time_offset:]

        if val_y[g] == 0 or (times_gt_g == 0).all():
            continue

        gs.append(g)

        fig, thr_prec, thr_rec, thr_f1, thr_base = threshold_based_detection(
            val_modes[g],
            times_gt_g,
            threshold=cfg.explainer.threshold,
            window_size=cfg.explainer.window_size,
            plot=cfg.explainer.plot,
        )

        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/time_gt/{g}_thr_{mode_idx}.pdf",
                bbox_inches="tight",
            )

        fig, mad_prec, mad_rec, mad_f1, mad_base = threshold_based_detection(
            val_modes[g],
            times_gt_g,
            threshold="mad",
            window_size=cfg.explainer.window_size,
            plot=cfg.explainer.plot,
        )

        # Baseline with PCA modes
        _, _, _, thr_pca_base, _ = threshold_based_detection(
            val_pca_modes[g, :, mode_idx],
            times_gt_g,  # first PC
            threshold=cfg.explainer.threshold,
            window_size=cfg.explainer.window_size,
            plot=False,
        )

        fig, win_prec, win_rec, win_f1, win_base = windowing_analysis(
            val_modes[g],
            times_gt_g,
            window_size=cfg.explainer.window_size,
            threshold=cfg.explainer.threshold,
            plot=cfg.explainer.plot,
        )
        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/time_gt/{g}_win_{mode_idx}.pdf",
                bbox_inches="tight",
            )

        # Baseline with PCA modes
        _, _, _, win_pca_base, _ = windowing_analysis(
            val_pca_modes[g, :, mode_idx],
            times_gt_g,  # first PC
            window_size=cfg.explainer.window_size,
            threshold=cfg.explainer.threshold,
            plot=False,
        )

        # Baseline with saliency map
        fig, sal_prec, sal_rec, sal_f1 = F1_baseline_saliency(
            sal_attr[val_idx[g]].cpu().numpy().sum(axis=-1),
            val_times_gt[g],
            window_size=cfg.explainer.window_size,
            plot=cfg.explainer.plot,
        )
        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/time_gt/{g}_sal_{mode_idx}.pdf",
                bbox_inches="tight",
            )

        fig, mw_p_value = mann_whitney_test(
            val_modes[g],
            times_gt_g,
            window_size=cfg.explainer.window_size,
            plot=cfg.explainer.plot,
        )
        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/time_gt/{g}_mw_{mode_idx}.pdf",
                bbox_inches="tight",
            )

        # Compare DMD modes and PCA
        _, ACR_distance = autocorrelation_distance(
            val_modes[g], val_pca_modes[g, :, mode_idx]
        )

        # Spatial explanation on nodes via DMD on modes from training dataset
        node_modes = change_basis(
            rearrange(val_nodes[g], "t n f -> n t f"), v, dmd.emb_engine
        )
        weights = node_modes[:, -1, mode_idx] - node_modes[:, -1, mode_idx].mean()
        fig, auc = auc_analysis_nodes(
            np.abs(weights),
            val_nodes_gt[g][-1],
            val_edge_indexes[g],
            plot=cfg.explainer.plot,
        )
        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/node_gt/global_dmd/{g}_mask_{mode_idx}.pdf",
                bbox_inches="tight",
            )

        # Spatial explanation on nodes via PCA only
        node_modes = change_basis(
            rearrange(val_nodes[g], "t n f -> n t f"), v_id, dmd.emb_engine
        )
        weights = node_modes[:, -1, 0] - node_modes[:, -1, 0].mean()  # first PC
        _, auc_pca = auc_analysis_nodes(
            np.abs(weights), val_nodes_gt[g][-1], val_edge_indexes[g], plot=False
        )

        plt.close("all")

        r_thr_prec.append(thr_prec)
        r_thr_rec.append(thr_rec)
        r_thr_f1.append(thr_f1)
        r_thr_base.append(thr_base)
        r_mad_prec.append(mad_prec)
        r_mad_rec.append(mad_rec)
        r_mad_f1.append(mad_f1)
        r_mad_base.append(mad_base)
        r_thr_pca_base.append(thr_pca_base)
        r_win_prec.append(win_prec)
        r_win_rec.append(win_rec)
        r_win_f1.append(win_f1)
        r_win_base.append(win_base)
        r_win_pca_base.append(win_pca_base)
        r_sal_prec.append(sal_prec)
        r_sal_rec.append(sal_rec)
        r_sal_f1.append(sal_f1)
        r_mann.append(mw_p_value)
        r_acr_distance.append(ACR_distance)
        aucs_nodes.append(auc)
        aucs_nodes_pca_val.append(auc_pca)

    # Mann-Whitney U test on whole dataset
    fig, mw_p_value_dt = mann_whitney_test_dataset(
        val_modes[val_y == 1],
        np.stack(val_times_gt)[val_y == 1, gt_time_offset:],
        window_size=cfg.explainer.window_size,
        plot=cfg.explainer.plot,
    )

    if fig is not None:
        fig.savefig(
            f"plots/{cfg.dataset.name}/time_gt/dataset_mw_{mode_idx}.pdf",
            bbox_inches="tight",
        )

    # Create a dataframe with the results
    results = pd.DataFrame(
        {
            "g": gs,
            "thr_precision": r_thr_prec,
            "thr_recall": r_thr_rec,
            "thr_f1_score": r_thr_f1,
            "thr_baseline_f1": r_thr_base,
            "mad_precision": r_mad_prec,
            "mad_recall": r_mad_rec,
            "mad_f1_score": r_mad_f1,
            "mad_baseline_f1": r_mad_base,
            "thr_pca_baseline_f1": r_thr_pca_base,
            "window_precision": r_win_prec,
            "window_recall": r_win_rec,
            "window_f1_score": r_win_f1,
            "window_baseline_f1": r_win_base,
            "window_pca_baseline_f1": r_win_pca_base,
            "saliency_precision": r_sal_prec,
            "saliency_recall": r_sal_rec,
            "saliency_f1_score": r_sal_f1,
            "acr_distance": r_acr_distance,
            "mw_p_value": r_mann,
            "mw_p_value_dt": mw_p_value_dt,
            "auc_tg": aucs_nodes,
            "auc_tg_pca_val": aucs_nodes_pca_val,
        }
    )

    # Log on wandb the averages
    for key in results.columns:
        wandb.log({f"{key}_avg": np.asarray(results[key]).mean()})

    # Save the dataframe to an Excel file in a new sheet
    writer = pd.ExcelWriter(path="results.xlsx", engine="xlsxwriter")
    results.to_excel(writer, sheet_name=f"time_gt_{cfg.dataset.name}", index=False)

    # Explanation on whole dataset
    gs = []
    aucs2, aucs3 = [], []
    aucs_nodes, aucs_nodes_tt = [], []
    aucs_nodes_sal_base, aucs_nodes_pca_base = [], []

    for g in tqdm(range(len(edges_gt)), desc="Whole dataset", leave=False):

        if graph_labels[g] == 0 or torch.sum(edges_gt[g]) == 0:
            continue

        gs.append(g)

        # Spatial explanation via SINDy on edges
        weights = get_weights_from_SINDy(
            edge_indexes[g],
            node_states[g],
            cfg.explainer.dim_red,
            add_self_dependency=cfg.explainer.add_self_dependency_sindy,
            degree=2,
            method=cfg.explainer.emb_method,
        )

        num_nodes = node_states[g].shape[1]
        fig, auc = auc_analysis_edges(
            weights, edge_indexes[g], edges_gt[g], num_nodes, plot=cfg.explainer.plot
        )
        aucs2.append(auc)

        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/edge_gt/deg2/{g}_mask.pdf",
                bbox_inches="tight",
            )

        plt.close("all")

        weights = get_weights_from_SINDy(
            edge_indexes[g],
            node_states[g],
            cfg.explainer.dim_red,
            add_self_dependency=cfg.explainer.add_self_dependency_sindy,
            degree=3,
            method=cfg.explainer.emb_method,
        )

        fig, auc = auc_analysis_edges(
            weights, edge_indexes[g], edges_gt[g], num_nodes, plot=cfg.explainer.plot
        )
        aucs3.append(auc)

        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/edge_gt/deg3/{g}_mask.pdf",
                bbox_inches="tight",
            )

        plt.close("all")

        # Spatial explanation on nodes via DMD on graph modes
        weights_t = get_weights_from_DMD(
            node_states[g],
            cfg.explainer.dim_red,
            mode_idx=mode_idx,
            method=cfg.explainer.emb_method,
        )
        weights = weights_t[:, -1]  # last time step
        weights = weights - weights.mean()
        fig, auc = auc_analysis_nodes(
            np.abs(weights), nodes_gt[g][-1], edge_indexes[g], plot=cfg.explainer.plot
        )
        aucs_nodes.append(auc)

        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/node_gt/local_dmd/{g}_mask.pdf",
                bbox_inches="tight",
            )

        # Spatial explanation on nodes via TT-DMD on graph modes
        weights_t = get_weights_from_TT_DMD(node_states[g], mode_idx=mode_idx)
        weights = weights_t[:, -1]  # last time step
        weights = weights - weights.mean()
        fig, auc = auc_analysis_nodes(
            np.abs(weights), nodes_gt[g][-1], edge_indexes[g], plot=cfg.explainer.plot
        )
        aucs_nodes_tt.append(auc)

        if fig is not None:
            fig.savefig(
                f"plots/{cfg.dataset.name}/node_gt/local_ttdmd/{g}_mask.pdf",
                bbox_inches="tight",
            )

        # Spatial explanation on nodes baseline via saliency
        weights_t = sal_attr[g].T.cpu().numpy()  # shape [nodes, times]
        weights = np.max(np.abs(weights_t), axis=1)  # max over time
        _, auc = auc_analysis_nodes(
            weights, nodes_gt[g][-1], edge_indexes[g], plot=cfg.explainer.plot
        )
        aucs_nodes_sal_base.append(auc)

        # Spatial explanation on nodes baseline via PCA only
        weights_t = get_weights_from_PCA(
            node_states[g], cfg.explainer.dim_red, method=cfg.explainer.emb_method
        )
        weights = weights_t[:, -1, 0]  # last time step, first PC
        weights = weights - weights.mean()
        _, auc = auc_analysis_nodes(
            np.abs(weights), nodes_gt[g][-1], edge_indexes[g], plot=False
        )
        aucs_nodes_pca_base.append(auc)

        plt.close("all")

    results = pd.DataFrame(
        {
            "g": gs,
            "auc_2": aucs2,
            "auc_3": aucs3,
            "auc_nodes": aucs_nodes,
            "auc_nodes_tt": aucs_nodes_tt,
            "auc_nodes_pca_base": aucs_nodes_pca_base,
            "auc_nodes_sal_base": aucs_nodes_sal_base,
        }
    )

    # Log on wandb the averages
    for key in results.columns:
        wandb.log({f"{key}_avg": np.asarray(results[key]).mean()})

    # Save the dataframe to an Excel file in a new sheet
    results.to_excel(writer, sheet_name=f"edge_gt_{cfg.dataset.name}", index=False)
    writer.close()


@hydra.main(version_base=None, config_path="configs", config_name="explain")
def explain(cfg: DictConfig) -> None:
    """Hydra entry point for explanation."""
    # Get run_id from config if provided
    run_id = (
        cfg.wandb.run_id if hasattr(cfg.wandb, "run_id") and cfg.wandb.run_id else None
    )

    # If not a sweep and no run_id, try to load saved config
    if not is_sweep() and not run_id:
        encoder_type = cfg.model.encoder_type
        dataset_name = cfg.dataset.name
        saved_config = load_saved_config(encoder_type, dataset_name)
        if saved_config:
            cli_overrides = get_cli_overrides()
            saved_cfg = OmegaConf.create(saved_config)
            # Preserve explainer config
            explainer_cfg = cfg.explainer
            # Apply saved config values (skip 'training' - not relevant for explain)
            for key, value in OmegaConf.to_container(saved_cfg, resolve=True).items():
                if key not in ["training", "explainer"]:
                    OmegaConf.update(cfg, key, value)
            # Apply CLI overrides on top
            for key in cli_overrides:
                OmegaConf.update(cfg, key, OmegaConf.select(cfg, key))
            cfg.explainer = explainer_cfg

    run_explanation(cfg, run_id)


if __name__ == "__main__":
    explain()