rc-foundry 0.1.1__py3-none-any.whl

rf3/model/RF3_structure.py

@@ -0,0 +1,303 @@
+import logging
+
+import torch
+import torch.nn as nn
+from rf3.model.layers.af3_diffusion_transformer import (
+    AtomAttentionEncoderDiffusion,
+    AtomTransformer,
+    DiffusionTransformer,
+)
+from rf3.model.layers.layer_utils import Transition, linearNoBias
+from rf3.model.layers.pairformer_layers import (
+    MSAModule,
+    PairformerBlock,
+    RelativePositionEncoding,
+    RF3TemplateEmbedder,
+)
+
+from foundry.model.layers.blocks import FourierEmbedding
+from foundry.training.checkpoint import activation_checkpointing
+
+logger = logging.getLogger(__name__)
+
+"""
+Glossary:
+    I: # tokens (coarse representation)
+    L: # atoms   (fine representation)
+    M: # msa
+    T: # templates
+    D: # diffusion structure batch dim
+"""
+
+
+class AtomAttentionDecoder(nn.Module):
+    def __init__(self, c_token, c_atom, c_atompair, atom_transformer):
+        super().__init__()
+        self.atom_transformer = AtomTransformer(
+            c_atom=c_atom, c_atompair=c_atompair, **atom_transformer
+        )
+        self.linear_1 = linearNoBias(c_token, c_atom)
+        self.to_r_update = nn.Sequential(
+            nn.LayerNorm((c_atom,)), linearNoBias(c_atom, 3)
+        )
+
+    def forward(
+        self,
+        f,
+        Ai,  # [L, C_token]
+        Ql_skip,  # [L, C_atom]
+        Cl_skip,  # [L, C_atom]
+        Plm_skip,  # [L, L, C_atompair]
+    ):
+        tok_idx = f["atom_to_token_map"]
+
+        @activation_checkpointing
+        def atom_decoder(Ai, Ql_skip, Cl_skip, Plm_skip, tok_idx):
+            # Broadcast per-token activiations to per-atom activations and add the skip connection
+            Ql = self.linear_1(Ai[..., tok_idx, :]) + Ql_skip
+
+            # Cross attention transformer.
+            Ql = self.atom_transformer(Ql, Cl_skip, Plm_skip)
+
+            # Map to positions update
+            Rl_update = self.to_r_update(Ql)
+
+            return Rl_update
+
+        return atom_decoder(Ai, Ql_skip, Cl_skip, Plm_skip, tok_idx)
+
+
+class DiffusionModule(nn.Module):
+    def __init__(
+        self,
+        sigma_data,
+        c_atom,
+        c_atompair,
+        c_token,
+        c_s,
+        c_z,
+        f_pred,
+        diffusion_conditioning,
+        atom_attention_encoder,
+        diffusion_transformer,
+        atom_attention_decoder,
+    ):
+        super().__init__()
+        self.sigma_data = sigma_data
+        self.c_atom = c_atom
+        self.c_atompair = c_atompair
+        self.c_token = c_token
+        self.c_s = c_s
+        self.f_pred = f_pred
+
+        self.diffusion_conditioning = DiffusionConditioning(
+            sigma_data=sigma_data, c_s=c_s, c_z=c_z, **diffusion_conditioning
+        )
+        self.atom_attention_encoder = AtomAttentionEncoderDiffusion(
+            c_token=c_token,
+            c_s=c_s,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            **atom_attention_encoder,
+        )
+        self.process_s = nn.Sequential(
+            nn.LayerNorm((c_s,)),
+            linearNoBias(c_s, c_token),
+        )
+        self.diffusion_transformer = DiffusionTransformer(
+            c_token=c_token, c_s=c_s, c_tokenpair=c_z, **diffusion_transformer
+        )
+        self.layer_norm_1 = nn.LayerNorm(c_token)
+        self.atom_attention_decoder = AtomAttentionDecoder(
+            c_token=c_token,
+            c_atom=c_atom,
+            c_atompair=c_atompair,
+            **atom_attention_decoder,
+        )
+
+    def forward(
+        self,
+        X_noisy_L,  # [B, L, 3]
+        t,  # [B] (0 is ground truth)
+        f,  # Dict (Input feature dictionary)
+        S_inputs_I,  # [B, I, C_S_input]
+        S_trunk_I,  # [B, I, C_S_trunk]
+        Z_trunk_II,  # [B, I, I, C_Z]
+    ):
+        # Conditioning
+        S_I, Z_II = self.diffusion_conditioning(
+            t, f, S_inputs_I.float(), S_trunk_I.float(), Z_trunk_II.float()
+        )
+
+        # Scale positions to dimensionless vectors with approximately unit variance
+        if self.f_pred == "edm":
+            R_noisy_L = X_noisy_L / torch.sqrt(
+                t[..., None, None] ** 2 + self.sigma_data**2
+            )
+        elif self.f_pred == "unconditioned":
+            R_noisy_L = torch.zeros_like(X_noisy_L)
+        elif self.f_pred == "noise_pred":
+            R_noisy_L = X_noisy_L
+        else:
+            raise Exception(f"{self.f_pred=} unrecognized")
+        # Sequence-local Atom Attention and aggregation to coarse-grained tokens
+        A_I, Q_skip_L, C_skip_L, P_skip_LL = self.atom_attention_encoder(
+            f, R_noisy_L, S_trunk_I.float(), Z_II
+        )
+        # Full self-attention on token level
+
+        A_I = A_I + self.process_s(S_I)
+        A_I = self.diffusion_transformer(A_I, S_I, Z_II, Beta_II=None)
+        A_I = self.layer_norm_1(A_I)
+
+        # Broadcast token activations to atoms and run Sequence-local Atom Attention
+        R_update_L = self.atom_attention_decoder(
+            f, A_I.float(), Q_skip_L, C_skip_L, P_skip_LL
+        )
+        # Rescale updates to positions and combine with input positions
+        if self.f_pred == "edm":
+            X_out_L = (self.sigma_data**2 / (self.sigma_data**2 + t**2))[
+                ..., None, None
+            ] * X_noisy_L + (self.sigma_data * t / (self.sigma_data**2 + t**2) ** 0.5)[
+                ..., None, None
+            ] * R_update_L
+        elif self.f_pred == "unconditioned":
+            X_out_L = R_update_L
+        elif self.f_pred == "noise_pred":
+            X_out_L = X_noisy_L + R_update_L
+        else:
+            raise Exception(f"{self.f_pred=} unrecognized")
+
+        return X_out_L
+
+
+class DiffusionConditioning(nn.Module):
+    def __init__(
+        self, sigma_data, c_z, c_s, c_s_inputs, c_t_embed, relative_position_encoding
+    ):
+        super().__init__()
+        self.sigma_data = sigma_data
+        self.relative_position_encoding = RelativePositionEncoding(
+            c_z=c_z, **relative_position_encoding
+        )
+        self.to_zii = nn.Sequential(
+            nn.LayerNorm(
+                c_z * 2
+            ),  # Operates on concatenated (z_ij_trunk: [..., c_z]), RelativePositionalEncoding: [..., c_z])
+            linearNoBias(c_z * 2, c_z),
+        )
+        self.transition_1 = nn.ModuleList(
+            [
+                Transition(c=c_z, n=2),
+                Transition(c=c_z, n=2),
+            ]
+        )
+        self.to_si = nn.Sequential(
+            nn.LayerNorm(c_s + c_s_inputs), linearNoBias(c_s + c_s_inputs, c_s)
+        )
+        c_t_embed = 256
+        self.fourier_embedding = FourierEmbedding(c_t_embed)
+        self.process_n = nn.Sequential(
+            nn.LayerNorm(c_t_embed), linearNoBias(c_t_embed, c_s)
+        )
+        self.transition_2 = nn.ModuleList(
+            [
+                Transition(c=c_s, n=2),
+                Transition(c=c_s, n=2),
+            ]
+        )
+
+    def forward(self, t, f, S_inputs_I, S_trunk_I, Z_trunk_II):
+        # Pair conditioning
+        Z_II = torch.cat([Z_trunk_II, self.relative_position_encoding(f)], dim=-1)
+
+        @activation_checkpointing
+        def _run_conditioning(Z_II, S_trunk_I, S_inputs_I):
+            Z_II = self.to_zii(Z_II)
+            for b in range(2):
+                Z_II = Z_II + self.transition_1[b](Z_II)
+
+            # Single conditioning
+            S_I = torch.cat([S_trunk_I, S_inputs_I], dim=-1)
+            S_I = self.to_si(S_I)
+            N_D = self.fourier_embedding(1 / 4 * torch.log(t / self.sigma_data))
+            S_I = self.process_n(N_D).unsqueeze(-2) + S_I
+            for b in range(2):
+                S_I = S_I + self.transition_2[b](S_I)
+
+            return S_I, Z_II
+
+        return _run_conditioning(Z_II, S_trunk_I, S_inputs_I)
+
+
+class DistogramHead(nn.Module):
+    def __init__(
+        self,
+        c_z,
+        bins,
+    ):
+        super().__init__()
+        self.predictor = nn.Linear(c_z, bins)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # initialize linear layer for final logit prediction
+        nn.init.zeros_(self.predictor.weight)
+        nn.init.zeros_(self.predictor.bias)
+
+    def forward(
+        self,
+        Z_II,
+    ):
+        return self.predictor(
+            Z_II + Z_II.transpose(-2, -3)  # symmetrize pair features
+        )
+
+
+class Recycler(nn.Module):
+    def __init__(
+        self,
+        c_s,
+        c_z,
+        template_embedder,
+        msa_module,
+        n_pairformer_blocks,
+        pairformer_block,
+    ):
+        super().__init__()
+        self.c_z = c_z
+        self.process_zh = nn.Sequential(
+            nn.LayerNorm(c_z),
+            linearNoBias(c_z, c_z),
+        )
+        self.template_embedder = RF3TemplateEmbedder(c_z=c_z, **template_embedder)
+        self.msa_module = MSAModule(**msa_module)
+        self.process_sh = nn.Sequential(
+            nn.LayerNorm(c_s),
+            linearNoBias(c_s, c_s),
+        )
+        self.pairformer_stack = nn.ModuleList(
+            [
+                PairformerBlock(c_s=c_s, c_z=c_z, **pairformer_block)
+                for _ in range(n_pairformer_blocks)
+            ]
+        )
+
+    def forward(
+        self,
+        f,
+        S_inputs_I,
+        S_init_I,
+        Z_init_II,
+        S_I,
+        Z_II,
+    ):
+        Z_II = Z_init_II + self.process_zh(Z_II)
+        Z_II = Z_II + self.template_embedder(f, Z_II)
+        # NOTE: Implementing bugfix from the Protenix Technical report, where residual-connecting the MSA module is redundant
+        # Reference: https://github.com/bytedance/Protenix/blob/main/Protenix_Technical_Report.pdf
+        Z_II = self.msa_module(f, Z_II, S_inputs_I)
+        S_I = S_init_I + self.process_sh(S_I)
+        for block in self.pairformer_stack:
+            S_I, Z_II = block(S_I, Z_II)
+        return S_I, Z_II

rf3/model/layers/af3_auxiliary_heads.py

@@ -0,0 +1,255 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from rf3.model.RF3_structure import PairformerBlock, linearNoBias
+
+# TODO: Get from RF2AA encoding instead
+CHEM_DATA_LEGACY = {"NHEAVY": 23, "aa2num": {"UNK": 20, "GLY": 7, "MAS": 21}}
+
+
+def discretize_distance_matrix(
+    distance_matrix, num_bins=38, min_distance=3.25, max_distance=50.75
+):
+    # Calculate the bin width
+    bin_width = (max_distance - min_distance) / num_bins
+    bins = (
+        torch.arange(num_bins, device=distance_matrix.device) * bin_width + min_distance
+    )
+
+    # Discretize distances into bins (bucketize automatically places out-of-range values in the last bin)
+    binned_distances = torch.bucketize(distance_matrix, bins)
+
+    return binned_distances
+
+
+class ConfidenceHead(nn.Module):
+    """Algorithm 31"""
+
+    def __init__(
+        self,
+        c_s,
+        c_z,
+        n_pairformer_layers,
+        pairformer,
+        n_bins_pae,
+        n_bins_pde,
+        n_bins_plddt,
+        n_bins_exp_resolved,
+        use_Cb_distances=False,
+        use_af3_style_binning_and_final_layer_norms=False,
+        symmetrize_Cb_logits=True,
+        layer_norm_along_feature_dimension=False,
+    ):
+        super(ConfidenceHead, self).__init__()
+        self.process_s_inputs_right = linearNoBias(449, c_z)
+        self.process_s_inputs_left = linearNoBias(449, c_z)
+        self.use_af3_style_binning_and_final_layer_norms = (
+            use_af3_style_binning_and_final_layer_norms
+        )
+        self.layer_norm_along_feature_dimension = layer_norm_along_feature_dimension
+        if self.use_af3_style_binning_and_final_layer_norms:
+            self.layernorm_pde = nn.LayerNorm(c_z)
+            self.layernorm_pae = nn.LayerNorm(c_z)
+            self.layernorm_plddt = nn.LayerNorm(c_s)
+            self.layernorm_exp_resolved = nn.LayerNorm(c_s)
+            self.process_pred_distances = linearNoBias(40, c_z)
+        else:
+            self.process_pred_distances = linearNoBias(11, c_z)
+
+        self.pairformer = nn.ModuleList(
+            [
+                PairformerBlock(c_s=c_s, c_z=c_z, **pairformer)
+                for _ in range(n_pairformer_layers)
+            ]
+        )
+
+        self.predict_pae = linearNoBias(c_z, n_bins_pae)
+        self.predict_pde = linearNoBias(c_z, n_bins_pde)
+        self.predict_plddt = linearNoBias(
+            c_s, CHEM_DATA_LEGACY["NHEAVY"] * n_bins_plddt
+        )
+        self.predict_exp_resolved = linearNoBias(
+            c_s, CHEM_DATA_LEGACY["NHEAVY"] * n_bins_exp_resolved
+        )
+        self.use_Cb_distances = use_Cb_distances
+        if self.use_Cb_distances:
+            self.process_Cb_distances = linearNoBias(25, c_z)
+        self.symmetrize_Cb_logits = symmetrize_Cb_logits
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(
+        self,
+        S_inputs_I,
+        S_trunk_I,
+        Z_trunk_II,
+        X_pred_L,
+        seq,
+        rep_atoms,
+        frame_atom_idxs=None,
+    ):
+        # stopgrad on S_trunk_I, Z_trunk_II, X_pred_L but not S_inputs_I (4.3.5)
+        S_trunk_I = S_trunk_I.detach().float()  # B, L, 384
+        Z_trunk_II = Z_trunk_II.detach().float()  # B, L, L, 128
+        if X_pred_L is not None:
+            X_pred_L = X_pred_L.detach().float()  # B, n_atoms, 3
+        S_inputs_I = S_inputs_I.detach().float()  # B, L, 384
+        seq = seq.detach()
+
+        if self.layer_norm_along_feature_dimension:
+            # do a layer norm on S_trunk_I
+            S_trunk_I = F.layer_norm(S_trunk_I, normalized_shape=(S_trunk_I.shape[-1]))
+            # do a layer norm on Z_trunk_II
+            Z_trunk_II = F.layer_norm(
+                Z_trunk_II, normalized_shape=(Z_trunk_II.shape[-1])
+            )
+            # do a layer norm on S_inputs_I
+            S_inputs_I = F.layer_norm(
+                S_inputs_I, normalized_shape=(S_inputs_I.shape[-1])
+            )
+        else:
+            S_trunk_I = F.layer_norm(S_trunk_I, normalized_shape=(S_trunk_I.shape))
+            Z_trunk_II = F.layer_norm(Z_trunk_II, normalized_shape=(Z_trunk_II.shape))
+            S_inputs_I = F.layer_norm(S_inputs_I, normalized_shape=(S_inputs_I.shape))
+
+        # embed S_inputs_I twice
+        S_inputs_I_right = self.process_s_inputs_right(S_inputs_I)
+        S_inputs_I_left = self.process_s_inputs_left(S_inputs_I)
+        # add outer product of two linear embeddings of S_inputs_I  to Z_II
+        # TODO: check the unsqueezed dimension is the correct one
+        Z_trunk_II = Z_trunk_II + (
+            S_inputs_I_right.unsqueeze(-2) + S_inputs_I_left.unsqueeze(-3)
+        )
+
+        # embed distances of representative atom from every token
+        #    in the pair representation
+        # if no coords are input, skip this connection
+        if X_pred_L is not None:
+            X_pred_rep_I = X_pred_L.index_select(1, rep_atoms)
+            dist = torch.cdist(X_pred_rep_I, X_pred_rep_I)
+            if not self.use_af3_style_binning_and_final_layer_norms:
+                # bins are 3.375 to 20.375 in 1.75 increments according to pseudocode
+                dist_one_hot = F.one_hot(
+                    discretize_distance_matrix(
+                        dist, min_distance=3.375, max_distance=20.875, num_bins=10
+                    ),
+                    num_classes=11,
+                )
+            else:
+                # published code is 3.25 to 50.75, with 39 bins
+                dist_one_hot = F.one_hot(
+                    discretize_distance_matrix(
+                        dist, min_distance=3.25, max_distance=50.75, num_bins=39
+                    ),
+                    num_classes=40,
+                )
+
+            Z_trunk_II = Z_trunk_II + self.process_pred_distances(dist_one_hot.float())
+
+            if self.use_Cb_distances:
+                # embed difference between observed cb and ideal cb positions
+                Cb_distances = calc_Cb_distances(
+                    X_pred_L, seq, rep_atoms, frame_atom_idxs
+                )
+                Cb_distances_one_hot = F.one_hot(
+                    discretize_distance_matrix(
+                        Cb_distances,
+                        min_distance=0.0001,
+                        max_distance=0.25,
+                        num_bins=24,
+                    ),
+                    num_classes=25,
+                )
+                Cb_logits = self.process_Cb_distances(Cb_distances_one_hot.float())
+                # symmetrize the logits
+                if self.symmetrize_Cb_logits:
+                    Cb_logits = Cb_logits[:, None, :, :] + Cb_logits[:, :, None, :]
+                else:
+                    Cb_logits = Cb_logits[:, None, :, :]
+
+                Z_trunk_II = Z_trunk_II + Cb_logits
+
+        if not self.use_af3_style_binning_and_final_layer_norms:
+            S_trunk_residual_I = S_trunk_I.clone()
+            Z_trunk_residual_II = Z_trunk_II.clone()
+
+        # process with pairformer stack
+        for n in range(len(self.pairformer)):
+            S_trunk_I, Z_trunk_II = self.pairformer[n](S_trunk_I, Z_trunk_II)
+
+        # despite doing so in their pseudocode, af3's published code does not add the residual back
+        if not self.use_af3_style_binning_and_final_layer_norms:
+            S_trunk_I = S_trunk_residual_I + S_trunk_I
+            Z_trunk_II = Z_trunk_residual_II + Z_trunk_II
+
+            # linearly project for each prediction task
+            pde_logits = self.predict_pde(
+                Z_trunk_II + Z_trunk_II.transpose(-2, -3)
+            )  # BUG: needs to be symmetrized correctly
+
+            pae_logits = self.predict_pae(Z_trunk_II)
+
+            plddt_logits = self.predict_plddt(S_trunk_I)
+            exp_resolved_logits = self.predict_exp_resolved(S_trunk_I)
+
+        # af3's published code does not add the residual back and has some additional layernorms before the linear projections
+        # they also do the pde slightly differently, adding the transpose after the linear projection
+        else:
+            left_distance_logits = self.predict_pde(self.layernorm_pde(Z_trunk_II))
+            right_distance_logits = left_distance_logits.transpose(-2, -3)
+            pde_logits = left_distance_logits + right_distance_logits
+
+            pae_logits = self.predict_pae(self.layernorm_pae(Z_trunk_II))
+            plddt_logits = self.predict_plddt(self.layernorm_plddt(S_trunk_I))
+            exp_resolved_logits = self.predict_exp_resolved(
+                self.layernorm_exp_resolved(S_trunk_I)
+            )
+
+        return dict(
+            pde_logits=pde_logits,
+            pae_logits=pae_logits,
+            plddt_logits=plddt_logits,
+            exp_resolved_logits=exp_resolved_logits,
+        )
+
+
+def calc_Cb_distances(X_pred_L, seq, rep_atoms, frame_atom_idxs):
+    frame_atom_idxs = frame_atom_idxs.unsqueeze(0).expand(X_pred_L.shape[0], -1, -1)
+
+    N = torch.gather(
+        X_pred_L, 1, frame_atom_idxs[..., 0].unsqueeze(-1).expand(-1, -1, 3)
+    )
+    Ca = torch.gather(
+        X_pred_L, 1, frame_atom_idxs[..., 1].unsqueeze(-1).expand(-1, -1, 3)
+    )
+    C = torch.gather(
+        X_pred_L, 1, frame_atom_idxs[..., 2].unsqueeze(-1).expand(-1, -1, 3)
+    )
+    Cb = X_pred_L.index_select(1, rep_atoms)
+
+    is_valid_Cb = (
+        (seq != CHEM_DATA_LEGACY.aa2num["UNK"])
+        & (seq != CHEM_DATA_LEGACY.aa2num["GLY"])
+        & (seq != CHEM_DATA_LEGACY.aa2num["MAS"])
+    )
+
+    def _legacy_is_protein(seq):
+        return (seq >= 0).all() & (seq < 20).all()
+
+    is_valid_Cb = is_valid_Cb & _legacy_is_protein(seq)
+
+    b = Ca - N
+    c = C - Ca
+    a = torch.cross(b, c, dim=-1)
+
+    ideal_Cb = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + Ca
+
+    Cb_distances = torch.norm(Cb - ideal_Cb, dim=-1)
+    Cb_distances[:, ~is_valid_Cb] = 0.0
+
+    return Cb_distances