model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv, GINConv
from torch_geometric.utils.scatter import scatter_add

# ------------------------------------------------------------
# Negative sampling (local helper, deduplicated)
# ------------------------------------------------------------
def negative_sampling(edge_index, num_nodes, ratio=1.0):
    """Generate negative edges excluding self-loops and existing positive edges."""
    E = edge_index.size(1)
    num_neg = int(E * ratio)
    if num_neg == 0:
        return torch.empty((2, 0), dtype=torch.long, device=edge_index.device)

    existing = set()
    for i in range(E):
        u = edge_index[0, i].item()
        v = edge_index[1, i].item()
        existing.add((u, v))
        existing.add((v, u))

    neg_edges = []
    max_tries = num_neg * 10
    tries = 0
    while len(neg_edges) < num_neg and tries < max_tries:
        u = torch.randint(0, num_nodes, (1,)).item()
        v = torch.randint(0, num_nodes, (1,)).item()
        tries += 1
        if u != v and (u, v) not in existing:
            neg_edges.append((u, v))
            existing.add((u, v))
            existing.add((v, u))

    while len(neg_edges) < num_neg:
        u = torch.randint(0, num_nodes, (1,)).item()
        v = torch.randint(0, num_nodes, (1,)).item()
        if u != v and (u, v) not in existing:
            neg_edges.append((u, v))
            existing.add((u, v))
            existing.add((v, u))

    neg = torch.tensor(neg_edges[:num_neg], dtype=torch.long).T.to(edge_index.device)
    return neg

# ------------------------------------------------------------
# LIF Neuron
# ------------------------------------------------------------
class LIFNeuron(nn.Module):
    """Leaky Integrate-and-Fire neuron with membrane potential state."""
    def __init__(self, thresh=1.0, tau_mem=5.0):
        super().__init__()
        self.threshold = thresh
        self.leak = 1 - 1.0 / tau_mem
        self.mem = None

    def reset(self, shape=None):
        if shape is not None:
            self.mem = torch.zeros(shape, device=self.mem.device if self.mem is not None else None)
        else:
            self.mem = None

    def forward(self, x):
        if self.mem is None or self.mem.shape != x.shape:
            self.mem = torch.zeros_like(x, device=x.device)
        self.mem = self.leak * self.mem + x
        spikes = (self.mem > self.threshold).float()
        self.mem = self.mem - self.threshold * spikes
        return spikes

# ------------------------------------------------------------
# Spiking GNN Base
# ------------------------------------------------------------
class SpikingGNNBase(nn.Module):
    """Abstract base for spiking GNN with LIF neurons and message passing."""
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.encoder = nn.Linear(config.node_feat_dim, config.hidden_dim)
        self.spike_gen = LIFNeuron(thresh=1.0, tau_mem=5.0)
        self.message_mlp = nn.Sequential(
            nn.Linear(2 * config.hidden_dim, config.hidden_dim),
            nn.ReLU(),
            nn.Linear(config.hidden_dim, config.hidden_dim)
        )
        self.readout = nn.Linear(config.hidden_dim, config.out_dim)
        self.T = config.num_time_steps
        self.num_layers = config.num_layers
        self.record_spikes = False
        self.pre_trace = {}
        self.post_trace = {}

    def compute_messages(self, h, edge_index):
        raise NotImplementedError

    def get_spike_count(self):
        return 0.0

    def get_mac_count(self):
        return 0.0

    def forward(self, x, edge_index_list, test_edges, record_spikes=False):
        # Ensure edge_index_list is a list of length T
        edge_index_list = None
        if not isinstance(edge_index_list, list):
            edge_index_list = [edge_index_list] * self.T
        # Pad if shorter
        while len(edge_index_list) < self.T:
            edge_index_list.append(edge_index_list[-1])

        x = self.encoder(x)
        self.spike_gen.reset(x.shape)

        if record_spikes:
            self.pre_trace = {}
            self.post_trace = {}
            for t in range(self.T):
                pre_layer = {}
                post_layer = {}
                for l in range(self.num_layers):
                    edge_idx = edge_index_list[t] if t < len(edge_index_list) else edge_index_list[-1]
                    pre_feats = x[edge_idx[0]].detach()
                    post_feats = x[edge_idx[1]].detach()
                    pre_layer[l] = pre_feats
                    post_layer[l] = post_feats
                self.pre_trace[t] = pre_layer
                self.post_trace[t] = post_layer

        for t in range(self.T):
            inp_spikes = self.spike_gen(x)
            h = inp_spikes
            for layer in range(self.num_layers):
                edge_index = edge_index_list[t] if t < len(edge_index_list) else edge_index_list[-1]
                messages = self.compute_messages(h, edge_index)
                h = self.spike_gen(messages)
        h_final = h

        u, v = test_edges[0], test_edges[1]
        edge_feats = torch.cat([h_final[u], h_final[v]], dim=-1)
        scores = torch.sigmoid(self.readout(edge_feats))
        return scores

# ------------------------------------------------------------
# Static Baselines
# ------------------------------------------------------------
class GCN(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.conv1 = GCNConv(config.node_feat_dim, config.hidden_dim)
        self.conv2 = GCNConv(config.hidden_dim, config.hidden_dim)
        self.readout = nn.Linear(config.hidden_dim, config.out_dim)
        self.dropout = config.dropout

    def forward(self, x, edge_index, test_edges):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.relu(self.conv2(x, edge_index))
        u, v = test_edges[0], test_edges[1]
        edge_feats = torch.cat([x[u], x[v]], dim=-1)
        scores = torch.sigmoid(self.readout(edge_feats))
        return scores

class GAT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.conv1 = GATConv(config.node_feat_dim, config.hidden_dim, heads=config.num_heads, concat=True)
        self.conv2 = GATConv(config.hidden_dim * config.num_heads, config.hidden_dim, heads=1, concat=False)
        self.readout = nn.Linear(config.hidden_dim, config.out_dim)
        self.dropout = config.dropout

    def forward(self, x, edge_index, test_edges):
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index)
        u, v = test_edges[0], test_edges[1]
        edge_feats = torch.cat([x[u], x[v]], dim=-1)
        scores = torch.sigmoid(self.readout(edge_feats))
        return scores

class GIN(nn.Module):
    def __init__(self, config):
        super().__init__()
        mlp1 = nn.Sequential(
            nn.Linear(config.node_feat_dim, config.hidden_dim),
            nn.BatchNorm1d(config.hidden_dim),
            nn.ReLU(),
            nn.Linear(config.hidden_dim, config.hidden_dim)
        )
        self.conv1 = GINConv(mlp1)
        mlp2 = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.BatchNorm1d(config.hidden_dim),
            nn.ReLU(),
            nn.Linear(config.hidden_dim, config.hidden_dim)
        )
        self.conv2 = GINConv(mlp2)
        self.readout = nn.Linear(config.hidden_dim, config.out_dim)
        self.dropout = config.dropout

    def forward(self, x, edge_index, test_edges):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.conv2(x, edge_index)
        u, v = test_edges[0], test_edges[1]
        edge_feats = torch.cat([x[u], x[v]], dim=-1)
        scores = torch.sigmoid(self.readout(edge_feats))
        return scores

# ------------------------------------------------------------
# Spiking GNN Variants
# ------------------------------------------------------------
class FullFinetuneSgnn(SpikingGNNBase):
    """Full fine-tuning: all parameters updated during adaptation."""
    def __init__(self, config):
        super().__init__(config)

    def compute_messages(self, h, edge_index):
        src, dst = h[edge_index[0]], h[edge_index[1]]
        messages = self.message_mlp(torch.cat([src, dst], dim=-1))
        out = scatter_add(messages, edge_index[1], dim=0, dim_size=h.size(0))
        return out

    def adapt(self, data, config):
        optimizer = torch.optim.Adam(self.parameters(), lr=config.lr)
        for epoch in range(config.epochs_adapt):
            for t in range(config.adapt_steps):
                edges = data['adapt_edges'][t]
                pos_scores = self.forward(data['x'], [edges], test_edges=data['pos_test_edges'], record_spikes=False)
                neg_edges = negative_sampling(edges, data['num_nodes'])
                neg_scores = self.forward(data['x'], [edges], test_edges=neg_edges, record_spikes=False)
                loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores)) \
                       + F.binary_cross_entropy(neg_scores, torch.zeros_like(neg_scores))
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                if loss.item() > 100:
                    print("FAIL: NaN/divergence detected")
                    return

class SilentSynapseAdaptation(SpikingGNNBase):
    """Mask learning via STDP-based weight binarization with per-synapse plasticity."""
    def __init__(self, config):
        super().__init__(config)
        # mask will be resized in adapt; start with large placeholder
        self.register_buffer('masks', torch.ones(10000) * config.mask_initial_ratio)

    def compute_messages(self, h, edge_index):
        mask_bin = torch.bernoulli(torch.sigmoid(self.masks[:edge_index.size(1)]))
        src, dst = h[edge_index[0]], h[edge_index[1]]
        messages = self.message_mlp(torch.cat([src, dst], dim=-1)) * mask_bin.unsqueeze(-1)
        out = scatter_add(messages, edge_index[1], dim=0, dim_size=h.size(0))
        return out

    def adapt(self, data, config):
        # Per-synapse STDP: iterate over each synapse in each adaptation snapshot
        for t in range(config.adapt_steps):
            edges = data['adapt_edges'][t]
            E = edges.size(1)
            # Ensure mask vector large enough
            if E > self.masks.size(0):
                new_masks = torch.ones(E, device=self.masks.device) * config.mask_initial_ratio
                self.masks = nn.Parameter(new_masks)  # replace; note: was register_buffer (we'll keep as buffer)
            else:
                # we need to set masks up to E (if smaller)
                pass

            # Forward pass with spike recording
            self.forward(data['x'], [edges], data['pos_test_edges'], record_spikes=True)

            # For each synapse (edge) compute local correlation
            with torch.no_grad():
                for i in range(E):
                    # Collect pre- and post-spike traces for all time steps for this edge
                    pre_sum = 0.0
                    post_sum = 0.0
                    for t_ in range(self.T):
                        if t_ in self.pre_trace and 0 in self.pre_trace[t_]:  # layer 0 traces
                            pre_spikes = self.pre_trace[t_][0][i]  # feature vector at edge i
                            post_spikes = self.post_trace[t_][0][i]
                            # use mean over feature dimensions as scalar trace
                            pre_val = pre_spikes.mean().item()
                            post_val = post_spikes.mean().item()
                            # Apply STDP window: exponential decay with time difference
                            for tau in range(config.stdp_window):
                                if t_ + tau < self.T:
                                    pre_sum += pre_val * math.exp(-tau / config.stdp_window)
                                if t_ - tau >= 0:
                                    post_sum += post_val * math.exp(-tau / config.stdp_window)
                    # Local correlation: product of pre and post (simplified)
                    corr = pre_sum * post_sum
                    if abs(corr) < 1e-12:
                        corr = 0.0
                    delta = config.stdp_learning_rate * (corr - config.stdp_decay)
                    new_val = self.masks[i] + delta
                    self.masks[i] = torch.clamp(new_val, 0, 1)
        # Binarize after adaptation
        with torch.no_grad():
            self.masks = (torch.sigmoid(self.masks) > 0.5).float()

class SilentSynapseSurrogate(SpikingGNNBase):
    """Mask learning via surrogate gradient (Gumbel-Softmax)."""
    def __init__(self, config):
        super().__init__(config)
        self.mask_logits = nn.Parameter(torch.zeros(10000))

    def compute_messages(self, h, edge_index):
        E = edge_index.size(1)
        mask_logits = self.mask_logits[:E]
        mask_soft = F.gumbel_softmax(
            torch.stack([-mask_logits, mask_logits], dim=-1),
            tau=self.config.gumbel_tau, hard=True
        )[:, 1]
        src, dst = h[edge_index[0]], h[edge_index[1]]
        messages = self.message_mlp(torch.cat([src, dst], dim=-1)) * mask_soft.unsqueeze(-1)
        out = scatter_add(messages, edge_index[1], dim=0, dim_size=h.size(0))
        return out

    def adapt(self, data, config):
        optimizer = torch.optim.Adam([self.mask_logits], lr=config.mask_lr)
        for epoch in range(config.epochs_adapt):
            for t in range(config.adapt_steps):
                edges = data['adapt_edges'][t]
                E = edges.size(1)
                if E > self.mask_logits.size(0):
                    self.mask_logits.data = torch.cat([
                        self.mask_logits.data,
                        torch.zeros(E - self.mask_logits.size(0), device=self.mask_logits.device)
                    ], dim=0)
                pos_scores = self.forward(data['x'], [edges], test_edges=data['pos_test_edges'])
                neg_edges = negative_sampling(edges, data['num_nodes'])
                neg_scores = self.forward(data['x'], [edges], test_edges=neg_edges)
                loss = F.binary_cross_entropy(pos_scores, torch.ones_like(pos_scores)) \
                       + F.binary_cross_entropy(neg_scores, torch.zeros_like(neg_scores))
                p = torch.sigmoid(self.mask_logits[:E])
                entropy = -p * torch.log(p + 1e-8) - (1 - p) * torch.log(1 - p + 1e-8)
                loss += config.entropy_reg_lambda * entropy.mean()
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()
                if loss.item() > 100:
                    print("FAIL: NaN/divergence detected")
                    return
        with torch.no_grad():
            self.register_buffer('masks', (torch.sigmoid(self.mask_logits) > 0.5).float())

class NoAdaptation(SpikingGNNBase):
    """No adaptation: frozen pre-trained model."""
    def __init__(self, config):
        super().__init__(config)

    def compute_messages(self, h, edge_index):
        src, dst = h[edge_index[0]], h[edge_index[1]]
        messages = self.message_mlp(torch.cat([src, dst], dim=-1))
        out = scatter_add(messages, edge_index[1], dim=0, dim_size=h.size(0))
        return out

    def adapt(self, data, config):
        pass

class RandomMaskBaseline(SpikingGNNBase):
    """Random fixed masks (50% active)."""
    def __init__(self, config):
        super().__init__(config)
        total_edges = 10000
        self.register_buffer('masks', torch.bernoulli(torch.full((total_edges,), 0.5)))

    def compute_messages(self, h, edge_index):
        new_masks = None
        if edge_index.size(1) > self.masks.size(0):
            new_masks = torch.bernoulli(torch.full((edge_index.size(1),), 0.5)).to(self.masks.device)
            self.masks = new_masks
        mask_bin = self.masks[:edge_index.size(1)]
        src, dst = h[edge_index[0]], h[edge_index[1]]
        messages = self.message_mlp(torch.cat([src, dst], dim=-1)) * mask_bin.unsqueeze(-1)
        out = scatter_add(messages, edge_index[1], dim=0, dim_size=h.size(0))
        return out

    def adapt(self, data, config):
        pass