model.py

# -*- coding: utf-8 -*-
"""
scGTN: Single-cell Graph Transformer Network
Dual-view clustering model
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from transformers.models.bert.modeling_bert import (
    BertPreTrainedModel, BertPooler, BertAttention,
    BertIntermediate, BertOutput
)
from torch_geometric.nn import GATConv
from torch_geometric.utils import dense_to_sparse

from layer import (
    ZINBLoss, ZINBDecoder,
    optimal_transport_clustering_loss, compute_soft_assignment,
    siamese_correlation_loss, graph_reconstruction_loss
)
from utils import construct_subgraph_features


######################################## GAT Encoder ########################################

class GATEncoder(nn.Module):
    """Graph Attention Network encoder"""

    def __init__(self, nfeat, nhid, nheads=4, dropout=0.1):
        super(GATEncoder, self).__init__()
        self.dropout = dropout
        head_dim = nhid // nheads

        self.gat_layer = GATConv(
            in_channels=nfeat,
            out_channels=head_dim,
            heads=nheads,
            dropout=dropout,
            concat=True
        )

        self.out_dim = head_dim * nheads
        self.residual_proj = nn.Linear(nfeat, self.out_dim) if nfeat != self.out_dim else nn.Identity()
        self.layer_norm = nn.LayerNorm(self.out_dim)
        self.activation = nn.ELU()

    def forward(self, x, adj):
        if adj.dim() == 2 and adj.shape[0] == adj.shape[1]:
            edge_index, _ = dense_to_sparse(adj)
        else:
            edge_index = adj

        x_in = x
        x = self.gat_layer(x, edge_index)
        x = x + self.residual_proj(x_in)
        x = self.layer_norm(x)
        x = self.activation(x)

        if self.training:
            x = F.dropout(x, p=self.dropout, training=True)
        return x


######################################## Transformer Components ########################################

class GraphBertEmbeddings(nn.Module):
    """Multi-modal embedding: Y = H + G + H_sp"""

    def __init__(self, config):
        super(GraphBertEmbeddings, self).__init__()
        self.raw_feature_embeddings = nn.Linear(config.hidden_size, config.hidden_size)
        self.inti_pos_embeddings = nn.Embedding(config.max_inti_pos_index, config.hidden_size)
        self.hop_dis_embeddings = nn.Embedding(config.max_hop_dis_index, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, raw_features, init_pos_ids, hop_dis_ids):
        raw_feature_embeds = self.raw_feature_embeddings(raw_features)
        position_embeddings = self.inti_pos_embeddings(init_pos_ids)
        hop_embeddings = self.hop_dis_embeddings(hop_dis_ids)

        embeddings = raw_feature_embeds + position_embeddings + hop_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class TransformerLayer(nn.Module):
    """Single Transformer layer"""

    def __init__(self, config):
        super().__init__()
        self.attention = BertAttention(config)
        self.is_decoder = config.is_decoder
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
        attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
        attention_output = attention_outputs[0]
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return (layer_output,) + attention_outputs[1:]


class TransformerEncoder(nn.Module):
    """Stacked Transformer encoder"""

    def __init__(self, config):
        super(TransformerEncoder, self).__init__()
        self.layer = nn.ModuleList([TransformerLayer(config) for _ in range(config.num_hidden_layers)])

    def forward(self, hidden_states, head_mask=None, residual_h=None):
        if head_mask is None:
            head_mask = [None] * len(self.layer)

        for i, layer_module in enumerate(self.layer):
            layer_outputs = layer_module(hidden_states, head_mask=head_mask[i])
            hidden_states = layer_outputs[0]
            if residual_h is not None:
                for idx in range(hidden_states.size(1)):
                    hidden_states[:, idx, :] += residual_h

        return (hidden_states,)


class SingleViewEncoder(BertPreTrainedModel):
    """Single-view Graph-Transformer encoder"""

    def __init__(self, config):
        super(SingleViewEncoder, self).__init__(config)
        self.config = config
        self.embeddings = GraphBertEmbeddings(config)
        self.encoder = TransformerEncoder(config)
        self.pooler = BertPooler(config)
        self.init_weights()

    def forward(self, raw_features, init_pos_ids, hop_dis_ids, head_mask=None):
        if head_mask is None:
            head_mask = [None] * self.config.num_hidden_layers

        embedding_output = self.embeddings(raw_features, init_pos_ids, hop_dis_ids)
        encoder_outputs = self.encoder(embedding_output, head_mask=head_mask)
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output)
        return (sequence_output, pooled_output)


######################################## Dual-View Clustering Model ########################################

class scGTN(BertPreTrainedModel):
    """
    scGTN: Dual-view Graph Transformer Network for single-cell clustering

    Architecture:
    1. Shared GAT encoder for both views
    2. Dual Transformer encoders (one per view)
    3. ZINB decoder for reconstruction
    4. Learnable fusion weight for view combination
    """

    def __init__(self, config):
        super(scGTN, self).__init__(config)
        self.config = config

        # Shared GAT encoder
        self.gat_encoder = GATEncoder(
            nfeat=config.x_size,
            nhid=config.hidden_size,
            nheads=4,
            dropout=0.1
        )

        # Dual Transformer encoders
        self.encoder1 = SingleViewEncoder(config)
        self.encoder2 = SingleViewEncoder(config)

        # ZINB decoder
        self.zinb_decoder = ZINBDecoder(config.hidden_size, config.x_size)
        self.zinb_loss_fn = ZINBLoss()

        self.init_weights()

        # Cluster centers
        self.cluster_layer = Parameter(torch.Tensor(config.y_size, config.hidden_size))
        torch.nn.init.xavier_normal_(self.cluster_layer.data)

        # Learnable fusion weight
        self.fusion_weight = Parameter(torch.tensor(0.0))

        # Data placeholder
        self.data = None

    def forward(self, raw_features, raw_features2, init_pos_ids, hop_dis_ids,
                init_pos_ids2, hop_dis_ids2, x_features=None, adj1=None, adj2=None,
                neighbor_indices=None, neighbor_indices2=None):
        """Dual-view forward pass"""
        # GAT encoding
        h1 = self.gat_encoder(x_features, adj1)
        h2 = self.gat_encoder(x_features, adj2)

        # Construct subgraph features
        if neighbor_indices is not None:
            raw_features = construct_subgraph_features(h1, neighbor_indices, self.config.k)
        if neighbor_indices2 is not None:
            raw_features2 = construct_subgraph_features(h2, neighbor_indices2, self.config.k)

        # Transformer encoding
        outputs1 = self.encoder1(raw_features, init_pos_ids, hop_dis_ids)
        outputs2 = self.encoder2(raw_features2, init_pos_ids2, hop_dis_ids2)

        # Mean pooling
        z1 = outputs1[0].mean(dim=1)
        z2 = outputs2[0].mean(dim=1)

        return z1, z2

    def compute_losses(self, z1, z2, x_raw, adj1, adj2,
                       use_otc=True, temperature=1.0, lambda_cor=0.005,
                       target_distribution=None):
        """Compute all losses"""
        losses = {}

        # Fused embedding
        alpha = torch.sigmoid(self.fusion_weight)
        z_fused = alpha * z1 + (1 - alpha) * z2

        # OTC loss
        if use_otc:
            q = compute_soft_assignment(z_fused, self.cluster_layer, temperature)
            loss_otc = optimal_transport_clustering_loss(
                q, self.cluster_layer, z_fused, temperature, 3, target_distribution
            )
            losses['loss_otc'] = loss_otc
        else:
            losses['loss_otc'] = torch.tensor(0.0, device=z1.device)

        # Correlation loss
        losses['loss_cor'] = siamese_correlation_loss(z1, z2, lambda_cor)

        # ZINB loss
        mean1, disp1, pi1 = self.zinb_decoder(z1)
        mean2, disp2, pi2 = self.zinb_decoder(z2)
        losses['loss_zinb'] = self.zinb_loss_fn(x_raw, mean1, disp1, pi1) + \
                              self.zinb_loss_fn(x_raw, mean2, disp2, pi2)

        # Graph reconstruction loss
        losses['loss_graph'] = graph_reconstruction_loss(z1, adj1) + \
                               graph_reconstruction_loss(z2, adj2)

        return losses

    def get_embeddings(self):
        """Get embeddings from stored data"""
        self.eval()
        with torch.no_grad():
            z1, z2 = self.forward(
                self.data['raw_embeddings'],
                self.data['raw_embeddings2'],
                self.data['int_embeddings'],
                self.data['hop_embeddings'],
                self.data['int_embeddings2'],
                self.data['hop_embeddings2'],
                x_features=self.data.get('X_features'),
                adj1=self.data.get('A1'),
                adj2=self.data.get('A2'),
                neighbor_indices=self.data.get('neighbor_indices'),
                neighbor_indices2=self.data.get('neighbor_indices2')
            )
        return z1, z2