CLCNet：面向植物基因组预测的对比学习与染色体感知网络¶

输入: 全基因组SNP矩阵
      │
      ├─> 染色体1编码器 ──> 染色体1嵌入
      ├─> 染色体2编码器 ──> 染色体2嵌入
      │   ...
      └─> 染色体k编码器 ──> 染色体k嵌入
                │
         [拼接/注意力融合]
                │
         个体最终嵌入向量
            /        \
      投影头（对比学习）  预测头（表型回归）
      对比损失 ←──┘        └──→ 预测损失

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

# ------------------------- 染色体感知编码器 -------------------------
class ChromosomeEncoder(nn.Module):
    """单条染色体的SNP特征编码器"""
    def __init__(self, input_dim, hidden_dim, embed_dim, dropout=0.2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.ReLU()
        )
    def forward(self, x):
        # x: (batch_size, input_dim)  某条染色体的SNP数据
        return self.net(x)

class CLCNet(nn.Module):
    """CLCNet主干：分染色体编码 + 注意力融合 + 投影头 + 预测头"""
    def __init__(self, chrom_snp_dims, hidden_dim=128, embed_dim=32, fuse_dim=64, proj_dim=32, dropout=0.2):
        """
        chrom_snp_dims: list[int]，每条染色体上的SNP数量
        """
        super().__init__()
        self.num_chrom = len(chrom_snp_dims)
        # 为每条染色体创建编码器
        self.encoders = nn.ModuleList([
            ChromosomeEncoder(in_dim, hidden_dim, embed_dim, dropout)
            for in_dim in chrom_snp_dims
        ])
        # 染色体注意力权重（可学习）
        self.chrom_att = nn.Parameter(torch.ones(self.num_chrom))
        # 融合层
        self.fuse_fc = nn.Linear(embed_dim, fuse_dim)
        # 预测头
        self.pred_head = nn.Sequential(
            nn.Linear(fuse_dim, 32),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(32, 1)   # 输出育种值
        )
        # 对比投影头
        self.proj_head = nn.Sequential(
            nn.Linear(fuse_dim, 32),
            nn.ReLU(),
            nn.Linear(32, proj_dim)
        )

    def forward(self, chrom_inputs):
        """
        chrom_inputs: list of Tensors, 每个形状为 (batch_size, chrom_snp_dim)
        返回: 预测值, 对比投影向量, 融合嵌入（可用于分析）
        """
        # 各染色体独立编码
        chrom_embeds = []
        for i, enc in enumerate(self.encoders):
            h = enc(chrom_inputs[i])          # (batch_size, embed_dim)
            chrom_embeds.append(h)
        # 堆叠为 (batch_size, num_chrom, embed_dim)
        stacked = torch.stack(chrom_embeds, dim=1)
        # 注意力融合
        att_weights = F.softmax(self.chrom_att, dim=0)  # (num_chrom,)
        fused = torch.sum(stacked * att_weights.view(1, -1, 1), dim=1)  # (batch_size, embed_dim)
        fused = self.fuse_fc(fused)                     # (batch_size, fuse_dim)

        # 预测值与投影向量
        pred = self.pred_head(fused).squeeze(-1)        # (batch_size,)
        proj = self.proj_head(fused)                    # (batch_size, proj_dim)
        return pred, proj, fused

# ------------------------- 对比损失 (NT-Xent) -------------------------
def nt_xent_loss(proj, labels, tau=0.5):
    """
    基于表型标签的对比损失（简化版）
    proj: (batch_size, proj_dim) 标准化投影
    labels: (batch_size,) 连续表型值，用于构建正负对（按分位数分组）
    """
    batch_size = proj.shape[0]
    # 将表型标签离散化为组（如按中位数分两组）
    groups = (labels > labels.median()).long()  # 高/低两组
    # 计算相似度矩阵
    proj = F.normalize(proj, dim=1)
    sim = torch.matmul(proj, proj.T) / tau   # (bs, bs)
    # 构建正样本mask：同组为1，不同组为0（忽略对角线）
    pos_mask = (groups.unsqueeze(0) == groups.unsqueeze(1)).float()
    pos_mask.fill_diagonal_(0)
    # 指数化
    exp_sim = torch.exp(sim)
    # 分母：所有非自身的exp和
    denom = exp_sim.sum(dim=1) - exp_sim.diag().sum(dim=0)
    # 分子：正样本对的exp和
    pos_sum = (exp_sim * pos_mask).sum(dim=1)
    loss = -torch.log(pos_sum / denom).mean()
    return loss

# ------------------------- 训练示例 -------------------------
def train_step(model, chrom_inputs, phenotypes, optimizer, lambda_con=0.1):
    model.train()
    optimizer.zero_grad()
    pred, proj, _ = model(chrom_inputs)
    # 预测损失
    pred_loss = F.mse_loss(pred, phenotypes)
    # 对比损失
    con_loss = nt_xent_loss(proj, phenotypes)
    total_loss = pred_loss + lambda_con * con_loss
    total_loss.backward()
    optimizer.step()
    return pred_loss.item(), con_loss.item(), total_loss.item()

# ------------------------- 使用示例 -------------------------
if __name__ == "__main__":
    # 假设5条染色体，SNP数量分别为2000, 1500, 1800, 2200, 1700
    chrom_dims = [2000, 1500, 1800, 2200, 1700]
    model = CLCNet(chrom_dims)
    batch_size = 64
    # 模拟一个batch的输入
    chrom_inputs = [torch.randn(batch_size, dim) for dim in chrom_dims]
    phenotypes = torch.randn(batch_size)  # 模拟表型
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(10):
        p_loss, c_loss, total = train_step(model, chrom_inputs, phenotypes, optimizer)
        print(f"Epoch {epoch+1}: Pred Loss={p_loss:.4f}, Con Loss={c_loss:.4f}, Total={total:.4f}")