boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/data/feature/featurizerv2.py

@@ -0,0 +1,2208 @@
+import math
+from typing import Optional
+
+import numba
+import numpy as np
+import numpy.typing as npt
+import rdkit.Chem.Descriptors
+import torch
+from numba import types
+from rdkit.Chem import Mol
+from scipy.spatial.distance import cdist
+from torch import Tensor, from_numpy
+from torch.nn.functional import one_hot
+
+from boltz.data import const
+from boltz.data.mol import (
+    get_amino_acids_symmetries,
+    get_chain_symmetries,
+    get_ligand_symmetries,
+    get_symmetries,
+)
+from boltz.data.pad import pad_dim
+from boltz.data.types import (
+    MSA,
+    MSADeletion,
+    MSAResidue,
+    MSASequence,
+    TemplateInfo,
+    Tokenized,
+)
+from boltz.model.modules.utils import center_random_augmentation
+
+####################################################################################################
+# HELPERS
+####################################################################################################
+
+
+def convert_atom_name(name: str) -> tuple[int, int, int, int]:
+    """Convert an atom name to a standard format.
+
+    Parameters
+    ----------
+    name : str
+        The atom name.
+
+    Returns
+    -------
+    tuple[int, int, int, int]
+        The converted atom name.
+
+    """
+    name = str(name).strip()
+    name = [ord(c) - 32 for c in name]
+    name = name + [0] * (4 - len(name))
+    return tuple(name)
+
+
+def sample_d(
+    min_d: float,
+    max_d: float,
+    n_samples: int,
+    random: np.random.Generator,
+) -> np.ndarray:
+    """Generate samples from a 1/d distribution between min_d and max_d.
+
+    Parameters
+    ----------
+    min_d : float
+        Minimum value of d
+    max_d : float
+        Maximum value of d
+    n_samples : int
+        Number of samples to generate
+    random : numpy.random.Generator
+        Random number generator
+
+    Returns
+    -------
+    numpy.ndarray
+        Array of samples drawn from the distribution
+
+    Notes
+    -----
+    The probability density function is:
+    f(d) = 1/(d * ln(max_d/min_d)) for d in [min_d, max_d]
+
+    The inverse CDF transform is:
+    d = min_d * (max_d/min_d)**u where u ~ Uniform(0,1)
+
+    """
+    # Generate n_samples uniform random numbers in [0, 1]
+    u = random.random(n_samples)
+    # Transform u using the inverse CDF
+    return min_d * (max_d / min_d) ** u
+
+
+def compute_frames_nonpolymer(
+    data: Tokenized,
+    coords,
+    resolved_mask,
+    atom_to_token,
+    frame_data: list,
+    resolved_frame_data: list,
+) -> tuple[list, list]:
+    """Get the frames for non-polymer tokens.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input data to the model.
+    frame_data : list
+        The frame data.
+    resolved_frame_data : list
+        The resolved frame data.
+
+    Returns
+    -------
+    tuple[list, list]
+        The frame data and resolved frame data.
+
+    """
+    frame_data = np.array(frame_data)
+    resolved_frame_data = np.array(resolved_frame_data)
+    asym_id_token = data.tokens["asym_id"]
+    asym_id_atom = data.tokens["asym_id"][atom_to_token]
+    token_idx = 0
+    atom_idx = 0
+    for id in np.unique(data.tokens["asym_id"]):
+        mask_chain_token = asym_id_token == id
+        mask_chain_atom = asym_id_atom == id
+        num_tokens = mask_chain_token.sum()
+        num_atoms = mask_chain_atom.sum()
+        if (
+            data.tokens[token_idx]["mol_type"] != const.chain_type_ids["NONPOLYMER"]
+            or num_atoms < 3  # noqa: PLR2004
+        ):
+            token_idx += num_tokens
+            atom_idx += num_atoms
+            continue
+        dist_mat = (
+            (
+                coords.reshape(-1, 3)[mask_chain_atom][:, None, :]
+                - coords.reshape(-1, 3)[mask_chain_atom][None, :, :]
+            )
+            ** 2
+        ).sum(-1) ** 0.5
+        resolved_pair = 1 - (
+            resolved_mask[mask_chain_atom][None, :]
+            * resolved_mask[mask_chain_atom][:, None]
+        ).astype(np.float32)
+        resolved_pair[resolved_pair == 1] = math.inf
+        indices = np.argsort(dist_mat + resolved_pair, axis=1)
+        frames = (
+            np.concatenate(
+                [
+                    indices[:, 1:2],
+                    indices[:, 0:1],
+                    indices[:, 2:3],
+                ],
+                axis=1,
+            )
+            + atom_idx
+        )
+        frame_data[token_idx : token_idx + num_atoms, :] = frames
+        resolved_frame_data[token_idx : token_idx + num_atoms] = resolved_mask[
+            frames
+        ].all(axis=1)
+        token_idx += num_tokens
+        atom_idx += num_atoms
+    frames_expanded = coords.reshape(-1, 3)[frame_data]
+
+    mask_collinear = compute_collinear_mask(
+        frames_expanded[:, 1] - frames_expanded[:, 0],
+        frames_expanded[:, 1] - frames_expanded[:, 2],
+    )
+    return frame_data, resolved_frame_data & mask_collinear
+
+
+def compute_collinear_mask(v1, v2):
+    norm1 = np.linalg.norm(v1, axis=1, keepdims=True)
+    norm2 = np.linalg.norm(v2, axis=1, keepdims=True)
+    v1 = v1 / (norm1 + 1e-6)
+    v2 = v2 / (norm2 + 1e-6)
+    mask_angle = np.abs(np.sum(v1 * v2, axis=1)) < 0.9063
+    mask_overlap1 = norm1.reshape(-1) > 1e-2
+    mask_overlap2 = norm2.reshape(-1) > 1e-2
+    return mask_angle & mask_overlap1 & mask_overlap2
+
+
+def dummy_msa(residues: np.ndarray) -> MSA:
+    """Create a dummy MSA for a chain.
+
+    Parameters
+    ----------
+    residues : np.ndarray
+        The residues for the chain.
+
+    Returns
+    -------
+    MSA
+        The dummy MSA.
+
+    """
+    residues = [res["res_type"] for res in residues]
+    deletions = []
+    sequences = [(0, -1, 0, len(residues), 0, 0)]
+    return MSA(
+        residues=np.array(residues, dtype=MSAResidue),
+        deletions=np.array(deletions, dtype=MSADeletion),
+        sequences=np.array(sequences, dtype=MSASequence),
+    )
+
+
+def construct_paired_msa(  # noqa: C901, PLR0915, PLR0912
+    data: Tokenized,
+    random: np.random.Generator,
+    max_seqs: int,
+    max_pairs: int = 8192,
+    max_total: int = 16384,
+    random_subset: bool = False,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Pair the MSA data.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input data to the model.
+
+    Returns
+    -------
+    Tensor
+        The MSA data.
+    Tensor
+        The deletion data.
+    Tensor
+        Mask indicating paired sequences.
+
+    """
+    # Get unique chains (ensuring monotonicity in the order)
+    assert np.all(np.diff(data.tokens["asym_id"], n=1) >= 0)
+    chain_ids = np.unique(data.tokens["asym_id"])
+
+    # Get relevant MSA, and create a dummy for chains without
+    msa: dict[int, MSA] = {}
+    for chain_id in chain_ids:
+        # Get input sequence
+        chain = data.structure.chains[chain_id]
+        res_start = chain["res_idx"]
+        res_end = res_start + chain["res_num"]
+        residues = data.structure.residues[res_start:res_end]
+
+        # Check if we have an MSA, and that the
+        # first sequence matches the input sequence
+        if chain_id in data.msa:
+            # Set the MSA
+            msa[chain_id] = data.msa[chain_id]
+
+            # Run length and residue type checks
+            first = data.msa[chain_id].sequences[0]
+            first_start = first["res_start"]
+            first_end = first["res_end"]
+            msa_residues = data.msa[chain_id].residues
+            first_residues = msa_residues[first_start:first_end]
+
+            warning = "Warning: MSA does not match input sequence, creating dummy."
+            if len(residues) == len(first_residues):
+                # If there is a mismatch, check if it is between MET & UNK
+                # If so, replace the first sequence with the input sequence.
+                # Otherwise, replace with a dummy MSA for this chain.
+                mismatches = residues["res_type"] != first_residues["res_type"]
+                if mismatches.sum().item():
+                    idx = np.where(mismatches)[0]
+                    is_met = residues["res_type"][idx] == const.token_ids["MET"]
+                    is_unk = residues["res_type"][idx] == const.token_ids["UNK"]
+                    is_msa_unk = (
+                        first_residues["res_type"][idx] == const.token_ids["UNK"]
+                    )
+                    if (np.all(is_met) and np.all(is_msa_unk)) or np.all(is_unk):
+                        msa_residues[first_start:first_end]["res_type"] = residues[
+                            "res_type"
+                        ]
+                    else:
+                        print(
+                            warning,
+                            "1",
+                            residues["res_type"],
+                            first_residues["res_type"],
+                            data.record.id,
+                        )
+                        msa[chain_id] = dummy_msa(residues)
+            else:
+                print(
+                    warning,
+                    "2",
+                    residues["res_type"],
+                    first_residues["res_type"],
+                    data.record.id,
+                )
+                msa[chain_id] = dummy_msa(residues)
+        else:
+            msa[chain_id] = dummy_msa(residues)
+
+    # Map taxonomies to (chain_id, seq_idx)
+    taxonomy_map: dict[str, list] = {}
+    for chain_id, chain_msa in msa.items():
+        sequences = chain_msa.sequences
+        sequences = sequences[sequences["taxonomy"] != -1]
+        for sequence in sequences:
+            seq_idx = sequence["seq_idx"]
+            taxon = sequence["taxonomy"]
+            taxonomy_map.setdefault(taxon, []).append((chain_id, seq_idx))
+
+    # Remove taxonomies with only one sequence and sort by the
+    # number of chain_id present in each of the taxonomies
+    taxonomy_map = {k: v for k, v in taxonomy_map.items() if len(v) > 1}
+    taxonomy_map = sorted(
+        taxonomy_map.items(),
+        key=lambda x: len({c for c, _ in x[1]}),
+        reverse=True,
+    )
+
+    # Keep track of the sequences available per chain, keeping the original
+    # order of the sequences in the MSA to favor the best matching sequences
+    visited = {(c, s) for c, items in taxonomy_map for s in items}
+    available = {}
+    for c in chain_ids:
+        available[c] = [
+            i for i in range(1, len(msa[c].sequences)) if (c, i) not in visited
+        ]
+
+    # Create sequence pairs
+    is_paired = []
+    pairing = []
+
+    # Start with the first sequence for each chain
+    is_paired.append({c: 1 for c in chain_ids})
+    pairing.append({c: 0 for c in chain_ids})
+
+    # Then add up to 8191 paired rows
+    for _, pairs in taxonomy_map:
+        # Group occurences by chain_id in case we have multiple
+        # sequences from the same chain and same taxonomy
+        chain_occurences = {}
+        for chain_id, seq_idx in pairs:
+            chain_occurences.setdefault(chain_id, []).append(seq_idx)
+
+        # We create as many pairings as the maximum number of occurences
+        max_occurences = max(len(v) for v in chain_occurences.values())
+        for i in range(max_occurences):
+            row_pairing = {}
+            row_is_paired = {}
+
+            # Add the chains present in the taxonomy
+            for chain_id, seq_idxs in chain_occurences.items():
+                # Roll over the sequence index to maximize diversity
+                idx = i % len(seq_idxs)
+                seq_idx = seq_idxs[idx]
+
+                # Add the sequence to the pairing
+                row_pairing[chain_id] = seq_idx
+                row_is_paired[chain_id] = 1
+
+            # Add any missing chains
+            for chain_id in chain_ids:
+                if chain_id not in row_pairing:
+                    row_is_paired[chain_id] = 0
+                    if available[chain_id]:
+                        # Add the next available sequence
+                        seq_idx = available[chain_id].pop(0)
+                        row_pairing[chain_id] = seq_idx
+                    else:
+                        # No more sequences available, we place a gap
+                        row_pairing[chain_id] = -1
+
+            pairing.append(row_pairing)
+            is_paired.append(row_is_paired)
+
+            # Break if we have enough pairs
+            if len(pairing) >= max_pairs:
+                break
+
+        # Break if we have enough pairs
+        if len(pairing) >= max_pairs:
+            break
+
+    # Now add up to 16384 unpaired rows total
+    max_left = max(len(v) for v in available.values())
+    for _ in range(min(max_total - len(pairing), max_left)):
+        row_pairing = {}
+        row_is_paired = {}
+        for chain_id in chain_ids:
+            row_is_paired[chain_id] = 0
+            if available[chain_id]:
+                # Add the next available sequence
+                seq_idx = available[chain_id].pop(0)
+                row_pairing[chain_id] = seq_idx
+            else:
+                # No more sequences available, we place a gap
+                row_pairing[chain_id] = -1
+
+        pairing.append(row_pairing)
+        is_paired.append(row_is_paired)
+
+        # Break if we have enough sequences
+        if len(pairing) >= max_total:
+            break
+
+    # Randomly sample a subset of the pairs
+    # ensuring the first row is always present
+    if random_subset:
+        num_seqs = len(pairing)
+        if num_seqs > max_seqs:
+            indices = random.choice(
+                np.arange(1, num_seqs), size=max_seqs - 1, replace=False
+            )  # noqa: NPY002
+            pairing = [pairing[0]] + [pairing[i] for i in indices]
+            is_paired = [is_paired[0]] + [is_paired[i] for i in indices]
+    else:
+        # Deterministic downsample to max_seqs
+        pairing = pairing[:max_seqs]
+        is_paired = is_paired[:max_seqs]
+
+    # Map (chain_id, seq_idx, res_idx) to deletion
+    deletions = {}
+    for chain_id, chain_msa in msa.items():
+        chain_deletions = chain_msa.deletions
+        for sequence in chain_msa.sequences:
+            del_start = sequence["del_start"]
+            del_end = sequence["del_end"]
+            chain_deletions = chain_msa.deletions[del_start:del_end]
+            for deletion_data in chain_deletions:
+                seq_idx = sequence["seq_idx"]
+                res_idx = deletion_data["res_idx"]
+                deletion = deletion_data["deletion"]
+                deletions[(chain_id, seq_idx, res_idx)] = deletion
+
+    # Add all the token MSA data
+    msa_data, del_data, paired_data = prepare_msa_arrays(
+        data.tokens, pairing, is_paired, deletions, msa
+    )
+
+    msa_data = torch.tensor(msa_data, dtype=torch.long)
+    del_data = torch.tensor(del_data, dtype=torch.float)
+    paired_data = torch.tensor(paired_data, dtype=torch.float)
+
+    return msa_data, del_data, paired_data
+
+
+def prepare_msa_arrays(
+    tokens,
+    pairing: list[dict[int, int]],
+    is_paired: list[dict[int, int]],
+    deletions: dict[tuple[int, int, int], int],
+    msa: dict[int, MSA],
+) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
+    """Reshape data to play nicely with numba jit."""
+    token_asym_ids_arr = np.array([t["asym_id"] for t in tokens], dtype=np.int64)
+    token_res_idxs_arr = np.array([t["res_idx"] for t in tokens], dtype=np.int64)
+
+    chain_ids = sorted(msa.keys())
+
+    # chain_ids are not necessarily contiguous (e.g. they might be 0, 24, 25).
+    # This allows us to look up a chain_id by it's index in the chain_ids list.
+    chain_id_to_idx = {chain_id: i for i, chain_id in enumerate(chain_ids)}
+    token_asym_ids_idx_arr = np.array(
+        [chain_id_to_idx[asym_id] for asym_id in token_asym_ids_arr], dtype=np.int64
+    )
+
+    pairing_arr = np.zeros((len(pairing), len(chain_ids)), dtype=np.int64)
+    is_paired_arr = np.zeros((len(is_paired), len(chain_ids)), dtype=np.int64)
+
+    for i, row_pairing in enumerate(pairing):
+        for chain_id in chain_ids:
+            pairing_arr[i, chain_id_to_idx[chain_id]] = row_pairing[chain_id]
+
+    for i, row_is_paired in enumerate(is_paired):
+        for chain_id in chain_ids:
+            is_paired_arr[i, chain_id_to_idx[chain_id]] = row_is_paired[chain_id]
+
+    max_seq_len = max(len(msa[chain_id].sequences) for chain_id in chain_ids)
+
+    # we want res_start from sequences
+    msa_sequences = np.full((len(chain_ids), max_seq_len), -1, dtype=np.int64)
+    for chain_id in chain_ids:
+        for i, seq in enumerate(msa[chain_id].sequences):
+            msa_sequences[chain_id_to_idx[chain_id], i] = seq["res_start"]
+
+    max_residues_len = max(len(msa[chain_id].residues) for chain_id in chain_ids)
+    msa_residues = np.full((len(chain_ids), max_residues_len), -1, dtype=np.int64)
+    for chain_id in chain_ids:
+        residues = msa[chain_id].residues.astype(np.int64)
+        idxs = np.arange(len(residues))
+        chain_idx = chain_id_to_idx[chain_id]
+        msa_residues[chain_idx, idxs] = residues
+
+    deletions_dict = numba.typed.Dict.empty(
+        key_type=numba.types.Tuple(
+            [numba.types.int64, numba.types.int64, numba.types.int64]
+        ),
+        value_type=numba.types.int64,
+    )
+    deletions_dict.update(deletions)
+
+    return _prepare_msa_arrays_inner(
+        token_asym_ids_arr,
+        token_res_idxs_arr,
+        token_asym_ids_idx_arr,
+        pairing_arr,
+        is_paired_arr,
+        deletions_dict,
+        msa_sequences,
+        msa_residues,
+        const.token_ids["-"],
+    )
+
+
+deletions_dict_type = types.DictType(types.UniTuple(types.int64, 3), types.int64)
+
+
+@numba.njit(
+    [
+        types.Tuple(
+            (
+                types.int64[:, ::1],  # msa_data
+                types.int64[:, ::1],  # del_data
+                types.int64[:, ::1],  # paired_data
+            )
+        )(
+            types.int64[::1],  # token_asym_ids
+            types.int64[::1],  # token_res_idxs
+            types.int64[::1],  # token_asym_ids_idx
+            types.int64[:, ::1],  # pairing
+            types.int64[:, ::1],  # is_paired
+            deletions_dict_type,  # deletions
+            types.int64[:, ::1],  # msa_sequences
+            types.int64[:, ::1],  # msa_residues
+            types.int64,  # gap_token
+        )
+    ],
+    cache=True,
+)
+def _prepare_msa_arrays_inner(
+    token_asym_ids: npt.NDArray[np.int64],
+    token_res_idxs: npt.NDArray[np.int64],
+    token_asym_ids_idx: npt.NDArray[np.int64],
+    pairing: npt.NDArray[np.int64],
+    is_paired: npt.NDArray[np.int64],
+    deletions: dict[tuple[int, int, int], int],
+    msa_sequences: npt.NDArray[np.int64],
+    msa_residues: npt.NDArray[np.int64],
+    gap_token: int,
+) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
+    n_tokens = len(token_asym_ids)
+    n_pairs = len(pairing)
+    msa_data = np.full((n_tokens, n_pairs), gap_token, dtype=np.int64)
+    paired_data = np.zeros((n_tokens, n_pairs), dtype=np.int64)
+    del_data = np.zeros((n_tokens, n_pairs), dtype=np.int64)
+
+    # Add all the token MSA data
+    for token_idx in range(n_tokens):
+        chain_id_idx = token_asym_ids_idx[token_idx]
+        chain_id = token_asym_ids[token_idx]
+        res_idx = token_res_idxs[token_idx]
+
+        for pair_idx in range(n_pairs):
+            seq_idx = pairing[pair_idx, chain_id_idx]
+            paired_data[token_idx, pair_idx] = is_paired[pair_idx, chain_id_idx]
+
+            # Add residue type
+            if seq_idx != -1:
+                res_start = msa_sequences[chain_id_idx, seq_idx]
+                res_type = msa_residues[chain_id_idx, res_start + res_idx]
+                k = (chain_id, seq_idx, res_idx)
+                if k in deletions:
+                    del_data[token_idx, pair_idx] = deletions[k]
+                msa_data[token_idx, pair_idx] = res_type
+
+    return msa_data, del_data, paired_data
+
+
+####################################################################################################
+# FEATURES
+####################################################################################################
+
+
+def select_subset_from_mask(mask, p, random: np.random.Generator) -> np.ndarray:
+    num_true = np.sum(mask)
+    v = random.geometric(p) + 1
+    k = min(v, num_true)
+
+    true_indices = np.where(mask)[0]
+
+    # Randomly select k indices from the true_indices
+    selected_indices = random.choice(true_indices, size=k, replace=False)
+
+    new_mask = np.zeros_like(mask)
+    new_mask[selected_indices] = 1
+
+    return new_mask
+
+
+def get_range_bin(value: float, range_dict: dict[tuple[float, float], int], default=0):
+    """Get the bin of a value given a range dictionary."""
+    value = float(value)
+    for k, idx in range_dict.items():
+        if k == "other":
+            continue
+        low, high = k
+        if low <= value < high:
+            return idx
+    return default
+
+
+def process_token_features(  # noqa: C901, PLR0915, PLR0912
+    data: Tokenized,
+    random: np.random.Generator,
+    max_tokens: Optional[int] = None,
+    binder_pocket_conditioned_prop: Optional[float] = 0.0,
+    contact_conditioned_prop: Optional[float] = 0.0,
+    binder_pocket_cutoff_min: Optional[float] = 4.0,
+    binder_pocket_cutoff_max: Optional[float] = 20.0,
+    binder_pocket_sampling_geometric_p: Optional[float] = 0.0,
+    only_ligand_binder_pocket: Optional[bool] = False,
+    only_pp_contact: Optional[bool] = False,
+    inference_pocket_constraints: Optional[bool] = False,
+    override_method: Optional[str] = None,
+) -> dict[str, Tensor]:
+    """Get the token features.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input data to the model.
+    max_tokens : int
+        The maximum number of tokens.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        The token features.
+
+    """
+    # Token data
+    token_data = data.tokens
+    token_bonds = data.bonds
+
+    # Token core features
+    token_index = torch.arange(len(token_data), dtype=torch.long)
+    residue_index = from_numpy(token_data["res_idx"]).long()
+    asym_id = from_numpy(token_data["asym_id"]).long()
+    entity_id = from_numpy(token_data["entity_id"]).long()
+    sym_id = from_numpy(token_data["sym_id"]).long()
+    mol_type = from_numpy(token_data["mol_type"]).long()
+    res_type = from_numpy(token_data["res_type"]).long()
+    res_type = one_hot(res_type, num_classes=const.num_tokens)
+    disto_center = from_numpy(token_data["disto_coords"])
+    modified = from_numpy(token_data["modified"]).long()  # float()
+    cyclic_period = from_numpy(token_data["cyclic_period"].copy())
+    affinity_mask = from_numpy(token_data["affinity_mask"]).float()
+
+    ## Conditioning features ##
+    method = (
+        np.zeros(len(token_data))
+        + const.method_types_ids[
+            (
+                "x-ray diffraction"
+                if override_method is None
+                else override_method.lower()
+            )
+        ]
+    )
+    if data.record is not None:
+        if (
+            override_method is None
+            and data.record.structure.method is not None
+            and data.record.structure.method.lower() in const.method_types_ids
+        ):
+            method = (method * 0) + const.method_types_ids[
+                data.record.structure.method.lower()
+            ]
+
+    method_feature = from_numpy(method).long()
+
+    # Token mask features
+    pad_mask = torch.ones(len(token_data), dtype=torch.float)
+    resolved_mask = from_numpy(token_data["resolved_mask"]).float()
+    disto_mask = from_numpy(token_data["disto_mask"]).float()
+
+    # Token bond features
+    if max_tokens is not None:
+        pad_len = max_tokens - len(token_data)
+        num_tokens = max_tokens if pad_len > 0 else len(token_data)
+    else:
+        num_tokens = len(token_data)
+
+    tok_to_idx = {tok["token_idx"]: idx for idx, tok in enumerate(token_data)}
+    bonds = torch.zeros(num_tokens, num_tokens, dtype=torch.float)
+    bonds_type = torch.zeros(num_tokens, num_tokens, dtype=torch.long)
+    for token_bond in token_bonds:
+        token_1 = tok_to_idx[token_bond["token_1"]]
+        token_2 = tok_to_idx[token_bond["token_2"]]
+        bonds[token_1, token_2] = 1
+        bonds[token_2, token_1] = 1
+        bond_type = token_bond["type"]
+        bonds_type[token_1, token_2] = bond_type
+        bonds_type[token_2, token_1] = bond_type
+
+    bonds = bonds.unsqueeze(-1)
+
+    # Pocket conditioned feature
+    contact_conditioning = (
+        np.zeros((len(token_data), len(token_data)))
+        + const.contact_conditioning_info["UNSELECTED"]
+    )
+    contact_threshold = np.zeros((len(token_data), len(token_data)))
+
+    if inference_pocket_constraints is not None:
+        for binder, contacts, max_distance in inference_pocket_constraints:
+            binder_mask = token_data["asym_id"] == binder
+
+            for idx, token in enumerate(token_data):
+                if (
+                    token["mol_type"] != const.chain_type_ids["NONPOLYMER"]
+                    and (token["asym_id"], token["res_idx"]) in contacts
+                ) or (
+                    token["mol_type"] == const.chain_type_ids["NONPOLYMER"]
+                    and (token["asym_id"], token["atom_idx"]) in contacts
+                ):
+                    contact_conditioning[binder_mask][:, idx] = (
+                        const.contact_conditioning_info["BINDER>POCKET"]
+                    )
+                    contact_conditioning[idx][binder_mask] = (
+                        const.contact_conditioning_info["POCKET>BINDER"]
+                    )
+                    contact_threshold[binder_mask][:, idx] = max_distance
+                    contact_threshold[idx][binder_mask] = max_distance
+
+    if binder_pocket_conditioned_prop > 0.0:
+        # choose as binder a random ligand in the crop, if there are no ligands select a protein chain
+        binder_asym_ids = np.unique(
+            token_data["asym_id"][
+                token_data["mol_type"] == const.chain_type_ids["NONPOLYMER"]
+            ]
+        )
+
+        if len(binder_asym_ids) == 0:
+            if not only_ligand_binder_pocket:
+                binder_asym_ids = np.unique(token_data["asym_id"])
+
+        while random.random() < binder_pocket_conditioned_prop:
+            if len(binder_asym_ids) == 0:
+                break
+
+            pocket_asym_id = random.choice(binder_asym_ids)
+            binder_asym_ids = binder_asym_ids[binder_asym_ids != pocket_asym_id]
+
+            binder_pocket_cutoff = sample_d(
+                min_d=binder_pocket_cutoff_min,
+                max_d=binder_pocket_cutoff_max,
+                n_samples=1,
+                random=random,
+            )
+
+            binder_mask = token_data["asym_id"] == pocket_asym_id
+
+            binder_coords = []
+            for token in token_data:
+                if token["asym_id"] == pocket_asym_id:
+                    _coords = data.structure.atoms["coords"][
+                        token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                    ]
+                    _is_present = data.structure.atoms["is_present"][
+                        token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                    ]
+                    binder_coords.append(_coords[_is_present])
+            binder_coords = np.concatenate(binder_coords, axis=0)
+
+            # find the tokens in the pocket
+            token_dist = np.zeros(len(token_data)) + 1000
+            for i, token in enumerate(token_data):
+                if (
+                    token["mol_type"] != const.chain_type_ids["NONPOLYMER"]
+                    and token["asym_id"] != pocket_asym_id
+                    and token["resolved_mask"] == 1
+                ):
+                    token_coords = data.structure.atoms["coords"][
+                        token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                    ]
+                    token_is_present = data.structure.atoms["is_present"][
+                        token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                    ]
+                    token_coords = token_coords[token_is_present]
+
+                    # find chain and apply chain transformation
+                    for chain in data.structure.chains:
+                        if chain["asym_id"] == token["asym_id"]:
+                            break
+
+                    token_dist[i] = np.min(
+                        np.linalg.norm(
+                            token_coords[:, None, :] - binder_coords[None, :, :],
+                            axis=-1,
+                        )
+                    )
+
+            pocket_mask = token_dist < binder_pocket_cutoff
+
+            if np.sum(pocket_mask) > 0:
+                if binder_pocket_sampling_geometric_p > 0.0:
+                    # select a subset of the pocket, according
+                    # to a geometric distribution with one as minimum
+                    pocket_mask = select_subset_from_mask(
+                        pocket_mask,
+                        binder_pocket_sampling_geometric_p,
+                        random,
+                    )
+
+                contact_conditioning[np.ix_(binder_mask, pocket_mask)] = (
+                    const.contact_conditioning_info["BINDER>POCKET"]
+                )
+                contact_conditioning[np.ix_(pocket_mask, binder_mask)] = (
+                    const.contact_conditioning_info["POCKET>BINDER"]
+                )
+                contact_threshold[np.ix_(binder_mask, pocket_mask)] = (
+                    binder_pocket_cutoff
+                )
+                contact_threshold[np.ix_(pocket_mask, binder_mask)] = (
+                    binder_pocket_cutoff
+                )
+
+    # Contact conditioning feature
+    if contact_conditioned_prop > 0.0:
+        while random.random() < contact_conditioned_prop:
+            contact_cutoff = sample_d(
+                min_d=binder_pocket_cutoff_min,
+                max_d=binder_pocket_cutoff_max,
+                n_samples=1,
+                random=random,
+            )
+            if only_pp_contact:
+                chain_asym_ids = np.unique(
+                    token_data["asym_id"][
+                        token_data["mol_type"] == const.chain_type_ids["PROTEIN"]
+                    ]
+                )
+            else:
+                chain_asym_ids = np.unique(token_data["asym_id"])
+
+            if len(chain_asym_ids) > 1:
+                chain_asym_id = random.choice(chain_asym_ids)
+
+                chain_coords = []
+                for token in token_data:
+                    if token["asym_id"] == chain_asym_id:
+                        _coords = data.structure.atoms["coords"][
+                            token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                        ]
+                        _is_present = data.structure.atoms["is_present"][
+                            token["atom_idx"] : token["atom_idx"] + token["atom_num"]
+                        ]
+                        chain_coords.append(_coords[_is_present])
+                chain_coords = np.concatenate(chain_coords, axis=0)
+
+                # find contacts in other chains
+                possible_other_chains = []
+                for other_chain_id in chain_asym_ids[chain_asym_ids != chain_asym_id]:
+                    for token in token_data:
+                        if token["asym_id"] == other_chain_id:
+                            _coords = data.structure.atoms["coords"][
+                                token["atom_idx"] : token["atom_idx"]
+                                + token["atom_num"]
+                            ]
+                            _is_present = data.structure.atoms["is_present"][
+                                token["atom_idx"] : token["atom_idx"]
+                                + token["atom_num"]
+                            ]
+                            if _is_present.sum() == 0:
+                                continue
+                            token_coords = _coords[_is_present]
+
+                            # check minimum distance
+                            if (
+                                np.min(cdist(chain_coords, token_coords))
+                                < contact_cutoff
+                            ):
+                                possible_other_chains.append(other_chain_id)
+                                break
+
+                if len(possible_other_chains) > 0:
+                    other_chain_id = random.choice(possible_other_chains)
+
+                    pairs = []
+                    for token_1 in token_data:
+                        if token_1["asym_id"] == chain_asym_id:
+                            _coords = data.structure.atoms["coords"][
+                                token_1["atom_idx"] : token_1["atom_idx"]
+                                + token_1["atom_num"]
+                            ]
+                            _is_present = data.structure.atoms["is_present"][
+                                token_1["atom_idx"] : token_1["atom_idx"]
+                                + token_1["atom_num"]
+                            ]
+                            if _is_present.sum() == 0:
+                                continue
+                            token_1_coords = _coords[_is_present]
+
+                            for token_2 in token_data:
+                                if token_2["asym_id"] == other_chain_id:
+                                    _coords = data.structure.atoms["coords"][
+                                        token_2["atom_idx"] : token_2["atom_idx"]
+                                        + token_2["atom_num"]
+                                    ]
+                                    _is_present = data.structure.atoms["is_present"][
+                                        token_2["atom_idx"] : token_2["atom_idx"]
+                                        + token_2["atom_num"]
+                                    ]
+                                    if _is_present.sum() == 0:
+                                        continue
+                                    token_2_coords = _coords[_is_present]
+
+                                    if (
+                                        np.min(cdist(token_1_coords, token_2_coords))
+                                        < contact_cutoff
+                                    ):
+                                        pairs.append(
+                                            (token_1["token_idx"], token_2["token_idx"])
+                                        )
+
+                    assert len(pairs) > 0
+
+                    pair = random.choice(pairs)
+                    token_1_mask = token_data["token_idx"] == pair[0]
+                    token_2_mask = token_data["token_idx"] == pair[1]
+
+                    contact_conditioning[np.ix_(token_1_mask, token_2_mask)] = (
+                        const.contact_conditioning_info["CONTACT"]
+                    )
+                    contact_conditioning[np.ix_(token_2_mask, token_1_mask)] = (
+                        const.contact_conditioning_info["CONTACT"]
+                    )
+
+            elif not only_pp_contact:
+                # only one chain, find contacts within the chain with minimum residue distance
+                pairs = []
+                for token_1 in token_data:
+                    _coords = data.structure.atoms["coords"][
+                        token_1["atom_idx"] : token_1["atom_idx"] + token_1["atom_num"]
+                    ]
+                    _is_present = data.structure.atoms["is_present"][
+                        token_1["atom_idx"] : token_1["atom_idx"] + token_1["atom_num"]
+                    ]
+                    if _is_present.sum() == 0:
+                        continue
+                    token_1_coords = _coords[_is_present]
+
+                    for token_2 in token_data:
+                        if np.abs(token_1["res_idx"] - token_2["res_idx"]) <= 8:
+                            continue
+
+                        _coords = data.structure.atoms["coords"][
+                            token_2["atom_idx"] : token_2["atom_idx"]
+                            + token_2["atom_num"]
+                        ]
+                        _is_present = data.structure.atoms["is_present"][
+                            token_2["atom_idx"] : token_2["atom_idx"]
+                            + token_2["atom_num"]
+                        ]
+                        if _is_present.sum() == 0:
+                            continue
+                        token_2_coords = _coords[_is_present]
+
+                        if (
+                            np.min(cdist(token_1_coords, token_2_coords))
+                            < contact_cutoff
+                        ):
+                            pairs.append((token_1["token_idx"], token_2["token_idx"]))
+
+                if len(pairs) > 0:
+                    pair = random.choice(pairs)
+                    token_1_mask = token_data["token_idx"] == pair[0]
+                    token_2_mask = token_data["token_idx"] == pair[1]
+
+                    contact_conditioning[np.ix_(token_1_mask, token_2_mask)] = (
+                        const.contact_conditioning_info["CONTACT"]
+                    )
+                    contact_conditioning[np.ix_(token_2_mask, token_1_mask)] = (
+                        const.contact_conditioning_info["CONTACT"]
+                    )
+
+    if np.all(contact_conditioning == const.contact_conditioning_info["UNSELECTED"]):
+        contact_conditioning = (
+            contact_conditioning
+            - const.contact_conditioning_info["UNSELECTED"]
+            + const.contact_conditioning_info["UNSPECIFIED"]
+        )
+    contact_conditioning = from_numpy(contact_conditioning).long()
+    contact_conditioning = one_hot(
+        contact_conditioning, num_classes=len(const.contact_conditioning_info)
+    )
+    contact_threshold = from_numpy(contact_threshold).float()
+
+    # compute cyclic polymer mask
+    cyclic_ids = {}
+    for idx_chain, asym_id_iter in enumerate(data.structure.chains["asym_id"]):
+        for connection in data.structure.bonds:
+            if (
+                idx_chain == connection["chain_1"] == connection["chain_2"]
+                and data.structure.chains[connection["chain_1"]]["res_num"] > 2
+                and connection["res_1"]
+                != connection["res_2"]  # Avoid same residue bonds!
+            ):
+                if (
+                    data.structure.chains[connection["chain_1"]]["res_num"]
+                    == (connection["res_2"] + 1)
+                    and connection["res_1"] == 0
+                ) or (
+                    data.structure.chains[connection["chain_1"]]["res_num"]
+                    == (connection["res_1"] + 1)
+                    and connection["res_2"] == 0
+                ):
+                    cyclic_ids[asym_id_iter] = data.structure.chains[
+                        connection["chain_1"]
+                    ]["res_num"]
+    cyclic = from_numpy(
+        np.array(
+            [
+                (cyclic_ids[asym_id_iter] if asym_id_iter in cyclic_ids else 0)
+                for asym_id_iter in token_data["asym_id"]
+            ]
+        )
+    ).float()
+
+    # cyclic period is either computed from the bonds or given as input flag
+    cyclic_period = torch.maximum(cyclic, cyclic_period)
+
+    # Pad to max tokens if given
+    if max_tokens is not None:
+        pad_len = max_tokens - len(token_data)
+        if pad_len > 0:
+            token_index = pad_dim(token_index, 0, pad_len)
+            residue_index = pad_dim(residue_index, 0, pad_len)
+            asym_id = pad_dim(asym_id, 0, pad_len)
+            entity_id = pad_dim(entity_id, 0, pad_len)
+            sym_id = pad_dim(sym_id, 0, pad_len)
+            mol_type = pad_dim(mol_type, 0, pad_len)
+            res_type = pad_dim(res_type, 0, pad_len)
+            disto_center = pad_dim(disto_center, 0, pad_len)
+            pad_mask = pad_dim(pad_mask, 0, pad_len)
+            resolved_mask = pad_dim(resolved_mask, 0, pad_len)
+            disto_mask = pad_dim(disto_mask, 0, pad_len)
+            contact_conditioning = pad_dim(contact_conditioning, 0, pad_len)
+            contact_conditioning = pad_dim(contact_conditioning, 1, pad_len)
+            contact_threshold = pad_dim(contact_threshold, 0, pad_len)
+            contact_threshold = pad_dim(contact_threshold, 1, pad_len)
+            method_feature = pad_dim(method_feature, 0, pad_len)
+            modified = pad_dim(modified, 0, pad_len)
+            cyclic_period = pad_dim(cyclic_period, 0, pad_len)
+            affinity_mask = pad_dim(affinity_mask, 0, pad_len)
+
+    token_features = {
+        "token_index": token_index,
+        "residue_index": residue_index,
+        "asym_id": asym_id,
+        "entity_id": entity_id,
+        "sym_id": sym_id,
+        "mol_type": mol_type,
+        "res_type": res_type,
+        "disto_center": disto_center,
+        "token_bonds": bonds,
+        "type_bonds": bonds_type,
+        "token_pad_mask": pad_mask,
+        "token_resolved_mask": resolved_mask,
+        "token_disto_mask": disto_mask,
+        "contact_conditioning": contact_conditioning,
+        "contact_threshold": contact_threshold,
+        "method_feature": method_feature,
+        "modified": modified,
+        "cyclic_period": cyclic_period,
+        "affinity_token_mask": affinity_mask,
+    }
+
+    return token_features
+
+
+def process_atom_features(
+    data: Tokenized,
+    random: np.random.Generator,
+    ensemble_features: dict,
+    molecules: dict[str, Mol],
+    atoms_per_window_queries: int = 32,
+    min_dist: float = 2.0,
+    max_dist: float = 22.0,
+    num_bins: int = 64,
+    max_atoms: Optional[int] = None,
+    max_tokens: Optional[int] = None,
+    disto_use_ensemble: Optional[bool] = False,
+    override_bfactor: bool = False,
+    compute_frames: bool = False,
+    override_coords: Optional[Tensor] = None,
+    bfactor_md_correction: bool = False,
+) -> dict[str, Tensor]:
+    """Get the atom features.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input to the model.
+    max_atoms : int, optional
+        The maximum number of atoms.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        The atom features.
+
+    """
+    # Filter to tokens' atoms
+    atom_data = []
+    atom_name = []
+    atom_element = []
+    atom_charge = []
+    atom_conformer = []
+    atom_chirality = []
+    ref_space_uid = []
+    coord_data = []
+    if compute_frames:
+        frame_data = []
+        resolved_frame_data = []
+    atom_to_token = []
+    token_to_rep_atom = []  # index on cropped atom table
+    r_set_to_rep_atom = []
+    disto_coords_ensemble = []
+    backbone_feat_index = []
+    token_to_center_atom = []
+
+    e_offsets = data.structure.ensemble["atom_coord_idx"]
+    atom_idx = 0
+
+    # Start atom idx in full atom table for structures chosen. Up to num_ensembles points.
+    ensemble_atom_starts = [
+        data.structure.ensemble[idx]["atom_coord_idx"]
+        for idx in ensemble_features["ensemble_ref_idxs"]
+    ]
+
+    # Set unk chirality id
+    unk_chirality = const.chirality_type_ids[const.unk_chirality_type]
+
+    chain_res_ids = {}
+    res_index_to_conf_id = {}
+    for token_id, token in enumerate(data.tokens):
+        # Get the chain residue ids
+        chain_idx, res_id = token["asym_id"], token["res_idx"]
+        chain = data.structure.chains[chain_idx]
+
+        if (chain_idx, res_id) not in chain_res_ids:
+            new_idx = len(chain_res_ids)
+            chain_res_ids[(chain_idx, res_id)] = new_idx
+        else:
+            new_idx = chain_res_ids[(chain_idx, res_id)]
+
+        # Get the molecule and conformer
+        mol = molecules[token["res_name"]]
+        atom_name_to_ref = {a.GetProp("name"): a for a in mol.GetAtoms()}
+
+        # Sample a random conformer
+        if (chain_idx, res_id) not in res_index_to_conf_id:
+            conf_ids = [int(conf.GetId()) for conf in mol.GetConformers()]
+            conf_id = int(random.choice(conf_ids))
+            res_index_to_conf_id[(chain_idx, res_id)] = conf_id
+
+        conf_id = res_index_to_conf_id[(chain_idx, res_id)]
+        conformer = mol.GetConformer(conf_id)
+
+        # Map atoms to token indices
+        ref_space_uid.extend([new_idx] * token["atom_num"])
+        atom_to_token.extend([token_id] * token["atom_num"])
+
+        # Add atom data
+        start = token["atom_idx"]
+        end = token["atom_idx"] + token["atom_num"]
+        token_atoms = data.structure.atoms[start:end]
+
+        # Add atom ref data
+        # element, charge, conformer, chirality
+        token_atom_name = np.array([convert_atom_name(a["name"]) for a in token_atoms])
+        token_atoms_ref = np.array([atom_name_to_ref[a["name"]] for a in token_atoms])
+        token_atoms_element = np.array([a.GetAtomicNum() for a in token_atoms_ref])
+        token_atoms_charge = np.array([a.GetFormalCharge() for a in token_atoms_ref])
+        token_atoms_conformer = np.array(
+            [
+                (
+                    conformer.GetAtomPosition(a.GetIdx()).x,
+                    conformer.GetAtomPosition(a.GetIdx()).y,
+                    conformer.GetAtomPosition(a.GetIdx()).z,
+                )
+                for a in token_atoms_ref
+            ]
+        )
+        token_atoms_chirality = np.array(
+            [
+                const.chirality_type_ids.get(a.GetChiralTag().name, unk_chirality)
+                for a in token_atoms_ref
+            ]
+        )
+
+        # Map token to representative atom
+        token_to_rep_atom.append(atom_idx + token["disto_idx"] - start)
+        token_to_center_atom.append(atom_idx + token["center_idx"] - start)
+        if (chain["mol_type"] != const.chain_type_ids["NONPOLYMER"]) and token[
+            "resolved_mask"
+        ]:
+            r_set_to_rep_atom.append(atom_idx + token["center_idx"] - start)
+
+        if chain["mol_type"] == const.chain_type_ids["PROTEIN"]:
+            backbone_index = [
+                (
+                    const.protein_backbone_atom_index[atom_name] + 1
+                    if atom_name in const.protein_backbone_atom_index
+                    else 0
+                )
+                for atom_name in token_atoms["name"]
+            ]
+        elif (
+            chain["mol_type"] == const.chain_type_ids["DNA"]
+            or chain["mol_type"] == const.chain_type_ids["RNA"]
+        ):
+            backbone_index = [
+                (
+                    const.nucleic_backbone_atom_index[atom_name]
+                    + 1
+                    + len(const.protein_backbone_atom_index)
+                    if atom_name in const.nucleic_backbone_atom_index
+                    else 0
+                )
+                for atom_name in token_atoms["name"]
+            ]
+        else:
+            backbone_index = [0] * token["atom_num"]
+        backbone_feat_index.extend(backbone_index)
+
+        # Get token coordinates across sampled ensembles  and apply transforms
+        token_coords = np.array(
+            [
+                data.structure.coords[
+                    ensemble_atom_start + start : ensemble_atom_start + end
+                ]["coords"]
+                for ensemble_atom_start in ensemble_atom_starts
+            ]
+        )
+        coord_data.append(token_coords)
+
+        if compute_frames:
+            # Get frame data
+            res_type = const.tokens[token["res_type"]]
+            res_name = str(token["res_name"])
+
+            if token["atom_num"] < 3 or res_type in ["PAD", "UNK", "-"]:
+                idx_frame_a, idx_frame_b, idx_frame_c = 0, 0, 0
+                mask_frame = False
+            elif (token["mol_type"] == const.chain_type_ids["PROTEIN"]) and (
+                res_name in const.ref_atoms
+            ):
+                idx_frame_a, idx_frame_b, idx_frame_c = (
+                    const.ref_atoms[res_name].index("N"),
+                    const.ref_atoms[res_name].index("CA"),
+                    const.ref_atoms[res_name].index("C"),
+                )
+                mask_frame = (
+                    token_atoms["is_present"][idx_frame_a]
+                    and token_atoms["is_present"][idx_frame_b]
+                    and token_atoms["is_present"][idx_frame_c]
+                )
+            elif (
+                token["mol_type"] == const.chain_type_ids["DNA"]
+                or token["mol_type"] == const.chain_type_ids["RNA"]
+            ) and (res_name in const.ref_atoms):
+                idx_frame_a, idx_frame_b, idx_frame_c = (
+                    const.ref_atoms[res_name].index("C1'"),
+                    const.ref_atoms[res_name].index("C3'"),
+                    const.ref_atoms[res_name].index("C4'"),
+                )
+                mask_frame = (
+                    token_atoms["is_present"][idx_frame_a]
+                    and token_atoms["is_present"][idx_frame_b]
+                    and token_atoms["is_present"][idx_frame_c]
+                )
+            elif token["mol_type"] == const.chain_type_ids["PROTEIN"]:
+                # Try to look for the atom nams in the modified residue
+                is_ca = token_atoms["name"] == "CA"
+                idx_frame_a = is_ca.argmax()
+                ca_present = (
+                    token_atoms[idx_frame_a]["is_present"] if is_ca.any() else False
+                )
+
+                is_n = token_atoms["name"] == "N"
+                idx_frame_b = is_n.argmax()
+                n_present = (
+                    token_atoms[idx_frame_b]["is_present"] if is_n.any() else False
+                )
+
+                is_c = token_atoms["name"] == "C"
+                idx_frame_c = is_c.argmax()
+                c_present = (
+                    token_atoms[idx_frame_c]["is_present"] if is_c.any() else False
+                )
+                mask_frame = ca_present and n_present and c_present
+
+            elif (token["mol_type"] == const.chain_type_ids["DNA"]) or (
+                token["mol_type"] == const.chain_type_ids["RNA"]
+            ):
+                # Try to look for the atom nams in the modified residue
+                is_c1 = token_atoms["name"] == "C1'"
+                idx_frame_a = is_c1.argmax()
+                c1_present = (
+                    token_atoms[idx_frame_a]["is_present"] if is_c1.any() else False
+                )
+
+                is_c3 = token_atoms["name"] == "C3'"
+                idx_frame_b = is_c3.argmax()
+                c3_present = (
+                    token_atoms[idx_frame_b]["is_present"] if is_c3.any() else False
+                )
+
+                is_c4 = token_atoms["name"] == "C4'"
+                idx_frame_c = is_c4.argmax()
+                c4_present = (
+                    token_atoms[idx_frame_c]["is_present"] if is_c4.any() else False
+                )
+                mask_frame = c1_present and c3_present and c4_present
+            else:
+                idx_frame_a, idx_frame_b, idx_frame_c = 0, 0, 0
+                mask_frame = False
+            frame_data.append(
+                [
+                    idx_frame_a + atom_idx,
+                    idx_frame_b + atom_idx,
+                    idx_frame_c + atom_idx,
+                ]
+            )
+            resolved_frame_data.append(mask_frame)
+
+        # Get distogram coordinates
+        disto_coords_ensemble_tok = data.structure.coords[
+            e_offsets + token["disto_idx"]
+        ]["coords"]
+        disto_coords_ensemble.append(disto_coords_ensemble_tok)
+
+        # Update atom data. This is technically never used again (we rely on coord_data),
+        # but we update for consistency and to make sure the Atom object has valid, transformed coordinates.
+        token_atoms = token_atoms.copy()
+        token_atoms["coords"] = token_coords[
+            0
+        ]  # atom has a copy of first coords in ensemble
+        atom_data.append(token_atoms)
+        atom_name.append(token_atom_name)
+        atom_element.append(token_atoms_element)
+        atom_charge.append(token_atoms_charge)
+        atom_conformer.append(token_atoms_conformer)
+        atom_chirality.append(token_atoms_chirality)
+        atom_idx += len(token_atoms)
+
+    disto_coords_ensemble = np.array(disto_coords_ensemble)  # (N_TOK, N_ENS, 3)
+
+    # Compute ensemble distogram
+    L = len(data.tokens)
+
+    if disto_use_ensemble:
+        # Use all available structures to create distogram
+        idx_list = range(disto_coords_ensemble.shape[1])
+    else:
+        # Only use a sampled structures to create distogram
+        idx_list = ensemble_features["ensemble_ref_idxs"]
+
+    # Create distogram
+    disto_target = torch.zeros(L, L, len(idx_list), num_bins)  # TODO1
+
+    # disto_target = torch.zeros(L, L, num_bins)
+    for i, e_idx in enumerate(idx_list):
+        t_center = torch.Tensor(disto_coords_ensemble[:, e_idx, :])
+        t_dists = torch.cdist(t_center, t_center)
+        boundaries = torch.linspace(min_dist, max_dist, num_bins - 1)
+        distogram = (t_dists.unsqueeze(-1) > boundaries).sum(dim=-1).long()
+        # disto_target += one_hot(distogram, num_classes=num_bins)
+        disto_target[:, :, i, :] = one_hot(distogram, num_classes=num_bins)  # TODO1
+
+    # Normalize distogram
+    # disto_target = disto_target / disto_target.sum(-1)[..., None]  # remove TODO1
+    atom_data = np.concatenate(atom_data)
+    atom_name = np.concatenate(atom_name)
+    atom_element = np.concatenate(atom_element)
+    atom_charge = np.concatenate(atom_charge)
+    atom_conformer = np.concatenate(atom_conformer)
+    atom_chirality = np.concatenate(atom_chirality)
+    coord_data = np.concatenate(coord_data, axis=1)
+    ref_space_uid = np.array(ref_space_uid)
+
+    # Compute features
+    disto_coords_ensemble = from_numpy(disto_coords_ensemble)
+    disto_coords_ensemble = disto_coords_ensemble[
+        :, ensemble_features["ensemble_ref_idxs"]
+    ].permute(1, 0, 2)
+    backbone_feat_index = from_numpy(np.asarray(backbone_feat_index)).long()
+    ref_atom_name_chars = from_numpy(atom_name).long()
+    ref_element = from_numpy(atom_element).long()
+    ref_charge = from_numpy(atom_charge).float()
+    ref_pos = from_numpy(atom_conformer).float()
+    ref_space_uid = from_numpy(ref_space_uid)
+    ref_chirality = from_numpy(atom_chirality).long()
+    coords = from_numpy(coord_data.copy())
+    resolved_mask = from_numpy(atom_data["is_present"])
+    pad_mask = torch.ones(len(atom_data), dtype=torch.float)
+    atom_to_token = torch.tensor(atom_to_token, dtype=torch.long)
+    token_to_rep_atom = torch.tensor(token_to_rep_atom, dtype=torch.long)
+    r_set_to_rep_atom = torch.tensor(r_set_to_rep_atom, dtype=torch.long)
+    token_to_center_atom = torch.tensor(token_to_center_atom, dtype=torch.long)
+    bfactor = from_numpy(atom_data["bfactor"].copy())
+    plddt = from_numpy(atom_data["plddt"].copy())
+    if override_bfactor:
+        bfactor = bfactor * 0.0
+
+    if bfactor_md_correction and data.record.structure.method.lower() == "md":
+        # MD bfactor was computed as RMSF
+        # Convert to b-factor
+        bfactor = 8 * (np.pi**2) * (bfactor**2)
+
+    # We compute frames within ensemble
+    if compute_frames:
+        frames = []
+        frame_resolved_mask = []
+        for i in range(coord_data.shape[0]):
+            frame_data_, resolved_frame_data_ = compute_frames_nonpolymer(
+                data,
+                coord_data[i],
+                atom_data["is_present"],
+                atom_to_token,
+                frame_data,
+                resolved_frame_data,
+            )  # Compute frames for NONPOLYMER tokens
+            frames.append(frame_data_.copy())
+            frame_resolved_mask.append(resolved_frame_data_.copy())
+        frames = from_numpy(np.stack(frames))  # (N_ENS, N_TOK, 3)
+        frame_resolved_mask = from_numpy(np.stack(frame_resolved_mask))
+
+    # Convert to one-hot
+    backbone_feat_index = one_hot(
+        backbone_feat_index,
+        num_classes=1
+        + len(const.protein_backbone_atom_index)
+        + len(const.nucleic_backbone_atom_index),
+    )
+    ref_atom_name_chars = one_hot(ref_atom_name_chars, num_classes=64)
+    ref_element = one_hot(ref_element, num_classes=const.num_elements)
+    atom_to_token = one_hot(atom_to_token, num_classes=token_id + 1)
+    token_to_rep_atom = one_hot(token_to_rep_atom, num_classes=len(atom_data))
+    r_set_to_rep_atom = one_hot(r_set_to_rep_atom, num_classes=len(atom_data))
+    token_to_center_atom = one_hot(token_to_center_atom, num_classes=len(atom_data))
+
+    # Center the ground truth coordinates
+    center = (coords * resolved_mask[None, :, None]).sum(dim=1)
+    center = center / resolved_mask.sum().clamp(min=1)
+    coords = coords - center[:, None]
+
+    if isinstance(override_coords, Tensor):
+        coords = override_coords.unsqueeze(0)
+
+    # Apply random roto-translation to the input conformers
+    for i in range(torch.max(ref_space_uid)):
+        included = ref_space_uid == i
+        if torch.sum(included) > 0 and torch.any(resolved_mask[included]):
+            ref_pos[included] = center_random_augmentation(
+                ref_pos[included][None], resolved_mask[included][None], centering=True
+            )[0]
+
+    # Compute padding and apply
+    if max_atoms is not None:
+        assert max_atoms % atoms_per_window_queries == 0
+        pad_len = max_atoms - len(atom_data)
+    else:
+        pad_len = (
+            (len(atom_data) - 1) // atoms_per_window_queries + 1
+        ) * atoms_per_window_queries - len(atom_data)
+
+    if pad_len > 0:
+        pad_mask = pad_dim(pad_mask, 0, pad_len)
+        ref_pos = pad_dim(ref_pos, 0, pad_len)
+        resolved_mask = pad_dim(resolved_mask, 0, pad_len)
+        ref_atom_name_chars = pad_dim(ref_atom_name_chars, 0, pad_len)
+        ref_element = pad_dim(ref_element, 0, pad_len)
+        ref_charge = pad_dim(ref_charge, 0, pad_len)
+        ref_chirality = pad_dim(ref_chirality, 0, pad_len)
+        backbone_feat_index = pad_dim(backbone_feat_index, 0, pad_len)
+        ref_space_uid = pad_dim(ref_space_uid, 0, pad_len)
+        coords = pad_dim(coords, 1, pad_len)
+        atom_to_token = pad_dim(atom_to_token, 0, pad_len)
+        token_to_rep_atom = pad_dim(token_to_rep_atom, 1, pad_len)
+        token_to_center_atom = pad_dim(token_to_center_atom, 1, pad_len)
+        r_set_to_rep_atom = pad_dim(r_set_to_rep_atom, 1, pad_len)
+        bfactor = pad_dim(bfactor, 0, pad_len)
+        plddt = pad_dim(plddt, 0, pad_len)
+
+    if max_tokens is not None:
+        pad_len = max_tokens - token_to_rep_atom.shape[0]
+        if pad_len > 0:
+            atom_to_token = pad_dim(atom_to_token, 1, pad_len)
+            token_to_rep_atom = pad_dim(token_to_rep_atom, 0, pad_len)
+            r_set_to_rep_atom = pad_dim(r_set_to_rep_atom, 0, pad_len)
+            token_to_center_atom = pad_dim(token_to_center_atom, 0, pad_len)
+            disto_target = pad_dim(pad_dim(disto_target, 0, pad_len), 1, pad_len)
+            disto_coords_ensemble = pad_dim(disto_coords_ensemble, 1, pad_len)
+
+            if compute_frames:
+                frames = pad_dim(frames, 1, pad_len)
+                frame_resolved_mask = pad_dim(frame_resolved_mask, 1, pad_len)
+
+    atom_features = {
+        "ref_pos": ref_pos,
+        "atom_resolved_mask": resolved_mask,
+        "ref_atom_name_chars": ref_atom_name_chars,
+        "ref_element": ref_element,
+        "ref_charge": ref_charge,
+        "ref_chirality": ref_chirality,
+        "atom_backbone_feat": backbone_feat_index,
+        "ref_space_uid": ref_space_uid,
+        "coords": coords,
+        "atom_pad_mask": pad_mask,
+        "atom_to_token": atom_to_token,
+        "token_to_rep_atom": token_to_rep_atom,
+        "r_set_to_rep_atom": r_set_to_rep_atom,
+        "token_to_center_atom": token_to_center_atom,
+        "disto_target": disto_target,
+        "disto_coords_ensemble": disto_coords_ensemble,
+        "bfactor": bfactor,
+        "plddt": plddt,
+    }
+
+    if compute_frames:
+        atom_features["frames_idx"] = frames
+        atom_features["frame_resolved_mask"] = frame_resolved_mask
+
+    return atom_features
+
+
+def process_msa_features(
+    data: Tokenized,
+    random: np.random.Generator,
+    max_seqs_batch: int,
+    max_seqs: int,
+    max_tokens: Optional[int] = None,
+    pad_to_max_seqs: bool = False,
+    msa_sampling: bool = False,
+    affinity: bool = False,
+) -> dict[str, Tensor]:
+    """Get the MSA features.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input to the model.
+    random : np.random.Generator
+        The random number generator.
+    max_seqs : int
+        The maximum number of MSA sequences.
+    max_tokens : int
+        The maximum number of tokens.
+    pad_to_max_seqs : bool
+        Whether to pad to the maximum number of sequences.
+    msa_sampling : bool
+        Whether to sample the MSA.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        The MSA features.
+
+    """
+    # Created paired MSA
+    msa, deletion, paired = construct_paired_msa(
+        data=data,
+        random=random,
+        max_seqs=max_seqs_batch,
+        random_subset=msa_sampling,
+    )
+    msa, deletion, paired = (
+        msa.transpose(1, 0),
+        deletion.transpose(1, 0),
+        paired.transpose(1, 0),
+    )  # (N_MSA, N_RES, N_AA)
+
+    # Prepare features
+    assert torch.all(msa >= 0) and torch.all(msa < const.num_tokens)
+    msa_one_hot = torch.nn.functional.one_hot(msa, num_classes=const.num_tokens)
+    msa_mask = torch.ones_like(msa)
+    profile = msa_one_hot.float().mean(dim=0)
+    has_deletion = deletion > 0
+    deletion = np.pi / 2 * np.arctan(deletion / 3)
+    deletion_mean = deletion.mean(axis=0)
+
+    # Pad in the MSA dimension (dim=0)
+    if pad_to_max_seqs:
+        pad_len = max_seqs - msa.shape[0]
+        if pad_len > 0:
+            msa = pad_dim(msa, 0, pad_len, const.token_ids["-"])
+            paired = pad_dim(paired, 0, pad_len)
+            msa_mask = pad_dim(msa_mask, 0, pad_len)
+            has_deletion = pad_dim(has_deletion, 0, pad_len)
+            deletion = pad_dim(deletion, 0, pad_len)
+
+    # Pad in the token dimension (dim=1)
+    if max_tokens is not None:
+        pad_len = max_tokens - msa.shape[1]
+        if pad_len > 0:
+            msa = pad_dim(msa, 1, pad_len, const.token_ids["-"])
+            paired = pad_dim(paired, 1, pad_len)
+            msa_mask = pad_dim(msa_mask, 1, pad_len)
+            has_deletion = pad_dim(has_deletion, 1, pad_len)
+            deletion = pad_dim(deletion, 1, pad_len)
+            profile = pad_dim(profile, 0, pad_len)
+            deletion_mean = pad_dim(deletion_mean, 0, pad_len)
+    if affinity:
+        return {
+            "deletion_mean_affinity": deletion_mean,
+            "profile_affinity": profile,
+        }
+    else:
+        return {
+            "msa": msa,
+            "msa_paired": paired,
+            "deletion_value": deletion,
+            "has_deletion": has_deletion,
+            "deletion_mean": deletion_mean,
+            "profile": profile,
+            "msa_mask": msa_mask,
+        }
+
+
+def load_dummy_templates_features(tdim: int, num_tokens: int) -> dict:
+    """Load dummy templates for v2."""
+    # Allocate features
+    res_type = np.zeros((tdim, num_tokens), dtype=np.int64)
+    frame_rot = np.zeros((tdim, num_tokens, 3, 3), dtype=np.float32)
+    frame_t = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    cb_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    ca_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    frame_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    cb_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    template_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    query_to_template = np.zeros((tdim, num_tokens), dtype=np.int64)
+    visibility_ids = np.zeros((tdim, num_tokens), dtype=np.float32)
+
+    # Convert to one-hot
+    res_type = torch.from_numpy(res_type)
+    res_type = one_hot(res_type, num_classes=const.num_tokens)
+
+    return {
+        "template_restype": res_type,
+        "template_frame_rot": torch.from_numpy(frame_rot),
+        "template_frame_t": torch.from_numpy(frame_t),
+        "template_cb": torch.from_numpy(cb_coords),
+        "template_ca": torch.from_numpy(ca_coords),
+        "template_mask_cb": torch.from_numpy(cb_mask),
+        "template_mask_frame": torch.from_numpy(frame_mask),
+        "template_mask": torch.from_numpy(template_mask),
+        "query_to_template": torch.from_numpy(query_to_template),
+        "visibility_ids": torch.from_numpy(visibility_ids),
+    }
+
+
+def compute_template_features(
+    query_tokens: Tokenized,
+    tmpl_tokens: list[dict],
+    num_tokens: int,
+) -> dict:
+    """Compute the template features."""
+    # Allocate features
+    res_type = np.zeros((num_tokens,), dtype=np.int64)
+    frame_rot = np.zeros((num_tokens, 3, 3), dtype=np.float32)
+    frame_t = np.zeros((num_tokens, 3), dtype=np.float32)
+    cb_coords = np.zeros((num_tokens, 3), dtype=np.float32)
+    ca_coords = np.zeros((num_tokens, 3), dtype=np.float32)
+    frame_mask = np.zeros((num_tokens,), dtype=np.float32)
+    cb_mask = np.zeros((num_tokens,), dtype=np.float32)
+    template_mask = np.zeros((num_tokens,), dtype=np.float32)
+    query_to_template = np.zeros((num_tokens,), dtype=np.int64)
+    visibility_ids = np.zeros((num_tokens,), dtype=np.float32)
+
+    # Now create features per token
+    asym_id_to_pdb_id = {}
+
+    for token_dict in tmpl_tokens:
+        idx = token_dict["q_idx"]
+        pdb_id = token_dict["pdb_id"]
+        token = token_dict["token"]
+        query_token = query_tokens.tokens[idx]
+        asym_id_to_pdb_id[query_token["asym_id"]] = pdb_id
+        res_type[idx] = token["res_type"]
+        frame_rot[idx] = token["frame_rot"].reshape(3, 3)
+        frame_t[idx] = token["frame_t"]
+        cb_coords[idx] = token["disto_coords"]
+        ca_coords[idx] = token["center_coords"]
+        cb_mask[idx] = token["disto_mask"]
+        frame_mask[idx] = token["frame_mask"]
+        template_mask[idx] = 1.0
+
+    # Set visibility_id for templated chains
+    for asym_id, pdb_id in asym_id_to_pdb_id.items():
+        indices = (query_tokens.tokens["asym_id"] == asym_id).nonzero()
+        visibility_ids[indices] = pdb_id
+
+    # Set visibility for non templated chain + olygomerics
+    for asym_id in np.unique(query_tokens.structure.chains["asym_id"]):
+        if asym_id not in asym_id_to_pdb_id:
+            # We hack the chain id to be negative to not overlap with the above
+            indices = (query_tokens.tokens["asym_id"] == asym_id).nonzero()
+            visibility_ids[indices] = -1 - asym_id
+
+    # Convert to one-hot
+    res_type = torch.from_numpy(res_type)
+    res_type = one_hot(res_type, num_classes=const.num_tokens)
+
+    return {
+        "template_restype": res_type,
+        "template_frame_rot": torch.from_numpy(frame_rot),
+        "template_frame_t": torch.from_numpy(frame_t),
+        "template_cb": torch.from_numpy(cb_coords),
+        "template_ca": torch.from_numpy(ca_coords),
+        "template_mask_cb": torch.from_numpy(cb_mask),
+        "template_mask_frame": torch.from_numpy(frame_mask),
+        "template_mask": torch.from_numpy(template_mask),
+        "query_to_template": torch.from_numpy(query_to_template),
+        "visibility_ids": torch.from_numpy(visibility_ids),
+    }
+
+
+def process_template_features(
+    data: Tokenized,
+    max_tokens: int,
+) -> dict[str, torch.Tensor]:
+    """Load the given input data.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input to the model.
+    max_tokens : int
+        The maximum number of tokens.
+
+    Returns
+    -------
+    dict[str, torch.Tensor]
+        The loaded template features.
+
+    """
+    # Group templates by name
+    name_to_templates: dict[str, list[TemplateInfo]] = {}
+    for template_info in data.record.templates:
+        name_to_templates.setdefault(template_info.name, []).append(template_info)
+
+    # Map chain name to asym_id
+    chain_name_to_asym_id = {}
+    for chain in data.structure.chains:
+        chain_name_to_asym_id[chain["name"]] = chain["asym_id"]
+
+    # Compute the offset
+    template_features = []
+    for template_id, (template_name, templates) in enumerate(name_to_templates.items()):
+        row_tokens = []
+        template_structure = data.templates[template_name]
+        template_tokens = data.template_tokens[template_name]
+        tmpl_chain_name_to_asym_id = {}
+        for chain in template_structure.chains:
+            tmpl_chain_name_to_asym_id[chain["name"]] = chain["asym_id"]
+
+        for template in templates:
+            offset = template.template_st - template.query_st
+
+            # Get query and template tokens to map residues
+            query_tokens = data.tokens
+            chain_id = chain_name_to_asym_id[template.query_chain]
+            q_tokens = query_tokens[query_tokens["asym_id"] == chain_id]
+            q_indices = dict(zip(q_tokens["res_idx"], q_tokens["token_idx"]))
+
+            # Get the template tokens at the query residues
+            chain_id = tmpl_chain_name_to_asym_id[template.template_chain]
+            toks = template_tokens[template_tokens["asym_id"] == chain_id]
+            toks = [t for t in toks if t["res_idx"] - offset in q_indices]
+            for t in toks:
+                q_idx = q_indices[t["res_idx"] - offset]
+                row_tokens.append(
+                    {
+                        "token": t,
+                        "pdb_id": template_id,
+                        "q_idx": q_idx,
+                    }
+                )
+
+        # Compute template features for each row
+        row_features = compute_template_features(data, row_tokens, max_tokens)
+        template_features.append(row_features)
+
+    # Stack each feature
+    out = {}
+    for k in template_features[0]:
+        out[k] = torch.stack([f[k] for f in template_features])
+
+    return out
+
+
+def process_symmetry_features(
+    cropped: Tokenized, symmetries: dict
+) -> dict[str, Tensor]:
+    """Get the symmetry features.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input to the model.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        The symmetry features.
+
+    """
+    features = get_chain_symmetries(cropped)
+    features.update(get_amino_acids_symmetries(cropped))
+    features.update(get_ligand_symmetries(cropped, symmetries))
+
+    return features
+
+
+def process_ensemble_features(
+    data: Tokenized,
+    random: np.random.Generator,
+    num_ensembles: int,
+    ensemble_sample_replacement: bool,
+    fix_single_ensemble: bool,
+) -> dict[str, Tensor]:
+    """Get the ensemble features.
+
+    Parameters
+    ----------
+    data : Tokenized
+        The input to the model.
+    random : np.random.Generator
+        The random number generator.
+    num_ensembles : int
+        The maximum number of ensembles to sample.
+    ensemble_sample_replacement : bool
+        Whether to sample with replacement.
+
+    Returns
+    -------
+    dict[str, Tensor]
+        The ensemble features.
+
+    """
+    assert num_ensembles > 0, "Number of conformers sampled must be greater than 0."
+
+    # Number of available conformers in the structure
+    # s_ensemble_num = min(len(cropped.structure.ensemble), 24)  # Limit to 24 conformers DEBUG: TODO: remove !
+    s_ensemble_num = len(data.structure.ensemble)
+
+    if fix_single_ensemble:
+        # Always take the first conformer for train and validation
+        assert (
+            num_ensembles == 1
+        ), "Number of conformers sampled must be 1 with fix_single_ensemble=True."
+        ensemble_ref_idxs = np.array([0])
+    else:
+        if ensemble_sample_replacement:
+            # Used in training
+            ensemble_ref_idxs = random.integers(0, s_ensemble_num, (num_ensembles,))
+        else:
+            # Used in validation
+            if s_ensemble_num < num_ensembles:
+                # Take all available conformers
+                ensemble_ref_idxs = np.arange(0, s_ensemble_num)
+            else:
+                # Sample without replacement
+                ensemble_ref_idxs = random.choice(
+                    s_ensemble_num, num_ensembles, replace=False
+                )
+
+    ensemble_features = {
+        "ensemble_ref_idxs": torch.Tensor(ensemble_ref_idxs).long(),
+    }
+
+    return ensemble_features
+
+
+def process_residue_constraint_features(data: Tokenized) -> dict[str, Tensor]:
+    residue_constraints = data.residue_constraints
+    if residue_constraints is not None:
+        rdkit_bounds_constraints = residue_constraints.rdkit_bounds_constraints
+        chiral_atom_constraints = residue_constraints.chiral_atom_constraints
+        stereo_bond_constraints = residue_constraints.stereo_bond_constraints
+        planar_bond_constraints = residue_constraints.planar_bond_constraints
+        planar_ring_5_constraints = residue_constraints.planar_ring_5_constraints
+        planar_ring_6_constraints = residue_constraints.planar_ring_6_constraints
+
+        rdkit_bounds_index = torch.tensor(
+            rdkit_bounds_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+        rdkit_bounds_bond_mask = torch.tensor(
+            rdkit_bounds_constraints["is_bond"].copy(), dtype=torch.bool
+        )
+        rdkit_bounds_angle_mask = torch.tensor(
+            rdkit_bounds_constraints["is_angle"].copy(), dtype=torch.bool
+        )
+        rdkit_upper_bounds = torch.tensor(
+            rdkit_bounds_constraints["upper_bound"].copy(), dtype=torch.float
+        )
+        rdkit_lower_bounds = torch.tensor(
+            rdkit_bounds_constraints["lower_bound"].copy(), dtype=torch.float
+        )
+
+        chiral_atom_index = torch.tensor(
+            chiral_atom_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+        chiral_reference_mask = torch.tensor(
+            chiral_atom_constraints["is_reference"].copy(), dtype=torch.bool
+        )
+        chiral_atom_orientations = torch.tensor(
+            chiral_atom_constraints["is_r"].copy(), dtype=torch.bool
+        )
+
+        stereo_bond_index = torch.tensor(
+            stereo_bond_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+        stereo_reference_mask = torch.tensor(
+            stereo_bond_constraints["is_reference"].copy(), dtype=torch.bool
+        )
+        stereo_bond_orientations = torch.tensor(
+            stereo_bond_constraints["is_e"].copy(), dtype=torch.bool
+        )
+
+        planar_bond_index = torch.tensor(
+            planar_bond_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+        planar_ring_5_index = torch.tensor(
+            planar_ring_5_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+        planar_ring_6_index = torch.tensor(
+            planar_ring_6_constraints["atom_idxs"].copy(), dtype=torch.long
+        ).T
+    else:
+        rdkit_bounds_index = torch.empty((2, 0), dtype=torch.long)
+        rdkit_bounds_bond_mask = torch.empty((0,), dtype=torch.bool)
+        rdkit_bounds_angle_mask = torch.empty((0,), dtype=torch.bool)
+        rdkit_upper_bounds = torch.empty((0,), dtype=torch.float)
+        rdkit_lower_bounds = torch.empty((0,), dtype=torch.float)
+        chiral_atom_index = torch.empty(
+            (
+                4,
+                0,
+            ),
+            dtype=torch.long,
+        )
+        chiral_reference_mask = torch.empty((0,), dtype=torch.bool)
+        chiral_atom_orientations = torch.empty((0,), dtype=torch.bool)
+        stereo_bond_index = torch.empty((4, 0), dtype=torch.long)
+        stereo_reference_mask = torch.empty((0,), dtype=torch.bool)
+        stereo_bond_orientations = torch.empty((0,), dtype=torch.bool)
+        planar_bond_index = torch.empty((6, 0), dtype=torch.long)
+        planar_ring_5_index = torch.empty((5, 0), dtype=torch.long)
+        planar_ring_6_index = torch.empty((6, 0), dtype=torch.long)
+
+    return {
+        "rdkit_bounds_index": rdkit_bounds_index,
+        "rdkit_bounds_bond_mask": rdkit_bounds_bond_mask,
+        "rdkit_bounds_angle_mask": rdkit_bounds_angle_mask,
+        "rdkit_upper_bounds": rdkit_upper_bounds,
+        "rdkit_lower_bounds": rdkit_lower_bounds,
+        "chiral_atom_index": chiral_atom_index,
+        "chiral_reference_mask": chiral_reference_mask,
+        "chiral_atom_orientations": chiral_atom_orientations,
+        "stereo_bond_index": stereo_bond_index,
+        "stereo_reference_mask": stereo_reference_mask,
+        "stereo_bond_orientations": stereo_bond_orientations,
+        "planar_bond_index": planar_bond_index,
+        "planar_ring_5_index": planar_ring_5_index,
+        "planar_ring_6_index": planar_ring_6_index,
+    }
+
+
+def process_chain_feature_constraints(data: Tokenized) -> dict[str, Tensor]:
+    structure = data.structure
+    if structure.bonds.shape[0] > 0:
+        connected_chain_index, connected_atom_index = [], []
+        for connection in structure.bonds:
+            if connection["chain_1"] == connection["chain_2"]:
+                continue
+            connected_chain_index.append([connection["chain_1"], connection["chain_2"]])
+            connected_atom_index.append([connection["atom_1"], connection["atom_2"]])
+        if len(connected_chain_index) > 0:
+            connected_chain_index = torch.tensor(
+                connected_chain_index, dtype=torch.long
+            ).T
+            connected_atom_index = torch.tensor(
+                connected_atom_index, dtype=torch.long
+            ).T
+        else:
+            connected_chain_index = torch.empty((2, 0), dtype=torch.long)
+            connected_atom_index = torch.empty((2, 0), dtype=torch.long)
+    else:
+        connected_chain_index = torch.empty((2, 0), dtype=torch.long)
+        connected_atom_index = torch.empty((2, 0), dtype=torch.long)
+
+    symmetric_chain_index = []
+    for i, chain_i in enumerate(structure.chains):
+        for j, chain_j in enumerate(structure.chains):
+            if j <= i:
+                continue
+            if chain_i["entity_id"] == chain_j["entity_id"]:
+                symmetric_chain_index.append([i, j])
+    if len(symmetric_chain_index) > 0:
+        symmetric_chain_index = torch.tensor(symmetric_chain_index, dtype=torch.long).T
+    else:
+        symmetric_chain_index = torch.empty((2, 0), dtype=torch.long)
+    return {
+        "connected_chain_index": connected_chain_index,
+        "connected_atom_index": connected_atom_index,
+        "symmetric_chain_index": symmetric_chain_index,
+    }
+
+
+class Boltz2Featurizer:
+    """Boltz2 featurizer."""
+
+    def process(
+        self,
+        data: Tokenized,
+        random: np.random.Generator,
+        molecules: dict[str, Mol],
+        training: bool,
+        max_seqs: int,
+        atoms_per_window_queries: int = 32,
+        min_dist: float = 2.0,
+        max_dist: float = 22.0,
+        num_bins: int = 64,
+        num_ensembles: int = 1,
+        ensemble_sample_replacement: bool = False,
+        disto_use_ensemble: Optional[bool] = False,
+        fix_single_ensemble: Optional[bool] = True,
+        max_tokens: Optional[int] = None,
+        max_atoms: Optional[int] = None,
+        pad_to_max_seqs: bool = False,
+        compute_symmetries: bool = False,
+        binder_pocket_conditioned_prop: Optional[float] = 0.0,
+        contact_conditioned_prop: Optional[float] = 0.0,
+        binder_pocket_cutoff_min: Optional[float] = 4.0,
+        binder_pocket_cutoff_max: Optional[float] = 20.0,
+        binder_pocket_sampling_geometric_p: Optional[float] = 0.0,
+        only_ligand_binder_pocket: Optional[bool] = False,
+        only_pp_contact: Optional[bool] = False,
+        pocket_constraints: Optional[list] = None,
+        single_sequence_prop: Optional[float] = 0.0,
+        msa_sampling: bool = False,
+        override_bfactor: float = False,
+        override_method: Optional[str] = None,
+        compute_frames: bool = False,
+        override_coords: Optional[Tensor] = None,
+        bfactor_md_correction: bool = False,
+        compute_constraint_features: bool = False,
+        inference_pocket_constraints: Optional[list] = None,
+        compute_affinity: bool = False,
+    ) -> dict[str, Tensor]:
+        """Compute features.
+
+        Parameters
+        ----------
+        data : Tokenized
+            The input to the model.
+        training : bool
+            Whether the model is in training mode.
+        max_tokens : int, optional
+            The maximum number of tokens.
+        max_atoms : int, optional
+            The maximum number of atoms
+        max_seqs : int, optional
+            The maximum number of sequences.
+
+        Returns
+        -------
+        dict[str, Tensor]
+            The features for model training.
+
+        """
+        # Compute random number of sequences
+        if training and max_seqs is not None:
+            if random.random() > single_sequence_prop:
+                max_seqs_batch = random.integers(1, max_seqs + 1)
+            else:
+                max_seqs_batch = 1
+        else:
+            max_seqs_batch = max_seqs
+
+        # Compute ensemble features
+        ensemble_features = process_ensemble_features(
+            data=data,
+            random=random,
+            num_ensembles=num_ensembles,
+            ensemble_sample_replacement=ensemble_sample_replacement,
+            fix_single_ensemble=fix_single_ensemble,
+        )
+
+        # Compute token features
+        token_features = process_token_features(
+            data=data,
+            random=random,
+            max_tokens=max_tokens,
+            binder_pocket_conditioned_prop=binder_pocket_conditioned_prop,
+            contact_conditioned_prop=contact_conditioned_prop,
+            binder_pocket_cutoff_min=binder_pocket_cutoff_min,
+            binder_pocket_cutoff_max=binder_pocket_cutoff_max,
+            binder_pocket_sampling_geometric_p=binder_pocket_sampling_geometric_p,
+            only_ligand_binder_pocket=only_ligand_binder_pocket,
+            only_pp_contact=only_pp_contact,
+            override_method=override_method,
+            inference_pocket_constraints=inference_pocket_constraints,
+        )
+
+        # Compute atom features
+        atom_features = process_atom_features(
+            data=data,
+            random=random,
+            molecules=molecules,
+            ensemble_features=ensemble_features,
+            atoms_per_window_queries=atoms_per_window_queries,
+            min_dist=min_dist,
+            max_dist=max_dist,
+            num_bins=num_bins,
+            max_atoms=max_atoms,
+            max_tokens=max_tokens,
+            disto_use_ensemble=disto_use_ensemble,
+            override_bfactor=override_bfactor,
+            compute_frames=compute_frames,
+            override_coords=override_coords,
+            bfactor_md_correction=bfactor_md_correction,
+        )
+
+        # Compute MSA features
+        msa_features = process_msa_features(
+            data=data,
+            random=random,
+            max_seqs_batch=max_seqs_batch,
+            max_seqs=max_seqs,
+            max_tokens=max_tokens,
+            pad_to_max_seqs=pad_to_max_seqs,
+            msa_sampling=training and msa_sampling,
+        )
+
+        # Compute MSA features
+        msa_features_affinity = {}
+        if compute_affinity:
+            msa_features_affinity = process_msa_features(
+                data=data,
+                random=random,
+                max_seqs_batch=1,
+                max_seqs=1,
+                max_tokens=max_tokens,
+                pad_to_max_seqs=pad_to_max_seqs,
+                msa_sampling=training and msa_sampling,
+                affinity=True,
+            )
+
+        # Compute affinity ligand Molecular Weight
+        ligand_to_mw = {}
+        if compute_affinity:
+            ligand_to_mw["affinity_mw"] = data.record.affinity.mw
+
+        # Compute template features
+        num_tokens = data.tokens.shape[0] if max_tokens is None else max_tokens
+        if data.templates:
+            template_features = process_template_features(
+                data=data,
+                max_tokens=num_tokens,
+            )
+        else:
+            template_features = load_dummy_templates_features(
+                tdim=1,
+                num_tokens=num_tokens,
+            )
+
+        # Compute symmetry features
+        symmetry_features = {}
+        if compute_symmetries:
+            symmetries = get_symmetries(molecules)
+            symmetry_features = process_symmetry_features(data, symmetries)
+
+        # Compute residue constraint features
+        residue_constraint_features = {}
+        if compute_constraint_features:
+            residue_constraint_features = process_residue_constraint_features(data)
+            chain_constraint_features = process_chain_feature_constraints(data)
+
+        return {
+            **token_features,
+            **atom_features,
+            **msa_features,
+            **msa_features_affinity,
+            **template_features,
+            **symmetry_features,
+            **ensemble_features,
+            **residue_constraint_features,
+            **chain_constraint_features,
+            **ligand_to_mw,
+        }