"""Bundled inverse-folding inference utilities for Evo-IF. This module contains only the model, structure featurization, checkpoint loading, sequence conversion, and iterative sampler used by ``infer.py``. """ from __future__ import annotations import json from collections import OrderedDict from pathlib import Path from typing import Literal, Mapping, Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils import torch.utils.checkpoint from .structure_io import CIFParser, PDBParser # ============================================================================ # Inverse-folding model # ============================================================================ def featurize(batch, polytype_to_int, restype_to_int, atom_dict, device): # Filter out failed structures (where b[0] is "pass" instead of a dict) # Also handle nested list from DataLoader valid_batch = [] for b in batch: # Skip if b[0] is a list (nested batch structure) if isinstance(b[0], list): continue # Skip if b[0] is not a dict (e.g., "pass" for failed structures) if not isinstance(b[0], dict): continue # Skip if b[1] is not a valid length (int or tensor) if not isinstance(b[1], (int, torch.Tensor)): continue valid_batch.append(b) batch = valid_batch B = len(batch) if B > 0: # Convert lengths to tensors if they are integers L_list = [] for b in batch: L = b[1] if isinstance(L, int): L = torch.tensor(L) L_list.append(L) L_stack = torch.stack(L_list) L_max = torch.max(L_stack) X = torch.zeros([B, L_max, len(atom_dict), 3], dtype=torch.float32) X_m = torch.zeros([B, L_max, len(atom_dict)], dtype=torch.int32) mask = torch.zeros([B, L_max], dtype=torch.int32) S = restype_to_int["PAD"] * torch.ones([B, L_max], dtype=torch.int64) R_idx = -100*torch.ones([B, L_max], dtype=torch.int32) chain_labels = -1*torch.ones([B, L_max], dtype=torch.int64) protein_mask = torch.zeros([B, L_max], dtype=torch.int32) dna_mask = torch.zeros([B, L_max], dtype=torch.int32) rna_mask = torch.zeros([B, L_max], dtype=torch.int32) R_polymer_type = polytype_to_int["PAD"] * torch.ones([B, L_max], dtype=torch.int64) interface_mask = torch.zeros([B, L_max], dtype=torch.int32) base_pair_mask = torch.zeros([B, L_max], dtype=torch.int32) base_pair_index = torch.zeros([B, L_max], dtype=torch.int64) canonical_base_pair_mask = torch.zeros([B, L_max], dtype=torch.int32) canonical_base_pair_index = torch.zeros([B, L_max], dtype=torch.int64) aligned_ppm = torch.zeros([B, L_max, len(restype_to_int)], dtype=torch.float64) ppm_mask = torch.zeros([B, L_max], dtype=torch.int32) structure_paths = [] assembly_ids = [] for i, b in enumerate(batch): out_dict = b[0] X[i,:L_stack[i]] = out_dict["X"][None,] X_m[i,:L_stack[i]] = out_dict["X_m"][None,] mask[i,:L_stack[i]] = torch.ones_like(out_dict["S"][None,], dtype=torch.int32) S[i,:L_stack[i]] = out_dict["S"][None,] R_idx[i,:L_stack[i]] = out_dict["R_idx"][None,] chain_labels[i,:L_stack[i]] = out_dict["chain_labels"][None,] protein_mask[i,:L_stack[i]] = out_dict["protein_mask"][None,] dna_mask[i,:L_stack[i]] = out_dict["dna_mask"][None,] rna_mask[i,:L_stack[i]] = out_dict["rna_mask"][None,] R_polymer_type[i,:L_stack[i]] = out_dict["R_polymer_type"][None,] interface_mask[i,:L_stack[i]] = out_dict["interface_mask"][None,] base_pair_mask[i,:L_stack[i]] = out_dict["base_pair_mask"][None,] base_pair_index[i,:L_stack[i]] = out_dict["base_pair_index"][None,] canonical_base_pair_mask[i,:L_stack[i]] = out_dict["canonical_base_pair_mask"][None,] canonical_base_pair_index[i,:L_stack[i]] = out_dict["canonical_base_pair_index"][None,] aligned_ppm[i,:L_stack[i]] = out_dict["aligned_ppm"][None,] ppm_mask[i,:L_stack[i]] = out_dict["ppm_mask"][None,] structure_paths.append(out_dict["structure_path"]) assembly_ids.append(out_dict["assembly_id"]) out_dict = {} out_dict["X"] = X.to(device) #[B, L, num_prot_atoms, 3] out_dict["X_m"] = X_m.to(device) #[B, L, num_prot_atoms] out_dict["mask"] = mask.to(device) #[B, L, num_prot_atoms] out_dict["S"] = S.long().to(device) #[B, L] out_dict["R_idx"] = R_idx.to(device) #[B, L] out_dict["chain_labels"] = chain_labels.to(device) #[B, L] out_dict["protein_mask"] = protein_mask.to(device) out_dict["dna_mask"] = dna_mask.to(device) out_dict["rna_mask"] = rna_mask.to(device) out_dict["R_polymer_type"] = R_polymer_type.to(device) out_dict["interface_mask"] = interface_mask.to(device) out_dict["base_pair_mask"] = base_pair_mask.to(device) out_dict["base_pair_index"] = base_pair_index.to(device) out_dict["canonical_base_pair_mask"] = canonical_base_pair_mask.to(device) out_dict["canonical_base_pair_index"] = canonical_base_pair_index.to(device) out_dict["aligned_ppm"] = aligned_ppm.to(device) out_dict["ppm_mask"] = ppm_mask.to(device) out_dict["structure_path"] = structure_paths out_dict["assembly_id"] = assembly_ids return out_dict else: return "pass" # The following gather functions def gather_edges(edges, neighbor_idx): # Features [B,N,N,C] at Neighbor indices [B,N,K] => Neighbor features [B,N,K,C] neighbors = neighbor_idx.unsqueeze(-1).expand(-1, -1, -1, edges.size(-1)) edge_features = torch.gather(edges, 2, neighbors) return edge_features def gather_nodes(nodes, neighbor_idx): # Features [B,N,C] at Neighbor indices [B,N,K] => [B,N,K,C] # Flatten and expand indices per batch [B,N,K] => [B,NK] => [B,NK,C] neighbors_flat = neighbor_idx.reshape((neighbor_idx.shape[0], -1)) neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2)) # Gather and re-pack neighbor_features = torch.gather(nodes, 1, neighbors_flat) neighbor_features = neighbor_features.view(list(neighbor_idx.shape)[:3] + [-1]) return neighbor_features def gather_nodes_t(nodes, neighbor_idx): # Features [B,N,C] at Neighbor index [B,K] => Neighbor features[B,K,C] idx_flat = neighbor_idx.unsqueeze(-1).expand(-1, -1, nodes.size(2)) neighbor_features = torch.gather(nodes, 1, idx_flat) return neighbor_features def cat_neighbors_nodes(h_nodes, h_neighbors, E_idx): h_nodes = gather_nodes(h_nodes, E_idx) h_nn = torch.cat([h_neighbors, h_nodes], -1) return h_nn class EncLayer(nn.Module): def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30): super(EncLayer, self).__init__() self.num_hidden = num_hidden self.num_in = num_in self.scale = scale self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(num_hidden) self.norm2 = nn.LayerNorm(num_hidden) self.norm3 = nn.LayerNorm(num_hidden) self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True) self.W2 = nn.Linear(num_hidden, num_hidden, bias=True) self.W3 = nn.Linear(num_hidden, num_hidden, bias=True) self.W11 = nn.Linear(num_hidden + num_in, num_hidden, bias=True) self.W12 = nn.Linear(num_hidden, num_hidden, bias=True) self.W13 = nn.Linear(num_hidden, num_hidden, bias=True) self.act = torch.nn.GELU() self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4) def forward(self, h_V, h_E, E_idx, mask_V=None, mask_attend=None): """ Parallel computation of full transformer layer """ h_EV = cat_neighbors_nodes(h_V, h_E, E_idx) h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1) h_EV = torch.cat([h_V_expand, h_EV], -1) h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV))))) if mask_attend is not None: h_message = mask_attend.unsqueeze(-1) * h_message dh = torch.sum(h_message, -2) / self.scale h_V = self.norm1(h_V + self.dropout1(dh)) dh = self.dense(h_V) h_V = self.norm2(h_V + self.dropout2(dh)) if mask_V is not None: mask_V = mask_V.unsqueeze(-1) h_V = mask_V * h_V h_EV = cat_neighbors_nodes(h_V, h_E, E_idx) h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_EV.size(-2),-1) h_EV = torch.cat([h_V_expand, h_EV], -1) h_message = self.W13(self.act(self.W12(self.act(self.W11(h_EV))))) h_E = self.norm3(h_E + self.dropout3(h_message)) return h_V, h_E class DecLayer(nn.Module): def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30): super(DecLayer, self).__init__() self.num_hidden = num_hidden self.num_in = num_in self.scale = scale self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(num_hidden) self.norm2 = nn.LayerNorm(num_hidden) self.W1 = nn.Linear(num_hidden + num_in, num_hidden, bias=True) self.W2 = nn.Linear(num_hidden, num_hidden, bias=True) self.W3 = nn.Linear(num_hidden, num_hidden, bias=True) self.act = torch.nn.GELU() self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4) def forward(self, h_V, h_E, mask_V=None, mask_attend=None): """ Parallel computation of full transformer layer """ # Concatenate h_V_i to h_E_ij h_V_expand = h_V.unsqueeze(-2).expand(-1,-1,h_E.size(-2),-1) h_EV = torch.cat([h_V_expand, h_E], -1) h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV))))) if mask_attend is not None: h_message = mask_attend.unsqueeze(-1) * h_message dh = torch.sum(h_message, -2) / self.scale h_V = self.norm1(h_V + self.dropout1(dh)) # Position-wise feedforward dh = self.dense(h_V) h_V = self.norm2(h_V + self.dropout2(dh)) if mask_V is not None: mask_V = mask_V.unsqueeze(-1) h_V = mask_V * h_V return h_V class PositionWiseFeedForward(nn.Module): def __init__(self, num_hidden, num_ff): super(PositionWiseFeedForward, self).__init__() self.W_in = nn.Linear(num_hidden, num_ff, bias=True) self.W_out = nn.Linear(num_ff, num_hidden, bias=True) self.act = torch.nn.GELU() def forward(self, h_V): h = self.act(self.W_in(h_V)) h = self.W_out(h) return h class PositionalEncodings(nn.Module): def __init__(self, num_embeddings, max_relative_feature=32): super(PositionalEncodings, self).__init__() self.num_embeddings = num_embeddings self.max_relative_feature = max_relative_feature self.linear = nn.Linear(2*max_relative_feature+1+1, num_embeddings) def forward(self, offset, mask): d = torch.clip(offset + self.max_relative_feature, 0, 2*self.max_relative_feature)*mask + (1-mask)*(2*self.max_relative_feature+1) d_onehot = torch.nn.functional.one_hot(d, 2*self.max_relative_feature+1+1) E = self.linear(d_onehot.float()) return E class ProteinFeatures(nn.Module): def __init__(self, edge_features, node_features, num_positional_embeddings=16, num_rbf=16, top_k=30, atom_dict=None, polytype_to_int=None, protein_augment_eps=0., dna_augment_eps=0., rna_augment_eps=0., na_ref_atom="C1'", include_pred_na_N=1, device=None): """ Extract protein features """ super(ProteinFeatures, self).__init__() if atom_dict is None: raise Exception("atom_dict is necessary for featurization!") if polytype_to_int is None: raise Exception("polytype_to_int is necessary for featurization!") self.N_idx = atom_dict["N"] self.CA_idx = atom_dict["CA"] self.C_idx = atom_dict["C"] self.O4prime_idx = atom_dict["O4'"] self.C1prime_idx = atom_dict["C1'"] self.C2prime_idx = atom_dict["C2'"] self.na_ref_atom_idx = atom_dict[na_ref_atom] self.edge_features = edge_features self.node_features = node_features self.top_k = top_k self.protein_augment_eps = protein_augment_eps self.dna_augment_eps = dna_augment_eps self.rna_augment_eps = rna_augment_eps self.num_rbf = num_rbf self.num_positional_embeddings = num_positional_embeddings self.embeddings = PositionalEncodings(num_positional_embeddings) self.num_polytypes = len(polytype_to_int) self.node_in = len(polytype_to_int) self.node_embedding = nn.Linear(self.node_in, node_features, bias=False) self.norm_nodes = nn.LayerNorm(node_features) total_atoms = len(atom_dict) + 1 self.include_pred_na_N = include_pred_na_N if self.include_pred_na_N: total_atoms = total_atoms + 1 self.edge_in = num_positional_embeddings + num_rbf*total_atoms*total_atoms self.edge_embedding = nn.Linear(self.edge_in, edge_features, bias=False) self.norm_edges = nn.LayerNorm(edge_features) def _dist(self, X, mask, eps = 1E-6): mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2) dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2) D = mask_2D * torch.sqrt(torch.sum(dX**2, 3) + eps) D_max, _ = torch.max(D, -1, keepdim=True) D_adjust = D + (1. - mask_2D) * D_max sampled_top_k = self.top_k D_neighbors, E_idx = torch.topk(D_adjust, np.minimum(self.top_k, X.shape[1]), dim=-1, largest=False) return D_neighbors, E_idx def _rbf(self, D): device = D.device D_min, D_max, D_count = 2., 22., self.num_rbf D_mu = torch.linspace(D_min, D_max, D_count, device=device) D_mu = D_mu.view([1,1,1,1,1,-1]) D_sigma = (D_max - D_min) / D_count D_expand = torch.unsqueeze(D, -1) RBF = torch.exp(-((D_expand - D_mu) / D_sigma)**2) return RBF def _get_all_rbf(self, X, E_idx, X_m): # [B,L,16,3] => [B,L,16*3] => [B,L,K,16*3] => [B,L,K,16,3] => [B,L,K,16,16] X_flat = X.reshape((X.shape[0], X.shape[1], -1)) X_flat_g = gather_nodes(X_flat, E_idx) #[B,L,K,16*3] X_g = X_flat_g.reshape(list(X_flat_g.shape)[:-1] + list(X.shape[-2:])) #[B,L,K,16,3] D = torch.sqrt(torch.sum((X[:,:,None,:,None,:] - X_g[:,:,:,None,:,:])**2,-1) + 1e-6) RBF_all = self._rbf(D) #[B, L, K, 16, 16, H] X_m_gathered = gather_nodes(X_m, E_idx) #[B,L,K,16] RBF_all = RBF_all*X_m[:,:,None,:,None,None]*X_m_gathered[:,:,:,None,:,None] #[B,L,K,16,16,H] RBF_all = RBF_all.view([X.shape[0], X.shape[1], E_idx.shape[2],-1]) #[B,L,K,16*16*H] return RBF_all def get_Cb(self, N, Ca, C, w_a, w_b, w_c): b = Ca - N c = C - Ca a = torch.cross(b, c, dim=-1) Cb = w_a * a + w_b * b + w_c * c + Ca #shift from CA return Cb def forward(self, feature_dict): X = feature_dict["X"] mask = feature_dict["mask"] R_idx = feature_dict["R_idx"] chain_labels = feature_dict["chain_labels"] X_m = feature_dict["X_m"] protein_mask = feature_dict["protein_mask"] dna_mask = feature_dict["dna_mask"] rna_mask = feature_dict["rna_mask"] R_polymer_type = feature_dict["R_polymer_type"] if self.training and (self.protein_augment_eps > 0 or \ self.dna_augment_eps > 0 or \ self.rna_augment_eps > 0): augment_eps = protein_mask * self.protein_augment_eps + \ dna_mask * self.dna_augment_eps + \ rna_mask * self.rna_augment_eps X = X + X_m[:,:,:,None] * augment_eps[:,:,None,None] * torch.randn_like(X) Ca = X[:,:,self.CA_idx,:] N = X[:,:,self.N_idx,:] C = X[:,:,self.C_idx,:] Cb = self.get_Cb(N, Ca, C, w_a = -0.58273431, w_b = 0.56802827, w_c = -0.54067466) na_ref_atom = X[:,:,self.na_ref_atom_idx,:] if self.include_pred_na_N: O4prime = X[:,:,self.O4prime_idx,:] C1prime = X[:,:,self.C1prime_idx,:] C2prime = X[:,:,self.C2prime_idx,:] N_na = self.get_Cb(O4prime, C1prime, C2prime, w_a = -0.56967352, w_b = 0.51055973, w_c = -0.53122153) augmented_X = (X, Cb[:,:,None,:], N_na[:,:,None,:]) augmented_X_m = (X_m, protein_mask[:,:,None], (rna_mask + dna_mask)[:,:,None]) else: augmented_X = (X, Cb[:,:,None,:]) augmented_X_m = (X_m, protein_mask[:,:,None]) augmented_X = torch.cat(augmented_X, -2) augmented_X_m = torch.cat(augmented_X_m, -1) # Ca + P because these vectors are disjoint. This sum represents the # center coordinates for all (protein or dna) residues. D_neighbors, E_idx = self._dist(Ca + na_ref_atom, mask) RBF_all = self._get_all_rbf(augmented_X, E_idx, augmented_X_m) offset = R_idx[:,:,None]-R_idx[:,None,:] offset = gather_edges(offset[:,:,:,None], E_idx)[:,:,:,0] #[B, L, K] d_chains = ((chain_labels[:, :, None] - chain_labels[:,None,:])==0).long() #find self vs non-self interaction E_chains = gather_edges(d_chains[:,:,:,None], E_idx)[:,:,:,0] E_positional = self.embeddings(offset.long(), E_chains) E = torch.cat((E_positional, RBF_all), -1) E = self.edge_embedding(E) E = self.norm_edges(E) R_polymer_type_one_hot = torch.nn.functional.one_hot(R_polymer_type, num_classes = self.num_polytypes).float() V = R_polymer_type_one_hot V = self.node_embedding(V) V = self.norm_nodes(V) return V, E, E_idx class ProteinMPNNDiffusion(nn.Module): """ NA-MPNN with Absorbing State Diffusion support. Key differences from ProteinMPNN: - Bidirectional attention in decoder (no causal masking) - Accepts masked sequences as input - Predicts all positions simultaneously (MLM-style) This enables iterative denoising for sequence generation, inspired by ProRefiner and DPLM. """ def __init__(self, node_features=128, edge_features=128, hidden_dim=128, num_encoder_layers=3, num_decoder_layers=3, atom_dict=None, restype_to_int=None, polytype_to_int=None, vocab=33, num_letters=33, k_neighbors=32, protein_augment_eps=0.1, dna_augment_eps=0.1, rna_augment_eps=0.1, dropout=0.1, na_ref_atom="C1'", include_pred_na_N=1, use_sequence_context=True, device=None): """ Args: use_sequence_context: If True, use sequence embeddings in decoder (allows model to see unmasked positions) """ super(ProteinMPNNDiffusion, self).__init__() # Hyperparameters self.node_features = node_features self.edge_features = edge_features self.vocab = vocab self.hidden_dim = hidden_dim self.use_sequence_context = use_sequence_context if restype_to_int is None: raise Exception("restype_to_int dictionary is necessary!") self.mask_token = restype_to_int["MAS"] self.features = ProteinFeatures(node_features, edge_features, top_k=k_neighbors, atom_dict=atom_dict, polytype_to_int=polytype_to_int, protein_augment_eps=protein_augment_eps, dna_augment_eps=dna_augment_eps, rna_augment_eps=rna_augment_eps, na_ref_atom=na_ref_atom, include_pred_na_N=include_pred_na_N, device=device) self.W_e = nn.Linear(edge_features, hidden_dim, bias=True) self.W_v = nn.Linear(node_features, hidden_dim, bias=True) self.W_s = nn.Embedding(vocab, hidden_dim) # Encoder layers self.encoder_layers = nn.ModuleList([ EncLayer(hidden_dim, hidden_dim*2, dropout=dropout) for _ in range(num_encoder_layers) ]) # Decoder layers (bidirectional, no causal masking) self.decoder_layers = nn.ModuleList([ DecLayer(hidden_dim, hidden_dim*3, dropout=dropout) for _ in range(num_decoder_layers) ]) self.W_out = nn.Linear(hidden_dim, num_letters, bias=True) for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward(self, feature_dict, return_embeddings=False): """ Forward pass with bidirectional attention. Unlike autoregressive ProteinMPNN, this uses full bidirectional attention in the decoder, allowing the model to see all context (both masked and unmasked positions) when making predictions. Args: feature_dict: Dictionary containing: - X: [B, L, num_atoms, 3] Coordinates - S: [B, L] Sequence (may contain MASK tokens) - mask: [B, L] Valid position mask - protein_mask, dna_mask, rna_mask: Polymer type masks - R_idx, chain_labels: Residue indices and chain labels - R_polymer_type: Polymer type for each position return_embeddings: If True, also return hidden embeddings Returns: log_probs: [B, L, V] Log probabilities for each position probs: [B, L, V] Probabilities for each position (optional) h_V: [B, L, H] Hidden embeddings """ X = feature_dict["X"] S = feature_dict["S"] # May contain MASK tokens mask = feature_dict["mask"] device = X.device # Prepare node and edge embeddings from structure V, E, E_idx = self.features(feature_dict) h_V = self.W_v(V) h_E = self.W_e(E) # Encoder: bidirectional self-attention over structure mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1) mask_attend = mask.unsqueeze(-1) * mask_attend for layer in self.encoder_layers: if self.training: h_V, h_E = torch.utils.checkpoint.checkpoint( layer, h_V, h_E, E_idx, mask, mask_attend, use_reentrant=False ) else: h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend) # Sequence embeddings (including MASK token embeddings) h_S = self.W_s(S) # Decoder: BIDIRECTIONAL attention (key difference from autoregressive) # All positions can attend to all other positions # This allows the model to use context from both directions if self.use_sequence_context: # Include sequence information in message passing h_ES = cat_neighbors_nodes(h_S, h_E, E_idx) else: # Structure-only context (useful for initial generation) h_ES = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx) # Build combined embeddings h_EXV_encoder = cat_neighbors_nodes(h_V, h_ES, E_idx) # Bidirectional attention mask (all valid positions attend to each other) # No causal masking - this is the key for diffusion! # Shape: [B, L, 1, 1] * [B, L, K, 1] -> [B, L, K, 1] mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1]) mask_attend_decoder = mask_1D * gather_nodes(mask.unsqueeze(-1), E_idx) for layer in self.decoder_layers: h_ESV = h_EXV_encoder * mask_attend_decoder if self.training: h_V = torch.utils.checkpoint.checkpoint( layer, h_V, h_ESV, mask, use_reentrant=False ) else: h_V = layer(h_V, h_ESV, mask) logits = self.W_out(h_V) log_probs = torch.nn.functional.log_softmax(logits, dim=-1) probs = torch.nn.functional.softmax(logits, dim=-1) if return_embeddings: return log_probs, probs, h_V return log_probs, probs # ============================================================================ # Structure feature dataset # ============================================================================ class PDBDataset(torch.utils.data.Dataset): """Parser metadata and structure assembly helpers needed at inference.""" def __init__( self, cif_parser, pdb_parser, atom_list_to_save=None, parse_protein=1, parse_dna=1, parse_rna=1, parse_rna_as_dna=0, na_shared_tokens=0, protein_backbone_occ_cutoff=0.8, protein_side_chain_occ_cutoff=0.5, dna_backbone_occ_cutoff=0.8, dna_side_chain_occ_cutoff=0.5, rna_backbone_occ_cutoff=0.8, rna_side_chain_occ_cutoff=0.5, crop_large_structures=0, batch_tokens=6000, na_ref_atom="C1'", parse_ppms=0, min_overlap_length=5, drop_protein_probability=0, na_only_as_uniform_ppm=0, protein_interface_residue_mutation_probability=0, mutate_base_pair_together=0, mutate_entire_side_chain_interface_probability=0, na_non_interface_as_uniform_ppm=0, ): # Keep the historical constructor surface so inference configurations # produced by earlier runs remain loadable. Training-only arguments are # deliberately accepted but unused. self.crop_large_structures = bool(crop_large_structures) del ( parse_ppms, min_overlap_length, drop_protein_probability, na_only_as_uniform_ppm, protein_interface_residue_mutation_probability, mutate_base_pair_together, mutate_entire_side_chain_interface_probability, na_non_interface_as_uniform_ppm, ) if atom_list_to_save is None: atom_list_to_save = [ "N", "CA", "C", "O", "OP1", "OP2", "P", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", ] self.cif_parser = cif_parser self.pdb_parser = pdb_parser self.atom_list_to_save = atom_list_to_save self.num_atoms_to_save = len(atom_list_to_save) self.atom_dict = dict(zip(atom_list_to_save, range(self.num_atoms_to_save))) self.parse_protein = parse_protein self.parse_dna = parse_dna self.parse_rna = parse_rna self.parse_rna_as_dna = parse_rna_as_dna self.na_shared_tokens = na_shared_tokens self.batch_tokens = batch_tokens self.na_ref_atom = na_ref_atom self.protein_backbone_occ_cutoff = protein_backbone_occ_cutoff self.protein_side_chain_occ_cutoff = protein_side_chain_occ_cutoff self.dna_backbone_occ_cutoff = dna_backbone_occ_cutoff self.dna_side_chain_occ_cutoff = dna_side_chain_occ_cutoff self.rna_backbone_occ_cutoff = rna_backbone_occ_cutoff self.rna_side_chain_occ_cutoff = rna_side_chain_occ_cutoff self.polytypes = ["PP", "DNA", "RNA", "UNK", "MAS", "PAD"] self.polytype_to_int = dict(zip(self.polytypes, range(len(self.polytypes)))) if self.parse_rna_as_dna: self.polytype_to_int["RNA"] = self.polytype_to_int["DNA"] self.restypes = [ "ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL", "UNK", "DA", "DC", "DG", "DT", "DX", "A", "C", "G", "U", "RX", "MAS", "PAD", ] self.protein_restypes = [ "ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS", "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL", "UNK", ] self.dna_restypes = ["DA", "DC", "DG", "DT", "DX"] self.rna_restypes = ["A", "C", "G", "U", "RX"] self.restype_3_to_1 = { "ALA": "A", "ARG": "R", "ASN": "N", "ASP": "D", "CYS": "C", "GLN": "Q", "GLU": "E", "GLY": "G", "HIS": "H", "ILE": "I", "LEU": "L", "LYS": "K", "MET": "M", "PHE": "F", "PRO": "P", "SER": "S", "THR": "T", "TRP": "W", "TYR": "Y", "VAL": "V", "UNK": "X", "DA": "a", "DC": "c", "DG": "g", "DT": "t", "DX": "x", "A": "b", "C": "d", "G": "h", "U": "u", "RX": "y", "MAS": "-", "PAD": "+", } self.restype_to_int = dict(zip(self.restypes, range(len(self.restypes)))) self.int_to_restype = dict(zip(range(len(self.restypes)), self.restypes)) if self.parse_rna_as_dna or self.na_shared_tokens: self.restype_to_int["A"] = self.restype_to_int["DA"] self.restype_to_int["C"] = self.restype_to_int["DC"] self.restype_to_int["G"] = self.restype_to_int["DG"] self.restype_to_int["U"] = self.restype_to_int["DT"] self.restype_to_int["RX"] = self.restype_to_int["DX"] self.protein_backbone_list = ["N", "CA", "C", "O"] self.dna_backbone_list = [ "OP1", "OP2", "P", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", ] self.rna_backbone_list = [ "OP1", "OP2", "P", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", ] self.protein_bb_idx_list = [ self.atom_dict[atom] for atom in self.atom_list_to_save if atom in self.protein_backbone_list ] self.dna_bb_idx_list = [ self.atom_dict[atom] for atom in self.atom_list_to_save if atom in self.dna_backbone_list ] self.rna_bb_idx_list = [ self.atom_dict[atom] for atom in self.atom_list_to_save if atom in self.rna_backbone_list ] def load_chains(self, chains): """Convert parser chain objects into arrays used by ``load_assembly``.""" supported_types = { "polypeptide(L)", "polydeoxyribonucleotide", "polyribonucleotide", "polydeoxyribonucleotide/polyribonucleotide hybrid", } macromolecule_chain_dict = {} for letter, chain in chains.items(): if chain.type not in supported_types: continue residues = OrderedDict() for atom_key in chain.atoms: _, residue_id, residue_name, _ = atom_key variants = residues.setdefault(residue_id, OrderedDict()) variants.setdefault(residue_name, []).append(atom_key) residue_count = len(residues) xyz = np.zeros( [residue_count, self.num_atoms_to_save, 3], dtype=np.float32 ) occ = np.zeros( [residue_count, self.num_atoms_to_save], dtype=np.float32 ) residue_idx = -100 * np.ones([residue_count], dtype=np.int32) raw_sequence = residue_count * ["UNK"] for residue_offset, (residue_id, variants) in enumerate(residues.items()): residue_name, residue_atoms = max( variants.items(), key=lambda item: sum( float(chain.atoms[atom_key].occ) for atom_key in item[1] ), ) for atom_key in residue_atoms: _, _, _, atom_name = atom_key if atom_name not in self.atom_dict: continue atom_index = self.atom_dict[atom_name] xyz[residue_offset, atom_index] = np.asarray( chain.atoms[atom_key].xyz ) occ[residue_offset, atom_index] = np.asarray( chain.atoms[atom_key].occ ) raw_sequence[residue_offset] = residue_name residue_idx[residue_offset] = int(residue_id) macromolecule_chain_dict[letter] = { "type": chain.type, "xyz": xyz, "occ": occ, "seq": raw_sequence, "residue_idx": residue_idx, } return macromolecule_chain_dict def load_assembly(self, macromolecule_chain_dict, asmb, assembly_id, ppms=None): """Apply biological-assembly transforms and select usable polymers. ``ppms`` is accepted for call compatibility. Evo-IF inference does not consume PPM inputs, so the returned PPM features are zero-filled. """ del ppms if assembly_id not in asmb: raise ValueError(f"Assembly {assembly_id!r} is not present in the structure.") assembly_transforms = asmb[assembly_id] if not assembly_transforms: raise ValueError(f"Assembly {assembly_id!r} contains no chain transforms.") X_list = [] X_occ_list = [] S_list = [] R_idx_list = [] chain_labels_list = [] protein_mask_list = [] dna_mask_list = [] rna_mask_list = [] for letter, transform_matrix in assembly_transforms: if letter not in macromolecule_chain_dict: continue chain = macromolecule_chain_dict[letter] xyz = chain["xyz"] transform_matrix = np.asarray(transform_matrix) rotation_matrix = transform_matrix[:3, :3] translation = transform_matrix[:3, 3] xyz = np.einsum("ij,raj->rai", rotation_matrix, xyz) xyz = xyz + translation[None, None, :] chain_length = len(chain["residue_idx"]) protein_mask = np.zeros(chain_length, dtype=np.int32) dna_mask = np.zeros(chain_length, dtype=np.int32) rna_mask = np.zeros(chain_length, dtype=np.int32) if chain["type"] == "polypeptide(L)": unknown_residue = "UNK" protein_mask[:] = 1 elif chain["type"] == "polydeoxyribonucleotide": unknown_residue = "DX" dna_mask[:] = 1 elif chain["type"] == "polyribonucleotide": unknown_residue = "RX" rna_mask[:] = 1 else: # Unknown residues in hybrid chains cannot be assigned to DNA # or RNA reliably; preserve the original zero-mask behaviour. unknown_residue = "DX" for residue_offset, residue_name in enumerate(chain["seq"]): if residue_name in self.dna_restypes: dna_mask[residue_offset] = 1 elif residue_name in self.rna_restypes: rna_mask[residue_offset] = 1 sequence = np.array( [ self.restype_to_int.get( residue_name, self.restype_to_int[unknown_residue] ) for residue_name in chain["seq"] ], dtype=np.int32, ) chain_label = len(chain_labels_list) X_list.append(xyz) X_occ_list.append(chain["occ"]) S_list.append(sequence) R_idx_list.append(chain["residue_idx"]) chain_labels_list.append( np.full(chain_length, chain_label, dtype=np.int32) ) protein_mask_list.append(protein_mask) dna_mask_list.append(dna_mask) rna_mask_list.append(rna_mask) if not X_list: available = ", ".join(map(str, macromolecule_chain_dict)) or "none" raise ValueError( f"Assembly {assembly_id!r} contains no supported macromolecular " f"chains (parsed chains: {available})." ) X = np.concatenate(X_list, axis=0) X_occ = np.concatenate(X_occ_list, axis=0) S = np.concatenate(S_list, axis=0) R_idx = np.concatenate(R_idx_list, axis=0) chain_labels = np.concatenate(chain_labels_list, axis=0) protein_mask = np.concatenate(protein_mask_list, axis=0) dna_mask = np.concatenate(dna_mask_list, axis=0) rna_mask = np.concatenate(rna_mask_list, axis=0) # PPMs are not an Evo-IF inference input. Preserve the historical field # shapes and dtypes expected by the model featurizer. aligned_ppm = np.zeros( (len(S), len(self.restype_to_int)), dtype=np.float64 ) ppm_mask = np.zeros(len(S), dtype=np.int32) R_polymer_type = ( protein_mask * self.polytype_to_int["PP"] + dna_mask * self.polytype_to_int["DNA"] + rna_mask * self.polytype_to_int["RNA"] + (1 - protein_mask - dna_mask - rna_mask) * self.polytype_to_int["UNK"] ) side_chain_occ_cutoff = ( protein_mask * self.protein_side_chain_occ_cutoff + dna_mask * self.dna_side_chain_occ_cutoff + rna_mask * self.rna_side_chain_occ_cutoff ) X_m = (X_occ > side_chain_occ_cutoff[:, None]).astype(np.int32) backbone_occ_cutoff = ( protein_mask * self.protein_backbone_occ_cutoff + dna_mask * self.dna_backbone_occ_cutoff + rna_mask * self.rna_backbone_occ_cutoff ) X_occ_mask = (X_occ > backbone_occ_cutoff[:, None]).astype(np.int32) protein_mask = protein_mask * np.prod( X_occ_mask[:, self.protein_bb_idx_list], axis=-1 ) dna_mask = dna_mask * np.prod( X_occ_mask[:, self.dna_bb_idx_list], axis=-1 ) rna_mask = rna_mask * np.prod( X_occ_mask[:, self.rna_bb_idx_list], axis=-1 ) if self.parse_rna_as_dna: dna_mask = np.bitwise_or(dna_mask, rna_mask) rna_mask = np.zeros_like(dna_mask) mask_for_output = np.zeros_like(protein_mask) out_dict = {} if self.parse_protein: mask_for_output = np.bitwise_or(mask_for_output, protein_mask) out_dict["protein_L"] = np.count_nonzero(protein_mask) else: out_dict["protein_L"] = 0 if self.parse_dna: mask_for_output = np.bitwise_or(mask_for_output, dna_mask) out_dict["dna_L"] = np.count_nonzero(dna_mask) else: out_dict["dna_L"] = 0 if self.parse_rna: mask_for_output = np.bitwise_or(mask_for_output, rna_mask) out_dict["rna_L"] = np.count_nonzero(rna_mask) else: out_dict["rna_L"] = 0 out_dict["macromolecule_L"] = np.count_nonzero(mask_for_output) mask_for_output = mask_for_output.astype(bool) for key, value in { "protein_mask": protein_mask, "dna_mask": dna_mask, "rna_mask": rna_mask, "X": X, "X_m": X_m, "S": S, "R_idx": R_idx, "chain_labels": chain_labels, "R_polymer_type": R_polymer_type, "aligned_ppm": aligned_ppm, "ppm_mask": ppm_mask, }.items(): out_dict[key] = value[mask_for_output] return out_dict def apply_crop_mask(self, out_dict, mask_to_keep): """Crop per-residue arrays and remap precomputed index features.""" mask_to_keep = np.asarray(mask_to_keep, dtype=np.bool_) for key in out_dict: if type(out_dict[key]) is np.ndarray: out_dict[key] = out_dict[key][mask_to_keep] mask_to_remove = np.logical_not(mask_to_keep) index_of_removed = np.where(mask_to_remove)[0] residues_removed_to_left = np.array( [0] + list(np.add.accumulate(mask_to_remove.astype(np.int32))[:-1]), dtype=np.int64, ) index_and_mask_key_pairs = [ ("base_pair_index", "base_pair_mask"), ("canonical_base_pair_index", "canonical_base_pair_mask"), ("nearest_protein_side_chain_index", "side_chain_interface_mask"), ] for index_key, mask_key in index_and_mask_key_pairs: index_in_removed = np.isin(out_dict[index_key], index_of_removed) out_dict[mask_key][index_in_removed] = 0 out_dict[index_key] = ( out_dict[index_key] - residues_removed_to_left[out_dict[index_key]] ) out_dict[index_key] = out_dict[index_key] * out_dict[mask_key] out_dict["protein_L"] = np.count_nonzero(out_dict["protein_mask"]) out_dict["dna_L"] = np.count_nonzero(out_dict["dna_mask"]) out_dict["rna_L"] = np.count_nonzero(out_dict["rna_mask"]) out_dict["macromolecule_L"] = ( out_dict["protein_L"] + out_dict["dna_L"] + out_dict["rna_L"] ) def random_crop_na(self, out_dict): """Crop to the nearest ``batch_tokens`` residues around a random NA.""" X = out_dict["X"] na_mask = out_dict["dna_mask"] + out_dict["rna_mask"] na_indices = np.where(na_mask == 1)[0] if len(na_indices) == 0: raise ValueError("NA-centred cropping requires at least one DNA/RNA residue.") if self.batch_tokens <= 0: raise ValueError("batch_tokens must be positive for NA-centred cropping.") protein_ref_atom_index = self.atom_dict["CA"] na_ref_atom_index = self.atom_dict[self.na_ref_atom] ref_atom_X = ( X[:, protein_ref_atom_index, :] + X[:, na_ref_atom_index, :] ) na_residue_index = np.random.choice(na_indices) distances = np.sqrt( np.sum((ref_atom_X - ref_atom_X[na_residue_index]) ** 2, axis=-1) ) indices_to_keep = np.argsort(distances)[: self.batch_tokens] mask_to_keep = np.zeros_like(out_dict["S"], dtype=np.bool_) mask_to_keep[indices_to_keep] = True self.apply_crop_mask(out_dict, mask_to_keep) # ============================================================================ # Iterative denoising # ============================================================================ def get_denoising_schedule( num_steps: int, schedule_type: Literal["linear", "cosine", "sqrt"] = "cosine", ) -> np.ndarray: """Return mask ratios from fully masked to fully revealed.""" t = np.linspace(0, 1, num_steps + 1) if schedule_type == "linear": return 1.0 - t if schedule_type == "cosine": return 0.5 * (1.0 + np.cos(t * np.pi)) if schedule_type == "sqrt": return 1.0 - np.sqrt(t) raise ValueError(f"Unknown schedule type: {schedule_type}") def select_positions_to_unmask( current_seq: torch.Tensor, logits: torch.Tensor, mask_token_idx: int, num_to_unmask: int, strategy: Literal["confidence", "random", "entropy"] = "confidence", temperature: float = 1.0, use_argmax: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """Reveal the requested number of currently masked positions.""" batch_size, length, _ = logits.shape device = logits.device is_masked = current_seq.eq(mask_token_idx) probs = F.softmax(logits / temperature, dim=-1) if strategy == "confidence": confidence = probs.max(dim=-1).values elif strategy == "entropy": entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1) confidence = -entropy elif strategy == "random": confidence = torch.rand(batch_size, length, device=device) else: raise ValueError(f"Unknown strategy: {strategy}") confidence = confidence.masked_fill(~is_masked, -float("inf")) new_seq = current_seq.clone() unmasked_positions = torch.zeros( batch_size, length, dtype=torch.bool, device=device ) for batch_index in range(batch_size): masked_indices = torch.where(is_masked[batch_index])[0] n_unmask = min(num_to_unmask, masked_indices.numel()) if n_unmask == 0: continue selected = confidence[batch_index, masked_indices].topk(n_unmask).indices positions = masked_indices[selected] for position in positions: if use_argmax: token = logits[batch_index, position].argmax() else: token = torch.multinomial( probs[batch_index, position], 1 ).squeeze() new_seq[batch_index, position] = token unmasked_positions[batch_index, position] = True return new_seq, unmasked_positions def rewalk_positions( current_seq: torch.Tensor, mask_token_idx: int, probs: torch.Tensor, rewalk_ratio: float = 0.1, min_confidence_threshold: float = 0.3, designable_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Re-mask a random subset of low-confidence designable positions.""" confidence = probs.max(dim=-1).values candidates = ( current_seq.ne(mask_token_idx) & confidence.lt(min_confidence_threshold) ) if designable_mask is not None: candidates &= designable_mask.gt(0) selected = torch.rand_like(confidence).lt(rewalk_ratio) & candidates rewalked_seq = current_seq.clone() rewalked_seq[selected] = mask_token_idx return rewalked_seq @torch.no_grad() def iterative_denoise( model: torch.nn.Module, feature_dict: dict, mask_token_idx: int, num_steps: int = 50, schedule_type: str = "cosine", selection_strategy: str = "confidence", temperature: float = 1.0, use_argmax: bool = True, rewalk_enabled: bool = True, rewalk_ratio: float = 0.1, rewalk_threshold: float = 0.3, designable_mask: Optional[torch.Tensor] = None, na_mask: Optional[torch.Tensor] = None, verbose: bool = False, ) -> tuple[torch.Tensor, list[torch.Tensor]]: """Generate a sequence with MaskGIT-style iterative unmasking.""" if designable_mask is None and na_mask is not None: designable_mask = na_mask model.eval() device = next(model.parameters()).device batch_size, length = feature_dict["S"].shape current_seq = feature_dict["S"].clone() if designable_mask is None: current_seq = torch.full( (batch_size, length), mask_token_idx, device=device, dtype=torch.long ) else: current_seq[designable_mask > 0] = mask_token_idx schedule = get_denoising_schedule(num_steps, schedule_type) trajectory = [current_seq.clone()] for step in range(num_steps): step_features = feature_dict.copy() step_features["S"] = current_seq log_probs, probs = model(step_features) if designable_mask is None: total_maskable = torch.full( (batch_size,), length, device=device, dtype=torch.float32 ) else: total_maskable = designable_mask.gt(0).sum(dim=1).float() masked_now = current_seq.eq(mask_token_idx).sum(dim=1).float() masked_target = (total_maskable * schedule[step + 1]).long() to_unmask = (masked_now - masked_target).clamp(min=0).long() average_to_unmask = int(to_unmask.float().mean().item()) if average_to_unmask > 0: current_seq, _ = select_positions_to_unmask( current_seq, log_probs, mask_token_idx, average_to_unmask, strategy=selection_strategy, temperature=temperature, use_argmax=use_argmax, ) if rewalk_enabled and step < num_steps - 1: current_seq = rewalk_positions( current_seq, mask_token_idx, probs, rewalk_ratio=rewalk_ratio, min_confidence_threshold=rewalk_threshold, designable_mask=designable_mask, ) trajectory.append(current_seq.clone()) if verbose and step % 10 == 0: n_masked = current_seq.eq(mask_token_idx).sum().item() print(f"Step {step}/{num_steps}: {n_masked} masked positions remaining") final_features = feature_dict.copy() final_features["S"] = current_seq log_probs, _ = model(final_features) remaining = current_seq.eq(mask_token_idx) if designable_mask is not None: remaining &= designable_mask.gt(0) if remaining.any(): predictions = log_probs.argmax(dim=-1) current_seq[remaining] = predictions[remaining] trajectory.append(current_seq.clone()) return current_seq, trajectory # ============================================================================ # Checkpoint and structure loading # ============================================================================ DEFAULT_PARAMS = { "ATOMS_TO_LOAD": "backbone", "PARSE_PROTEIN": 1, "PARSE_DNA": 1, "PARSE_RNA": 1, "PARSE_RNA_AS_DNA": 0, "NA_SHARED_TOKENS": 1, "PROTEIN_BACKBONE_OCC_CUTOFF": 0.8, "PROTEIN_SIDE_CHAIN_OCC_CUTOFF": 0.5, "DNA_BACKBONE_OCC_CUTOFF": 0.8, "DNA_SIDE_CHAIN_OCC_CUTOFF": 0.5, "RNA_BACKBONE_OCC_CUTOFF": 0.8, "RNA_SIDE_CHAIN_OCC_CUTOFF": 0.5, "CROP_LARGE_STRUCTURES": 1, "BATCH_TOKENS": 6000, "NA_REF_ATOM": "C1'", "EXCLUDE_RES": ["HOH", "NA", "CL", "K", "BR"], "RANDOMIZE_NMR_MODEL": 0, "HIDDEN_DIM": 128, "NUM_ENCODER_LAYERS": 3, "NUM_DECODER_LAYERS": 3, "NUM_NEIGHBORS": 32, "INCLUDE_PRED_NA_N": 1, "USE_SEQUENCE_CONTEXT": True, "VOCAB_SIZE": 33, "NUM_LETTERS": 33, } def _torch_load(path: str | Path, map_location): try: return torch.load(path, map_location=map_location, weights_only=True) except TypeError: return torch.load(path, map_location=map_location) def load_model_and_dataset( checkpoint_path: Optional[str] = None, config_path: Optional[str] = None, model_state_dict: Optional[Mapping[str, torch.Tensor]] = None, device: Optional[torch.device] = None, ): """Create the diffusion model and dataset from a checkpoint or embedded state.""" if device is None: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if model_state_dict is not None: checkpoint = {"model_state_dict": model_state_dict} elif checkpoint_path: checkpoint = _torch_load(checkpoint_path, map_location=device) else: raise ValueError("checkpoint_path or model_state_dict is required") params = DEFAULT_PARAMS.copy() if config_path: with Path(config_path).expanduser().open(encoding="utf-8") as handle: params.update(json.load(handle)) # Atom list if params["ATOMS_TO_LOAD"] == "backbone": atom_list_to_save = [ 'N', 'CA', 'C', 'O', 'OP1', 'OP2', 'P', "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'" ] else: atom_list_to_save = [ 'N', 'CA', 'C', 'CB', 'O', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD', 'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3', 'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2', 'CZ3', 'NZ', 'OXT', 'OP1', 'OP2', 'P', "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", 'N9', 'C8', 'C7', 'N7', 'C6', 'N6', 'O6', 'C5', 'C4', 'N4', 'O4', 'N3', 'C2', 'N2', 'O2', 'N1' ] # Create parsers and dataset cif_parser = CIFParser( skip_res=params.get("EXCLUDE_RES", []), randomize_nmr_model=params.get("RANDOMIZE_NMR_MODEL", 0) ) pdb_parser = PDBParser() pdb_dataset = PDBDataset( cif_parser=cif_parser, pdb_parser=pdb_parser, atom_list_to_save=atom_list_to_save, parse_protein=params["PARSE_PROTEIN"], parse_dna=params["PARSE_DNA"], parse_rna=params["PARSE_RNA"], parse_rna_as_dna=params["PARSE_RNA_AS_DNA"], na_shared_tokens=params["NA_SHARED_TOKENS"], protein_backbone_occ_cutoff=params["PROTEIN_BACKBONE_OCC_CUTOFF"], protein_side_chain_occ_cutoff=params["PROTEIN_SIDE_CHAIN_OCC_CUTOFF"], dna_backbone_occ_cutoff=params["DNA_BACKBONE_OCC_CUTOFF"], dna_side_chain_occ_cutoff=params["DNA_SIDE_CHAIN_OCC_CUTOFF"], rna_backbone_occ_cutoff=params["RNA_BACKBONE_OCC_CUTOFF"], rna_side_chain_occ_cutoff=params["RNA_SIDE_CHAIN_OCC_CUTOFF"], crop_large_structures=params["CROP_LARGE_STRUCTURES"], batch_tokens=params["BATCH_TOKENS"], na_ref_atom=params["NA_REF_ATOM"], ) # Create model model = ProteinMPNNDiffusion( node_features=params["HIDDEN_DIM"], edge_features=params["HIDDEN_DIM"], hidden_dim=params["HIDDEN_DIM"], num_encoder_layers=params["NUM_ENCODER_LAYERS"], num_decoder_layers=params["NUM_DECODER_LAYERS"], k_neighbors=params["NUM_NEIGHBORS"], dropout=0.0, # No dropout during inference atom_dict=pdb_dataset.atom_dict, restype_to_int=pdb_dataset.restype_to_int, polytype_to_int=pdb_dataset.polytype_to_int, protein_augment_eps=0.0, dna_augment_eps=0.0, rna_augment_eps=0.0, na_ref_atom=params["NA_REF_ATOM"], include_pred_na_N=params["INCLUDE_PRED_NA_N"], use_sequence_context=params.get("USE_SEQUENCE_CONTEXT", True), device=device, vocab=params["VOCAB_SIZE"], num_letters=params["NUM_LETTERS"] ) # Load weights model.load_state_dict(checkpoint['model_state_dict']) model.to(device) model.eval() return model, pdb_dataset, params, device def load_structure( pdb_path: str, pdb_dataset: PDBDataset, device: torch.device, max_tokens: Optional[int] = None, ): """Load and featurize a structure, optionally using NAIAD's NA-centered crop.""" lower_path = str(pdb_path).lower() if lower_path.endswith((".pdb", ".pdb.gz")): chains, asmb, _covalei, _meta = pdb_dataset.pdb_parser.parse(pdb_path) elif lower_path.endswith((".cif", ".cif.gz", ".mmcif", ".mmcif.gz")): chains, asmb, _covalei, _meta = pdb_dataset.cif_parser.parse(pdb_path) else: raise ValueError(f"Unsupported structure format: {pdb_path}") # Load chains macromolecule_chain_dict = pdb_dataset.load_chains(chains) # Get first assembly if not asmb: raise ValueError(f"No usable assembly found in {pdb_path}") assembly_id = next(iter(asmb)) # Load assembly out_dict = pdb_dataset.load_assembly( macromolecule_chain_dict, asmb, assembly_id, ppms=[] ) # Create dummy preprocessed data (zeros for inference) L = out_dict["macromolecule_L"] out_dict["original_macromolecule_L"] = int(L) out_dict["interface_mask"] = np.zeros(L, dtype=np.int32) out_dict["side_chain_interface_mask"] = np.zeros(L, dtype=np.int32) out_dict["nearest_protein_side_chain_index"] = np.zeros(L, dtype=np.int64) out_dict["base_pair_mask"] = np.zeros(L, dtype=np.int32) out_dict["base_pair_index"] = np.zeros(L, dtype=np.int64) out_dict["canonical_base_pair_mask"] = np.zeros(L, dtype=np.int32) out_dict["canonical_base_pair_index"] = np.zeros(L, dtype=np.int64) if max_tokens is None and pdb_dataset.crop_large_structures: max_tokens = pdb_dataset.batch_tokens out_dict["crop_applied"] = False if max_tokens and L > max_tokens and (out_dict["dna_L"] + out_dict["rna_L"]) > 0: original_batch_tokens = pdb_dataset.batch_tokens try: pdb_dataset.batch_tokens = int(max_tokens) pdb_dataset.random_crop_na(out_dict) finally: pdb_dataset.batch_tokens = original_batch_tokens L = out_dict["macromolecule_L"] out_dict["crop_applied"] = True out_dict["structure_path"] = pdb_path out_dict["assembly_id"] = assembly_id out_dict["ppm_paths"] = "[]" out_dict["ppm_paths_chosen"] = [] for key, value in list(out_dict.items()): if isinstance(value, np.ndarray): out_dict[key] = torch.from_numpy(value) # Create batch batch = [[(out_dict, torch.tensor(L, dtype=torch.long))]] # Featurize feature_dict = featurize( batch[0], pdb_dataset.polytype_to_int, pdb_dataset.restype_to_int, pdb_dataset.atom_dict, device ) return feature_dict, out_dict def sequence_to_string( seq_tensor: torch.Tensor, pdb_dataset: PDBDataset, mask: Optional[torch.Tensor] = None, dna_mask: Optional[torch.Tensor] = None, rna_mask: Optional[torch.Tensor] = None, chain_labels: Optional[torch.Tensor] = None, chain_break_character: str = "/", ) -> str: """Convert sequence tensor to string representation.""" seq = seq_tensor.cpu().numpy() if mask is not None: mask = mask.cpu().numpy() if dna_mask is not None: dna_mask = dna_mask.cpu().numpy() if rna_mask is not None: rna_mask = rna_mask.cpu().numpy() if chain_labels is not None: chain_labels = chain_labels.cpu().numpy() dna_chars_by_restype = { "DA": "a", "A": "a", "DC": "c", "C": "c", "DG": "g", "G": "g", "DT": "t", "U": "t", "DX": "x", "RX": "x", } rna_chars_by_restype = { "DA": "b", "A": "b", "DC": "d", "C": "d", "DG": "h", "G": "h", "DT": "u", "U": "u", "DX": "y", "RX": "y", } result = [] previous_chain_label = None for i, token_idx in enumerate(seq): if mask is not None and mask[i] == 0: continue restype = pdb_dataset.int_to_restype.get(token_idx, 'X') if rna_mask is not None and rna_mask[i] > 0: char = rna_chars_by_restype.get(restype, 'X') elif dna_mask is not None and dna_mask[i] > 0: char = dna_chars_by_restype.get(restype, 'X') else: char = pdb_dataset.restype_3_to_1.get(restype, 'X') if chain_labels is not None: chain_label = chain_labels[i] if ( previous_chain_label is not None and chain_label != previous_chain_label ): result.append(chain_break_character) previous_chain_label = chain_label result.append(char) return ''.join(result) __all__ = [ "PDBDataset", "ProteinMPNNDiffusion", "featurize", "iterative_denoise", "load_model_and_dataset", "load_structure", "sequence_to_string", ]