script/inference_utils.py

"""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",
]