GNN LRP: Graph Motif Detection

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import igraph
import random

def create_house_motif():
    # House: Base square (0-1-2-3) + Roof (4 connected to 2,3)
    # 5 nodes
    adj = np.zeros((5, 5))
    edges = [
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 0),  # Square base
        (2, 4),
        (3, 4),
    ]  # Roof
    for i, j in edges:
        adj[i, j] = adj[j, i] = 1
    return adj


def create_dataset_graph(nodes_nr=15, has_motif=False):
    # 1. Base: Erdos-Renyi random graph
    p = 0.2
    adj = np.zeros((nodes_nr, nodes_nr))
    for i in range(nodes_nr):
        for j in range(i + 1, nodes_nr):
            if random.random() < p:
                adj[i, j] = adj[j, i] = 1

    # 2. Insert Motif if needed
    if has_motif:
        motif = create_house_motif()
        m_size = len(motif)
        # Pick random start node (ensure fit)
        start = random.randint(0, nodes_nr - m_size)

        # Overlay motif edges
        # We overwrite existing edges to ensure the shape exists
        for i in range(m_size):
            for j in range(m_size):
                if motif[i, j] == 1:
                    adj[start + i, start + j] = adj[start + j, start + i] = 1

    # 3. Add self-loops (GNN standard)
    adj = adj + np.eye(nodes_nr)

    # 4. Compute Laplacian (Normalization)
    D = adj.sum(axis=1)
    with np.errstate(divide="ignore"):
        D_inv_sqrt = np.power(D, -0.5)
    D_inv_sqrt[np.isinf(D_inv_sqrt)] = 0
    D_mat = np.diag(D_inv_sqrt)
    laplacian = torch.FloatTensor(D_mat @ adj @ D_mat)

    return {
        "adjacency": torch.FloatTensor(adj),
        "laplacian": laplacian,
        "target": 1 if has_motif else 0,
        "walks": compute_walks(adj),  # For visualization
    }


def compute_walks(adj):
    # Find all walks of length 3 (v1-v2-v3) for LRP visualization
    # Limit to reasonable number to prevent lag
    w = []
    nodes = len(adj)
    for v1 in range(nodes):
        neighbors_v1 = np.where(adj[v1] > 0)[0]
        for v2 in neighbors_v1:
            if v1 == v2:
                continue  # Skip self-loops for walks
            neighbors_v2 = np.where(adj[v2] > 0)[0]
            for v3 in neighbors_v2:
                if v2 == v3:
                    continue
                w.append((v1, v2, v3))
    return w

class GraphNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        # Since we have no node features, we use Identity as input features (One-Hot)
        # So input_dim = num_nodes
        self.W1 = nn.Parameter(torch.randn(input_dim, hidden_dim) * 0.1)
        self.W2 = nn.Parameter(torch.randn(hidden_dim, hidden_dim) * 0.1)
        self.W3 = nn.Parameter(torch.randn(hidden_dim, output_dim) * 0.1)
        self.params = [self.W1, self.W2, self.W3]

    def forward(self, laplacian):
        # Input features X = Identity
        X = torch.eye(len(laplacian))

        # Layer 1: H1 = ReLU(L * X * W1)
        H1 = torch.relu(laplacian @ X @ self.W1)

        # Layer 2: H2 = ReLU(L * H1 * W2)
        H2 = torch.relu(laplacian @ H1 @ self.W2)

        # Layer 3: Out = Mean(ReLU(L * H2 * W3))
        # Global Pooling (Mean)
        H3 = torch.relu(laplacian @ H2 @ self.W3)
        Out = H3.mean(dim=0)

        return Out

    def lrp(self, laplacian, target_class, gamma=0.1):
        # Implementation of GNN-LRP (simplified)
        # This is a complex backward pass similar to the original script

        # Re-run forward to get activations
        X = torch.eye(len(laplacian)).requires_grad_(True)

        # Weights with Gamma rule (enhance positive weights)
        W1p = self.W1 + gamma * self.W1.clamp(min=0)
        W2p = self.W2 + gamma * self.W2.clamp(min=0)
        W3p = self.W3 + gamma * self.W3.clamp(min=0)

        # Forward with Gamma weights to compute Relevance proportions
        H1 = (laplacian @ X @ self.W1).clamp(min=0)
        H2 = (laplacian @ H1 @ self.W2).clamp(min=0)
        H3 = (laplacian @ H2 @ self.W3).clamp(min=0)

        # Final output for target class
        score = H3.mean(dim=0)[target_class]

        # Backward (Gradient * Input) as approximation for LRP in deep nets
        # This is "Gradient x Input" which is related to LRP-z rule
        score.backward()

        # Relevance = Input * Gradient
        # We aggregate relevance on the input Identity matrix (Node importance)
        relevance = X.data * X.grad
        node_relevance = relevance.sum(dim=1)  # Sum over feature dimension

        return node_relevance

nodes_nr = 15
model = GraphNet(nodes_nr, 32, 2)
optimizer = torch.optim.Adam(model.params, lr=0.01)

losses = []
for i in range(500):
    # Generate batch (1 sample at a time for simplicity)
    has_motif = i % 2 == 0
    data = create_dataset_graph(nodes_nr, has_motif)

    # Forward
    out = model(data["laplacian"])

    # Loss (MSE for simplicity on One-Hot target)
    target = torch.zeros(2)
    target[data["target"]] = 1
    loss = ((out - target) ** 2).mean()

    # Backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if i % 100 == 0:
        print(f"Iter {i}, Loss: {loss.item():.4f}")

test_data = create_dataset_graph(nodes_nr, has_motif=True)
relevance = model.lrp(test_data["laplacian"], target_class=1)

print("Node Relevance Scores:", relevance)

# Visualization
adj = test_data["adjacency"].numpy()
g = igraph.Graph.Adjacency((adj > 0).tolist(), mode="undirected")
layout = g.layout_kamada_kawai()

# Normalize relevance for color mapping
r = relevance.numpy()
r = (r - r.min()) / (r.max() - r.min() + 1e-9)

plt.figure(figsize=(8, 8))
# Plot Edges
for i in range(nodes_nr):
    for j in range(i + 1, nodes_nr):
        if adj[i, j] > 0:
            # Color edge based on avg relevance of connected nodes
            # High relevance -> Red, Low -> Gray
            avg_rel = (r[i] + r[j]) / 2
            color = plt.cm.Reds(avg_rel) if avg_rel > 0.5 else "gray"
            width = 3 if avg_rel > 0.5 else 1

            p1 = layout[i]
            p2 = layout[j]
            plt.plot(
                [p1[0], p2[0]], [p1[1], p2[1]], color=color, linewidth=width, alpha=0.8
            )

# Plot Nodes
for i in range(nodes_nr):
    plt.scatter(
        layout[i][0],
        layout[i][1],
        c=r[i],
        cmap="Reds",
        s=200,
        edgecolor="black",
        zorder=10,
    )
    plt.text(
        layout[i][0],
        layout[i][1],
        str(i),
        fontsize=12,
        color="white",
        ha="center",
        va="center",
        zorder=11,
    )

plt.title("GNN-LRP: Detecting House Motif (Red = High Relevance)")
plt.axis("off")
plt.show()