src/models/gnn.py

"""
Modular Graph Neural Network (GNN) Architectures for Molecular Property Prediction.
Includes GCN, GAT, GIN, and MPNN (NNConv) with dual global pooling and residual connections.
"""
from typing import Optional, List, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import (
    GCNConv,
    GATConv,
    GINConv,
    NNConv,
    global_mean_pool,
    global_add_pool
)


class BaseGNN(nn.Module):
    """
    Base class for molecular GNNs.
    Implements multi-layer messaging, batch norm, residual connections,
    dual pooling ([mean_pool || add_pool]), and task prediction head.
    """
    def __init__(
        self,
        in_channels: int,
        hidden_dim: int = 64,
        num_layers: int = 3,
        num_tasks: int = 1,
        dropout: float = 0.2,
        residual: bool = True
    ):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_tasks = num_tasks
        self.dropout_rate = dropout
        self.residual = residual

        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()

        # Dual pooling produces 2 * hidden_dim features
        self.head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_tasks)
        )

    def forward(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None,
        batch: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Forward pass through GNN backbone, dual pooling, and task head."""
        h = self.get_node_representations(x, edge_index, edge_attr)

        if batch is None:
            batch = torch.zeros(h.size(0), dtype=torch.long, device=h.device)

        # Dual pooling: concatenation of mean and sum pooling
        mean_p = global_mean_pool(h, batch)
        sum_p = global_add_pool(h, batch)
        pooled = torch.cat([mean_p, sum_p], dim=1)

        return self.head(pooled)

    def get_embedding(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None,
        batch: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Extract graph-level latent representations (used for t-SNE / analysis)."""
        h = self.get_node_representations(x, edge_index, edge_attr)
        if batch is None:
            batch = torch.zeros(h.size(0), dtype=torch.long, device=h.device)

        mean_p = global_mean_pool(h, batch)
        sum_p = global_add_pool(h, batch)
        return torch.cat([mean_p, sum_p], dim=1)

    def get_node_representations(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Subclasses implement convolution message passing over nodes."""
        raise NotImplementedError


class GCN(BaseGNN):
    """Graph Convolutional Network (Kipf & Welling)."""
    def __init__(
        self,
        in_channels: int,
        hidden_dim: int = 64,
        num_layers: int = 3,
        num_tasks: int = 1,
        dropout: float = 0.2,
        residual: bool = True,
        **kwargs
    ):
        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
        for i in range(num_layers):
            dim_in = in_channels if i == 0 else hidden_dim
            self.convs.append(GCNConv(dim_in, hidden_dim))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

    def get_node_representations(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        h = x
        for i in range(self.num_layers):
            h_new = self.convs[i](h, edge_index)
            h_new = self.bns[i](h_new)
            h_new = F.relu(h_new)
            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
            if self.residual and h.shape == h_new.shape:
                h = h + h_new
            else:
                h = h_new
        return h


class GAT(BaseGNN):
    """Graph Attention Network with Multi-Head Attention (Veličković et al.)."""
    def __init__(
        self,
        in_channels: int,
        hidden_dim: int = 64,
        num_layers: int = 3,
        num_tasks: int = 1,
        dropout: float = 0.2,
        residual: bool = True,
        heads: int = 4,
        **kwargs
    ):
        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
        self.heads = heads
        assert hidden_dim % heads == 0, f"hidden_dim ({hidden_dim}) must be divisible by heads ({heads})"
        head_dim = hidden_dim // heads

        for i in range(num_layers):
            dim_in = in_channels if i == 0 else hidden_dim
            self.convs.append(GATConv(dim_in, head_dim, heads=heads, concat=True, dropout=dropout))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

    def get_node_representations(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        h = x
        for i in range(self.num_layers):
            h_new = self.convs[i](h, edge_index)
            h_new = self.bns[i](h_new)
            h_new = F.elu(h_new)
            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
            if self.residual and h.shape == h_new.shape:
                h = h + h_new
            else:
                h = h_new
        return h


class GIN(BaseGNN):
    """Graph Isomorphism Network with MLP aggregators (Xu et al.)."""
    def __init__(
        self,
        in_channels: int,
        hidden_dim: int = 64,
        num_layers: int = 3,
        num_tasks: int = 1,
        dropout: float = 0.2,
        residual: bool = True,
        train_eps: bool = True,
        **kwargs
    ):
        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
        for i in range(num_layers):
            dim_in = in_channels if i == 0 else hidden_dim
            mlp = nn.Sequential(
                nn.Linear(dim_in, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim)
            )
            self.convs.append(GINConv(mlp, train_eps=train_eps))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

    def get_node_representations(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        h = x
        for i in range(self.num_layers):
            h_new = self.convs[i](h, edge_index)
            h_new = self.bns[i](h_new)
            h_new = F.relu(h_new)
            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
            if self.residual and h.shape == h_new.shape:
                h = h + h_new
            else:
                h = h_new
        return h


class MPNN(BaseGNN):
    """
    Message Passing Neural Network using Edge Attributes (Gilmer et al.).
    Utilizes NNConv where message passing weights are conditioned on edge features.
    """
    def __init__(
        self,
        in_channels: int,
        edge_dim: int,
        hidden_dim: int = 64,
        num_layers: int = 3,
        num_tasks: int = 1,
        dropout: float = 0.2,
        residual: bool = True,
        **kwargs
    ):
        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
        self.edge_dim = edge_dim
        self.edge_nns = nn.ModuleList()

        for i in range(num_layers):
            dim_in = in_channels if i == 0 else hidden_dim
            edge_nn = nn.Sequential(
                nn.Linear(edge_dim, 32),
                nn.ReLU(),
                nn.Linear(32, dim_in * hidden_dim)
            )
            self.edge_nns.append(edge_nn)
            self.convs.append(NNConv(dim_in, hidden_dim, edge_nn, aggr='mean'))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

    def get_node_representations(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        # Fallback for empty edges or missing edge attributes
        if edge_attr is None or edge_attr.size(0) == 0:
            edge_attr = torch.zeros((edge_index.size(1), self.edge_dim), dtype=torch.float, device=x.device)

        h = x
        for i in range(self.num_layers):
            h_new = self.convs[i](h, edge_index, edge_attr)
            h_new = self.bns[i](h_new)
            h_new = F.relu(h_new)
            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
            if self.residual and h.shape == h_new.shape:
                h = h + h_new
            else:
                h = h_new
        return h


def build_gnn_model(
    model_name: str,
    in_channels: int,
    edge_dim: int = 10,
    num_tasks: int = 1,
    hidden_dim: int = 64,
    num_layers: int = 3,
    dropout: float = 0.2,
    residual: bool = True,
    **kwargs
) -> BaseGNN:
    """Factory function for creating GNN models."""
    name = model_name.lower()
    if name == 'gcn':
        return GCN(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
    elif name == 'gat':
        return GAT(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
    elif name == 'gin':
        return GIN(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
    elif name == 'mpnn':
        return MPNN(in_channels, edge_dim, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
    else:
        raise ValueError(f"Unknown GNN architecture '{model_name}'. Choose from: gcn, gat, gin, mpnn")