src/dst_gnn/explain.py

"""Explainability and ranking helpers for DST-GNN outputs."""

from __future__ import annotations

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch_geometric.data import Data
from torch_geometric.explain import Explainer, GNNExplainer

from .model import DSTGNN


def rank_nodes_by_embedding_drift(
    node_names: list[str],
    start_embeddings: torch.Tensor,
    end_embeddings: torch.Tensor,
) -> pd.DataFrame:
    """Rank nodes by embedding change magnitude between two stages."""

    drift = torch.norm(end_embeddings - start_embeddings, dim=1).detach().cpu().numpy()
    frame = pd.DataFrame({"node": node_names, "drift_score": drift})
    return frame.sort_values("drift_score", ascending=False).reset_index(drop=True)


def rank_edges_by_stage_change(
    node_names: list[str],
    start_matrix: np.ndarray,
    end_matrix: np.ndarray,
) -> pd.DataFrame:
    """Rank directed edges by absolute SSS change between two stages."""

    records: list[dict[str, float | str]] = []
    num_nodes = len(node_names)
    for source in range(num_nodes):
        for target in range(num_nodes):
            if source == target:
                continue
            delta = float(end_matrix[source, target] - start_matrix[source, target])
            records.append(
                {
                    "source": node_names[source],
                    "target": node_names[target],
                    "delta_sss": delta,
                    "abs_delta_sss": abs(delta),
                }
            )

    frame = pd.DataFrame.from_records(records)
    return frame.sort_values("abs_delta_sss", ascending=False).reset_index(drop=True)


class _TransitionReadout(nn.Module):
    """Wrapper that exposes next-stage adjacency predictions for GNNExplainer."""

    def __init__(self, model: DSTGNN, prev_hidden: torch.Tensor | None = None) -> None:
        super().__init__()
        self.model = model
        if prev_hidden is not None:
            self.register_buffer("prev_hidden", prev_hidden.detach().clone())
        else:
            self.prev_hidden = None

    def forward(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_weight: torch.Tensor,
    ) -> torch.Tensor:
        hidden = self.model.encode_step(x, edge_index, edge_weight, self.prev_hidden)
        return self.model.decode_adjacency(hidden)


def explain_node_transition(
    model: DSTGNN,
    graph: Data,
    target_node: int,
    prev_hidden: torch.Tensor | None = None,
    epochs: int = 200,
):
    """Explain a node's predicted next-stage connectivity profile."""

    wrapper = _TransitionReadout(model, prev_hidden=prev_hidden)
    explainer = Explainer(
        model=wrapper,
        algorithm=GNNExplainer(epochs=epochs),
        explanation_type="model",
        node_mask_type="attributes",
        edge_mask_type="object",
        model_config={
            "mode": "regression",
            "task_level": "node",
            "return_type": "raw",
        },
    )
    return explainer(
        graph.x,
        graph.edge_index,
        edge_weight=graph.edge_weight,
        index=target_node,
    )


def edge_mask_to_frame(
    graph: Data,
    node_names: list[str],
    edge_mask: torch.Tensor,
) -> pd.DataFrame:
    """Convert a learned edge mask to a sortable DataFrame."""

    source_nodes = graph.edge_index[0].detach().cpu().numpy()
    target_nodes = graph.edge_index[1].detach().cpu().numpy()
    weights = graph.edge_weight.detach().cpu().numpy()
    scores = edge_mask.detach().cpu().numpy()

    frame = pd.DataFrame(
        {
            "source": [node_names[index] for index in source_nodes],
            "target": [node_names[index] for index in target_nodes],
            "message_passing_weight": weights,
            "explanation_score": scores,
        }
    )
    return frame.sort_values("explanation_score", ascending=False).reset_index(drop=True)