变分自编码器 VAE¶

L = E[log p(x|z)] - KL(q(z|x) || p(z))
      ↑ 重建质量         ↑ 潜空间正则化

z = μ + ε × σ，其中 ε ~ N(0,1)

import torch
import torch.nn as nn
import torch.nn.functional as F

# ===== 单细胞基因表达 VAE（参考 scVI 架构）=====

class VAE(nn.Module):
    """
    变分自编码器用于单细胞基因表达数据
    输入：细胞 × 基因 矩阵（归一化后的基因表达值）
    潜空间：低维连续表示，可用于聚类/可视化
    """
    def __init__(self, input_dim=3000, latent_dim=10, hidden_dim=128):
        super().__init__()
        # === 编码器 ===
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),   # 基因维度 -> 隐藏层
            nn.BatchNorm1d(hidden_dim),          # 批归一化
            nn.ReLU()
        )
        # 编码为均值和对数方差（各一个线性层）
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)      # 潜变量均值 μ
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)  # 对数方差 log(σ²)

        # === 解码器 ===
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),  # 潜变量 -> 隐藏层
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),   # 隐藏层 -> 基因维度
            nn.Sigmoid()                         # 输出 [0,1]（归一化表达值）
        )

    def encode(self, x):
        """编码：输入基因表达，输出均值和对数方差"""
        h = self.encoder(x)
        mu = self.fc_mu(h)          # 均值向量
        logvar = self.fc_logvar(h)  # 对数方差向量
        return mu, logvar

    def reparameterize(self, mu, logvar):
        """重参数化技巧：z = μ + ε × σ，保证梯度可传播"""
        if self.training:
            std = torch.exp(0.5 * logvar)   # σ = exp(0.5 * logvar)
            eps = torch.randn_like(std)      # ε ~ N(0,1) 标准正态随机数
            return mu + eps * std            # 采样 z
        else:
            return mu  # 推断时直接用均值（确定性）

    def decode(self, z):
        """解码：从潜变量重建基因表达"""
        return self.decoder(z)

    def forward(self, x):
        mu, logvar = self.encode(x)          # 编码
        z = self.reparameterize(mu, logvar)  # 采样潜变量
        x_recon = self.decode(z)             # 解码重建
        return x_recon, mu, logvar

# ===== 损失函数 =====
def vae_loss(x_recon, x, mu, logvar, beta=1.0):
    """
    VAE 损失 = 重建损失 + β × KL 散度
    β=1：标准 VAE；β>1：β-VAE，潜空间更解耦
    """
    # 重建损失（均方误差，适合连续值）
    recon_loss = F.mse_loss(x_recon, x, reduction='sum')

    # KL 散度：让 q(z|x) 接近 N(0,1)
    # 公式：-0.5 * sum(1 + logvar - μ² - e^logvar)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon_loss + beta * kl_loss

# ===== 训练示例 =====
input_dim = 3000   # 高变基因数量
latent_dim = 10    # 潜空间维度（降维后的表示）
batch_size = 32

model = VAE(input_dim=input_dim, latent_dim=latent_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 模拟单细胞数据（实际替换为 scanpy 加载的归一化矩阵）
fake_data = torch.rand(500, input_dim)  # 500 个细胞

for epoch in range(30):
    idx = torch.randint(0, 500, (batch_size,))
    x = fake_data[idx]

    x_recon, mu, logvar = model(x)               # 前向传播
    loss = vae_loss(x_recon, x, mu, logvar)       # 计算损失

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.1f}')

# ===== 获取细胞嵌入（用于聚类/UMAP）=====
model.eval()
with torch.no_grad():
    mu, _ = model.encode(fake_data)  # 取均值作为细胞的低维表示
    # mu 形状: [500, 10]，可直接用于 UMAP、KMeans 等
    print(f'细胞嵌入维度: {mu.shape}')