layer.py

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


######################################## ZINB Loss ########################################

class ZINBLoss(nn.Module):
    """Zero-Inflated Negative Binomial Loss for scRNA-seq data"""

    def __init__(self, eps=1e-10, ridge_lambda=0.0, reduction="mean"):
        super(ZINBLoss, self).__init__()
        self.eps = eps
        self.ridge_lambda = ridge_lambda
        self.reduction = reduction

    def forward(self, x, mean, disp, pi, size_factors=None):
        if size_factors is not None:
            size_factors = size_factors.unsqueeze(1)
            mean = mean * size_factors

        mean = mean + self.eps
        disp = disp + self.eps
        pi = torch.clamp(pi, self.eps, 1 - self.eps)

        t1 = torch.lgamma(disp) + torch.lgamma(x + 1.0) - torch.lgamma(x + disp)
        t2 = (disp + x) * torch.log(1.0 + mean / disp) + x * (torch.log(disp) - torch.log(mean))
        nb_log_likelihood = t1 + t2

        nb_case = nb_log_likelihood - torch.log(1.0 - pi)
        zero_nb = torch.pow(disp / (disp + mean), disp)
        zero_case = -torch.log(pi + (1.0 - pi) * zero_nb)

        zinb_loss = torch.where(x <= 1e-8, zero_case, nb_case)

        if self.ridge_lambda > 0:
            ridge_penalty = self.ridge_lambda * torch.square(pi)
            zinb_loss = zinb_loss + ridge_penalty

        if self.reduction == "mean":
            return torch.mean(zinb_loss)
        elif self.reduction == "sum":
            return torch.sum(zinb_loss)
        return zinb_loss


class ZINBDecoder(nn.Module):
    """ZINB decoder: generates ZINB parameters from embeddings"""

    def __init__(self, input_dim, output_dim):
        super(ZINBDecoder, self).__init__()
        self.mean_decoder = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.Softplus()
        )
        self.disp_decoder = nn.Parameter(torch.randn(output_dim))
        self.pi_decoder = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.Sigmoid()
        )

    def forward(self, z):
        mean = self.mean_decoder(z)
        disp = torch.exp(self.disp_decoder)
        pi = self.pi_decoder(z)
        return mean, disp, pi


######################################## Optimal Transport Loss ########################################

def sinkhorn_knopp(Q, n_iterations=3, target_distribution=None):
    """Sinkhorn-Knopp algorithm for balanced cluster assignment"""
    N, K = Q.shape
    if target_distribution is None:
        target_distribution = torch.ones(K, device=Q.device) / K

    Q = torch.clamp(Q, min=1e-8)

    for _ in range(n_iterations):
        Q = Q / Q.sum(dim=1, keepdim=True)
        Q = Q / Q.sum(dim=0, keepdim=True)
        Q = Q * (N * target_distribution).unsqueeze(0)

    Q = Q / Q.sum(dim=1, keepdim=True)
    return Q


def compute_soft_assignment(embeddings, cluster_centers, temperature=1.0, alpha=1.0):
    """Compute soft cluster assignment using Student t-distribution"""
    distances = torch.cdist(embeddings, cluster_centers, p=2).pow(2)
    q = 1.0 / (1.0 + distances / alpha)
    q = q.pow((alpha + 1.0) / 2.0)

    if temperature != 1.0:
        q = q.pow(1.0 / temperature)

    q = q / q.sum(dim=1, keepdim=True)
    return q


def optimal_transport_clustering_loss(
    Q, cluster_centers, embeddings,
    temperature=1.0, sinkhorn_iterations=3, target_distribution=None
):
    """OTC Loss: KL divergence between Q and Sinkhorn-balanced P_hat"""
    if Q.shape != (embeddings.shape[0], cluster_centers.shape[0]):
        Q = compute_soft_assignment(embeddings, cluster_centers, temperature)

    with torch.no_grad():
        P_hat = sinkhorn_knopp(Q.detach(), sinkhorn_iterations, target_distribution)

    eps = 1e-8
    Q_safe = torch.clamp(Q, min=eps, max=1.0)
    P_hat_safe = torch.clamp(P_hat, min=eps, max=1.0)

    loss_otc = F.kl_div(Q_safe.log(), P_hat_safe, reduction='batchmean')
    return loss_otc


######################################## Siamese Correlation Loss ########################################

def siamese_correlation_loss(z1, z2, lambda_off_diagonal=0.005):
    """Barlow Twins style correlation loss for dual-view alignment"""
    N, D = z1.shape

    z1 = (z1 - z1.mean(dim=0)) / (z1.std(dim=0) + 1e-8)
    z2 = (z2 - z2.mean(dim=0)) / (z2.std(dim=0) + 1e-8)

    C = torch.matmul(z1.T, z2) / N

    diagonal_loss = torch.pow(1.0 - torch.diagonal(C), 2).sum()
    off_diagonal_mask = ~torch.eye(D, dtype=bool, device=z1.device)
    off_diagonal_loss = torch.pow(C[off_diagonal_mask], 2).sum()

    return diagonal_loss + lambda_off_diagonal * off_diagonal_loss


######################################## Graph Reconstruction Loss ########################################

def graph_reconstruction_loss(embeddings, adj_matrix, normalize=True):
    """Graph structure reconstruction loss"""
    embeddings_norm = F.normalize(embeddings, p=2, dim=1)
    cosine_sim = torch.matmul(embeddings_norm, embeddings_norm.T)

    if normalize:
        cosine_sim = (cosine_sim + 1.0) / 2.0

    if adj_matrix.is_sparse:
        adj_dense = adj_matrix.to_dense()
    else:
        adj_dense = adj_matrix

    return F.mse_loss(cosine_sim, adj_dense)