fastplms/models/boltz/vb_modules_trunkv2.py

"""Input, template, MSA, and prediction-head modules for the Boltz2 trunk."""

from __future__ import annotations

import torch
from typing import Any, cast
from torch import Tensor, nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint

from . import vb_const as const
from .vb_layers_dropout import get_dropout_mask
from .vb_layers_outer_product_mean import OuterProductMean
from .vb_layers_pair_averaging import PairWeightedAveraging
from .vb_layers_pairformer import PairformerNoSeqLayer, PairformerNoSeqModule
from .vb_layers_transition import Transition
from .vb_modules_encodersv2 import (
    AtomAttentionEncoder,
    AtomEncoder,
    FourierEmbedding,
)


class ContactConditioning(nn.Module):
    """Encode selected, unselected, and distance-threshold contact constraints."""

    def __init__(self, token_z: int, cutoff_min: float, cutoff_max: float) -> None:
        super().__init__()
        self.fourier_embedding = FourierEmbedding(token_z)
        input_width = token_z + len(const.contact_conditioning_info) - 1
        self.encoder = nn.Linear(input_width, token_z)
        self.encoding_unspecified = nn.Parameter(torch.zeros(token_z))  # (d_z,)
        self.encoding_unselected = nn.Parameter(torch.zeros(token_z))  # (d_z,)
        self.cutoff_min = cutoff_min
        self.cutoff_max = cutoff_max

    def forward(self, feats: dict[str, Tensor]) -> Tensor:
        """Return contact tensor C with shape ``(b, l, l, d_z)``."""

        if const.contact_conditioning_info["UNSPECIFIED"] != 0:
            raise ValueError("UNSPECIFIED contact conditioning must use channel zero")
        if const.contact_conditioning_info["UNSELECTED"] != 1:
            raise ValueError("UNSELECTED contact conditioning must use channel one")

        # c_contact is the number of contact-conditioning categories.
        categories = feats["contact_conditioning"]  # (b, l, l, c_contact)
        threshold = feats["contact_threshold"]  # (b, l, l)
        # (b, l, l)
        normalized = (threshold - self.cutoff_min) / (self.cutoff_max - self.cutoff_min)
        # (b, l, l, d_z)
        fourier = self.fourier_embedding(normalized.flatten()).reshape((*normalized.shape, -1))
        selected_features = torch.cat(  # (b, l, l, c_contact - 1 + d_z)
            [categories[..., 2:], normalized.unsqueeze(-1), fourier],
            dim=-1,
        )
        selected = self.encoder(selected_features)  # (b, l, l, d_z)
        special = categories[..., :2]  # (b, l, l, 2)
        return cast(
            Tensor,
            selected * (1 - special.sum(dim=-1, keepdim=True))
            + self.encoding_unspecified * special[..., 0:1]
            + self.encoding_unselected * special[..., 1:2],
        )  # (b, l, l, d_z)


class InputEmbedder(nn.Module):
    """Combine atom, residue, profile, and optional experimental features."""

    def __init__(
        self,
        atom_s: int,
        atom_z: int,
        token_s: int,
        token_z: int,
        atoms_per_window_queries: int,
        atoms_per_window_keys: int,
        atom_feature_dim: int,
        atom_encoder_depth: int,
        atom_encoder_heads: int,
        activation_checkpointing: bool = False,
        add_method_conditioning: bool = False,
        add_modified_flag: bool = False,
        add_cyclic_flag: bool = False,
        add_mol_type_feat: bool = False,
        use_no_atom_char: bool = False,
        use_atom_backbone_feat: bool = False,
        use_residue_feats_atoms: bool = False,
    ) -> None:
        super().__init__()
        self.token_s = token_s
        self.add_method_conditioning = add_method_conditioning
        self.add_modified_flag = add_modified_flag
        self.add_cyclic_flag = add_cyclic_flag
        self.add_mol_type_feat = add_mol_type_feat
        self.atom_encoder = AtomEncoder(
            atom_s=atom_s,
            atom_z=atom_z,
            token_s=token_s,
            token_z=token_z,
            atoms_per_window_queries=atoms_per_window_queries,
            atoms_per_window_keys=atoms_per_window_keys,
            atom_feature_dim=atom_feature_dim,
            structure_prediction=False,
            use_no_atom_char=use_no_atom_char,
            use_atom_backbone_feat=use_atom_backbone_feat,
            use_residue_feats_atoms=use_residue_feats_atoms,
        )
        self.atom_enc_proj_z = nn.Sequential(
            nn.LayerNorm(atom_z),
            nn.Linear(atom_z, atom_encoder_depth * atom_encoder_heads, bias=False),
        )
        self.atom_attention_encoder = AtomAttentionEncoder(
            atom_s=atom_s,
            token_s=token_s,
            atoms_per_window_queries=atoms_per_window_queries,
            atoms_per_window_keys=atoms_per_window_keys,
            atom_encoder_depth=atom_encoder_depth,
            atom_encoder_heads=atom_encoder_heads,
            structure_prediction=False,
            activation_checkpointing=activation_checkpointing,
        )
        self.res_type_encoding = nn.Linear(const.num_tokens, token_s, bias=False)
        self.msa_profile_encoding = nn.Linear(const.num_tokens + 1, token_s, bias=False)

        if add_method_conditioning:
            self.method_conditioning_init = nn.Embedding(const.num_method_types, token_s)
            self.method_conditioning_init.weight.data.fill_(0)  # (n_method, d_s)
        if add_modified_flag:
            self.modified_conditioning_init = nn.Embedding(2, token_s)
            self.modified_conditioning_init.weight.data.fill_(0)  # (2, d_s)
        if add_cyclic_flag:
            self.cyclic_conditioning_init = nn.Linear(1, token_s, bias=False)
            self.cyclic_conditioning_init.weight.data.fill_(0)  # (d_s, 1)
        if add_mol_type_feat:
            self.mol_type_conditioning_init = nn.Embedding(
                len(const.chain_type_ids),
                token_s,
            )
            self.mol_type_conditioning_init.weight.data.fill_(0)  # (n_mol_type, d_s)

    def forward(self, feats: dict[str, Tensor], affinity: bool = False) -> Tensor:
        """Return embedded sequence tensor S with shape ``(b, l, d_s)``."""

        # n_atom is the padded atom count; k is the number of atom windows.
        residue_type = feats["res_type"].float()  # (b, l, n_token_type)
        suffix = "_affinity" if affinity else ""
        profile = feats[f"profile{suffix}"]  # (b, l, n_token_type)
        deletion_mean = feats[f"deletion_mean{suffix}"].unsqueeze(-1)  # (b, l, 1)
        # (b, n_atom, d_a), (b, n_atom, d_a), (b, k, w_q, w_k, d_az), callable
        atom_queries, atom_conditioning, atom_pairs, to_keys = self.atom_encoder(feats)
        atom_bias = self.atom_enc_proj_z(atom_pairs)  # (b, k, w_q, w_k, n_layer * h)
        atom_output, _, _, _ = self.atom_attention_encoder(
            feats=feats,
            q=atom_queries,
            c=atom_conditioning,
            atom_enc_bias=atom_bias,
            to_keys=to_keys,
        )  # (b, l, d_s), (b, n_atom, d_a), (b, n_atom, d_a), callable
        output = (  # (b, l, d_s)
            atom_output
            + self.res_type_encoding(residue_type)
            + self.msa_profile_encoding(torch.cat([profile, deletion_mean], dim=-1))
        )
        if self.add_method_conditioning:
            # method_feature: (b, l); output: (b, l, d_s).
            output = output + self.method_conditioning_init(feats["method_feature"])
        if self.add_modified_flag:
            # modified: (b, l); output: (b, l, d_s).
            output = output + self.modified_conditioning_init(feats["modified"])
        if self.add_cyclic_flag:
            cyclic = feats["cyclic_period"].clamp(max=1.0).unsqueeze(-1)  # (b, l, 1)
            output = output + self.cyclic_conditioning_init(cyclic)  # (b, l, d_s)
        if self.add_mol_type_feat:
            # mol_type: (b, l); output: (b, l, d_s).
            output = output + self.mol_type_conditioning_init(feats["mol_type"])
        return cast(Tensor, output)  # (b, l, d_s)


class _TemplateBase(nn.Module):
    def __init__(
        self,
        token_z: int,
        template_dim: int,
        template_blocks: int,
        dropout: float,
        pairwise_head_width: int,
        pairwise_num_heads: int,
        post_layer_norm: bool,
        activation_checkpointing: bool,
        min_dist: float,
        max_dist: float,
        num_bins: int,
    ) -> None:
        super().__init__()
        self.min_dist = min_dist
        self.max_dist = max_dist
        self.num_bins = num_bins
        self.relu = nn.ReLU()
        self.z_norm = nn.LayerNorm(token_z)
        self.v_norm = nn.LayerNorm(template_dim)
        self.z_proj = nn.Linear(token_z, template_dim, bias=False)
        feature_width = const.num_tokens * 2 + num_bins + 5
        self.a_proj = nn.Linear(feature_width, template_dim, bias=False)
        self.u_proj = nn.Linear(template_dim, token_z, bias=False)
        self.pairformer = PairformerNoSeqModule(
            template_dim,
            num_blocks=template_blocks,
            dropout=dropout,
            pairwise_head_width=pairwise_head_width,
            pairwise_num_heads=pairwise_num_heads,
            post_layer_norm=post_layer_norm,
            activation_checkpointing=activation_checkpointing,
        )

    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
        raise NotImplementedError

    def _template_features(
        self,
        feats: dict[str, Tensor],
        template_pair_mask: Tensor,
    ) -> Tensor:
        # residue_type: (b, t, l, n_token_type); template_pair_mask: (b, t, l, l).
        residue_type = feats["template_restype"]  # (b, t, l, n_token_type)
        cb_mask = feats["template_mask_cb"]  # (b, t, l)
        frame_mask = feats["template_mask_frame"]  # (b, t, l)
        # (b, t, l, l, 1)
        cb_pair_mask = (cb_mask[..., :, None] * cb_mask[..., None, :]).unsqueeze(-1)
        # (b, t, l, l, 1)
        frame_pair_mask = (frame_mask[..., :, None] * frame_mask[..., None, :]).unsqueeze(-1)
        with torch.autocast(device_type="cuda", enabled=False):
            # template_cb: (b, t, l, 3).
            # (b, t, l, l)
            cb_distances = torch.cdist(feats["template_cb"], feats["template_cb"])
            boundaries = torch.linspace(  # (n_bin - 1,)
                self.min_dist,
                self.max_dist,
                self.num_bins - 1,
            ).to(cb_distances.device)
            bins = (cb_distances[..., None] > boundaries).sum(dim=-1).long()  # (b, t, l, l)
            distogram = F.one_hot(bins, num_classes=self.num_bins)  # (b, t, l, l, n_bin)

            # (b, t, 1, l, 3, 3)
            rotations = feats["template_frame_rot"].unsqueeze(2).transpose(-1, -2)
            # (b, t, 1, l, 3, 1)
            translations = feats["template_frame_t"].unsqueeze(2).unsqueeze(-1)
            # (b, t, l, 1, 3, 1)
            ca_coordinates = feats["template_ca"].unsqueeze(3).unsqueeze(-1)
            # (b, t, l, l, 3, 1)
            vectors = torch.matmul(rotations, ca_coordinates - translations)
            norms = torch.norm(vectors, dim=-1, keepdim=True)  # (b, t, l, l, 3, 1)
            unit_vectors = torch.where(  # (b, t, l, l, 3)
                norms > 0,
                vectors / norms,
                torch.zeros_like(vectors),
            ).squeeze(-1)
            pair_features = torch.cat(  # (b, t, l, l, n_bin + 5)
                [distogram, cb_pair_mask, unit_vectors, frame_pair_mask],
                dim=-1,
            )
            # (b, t, l, l, n_bin + 5)
            pair_features = pair_features * template_pair_mask.unsqueeze(-1)
            residue_i = residue_type[:, :, :, None].expand(  # (b, t, l, l, n_token_type)
                -1,
                -1,
                -1,
                residue_type.size(2),
                -1,
            )
            residue_j = residue_type[:, :, None, :].expand(  # (b, t, l, l, n_token_type)
                -1,
                -1,
                residue_type.size(2),
                -1,
                -1,
            )
            return cast(
                Tensor,
                self.a_proj(torch.cat([pair_features, residue_i, residue_j], dim=-1)),
            )  # (b, t, l, l, d_t)

    def forward(
        self,
        z: Tensor,
        feats: dict[str, Tensor],
        pair_mask: Tensor,
        use_kernels: bool = False,
    ) -> Tensor:
        """Aggregate template pair tensor V into trunk update U."""

        # z: (b, l, l, d_z); pair_mask: (b, l, l).
        residue_type = feats["template_restype"]  # (b, t, l, n_token_type)
        batch_size, template_count = residue_type.shape[:2]
        template_present = feats["template_mask"].any(dim=2).float()  # (b, t)
        present_count = template_present.sum(dim=1).clamp(min=1)  # (b,)
        features = self._template_features(  # (b, t, l, l, d_t)
            feats,
            self._template_pair_mask(feats, template_count),
        )
        expanded_mask = pair_mask[:, None].expand(-1, template_count, -1, -1)  # (b, t, l, l)
        expanded_mask = expanded_mask.reshape(  # (b * t, l, l)
            batch_size * template_count,
            *expanded_mask.shape[2:],
        )
        template_states = self.z_proj(self.z_norm(z[:, None])) + features  # (b, t, l, l, d_t)
        template_states = template_states.view(  # (b * t, l, l, d_t)
            batch_size * template_count,
            *template_states.shape[2:],
        )
        template_states = template_states + self.pairformer(  # (b * t, l, l, d_t)
            template_states,
            expanded_mask,
            use_kernels=use_kernels,
        )
        template_states = self.v_norm(template_states).view(  # (b, t, l, l, d_t)
            batch_size,
            template_count,
            *template_states.shape[1:],
        )
        weights = template_present[:, :, None, None, None]  # (b, t, 1, 1, 1)
        aggregate = (template_states * weights).sum(dim=1)  # (b, l, l, d_t)
        # (b, l, l, d_t)
        aggregate = aggregate / present_count[:, None, None, None].to(template_states)
        return cast(Tensor, self.u_proj(self.relu(aggregate)))  # (b, l, l, d_z)


class TemplateModule(_TemplateBase):
    """Aggregate templates while restricting features to the same chain."""

    def __init__(
        self,
        token_z: int,
        template_dim: int,
        template_blocks: int,
        dropout: float = 0.25,
        pairwise_head_width: int = 32,
        pairwise_num_heads: int = 4,
        post_layer_norm: bool = False,
        activation_checkpointing: bool = False,
        min_dist: float = 3.25,
        max_dist: float = 50.75,
        num_bins: int = 38,
        **kwargs: Any,
    ) -> None:
        del kwargs
        super().__init__(
            token_z,
            template_dim,
            template_blocks,
            dropout,
            pairwise_head_width,
            pairwise_num_heads,
            post_layer_norm,
            activation_checkpointing,
            min_dist,
            max_dist,
            num_bins,
        )

    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
        asym_id = feats["asym_id"]  # (b, l)
        same_chain = (asym_id[:, :, None] == asym_id[:, None, :]).float()  # (b, l, l)
        return same_chain[:, None].expand(-1, count, -1, -1)  # (b, t, l, l)


class TemplateV2Module(_TemplateBase):
    """Aggregate templates under per-template visibility groups."""

    def __init__(
        self,
        token_z: int,
        template_dim: int,
        template_blocks: int,
        dropout: float = 0.25,
        pairwise_head_width: int = 32,
        pairwise_num_heads: int = 4,
        post_layer_norm: bool = False,
        activation_checkpointing: bool = False,
        min_dist: float = 3.25,
        max_dist: float = 50.75,
        num_bins: int = 38,
        **kwargs: Any,
    ) -> None:
        del kwargs
        super().__init__(
            token_z,
            template_dim,
            template_blocks,
            dropout,
            pairwise_head_width,
            pairwise_num_heads,
            post_layer_norm,
            activation_checkpointing,
            min_dist,
            max_dist,
            num_bins,
        )

    def _template_pair_mask(self, feats: dict[str, Tensor], count: int) -> Tensor:
        del count
        visibility = feats["visibility_ids"]  # (b, t, l)
        return (visibility[..., :, None] == visibility[..., None, :]).float()  # (b, t, l, l)


class MSAModule(nn.Module):
    """Embed and update an MSA before returning its accumulated pair update."""

    def __init__(
        self,
        msa_s: int,
        token_z: int,
        token_s: int,
        msa_blocks: int,
        msa_dropout: float,
        z_dropout: float,
        pairwise_head_width: int = 32,
        pairwise_num_heads: int = 4,
        activation_checkpointing: bool = False,
        use_paired_feature: bool = True,
        subsample_msa: bool = False,
        num_subsampled_msa: int = 1024,
        **kwargs: Any,
    ) -> None:
        del kwargs
        super().__init__()
        self.msa_blocks = msa_blocks
        self.msa_dropout = msa_dropout
        self.z_dropout = z_dropout
        self.use_paired_feature = use_paired_feature
        self.activation_checkpointing = activation_checkpointing
        self.subsample_msa = subsample_msa
        self.num_subsampled_msa = num_subsampled_msa
        self.s_proj = nn.Linear(token_s, msa_s, bias=False)
        input_width = const.num_tokens + 2 + int(use_paired_feature)
        self.msa_proj = nn.Linear(input_width, msa_s, bias=False)
        self.layers = nn.ModuleList(
            [
                MSALayer(
                    msa_s,
                    token_z,
                    msa_dropout,
                    z_dropout,
                    pairwise_head_width,
                    pairwise_num_heads,
                )
                for _ in range(msa_blocks)
            ]
        )

    @staticmethod
    def _chunk_configuration(
        pair_states: Tensor,
        training: bool,
    ) -> tuple[bool, int | None, int | None, int | None, int | None]:
        # pair_states: (b, l, l, d_z).
        if training:
            return False, None, None, None, None
        if pair_states.shape[1] > const.chunk_size_threshold:
            return True, 64, 32, 4, 128
        return False, None, None, None, 512

    def forward(
        self,
        z: Tensor,
        emb: Tensor,
        feats: dict[str, Tensor],
        use_kernels: bool = False,
    ) -> Tensor:
        """Return updated pair tensor Z after every MSA block."""

        chunking = self._chunk_configuration(z, self.training)
        # z: (b, l, l, d_z); emb: (b, l, d_s).
        # s is MSA depth; n_token_type is the residue vocabulary size.
        msa = feats["msa"]  # (b, s, l) or (b, s, l, n_token_type)
        if msa.dtype in (torch.long, torch.int32, torch.int64):
            msa = F.one_hot(msa, num_classes=const.num_tokens).float()  # (b, s, l, n_token_type)
        msa_mask = feats["msa_mask"]  # (b, s, l)
        components = [  # (b, s, l, n_token_type), then two (b, s, l, 1) tensors
            msa,
            feats["has_deletion"].unsqueeze(-1),  # (b, s, l, 1)
            feats["deletion_value"].unsqueeze(-1),  # (b, s, l, 1)
        ]
        if self.use_paired_feature:
            components.append(feats["msa_paired"].unsqueeze(-1))  # (b, s, l, 1)
        msa_input = torch.cat(components, dim=-1)  # (b, s, l, n_token_type + 2 or 3)
        if self.subsample_msa:
            indices = torch.randperm(msa.shape[1])[: self.num_subsampled_msa]  # (s_sub,)
            msa_input = msa_input[:, indices]  # (b, s_sub, l, n_token_type + 2 or 3)
            msa_mask = msa_mask[:, indices]  # (b, s_sub, l)

        msa_states = self.msa_proj(msa_input) + self.s_proj(emb).unsqueeze(1)  # (b, s, l, d_m)
        token_mask = feats["token_pad_mask"].float()  # (b, l)
        pair_mask = token_mask[:, :, None] * token_mask[:, None, :]  # (b, l, l)
        pair_states = z  # (b, l, l, d_z)
        for layer in self.layers:
            # Tensor arguments: pair_states (b, l, l, d_z), msa_states (b, s, l, d_m),
            # pair_mask (b, l, l), msa_mask (b, s, l).
            arguments = (
                pair_states,
                msa_states,
                pair_mask,
                msa_mask,
                *chunking,
                use_kernels,
            )
            if self.activation_checkpointing and self.training:
                pair_states, msa_states = checkpoint(  # (b, l, l, d_z), (b, s, l, d_m)
                    layer,
                    *arguments,
                )
            else:
                pair_states, msa_states = layer(*arguments)  # (b, l, l, d_z), (b, s, l, d_m)
        return pair_states  # (b, l, l, d_z)


class MSALayer(nn.Module):
    """Exchange information between MSA tensor M and pair tensor Z."""

    def __init__(
        self,
        msa_s: int,
        token_z: int,
        msa_dropout: float,
        z_dropout: float,
        pairwise_head_width: int = 32,
        pairwise_num_heads: int = 4,
    ) -> None:
        super().__init__()
        self.msa_dropout = msa_dropout
        self.msa_transition = Transition(msa_s, msa_s * 4)
        self.pair_weighted_averaging = PairWeightedAveraging(
            c_m=msa_s,
            c_z=token_z,
            c_h=32,
            num_heads=8,
        )
        self.pairformer_layer = PairformerNoSeqLayer(
            token_z=token_z,
            dropout=z_dropout,
            pairwise_head_width=pairwise_head_width,
            pairwise_num_heads=pairwise_num_heads,
        )
        self.outer_product_mean = OuterProductMean(msa_s, 32, token_z)

    def forward(
        self,
        z: Tensor,
        m: Tensor,
        token_mask: Tensor,
        msa_mask: Tensor,
        chunk_heads_pwa: bool = False,
        chunk_size_transition_z: int | None = None,
        chunk_size_transition_msa: int | None = None,
        chunk_size_outer_product: int | None = None,
        chunk_size_tri_attn: int | None = None,
        use_kernels: bool = False,
    ) -> tuple[Tensor, Tensor]:
        """Return updated Z and M tensors."""

        del chunk_size_transition_z
        # z: (b, l, l, d_z); m: (b, s, l, d_m).
        # token_mask: (b, l, l); msa_mask: (b, s, l).
        dropout = get_dropout_mask(self.msa_dropout, m, self.training)  # (b, s, 1, 1)
        msa_states = m + dropout * self.pair_weighted_averaging(  # (b, s, l, d_m)
            m,
            z,
            token_mask,
            chunk_heads_pwa,
        )
        msa_states = msa_states + self.msa_transition(  # (b, s, l, d_m)
            msa_states,
            chunk_size_transition_msa,
        )
        pair_states = z + self.outer_product_mean(  # (b, l, l, d_z)
            msa_states,
            msa_mask,
            chunk_size_outer_product,
        )
        pair_states = self.pairformer_layer(  # (b, l, l, d_z)
            pair_states,
            token_mask,
            chunk_size_tri_attn,
            use_kernels=use_kernels,
        )
        return pair_states, msa_states  # (b, l, l, d_z), (b, s, l, d_m)


class BFactorModule(nn.Module):
    """Predict a per-token B-factor histogram."""

    def __init__(self, token_s: int, num_bins: int) -> None:
        super().__init__()
        self.bfactor = nn.Linear(token_s, num_bins)
        self.num_bins = num_bins

    def forward(self, s: Tensor) -> Tensor:
        # s: (..., d_s).
        return cast(Tensor, self.bfactor(s))  # (..., n_bin)


class DistogramModule(nn.Module):
    """Predict symmetric residue-pair distance histograms."""

    def __init__(self, token_z: int, num_bins: int, num_distograms: int = 1) -> None:
        super().__init__()
        self.distogram = nn.Linear(token_z, num_distograms * num_bins)
        self.num_distograms = num_distograms
        self.num_bins = num_bins

    def forward(self, z: Tensor) -> Tensor:
        # z: (b, l, l, d_z).
        symmetric = z + z.transpose(1, 2)  # (b, l, l, d_z)
        logits = self.distogram(symmetric)  # (b, l, l, n_distogram * n_bin)
        return cast(
            Tensor,
            logits.reshape(
                symmetric.shape[0],
                symmetric.shape[1],
                symmetric.shape[2],
                self.num_distograms,
                self.num_bins,
            ),
        )  # (b, l, l, n_distogram, n_bin)