graph-explain 0.7.0__py3-none-any.whl

graph_explain/methods/gradient/integrated_gradients.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ...core.explanation import Explanation
+from ...core.registry import register
+from ..base import ExplanationAlgorithm
+
+
+@register("integrated_gradients", "ig")
+class IntegratedGradients(ExplanationAlgorithm):
+    graph_level = True
+
+    def __init__(
+        self,
+        steps: int = 50,
+        method: str = "riemann",
+        edge_grads: bool = True,
+        **kwargs,
+    ):
+        self.steps = steps
+        self.method = method
+        self.edge_grads = edge_grads
+
+    def explain(
+        self,
+        backend: Any,
+        model: Any,
+        data: Any,
+        index: int | list[int] | torch.Tensor,
+        target_class: int | None = None,
+        **kwargs,
+    ) -> Explanation:
+        model.eval()
+        x = backend.node_features(data)
+        edge_index = backend.edge_index(data)
+
+        if index is None:
+            idx = torch.zeros(1, dtype=torch.long, device=x.device)
+        elif isinstance(index, int):
+            idx = torch.tensor([index], device=x.device)
+        else:
+            idx = torch.as_tensor(index, device=x.device)
+
+        with torch.no_grad():
+            logits = backend.forward(model, x, edge_index)
+        if target_class is None and logits.dim() == 2:
+            target_class = int(logits[idx[0]].argmax().item())
+
+        baseline = torch.zeros_like(x)
+        alphas, weights = self._alphas(self.steps, self.method, device=x.device)
+
+        ig_feat = torch.zeros_like(x)
+        ig_edge = None
+        compute_edge = self.edge_grads and backend.supports_edge_weight(model)
+        if compute_edge:
+            ig_edge = torch.zeros(
+                edge_index.size(1), dtype=torch.float32, device=x.device
+            )
+
+        for alpha, w in zip(alphas, weights):
+            x_step = (baseline + alpha * (x - baseline)).requires_grad_(True)
+            ew = None
+            if compute_edge:
+                ew = (
+                    torch.ones(edge_index.size(1), device=x.device) * (1.0 - alpha)
+                ).requires_grad_(True)
+            out = backend.forward(model, x_step, edge_index, edge_weight=ew)
+            if out.dim() == 2:
+                score = out[idx, target_class].sum()
+            else:
+                score = out[idx].sum()
+            model.zero_grad()
+            score.backward()
+            ig_feat = ig_feat + w * x_step.grad
+            if compute_edge and ew is not None and ew.grad is not None:
+                ig_edge = ig_edge + w * ew.grad
+
+        ig_feat = ig_feat * (x - baseline)
+        if compute_edge and ig_edge is not None:
+            ig_edge = ig_edge * (-1.0)
+
+        grad = ig_feat.detach()
+        node_importance = grad.abs().sum(dim=-1)
+        feature_importance = grad
+
+        return Explanation(
+            node_importance=node_importance.cpu(),
+            edge_importance=(
+                ig_edge.detach().abs().cpu() if ig_edge is not None else None
+            ),
+            feature_importance=feature_importance.cpu(),
+            prediction_original=logits[idx[0]].detach().reshape(1, -1).cpu(),
+            prediction_explanation=None,
+            node_idx=int(idx[0].item()) if isinstance(index, int) else index,
+            target_class=target_class,
+        )
+
+    def _alphas(self, steps: int, method: str, device):
+        if method in ("riemann", "left"):
+            alphas = torch.arange(0.0, 1.0, 1.0 / steps, device=device)
+            return alphas, torch.full_like(alphas, 1.0 / steps)
+        if method == "right":
+            alphas = torch.arange(1.0 / steps, 1.0 + 1e-6, 1.0 / steps, device=device)
+            return alphas, torch.full_like(alphas, 1.0 / steps)
+        if method == "gausslegendre":
+            from numpy.polynomial.legendre import leggauss
+
+            xs, ws = leggauss(steps)
+            alphas = torch.as_tensor((xs + 1) / 2, device=device, dtype=torch.float32)
+            weights = torch.as_tensor(ws / 2, device=device, dtype=torch.float32)
+            return alphas, weights
+        raise ValueError(f"method desconocido: {method}")

graph_explain/methods/gradient/saliency.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from ...core.explanation import Explanation
+from ...core.registry import register
+from ..base import ExplanationAlgorithm
+
+
+@register("saliency", "gradient", "grad")
+class Saliency(ExplanationAlgorithm):
+    graph_level = True
+
+    def __init__(
+        self,
+        absolute: bool = True,
+        aggregate: str = "sum",
+        node_mask_type: str | None = None,
+    ):
+        self.absolute = absolute
+        self.aggregate = aggregate
+        self.node_mask_type = node_mask_type
+
+    def explain(
+        self,
+        backend: Any,
+        model: Any,
+        data: Any,
+        index: int | list[int] | torch.Tensor,
+        target_class: int | None = None,
+        **kwargs,
+    ) -> Explanation:
+        model.eval()
+        x = backend.node_features(data).detach().clone().requires_grad_(True)
+        edge_index = backend.edge_index(data)
+        edge_weight = backend.edge_weight(data)
+
+        out = backend.forward(model, x, edge_index, edge_weight=edge_weight)
+        logits = out
+
+        if index is None:
+            idx = torch.zeros(1, dtype=torch.long, device=x.device)
+        elif isinstance(index, int):
+            idx = torch.tensor([index], device=x.device)
+        else:
+            idx = torch.as_tensor(index, device=x.device)
+
+        if target_class is None and logits.dim() == 2:
+            target_class = int(logits[idx].argmax(dim=-1)[0].item())
+
+        if logits.dim() == 2:
+            score = logits[idx, target_class].sum()
+        else:
+            score = logits[idx].sum()
+
+        model.zero_grad()
+        score.backward()
+
+        grad = x.grad
+        if grad is None:
+            raise RuntimeError("No se obtuvieron gradientes del modelo.")
+
+        if self.absolute:
+            grad = grad.abs()
+        if self.aggregate == "sum":
+            node_importance = grad.sum(dim=-1)
+        elif self.aggregate == "mean":
+            node_importance = grad.mean(dim=-1)
+        elif self.aggregate == "max":
+            node_importance = grad.max(dim=-1).values
+        else:
+            raise ValueError(f"aggregate desconocido: {self.aggregate}")
+
+        prediction_original = logits[idx].detach()
+        feature_importance = (
+            grad.detach() if self.node_mask_type == "attributes" else None
+        )
+
+        return Explanation(
+            node_importance=node_importance.detach().cpu(),
+            feature_importance=(
+                feature_importance.detach().cpu()
+                if feature_importance is not None
+                else None
+            ),
+            edge_importance=None,
+            prediction_original=prediction_original.cpu(),
+            prediction_explanation=None,
+            node_idx=int(idx[0].item()) if isinstance(index, int) else index,
+            target_class=target_class,
+        )

graph_explain/methods/perturbation/gnn_explainer.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from ...core.explanation import Explanation
+from ...core.registry import register
+from ..base import ExplanationAlgorithm
+
+
+@register("gnn_explainer", "gnnexplainer")
+class GNNExplainer(ExplanationAlgorithm):
+    graph_level = True
+
+    def __init__(
+        self,
+        epochs: int = 200,
+        lr: float = 0.01,
+        edge_entropy: float = 0.001,
+        node_entropy: float = 0.001,
+        node_mask_type: str | None = "attributes",
+        edge_mask_type: str | None = "object",
+        prints: int = 20,
+        **kwargs,
+    ):
+        self.epochs = epochs
+        self.lr = lr
+        self.edge_entropy = edge_entropy
+        self.node_entropy = node_entropy
+        self.node_mask_type = node_mask_type
+        self.edge_mask_type = edge_mask_type
+        self.prints = prints
+
+    def explain(
+        self,
+        backend: Any,
+        model: Any,
+        data: Any,
+        index: int | torch.Tensor,
+        target_class: int | None = None,
+        **kwargs,
+    ) -> Explanation:
+        node_mask_type = kwargs.get("node_mask_type", self.node_mask_type)
+        edge_mask_type = kwargs.get("edge_mask_type", self.edge_mask_type)
+
+        model.eval()
+        x = backend.node_features(data)
+        edge_index = backend.edge_index(data)
+        num_nodes = backend.num_nodes(data)
+
+        sub_nodes, sub_edge_index, mapping, sub_edge_mask = self._extract_subgraph(
+            backend, data, index, edge_index
+        )
+        device = x.device
+        x_sub = x[sub_nodes].to(device)
+        sub_graph_level = index is None
+
+        node_mask = None
+        if node_mask_type is not None and not sub_graph_level:
+            node_mask = torch.nn.Parameter(torch.randn(x_sub.size(0), device=device))
+
+        edge_mask = None
+        if edge_mask_type is not None:
+            edge_mask = torch.nn.Parameter(
+                torch.randn(sub_edge_index.size(1), device=device)
+            )
+
+        params = [p for p in (node_mask, edge_mask) if p is not None]
+        optimizer = torch.optim.Adam(params, lr=self.lr)
+
+        with torch.no_grad():
+            orig_logits = backend.forward(model, x_sub, sub_edge_index)
+        if sub_graph_level:
+            if target_class is None:
+                target_class = (
+                    int(orig_logits[0].argmax().item()) if orig_logits.dim() == 2 else 0
+                )
+            tgt_idx = 0
+        else:
+            ni = int(mapping.item() if torch.is_tensor(mapping) else mapping)
+            if target_class is None and orig_logits.dim() == 2:
+                target_class = int(orig_logits[ni].argmax().item())
+            tgt_idx = ni
+
+        for epoch in range(self.epochs):
+            optimizer.zero_grad()
+            mask_node = None
+            if node_mask is not None:
+                mask_node = torch.sigmoid(node_mask)
+            eweight = None
+            if edge_mask is not None:
+                eweight = torch.sigmoid(edge_mask)
+
+            pred = backend.forward(
+                model, x_sub, sub_edge_index, edge_weight=eweight, node_mask=mask_node
+            )
+            loss = self._loss(
+                pred, tgt_idx, sub_graph_level, target_class, node_mask, edge_mask
+            )
+            loss.backward()
+            optimizer.step()
+
+        logits = backend.forward(model, x_sub, sub_edge_index)
+
+        if sub_graph_level:
+            target_class = target_class or (
+                int(logits.argmax(-1)[0].item()) if logits.dim() == 2 else 0
+            )
+            ni = 0
+            pred_orig = logits[0].detach()
+            final_mask_node = None
+            final_mask_edge = None
+            if node_mask is not None:
+                final_mask_node = torch.sigmoid(node_mask)
+            if edge_mask is not None:
+                final_mask_edge = torch.sigmoid(edge_mask)
+            pred_masked = backend.forward(
+                model,
+                x_sub,
+                sub_edge_index,
+                edge_weight=final_mask_edge,
+                node_mask=final_mask_node,
+            )[0].detach()
+        else:
+            ni = int(mapping.item() if torch.is_tensor(mapping) else mapping)
+            if target_class is None and logits.dim() == 2:
+                target_class = int(logits[ni].argmax().item())
+            pred_orig = logits[ni].detach()
+            final_mask_node = None
+            final_mask_edge = None
+            if node_mask is not None:
+                final_mask_node = torch.sigmoid(node_mask)
+            if edge_mask is not None:
+                final_mask_edge = torch.sigmoid(edge_mask)
+            pred_masked = backend.forward(
+                model,
+                x_sub,
+                sub_edge_index,
+                edge_weight=final_mask_edge,
+                node_mask=final_mask_node,
+            )[ni].detach()
+
+        full_num_nodes = num_nodes
+        full_edge_count = edge_index.size(1)
+
+        node_full = self._scatter_node(node_mask, sub_nodes, full_num_nodes)
+        edge_full = self._scatter_edge(edge_mask, sub_edge_mask, full_edge_count)
+
+        return Explanation(
+            node_importance=node_full,
+            edge_importance=edge_full,
+            feature_importance=None,
+            prediction_original=pred_orig.cpu(),
+            prediction_explanation=pred_masked.cpu(),
+            node_idx=None
+            if sub_graph_level
+            else (int(index[0]) if torch.is_tensor(index) else int(index)),
+            target_class=target_class,
+            metadata={
+                "sub_nodes": sub_nodes,
+                "sub_edge_index": sub_edge_index,
+                "sub_edge_mask": sub_edge_mask,
+            },
+        )
+
+    @staticmethod
+    def _scatter_node(
+        mask: torch.nn.Parameter | None,
+        sub_nodes: torch.Tensor,
+        num_nodes: int,
+    ) -> torch.Tensor | None:
+        if mask is None:
+            return None
+        full = torch.zeros(num_nodes, dtype=torch.float32)
+        vals = torch.sigmoid(mask).detach().cpu()
+        full[sub_nodes.cpu()] = vals
+        return full
+
+    @staticmethod
+    def _scatter_edge(
+        mask: torch.nn.Parameter | None,
+        sub_edge_mask: torch.Tensor,
+        num_edges: int,
+    ) -> torch.Tensor | None:
+        if mask is None:
+            return None
+        full = torch.zeros(num_edges, dtype=torch.float32)
+        idx = sub_edge_mask.nonzero(as_tuple=False).view(-1)
+        full[idx.cpu()] = torch.sigmoid(mask).detach().cpu()
+        return full
+
+    def _loss(
+        self,
+        pred: torch.Tensor,
+        node_idx: int | torch.Tensor,
+        sub_graph_level: bool,
+        target_class: int | None,
+        node_mask: torch.nn.Parameter | None,
+        edge_mask: torch.nn.Parameter | None,
+    ) -> torch.Tensor:
+        if pred.dim() == 2:
+            log_logits = pred.log_softmax(dim=-1)
+        else:
+            log_logits = pred
+
+        if sub_graph_level:
+            idx = torch.zeros(1, dtype=torch.long, device=pred.device)
+            if target_class is None:
+                target_class = int(pred[0].argmax().item())
+            loss = F.nll_loss(
+                log_logits[0].unsqueeze(0),
+                torch.tensor([target_class], device=pred.device),
+            )
+        else:
+            if isinstance(node_idx, torch.Tensor) and node_idx.dim() == 0:
+                idx = node_idx.unsqueeze(0)
+            else:
+                idx = (
+                    torch.as_tensor([node_idx], device=pred.device)
+                    if not torch.is_tensor(node_idx)
+                    else node_idx.reshape(-1)
+                )
+            if target_class is None:
+                target_class = int(pred[idx].argmax(dim=-1)[0].item())
+            loss = F.nll_loss(
+                log_logits[idx], torch.tensor([target_class], device=pred.device)
+            )
+
+        if edge_mask is not None and self.edge_entropy > 0:
+            loss += self.edge_entropy * self._entropy(torch.sigmoid(edge_mask))
+        if node_mask is not None and self.node_entropy > 0:
+            loss += self.node_entropy * self._entropy(torch.sigmoid(node_mask))
+        return loss
+
+    @staticmethod
+    def _entropy(p: torch.Tensor) -> torch.Tensor:
+        eps = 1e-8
+        return -(p * torch.log(p + eps) + (1 - p) * torch.log(1 - p + eps)).mean()
+
+    def _extract_subgraph(
+        self,
+        backend: Any,
+        data: Any,
+        index: int | torch.Tensor | None,
+        edge_index: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        from torch_geometric.utils import k_hop_subgraph
+
+        if index is None:
+            n = edge_index.max().item() + 1
+            node_idx = torch.arange(n, device=edge_index.device)
+            return (
+                node_idx,
+                edge_index,
+                node_idx,
+                torch.ones(
+                    edge_index.size(1), dtype=torch.bool, device=edge_index.device
+                ),
+            )
+        node_idx = torch.as_tensor([index], device=edge_index.device).reshape(-1)
+        return k_hop_subgraph(
+            node_idx, num_hops=3, edge_index=edge_index, relabel_nodes=True
+        )

graph_explain/methods/perturbation/node_mask.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from ...core.explanation import Explanation
+from ...core.registry import register
+from ..base import ExplanationAlgorithm
+
+
+@register("node_mask", "nodemask", "nm")
+class NodeMask(ExplanationAlgorithm):
+    """NodeMask: node mask learned by optimization.
+
+    Optimizes a (sigmoid) mask over the nodes of the target node's k-hop
+    subgraph so the model keeps its prediction, with an entropy regularizer to
+    force sparsity. The resulting node importance is re-projected onto the full
+    graph (0 outside the neighborhood).
+    """
+
+    def __init__(
+        self,
+        epochs: int = 200,
+        lr: float = 0.05,
+        hops: int = 3,
+        suppress_ratio: float = 0.8,
+        entropy: float = 0.05,
+        **kwargs,
+    ):
+        self.epochs = epochs
+        self.lr = lr
+        self.hops = hops
+        self.suppress_ratio = suppress_ratio
+        self.entropy = entropy
+
+    def explain(
+        self,
+        backend: Any,
+        model: Any,
+        data: Any,
+        index: int | torch.Tensor | None = None,
+        target_class: int | None = None,
+        **kwargs,
+    ) -> Explanation:
+        model.eval()
+        x = backend.node_features(data)
+        edge_index = backend.edge_index(data)
+        num_nodes = backend.num_nodes(data)
+
+        nodes = self._to_node_ids(index, num_nodes)
+        root = int(nodes[0])
+
+        from torch_geometric.utils import k_hop_subgraph
+
+        device = x.device
+        sub_nodes, sub_edge_index, mapping, _ = k_hop_subgraph(
+            [root],
+            num_hops=self.hops,
+            edge_index=edge_index,
+            relabel_nodes=True,
+        )
+        x_sub = x[sub_nodes].to(device)
+
+        with torch.no_grad():
+            orig = backend.forward(model, x_sub, sub_edge_index)
+        ni = int(mapping.item() if torch.is_tensor(mapping) else mapping)
+        if target_class is None and orig.dim() == 2:
+            target_class = int(orig[ni].argmax().item())
+        target = target_class if target_class is not None else 0
+
+        if self.suppress_ratio > 0:
+            k = int(self.suppress_ratio * sub_nodes.numel())
+            k = max(0, k)
+        else:
+            k = max(0, sub_nodes.numel() - 1)
+
+        mask = torch.nn.Parameter(torch.zeros(sub_nodes.numel(), device=device))
+        optimizer = torch.optim.Adam([mask], lr=self.lr)
+
+        for _ in range(self.epochs):
+            optimizer.zero_grad()
+            node_mask = torch.sigmoid(mask)
+            pred = backend.forward(model, x_sub, sub_edge_index, node_mask=node_mask)
+            topk = node_mask.topk(max(k, 1)).values.min()
+            loss = self._loss(pred, ni, target, node_mask, topk)
+            loss.backward()
+            optimizer.step()
+
+        final = torch.sigmoid(mask).detach()
+        full = torch.zeros(num_nodes, dtype=torch.float32)
+        full[sub_nodes.cpu()] = final.cpu()
+
+        with torch.no_grad():
+            pred_masked = backend.forward(
+                model, x_sub, sub_edge_index, node_mask=final
+            )[ni]
+
+        return Explanation(
+            node_importance=full,
+            edge_importance=None,
+            feature_importance=None,
+            prediction_original=orig[ni].detach().cpu(),
+            prediction_explanation=pred_masked.detach().cpu(),
+            node_idx=root,
+            target_class=target,
+        )
+
+    def _loss(self, pred, node_idx, target_class, node_mask, topk) -> torch.Tensor:
+        if pred.dim() == 2:
+            log_logits = pred.log_softmax(dim=-1)
+        else:
+            log_logits = pred
+        idx = torch.as_tensor([node_idx], device=pred.device).reshape(-1)
+        loss = F.nll_loss(
+            log_logits[idx],
+            torch.tensor([target_class], device=pred.device),
+        )
+        loss += self.entropy * self._entropy(node_mask)
+        loss += (node_mask - topk.detach()).relu().mean()
+        return loss
+
+    @staticmethod
+    def _entropy(p: torch.Tensor) -> torch.Tensor:
+        eps = 1e-8
+        return -(p * torch.log(p + eps) + (1 - p) * torch.log(1 - p + eps)).mean()
+
+    @staticmethod
+    def _to_node_ids(index, num_nodes: int) -> torch.Tensor:
+        if index is None:
+            return torch.zeros(1, dtype=torch.long)
+        if isinstance(index, int):
+            return torch.tensor([index], dtype=torch.long)
+        idx = torch.as_tensor(index, dtype=torch.long)
+        return idx.reshape(-1) if idx.numel() else torch.zeros(1, dtype=torch.long)