AlphaFold 蛋白质结构预测

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

class MSARowAttention(nn.Module):
    """MSA 行注意力：沿序列维度做注意力"""
    def __init__(self, d_model=64, n_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, msa):
        # msa: (batch, N_seq, N_res, d_model)
        B, S, R, D = msa.shape
        # 对每个 MSA 序列，沿残基维度做注意力
        qkv = self.qkv(msa).reshape(B, S, R, 3, self.n_heads, self.d_k)
        q, k, v = qkv.unbind(3)  # 每个 (B, S, R, heads, d_k)

        # 沿残基维度做注意力 (R 维度)
        scores = torch.einsum('bsihd,bsjhd->bsij', q, k) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        out = torch.einsum('bsij,bsjhd->bsihd', attn, v)
        return self.out(out.reshape(B, S, R, D))


class MSAColumnAttention(nn.Module):
    """MSA 列注意力：沿 MSA 维度做注意力"""
    def __init__(self, d_model=64, n_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, msa):
        # msa: (batch, N_seq, N_res, d_model)
        B, S, R, D = msa.shape
        msa_t = msa.transpose(1, 2)  # (B, R, S, D) — 换轴
        qkv = self.qkv(msa_t).reshape(B, R, S, 3, self.n_heads, self.d_k)
        q, k, v = qkv.unbind(3)

        # 沿 MSA 序列维度做注意力 (S 维度)
        scores = torch.einsum('brihd,brjhd->brij', q, k) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        out = torch.einsum('brij,brjhd->brihd', attn, v)
        return self.out(out.reshape(B, R, S, D)).transpose(1, 2)


class OuterProductMean(nn.Module):
    """MSA → 配对表示：外积均值"""
    def __init__(self, d_msa=64, d_pair=128, d_hidden=32):
        super().__init__()
        self.linear_a = nn.Linear(d_msa, d_hidden)
        self.linear_b = nn.Linear(d_msa, d_hidden)
        self.linear_out = nn.Linear(d_hidden * d_hidden, d_pair)

    def forward(self, msa):
        # msa: (B, S, R, d_msa)
        a = self.linear_a(msa)  # (B, S, R, d_hidden)
        b = self.linear_b(msa)  # (B, S, R, d_hidden)
        # 外积沿 MSA 维度求均值 → (B, R, R, d_h, d_h)
        outer = torch.einsum('bsid,bsjd->bijcd', a, b)
        outer = outer.mean(dim=1)  # 平均 MSA 维度
        B, R1, R2, H1, H2 = outer.shape
        return self.linear_out(outer.reshape(B, R1, R2, H1 * H2))


class EvoformerBlock(nn.Module):
    """一个 Evoformer Block"""
    def __init__(self, d_msa=64, d_pair=128, n_heads=8):
        super().__init__()
        self.msa_row_attn = MSARowAttention(d_msa, n_heads)
        self.msa_col_attn = MSAColumnAttention(d_msa, n_heads)
        self.outer_product = OuterProductMean(d_msa, d_pair)
        self.pair_update = nn.Sequential(
            nn.Linear(d_pair, d_pair),
            nn.ReLU(),
            nn.Linear(d_pair, d_pair),
        )
        self.norm_msa = nn.LayerNorm(d_msa)
        self.norm_pair = nn.LayerNorm(d_pair)

    def forward(self, msa_repr, pair_repr):
        # MSA 处理
        msa_repr = msa_repr + self.msa_row_attn(msa_repr)
        msa_repr = msa_repr + self.msa_col_attn(msa_repr)
        msa_repr = self.norm_msa(msa_repr)

        # MSA → 配对表示
        pair_repr = pair_repr + self.outer_product(msa_repr)
        pair_repr = pair_repr + self.pair_update(pair_repr)
        pair_repr = self.norm_pair(pair_repr)

        return msa_repr, pair_repr


# 使用示例
block = EvoformerBlock(d_msa=64, d_pair=128)
msa = torch.randn(1, 16, 50, 64)   # 16 条序列, 50 个残基
pair = torch.randn(1, 50, 50, 128) # 50×50 配对矩阵
msa_out, pair_out = block(msa, pair)
print(f"MSA: {msa_out.shape}, Pair: {pair_out.shape}")