src/dst_gnn/model.py

"""Model and training utilities for DST-GNN."""

from __future__ import annotations

from typing import Iterable

import torch
import torch.nn.functional as F
from torch import nn
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv


class DSTGNN(nn.Module):
    """A simple spatio-temporal GNN for sequential COSTE graphs.

    The model follows the manuscript description with:
    - one-hot node identity features
    - diffusion-style graph convolutions over stage graphs
    - explicit temporal state propagation across T1 -> T2 -> T3
    - an MSE objective on later-stage SSS reconstruction
    """

    def __init__(
        self,
        num_nodes: int,
        hidden_channels: int = 32,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.num_nodes = num_nodes
        self.hidden_channels = hidden_channels
        self.dropout = dropout

        self.input_projection = nn.Linear(num_nodes, hidden_channels)
        self.conv1 = GCNConv(hidden_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, hidden_channels)
        self.temporal_update = nn.GRUCell(hidden_channels, hidden_channels)
        self.edge_decoder = nn.Sequential(
            nn.Linear(hidden_channels * 2, hidden_channels),
            nn.ReLU(),
            nn.Linear(hidden_channels, 1),
            nn.Sigmoid(),
        )

    def encode_step(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_weight: torch.Tensor,
        prev_hidden: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Encode one stage graph while propagating temporal state."""

        base_state = self.input_projection(x)
        if prev_hidden is None:
            prev_hidden = torch.zeros_like(base_state)

        hidden = base_state + prev_hidden
        hidden = self.conv1(hidden, edge_index, edge_weight=edge_weight)
        hidden = F.relu(hidden)
        hidden = F.dropout(hidden, p=self.dropout, training=self.training)
        hidden = self.conv2(hidden, edge_index, edge_weight=edge_weight)
        hidden = F.relu(hidden)
        hidden = self.temporal_update(hidden, prev_hidden)
        return hidden

    def decode_adjacency(self, node_embeddings: torch.Tensor) -> torch.Tensor:
        """Decode directed SSS predictions for all ordered cell-type pairs."""

        num_nodes = node_embeddings.size(0)
        left = node_embeddings.unsqueeze(1).expand(num_nodes, num_nodes, -1)
        right = node_embeddings.unsqueeze(0).expand(num_nodes, num_nodes, -1)
        logits = self.edge_decoder(torch.cat([left, right], dim=-1)).squeeze(-1)
        logits = logits.clamp(0.0, 1.0)
        logits.fill_diagonal_(0.0)
        return logits

    def forecast_sequence(self, graphs: Iterable[Data]) -> list[torch.Tensor]:
        """Forecast later-stage SSS matrices from earlier observed graphs."""

        predictions: list[torch.Tensor] = []
        hidden: torch.Tensor | None = None
        graph_list = list(graphs)
        for graph in graph_list[:-1]:
            hidden = self.encode_step(graph.x, graph.edge_index, graph.edge_weight, hidden)
            predictions.append(self.decode_adjacency(hidden))
        return predictions

    def encode_observed_sequence(self, graphs: Iterable[Data]) -> list[torch.Tensor]:
        """Return node embeddings after each observed stage graph."""

        embeddings: list[torch.Tensor] = []
        hidden: torch.Tensor | None = None
        for graph in graphs:
            hidden = self.encode_step(graph.x, graph.edge_index, graph.edge_weight, hidden)
            embeddings.append(hidden)
        return embeddings


def compute_transition_loss(model: DSTGNN, graphs: list[Data]) -> torch.Tensor:
    """Average MSE between predicted and observed later-stage SSS matrices."""

    predictions = model.forecast_sequence(graphs)
    losses = [
        F.mse_loss(prediction, graphs[index + 1].y)
        for index, prediction in enumerate(predictions)
    ]
    if not losses:
        raise ValueError("At least two stage graphs are required for temporal training.")
    return torch.stack(losses).mean()


def train_model(
    model: DSTGNN,
    graphs: list[Data],
    epochs: int = 400,
    lr: float = 0.01,
    weight_decay: float = 5e-4,
    verbose_every: int = 50,
) -> list[dict[str, float]]:
    """Train DST-GNN with Adam and an MSE objective."""

    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    history: list[dict[str, float]] = []

    for epoch in range(1, epochs + 1):
        model.train()
        optimizer.zero_grad()
        loss = compute_transition_loss(model, graphs)
        loss.backward()
        optimizer.step()

        record = {"epoch": float(epoch), "loss": float(loss.item())}
        history.append(record)

        if verbose_every and (epoch == 1 or epoch % verbose_every == 0 or epoch == epochs):
            print(f"Epoch {epoch:04d} | loss={loss.item():.6f}")

    return history