Graph Neural Network (GNN)

graph TD
    A[Node 1] -->|Edge| B[Node 2]
    A -->|Edge| C[Node 3]
    B -->|Edge| C
    C -->|Edge| D[Node 4]
    B -->|Edge| D
    style A fill:#f9f,stroke:#333
    style B fill:#bbf,stroke:#333
    style C fill:#f96,stroke:#333
    style D fill:#6f9,stroke:#333

h_v^(k+1) = UPDATE(h_v^k, AGGREGATE({h_u^k | u ∈ N(v)}))

# GCN layer implementation
import torch
import torch.nn as nn
import torch.nn.functional as F

class GCNLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super(GCNLayer, self).__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x, adjacency):
        """
        x: Node features (n_nodes, in_features)
        adjacency: Adjacency matrix (n_nodes, n_nodes)
        """
        # Normalize adjacency matrix
        degree = torch.diag(1 / torch.sqrt(torch.sum(adjacency, dim=1)))
        norm_adj = degree @ adjacency @ degree

        # Message passing
        x = norm_adj @ x

        # Linear transformation
        x = self.linear(x)

        return F.relu(x)

# GAT layer implementation
class GATLayer(nn.Module):
    def __init__(self, in_features, out_features, n_heads=8):
        super(GATLayer, self).__init__()
        self.n_heads = n_heads
        self.out_features = out_features

        # Attention parameters
        self.W = nn.Linear(in_features, n_heads * out_features)
        self.a = nn.Linear(2 * out_features, n_heads)

    def forward(self, x, adjacency):
        """
        x: Node features (n_nodes, in_features)
        adjacency: Adjacency matrix (n_nodes, n_nodes)
        """
        n_nodes = x.size(0)

        # Linear transformation
        h = self.W(x)  # (n_nodes, n_heads * out_features)
        h = h.view(n_nodes, self.n_heads, self.out_features)

        # Attention mechanism
        attention_scores = []
        for i in range(n_nodes):
            for j in range(n_nodes):
                if adjacency[i, j] > 0:  # If nodes are connected
                    # Concatenate features
                    concat = torch.cat([h[i], h[j]], dim=1)
                    # Compute attention score
                    score = torch.softmax(self.a(concat), dim=0)
                    attention_scores.append(score)

        # Apply attention
        output = torch.zeros_like(h)
        for i in range(n_nodes):
            neighbor_indices = torch.nonzero(adjacency[i], as_tuple=True)[0]
            for j in neighbor_indices:
                output[i] += attention_scores[i * n_nodes + j] * h[j]

        return output.mean(dim=1)  # Average over heads

# GraphSAGE layer implementation
class GraphSAGELayer(nn.Module):
    def __init__(self, in_features, out_features, aggregator='mean'):
        super(GraphSAGELayer, self).__init__()
        self.aggregator = aggregator
        self.linear = nn.Linear(2 * in_features, out_features)

    def forward(self, x, adjacency, sample_size=2):
        """
        x: Node features (n_nodes, in_features)
        adjacency: Adjacency matrix (n_nodes, n_nodes)
        """
        n_nodes = x.size(0)
        output = torch.zeros(n_nodes, self.linear.in_features)

        for i in range(n_nodes):
            # Get neighbors
            neighbor_indices = torch.nonzero(adjacency[i], as_tuple=True)[0]

            # Sample neighbors if too many
            if len(neighbor_indices) > sample_size:
                neighbor_indices = neighbor_indices[torch.randperm(len(neighbor_indices))[:sample_size]]

            # Aggregate neighbor features
            if self.aggregator == 'mean':
                neighbor_features = x[neighbor_indices].mean(dim=0)
            elif self.aggregator == 'max':
                neighbor_features = x[neighbor_indices].max(dim=0)[0]
            elif self.aggregator == 'sum':
                neighbor_features = x[neighbor_indices].sum(dim=0)

            # Concatenate self and neighbor features
            output[i] = torch.cat([x[i], neighbor_features])

        return F.relu(self.linear(output))

# Graph autoencoder implementation
class GraphAutoencoder(nn.Module):
    def __init__(self, in_features, hidden_dim):
        super(GraphAutoencoder, self).__init__()
        # Encoder
        self.encoder = nn.Sequential(
            GCNLayer(in_features, hidden_dim),
            GCNLayer(hidden_dim, hidden_dim)
        )
        # Decoder
        self.decoder = nn.Sequential(
            GCNLayer(hidden_dim, hidden_dim),
            GCNLayer(hidden_dim, in_features)
        )

    def forward(self, x, adjacency):
        # Encode
        z = self.encoder(x, adjacency)
        # Decode
        x_recon = self.decoder(z, adjacency)
        return x_recon, z

# Graph U-Net implementation (conceptual)
class GraphUNet(nn.Module):
    def __init__(self, in_features, hidden_dim, depth=3):
        super(GraphUNet, self).__init__()
        self.depth = depth

        # Downsampling layers
        self.down_layers = nn.ModuleList()
        for i in range(depth):
            self.down_layers.append(GCNLayer(
                in_features if i == 0 else hidden_dim,
                hidden_dim
            ))

        # Upsampling layers
        self.up_layers = nn.ModuleList()
        for i in range(depth):
            self.up_layers.append(GCNLayer(
                hidden_dim * 2 if i == 0 else hidden_dim,
                hidden_dim if i < depth - 1 else in_features
            ))

    def forward(self, x, adjacency):
        # Downsampling
        down_outputs = []
        for layer in self.down_layers:
            x = layer(x, adjacency)
            down_outputs.append(x)

        # Upsampling
        for i, layer in enumerate(self.up_layers):
            if i > 0:
                x = torch.cat([x, down_outputs[self.depth - i - 1]], dim=1)
            x = layer(x, adjacency)

        return x

# Common loss functions for GNNs
def node_classification_loss(predictions, labels):
    """Cross-entropy loss for node classification"""
    return F.cross_entropy(predictions, labels)

def link_prediction_loss(predictions, adjacency):
    """Binary cross-entropy for link prediction"""
    # Create negative samples
    n_nodes = adjacency.size(0)
    neg_adj = torch.ones_like(adjacency) - adjacency - torch.eye(n_nodes)

    # Sample negative edges
    neg_edges = torch.nonzero(neg_adj)
    neg_samples = neg_edges[torch.randperm(len(neg_edges))[:len(predictions)]]

    # Create labels
    pos_labels = torch.ones(len(predictions))
    neg_labels = torch.zeros(len(neg_samples))

    # Combine
    all_predictions = torch.cat([predictions, predictions[neg_samples[:, 0], neg_samples[:, 1]]])
    all_labels = torch.cat([pos_labels, neg_labels])

    return F.binary_cross_entropy_with_logits(all_predictions, all_labels)

def graph_classification_loss(predictions, labels):
    """Cross-entropy for graph classification"""
    return F.cross_entropy(predictions, labels)

# GNN training loop
def train_gnn(model, optimizer, data_loader, epochs, task='node_classification'):
    model.train()

    for epoch in range(epochs):
        total_loss = 0

        for batch in data_loader:
            optimizer.zero_grad()

            # Forward pass
            if task == 'node_classification':
                predictions = model(batch.x, batch.edge_index)
                loss = node_classification_loss(predictions[batch.train_mask], batch.y[batch.train_mask])
            elif task == 'link_prediction':
                predictions = model(batch.x, batch.edge_index)
                loss = link_prediction_loss(predictions[batch.edge_index[0], batch.edge_index[1]], batch.adjacency)
            elif task == 'graph_classification':
                predictions = model(batch.x, batch.edge_index)
                loss = graph_classification_loss(predictions, batch.y)

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

            total_loss += loss.item()

        print(f'Epoch {epoch+1}, Loss: {total_loss / len(data_loader):.4f}')

# Node classification with GNN
class NodeClassifier(nn.Module):
    def __init__(self, in_features, hidden_dim, n_classes):
        super(NodeClassifier, self).__init__()
        self.gcn1 = GCNLayer(in_features, hidden_dim)
        self.gcn2 = GCNLayer(hidden_dim, hidden_dim)
        self.classifier = nn.Linear(hidden_dim, n_classes)

    def forward(self, x, edge_index):
        # Convert edge_index to adjacency matrix
        n_nodes = x.size(0)
        adjacency = torch.zeros(n_nodes, n_nodes)
        adjacency[edge_index[0], edge_index[1]] = 1

        # GNN layers
        x = self.gcn1(x, adjacency)
        x = self.gcn2(x, adjacency)

        # Classification
        return self.classifier(x)

# Link prediction with GNN
class LinkPredictor(nn.Module):
    def __init__(self, in_features, hidden_dim):
        super(LinkPredictor, self).__init__()
        self.gcn1 = GCNLayer(in_features, hidden_dim)
        self.gcn2 = GCNLayer(hidden_dim, hidden_dim)
        self.predictor = nn.Linear(2 * hidden_dim, 1)

    def forward(self, x, edge_index):
        # Convert edge_index to adjacency matrix
        n_nodes = x.size(0)
        adjacency = torch.zeros(n_nodes, n_nodes)
        adjacency[edge_index[0], edge_index[1]] = 1

        # GNN layers
        x = self.gcn1(x, adjacency)
        x = self.gcn2(x, adjacency)

        # Create edge features by concatenating node features
        edge_features = torch.cat([
            x[edge_index[0]],
            x[edge_index[1]]
        ], dim=1)

        return torch.sigmoid(self.predictor(edge_features)).squeeze()

# Graph classification with GNN
class GraphClassifier(nn.Module):
    def __init__(self, in_features, hidden_dim, n_classes):
        super(GraphClassifier, self).__init__()
        self.gcn1 = GCNLayer(in_features, hidden_dim)
        self.gcn2 = GCNLayer(hidden_dim, hidden_dim)
        self.pool = nn.AdaptiveAvgPool1d(1)  # Global pooling
        self.classifier = nn.Linear(hidden_dim, n_classes)

    def forward(self, x, edge_index):
        # Convert edge_index to adjacency matrix
        n_nodes = x.size(0)
        adjacency = torch.zeros(n_nodes, n_nodes)
        adjacency[edge_index[0], edge_index[1]] = 1

        # GNN layers
        x = self.gcn1(x, adjacency)
        x = self.gcn2(x, adjacency)

        # Global pooling
        x = x.unsqueeze(0)  # Add batch dimension
        x = self.pool(x.transpose(1, 2)).squeeze()

        # Classification
        return self.classifier(x)

# Molecular property prediction with GNN
class MolecularGNN(nn.Module):
    def __init__(self, node_features, edge_features, hidden_dim, n_classes):
        super(MolecularGNN, self).__init__()
        # Node and edge feature encoders
        self.node_encoder = nn.Linear(node_features, hidden_dim)
        self.edge_encoder = nn.Linear(edge_features, hidden_dim)

        # GNN layers
        self.gcn1 = GCNLayer(hidden_dim, hidden_dim)
        self.gcn2 = GCNLayer(hidden_dim, hidden_dim)

        # Readout
        self.readout = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, n_classes)
        )

    def forward(self, node_features, edge_index, edge_features):
        # Encode features
        x = self.node_encoder(node_features)

        # Convert edge_index to adjacency matrix
        n_nodes = x.size(0)
        adjacency = torch.zeros(n_nodes, n_nodes)
        adjacency[edge_index[0], edge_index[1]] = 1

        # Incorporate edge features
        for i in range(edge_index.size(1)):
            src, dst = edge_index[:, i]
            adjacency[src, dst] *= self.edge_encoder(edge_features[i])

        # GNN layers
        x = self.gcn1(x, adjacency)
        x = self.gcn2(x, adjacency)

        # Global pooling
        x = x.mean(dim=0)

        # Readout
        return self.readout(x)