boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/data/parse/yaml.py

@@ -0,0 +1,68 @@
+from pathlib import Path
+
+import yaml
+from rdkit.Chem.rdchem import Mol
+
+from boltz.data.parse.schema import parse_boltz_schema
+from boltz.data.types import Target
+
+
+def parse_yaml(
+    path: Path,
+    ccd: dict[str, Mol],
+    mol_dir: Path,
+    boltz2: bool = False,
+) -> Target:
+    """Parse a Boltz input yaml / json.
+
+    The input file should be a yaml file with the following format:
+
+    sequences:
+        - protein:
+            id: A
+            sequence: "MADQLTEEQIAEFKEAFSLF"
+        - protein:
+            id: [B, C]
+            sequence: "AKLSILPWGHC"
+        - rna:
+            id: D
+            sequence: "GCAUAGC"
+        - ligand:
+            id: E
+            smiles: "CC1=CC=CC=C1"
+        - ligand:
+            id: [F, G]
+            ccd: []
+    constraints:
+        - bond:
+            atom1: [A, 1, CA]
+            atom2: [A, 2, N]
+        - pocket:
+            binder: E
+            contacts: [[B, 1], [B, 2]]
+    templates:
+        - path: /path/to/template.pdb
+          ids: [A] # optional, specify which chains to template
+
+    version: 1
+
+    Parameters
+    ----------
+    path : Path
+        Path to the YAML input format.
+    components : Dict
+        Dictionary of CCD components.
+    boltz2 : bool
+        Whether to parse the input for Boltz2.
+
+    Returns
+    -------
+    Target
+        The parsed target.
+
+    """
+    with path.open("r") as file:
+        data = yaml.safe_load(file)
+
+    name = path.stem
+    return parse_boltz_schema(name, data, ccd, mol_dir, boltz2)

boltz/data/sample/cluster.py

@@ -0,0 +1,283 @@
+from typing import Dict, Iterator, List
+
+import numpy as np
+from numpy.random import RandomState
+
+from boltz.data import const
+from boltz.data.types import ChainInfo, InterfaceInfo, Record
+from boltz.data.sample.sampler import Sample, Sampler
+
+
+def get_chain_cluster(chain: ChainInfo, record: Record) -> str:  # noqa: ARG001
+    """Get the cluster id for a chain.
+
+    Parameters
+    ----------
+    chain : ChainInfo
+        The chain id to get the cluster id for.
+    record : Record
+        The record the interface is part of.
+
+    Returns
+    -------
+    str
+        The cluster id of the chain.
+
+    """
+    return chain.cluster_id
+
+
+def get_interface_cluster(interface: InterfaceInfo, record: Record) -> str:
+    """Get the cluster id for an interface.
+
+    Parameters
+    ----------
+    interface : InterfaceInfo
+        The interface to get the cluster id for.
+    record : Record
+        The record the interface is part of.
+
+    Returns
+    -------
+    str
+        The cluster id of the interface.
+
+    """
+    chain1 = record.chains[interface.chain_1]
+    chain2 = record.chains[interface.chain_2]
+
+    cluster_1 = str(chain1.cluster_id)
+    cluster_2 = str(chain2.cluster_id)
+
+    cluster_id = (cluster_1, cluster_2)
+    cluster_id = tuple(sorted(cluster_id))
+
+    return cluster_id
+
+
+def get_chain_weight(
+    chain: ChainInfo,
+    record: Record,  # noqa: ARG001
+    clusters: Dict[str, int],
+    beta_chain: float,
+    alpha_prot: float,
+    alpha_nucl: float,
+    alpha_ligand: float,
+) -> float:
+    """Get the weight of a chain.
+
+    Parameters
+    ----------
+    chain : ChainInfo
+        The chain to get the weight for.
+    record : Record
+        The record the chain is part of.
+    clusters : Dict[str, int]
+        The cluster sizes.
+    beta_chain : float
+        The beta value for chains.
+    alpha_prot : float
+        The alpha value for proteins.
+    alpha_nucl : float
+        The alpha value for nucleic acids.
+    alpha_ligand : float
+        The alpha value for ligands.
+
+    Returns
+    -------
+    float
+        The weight of the chain.
+
+    """
+    prot_id = const.chain_type_ids["PROTEIN"]
+    rna_id = const.chain_type_ids["RNA"]
+    dna_id = const.chain_type_ids["DNA"]
+    ligand_id = const.chain_type_ids["NONPOLYMER"]
+
+    weight = beta_chain / clusters[chain.cluster_id]
+    if chain.mol_type == prot_id:
+        weight *= alpha_prot
+    elif chain.mol_type in [rna_id, dna_id]:
+        weight *= alpha_nucl
+    elif chain.mol_type == ligand_id:
+        weight *= alpha_ligand
+
+    return weight
+
+
+def get_interface_weight(
+    interface: InterfaceInfo,
+    record: Record,
+    clusters: Dict[str, int],
+    beta_interface: float,
+    alpha_prot: float,
+    alpha_nucl: float,
+    alpha_ligand: float,
+) -> float:
+    """Get the weight of an interface.
+
+    Parameters
+    ----------
+    interface : InterfaceInfo
+        The interface to get the weight for.
+    record : Record
+        The record the interface is part of.
+    clusters : Dict[str, int]
+        The cluster sizes.
+    beta_interface : float
+        The beta value for interfaces.
+    alpha_prot : float
+        The alpha value for proteins.
+    alpha_nucl : float
+        The alpha value for nucleic acids.
+    alpha_ligand : float
+        The alpha value for ligands.
+
+    Returns
+    -------
+    float
+        The weight of the interface.
+
+    """
+    prot_id = const.chain_type_ids["PROTEIN"]
+    rna_id = const.chain_type_ids["RNA"]
+    dna_id = const.chain_type_ids["DNA"]
+    ligand_id = const.chain_type_ids["NONPOLYMER"]
+
+    chain1 = record.chains[interface.chain_1]
+    chain2 = record.chains[interface.chain_2]
+
+    n_prot = (chain1.mol_type) == prot_id
+    n_nuc = chain1.mol_type in [rna_id, dna_id]
+    n_ligand = chain1.mol_type == ligand_id
+
+    n_prot += chain2.mol_type == prot_id
+    n_nuc += chain2.mol_type in [rna_id, dna_id]
+    n_ligand += chain2.mol_type == ligand_id
+
+    weight = beta_interface / clusters[get_interface_cluster(interface, record)]
+    weight *= alpha_prot * n_prot + alpha_nucl * n_nuc + alpha_ligand * n_ligand
+    return weight
+
+
+class ClusterSampler(Sampler):
+    """The weighted sampling approach, as described in AF3.
+
+    Each chain / interface is given a weight according
+    to the following formula, and sampled accordingly:
+
+    w = b / n_clust *(a_prot * n_prot + a_nuc * n_nuc
+        + a_ligand * n_ligand)
+
+    """
+
+    def __init__(
+        self,
+        alpha_prot: float = 3.0,
+        alpha_nucl: float = 3.0,
+        alpha_ligand: float = 1.0,
+        beta_chain: float = 0.5,
+        beta_interface: float = 1.0,
+    ) -> None:
+        """Initialize the sampler.
+
+        Parameters
+        ----------
+        alpha_prot : float, optional
+            The alpha value for proteins.
+        alpha_nucl : float, optional
+            The alpha value for nucleic acids.
+        alpha_ligand : float, optional
+            The alpha value for ligands.
+        beta_chain : float, optional
+            The beta value for chains.
+        beta_interface : float, optional
+            The beta value for interfaces.
+
+        """
+        self.alpha_prot = alpha_prot
+        self.alpha_nucl = alpha_nucl
+        self.alpha_ligand = alpha_ligand
+        self.beta_chain = beta_chain
+        self.beta_interface = beta_interface
+
+    def sample(self, records: List[Record], random: RandomState) -> Iterator[Sample]:  # noqa: C901, PLR0912
+        """Sample a structure from the dataset infinitely.
+
+        Parameters
+        ----------
+        records : List[Record]
+            The records to sample from.
+        random : RandomState
+            The random state for reproducibility.
+
+        Yields
+        ------
+        Sample
+            A data sample.
+
+        """
+        # Compute chain cluster sizes
+        chain_clusters: Dict[str, int] = {}
+        for record in records:
+            for chain in record.chains:
+                if not chain.valid:
+                    continue
+                cluster_id = get_chain_cluster(chain, record)
+                if cluster_id not in chain_clusters:
+                    chain_clusters[cluster_id] = 0
+                chain_clusters[cluster_id] += 1
+
+        # Compute interface clusters sizes
+        interface_clusters: Dict[str, int] = {}
+        for record in records:
+            for interface in record.interfaces:
+                if not interface.valid:
+                    continue
+                cluster_id = get_interface_cluster(interface, record)
+                if cluster_id not in interface_clusters:
+                    interface_clusters[cluster_id] = 0
+                interface_clusters[cluster_id] += 1
+
+        # Compute weights
+        items, weights = [], []
+        for record in records:
+            for chain_id, chain in enumerate(record.chains):
+                if not chain.valid:
+                    continue
+                weight = get_chain_weight(
+                    chain,
+                    record,
+                    chain_clusters,
+                    self.beta_chain,
+                    self.alpha_prot,
+                    self.alpha_nucl,
+                    self.alpha_ligand,
+                )
+                items.append((record, 0, chain_id))
+                weights.append(weight)
+
+            for int_id, interface in enumerate(record.interfaces):
+                if not interface.valid:
+                    continue
+                weight = get_interface_weight(
+                    interface,
+                    record,
+                    interface_clusters,
+                    self.beta_interface,
+                    self.alpha_prot,
+                    self.alpha_nucl,
+                    self.alpha_ligand,
+                )
+                items.append((record, 1, int_id))
+                weights.append(weight)
+
+        # Sample infinitely
+        weights = np.array(weights) / np.sum(weights)
+        while True:
+            item_idx = random.choice(len(items), p=weights)
+            record, kind, index = items[item_idx]
+            if kind == 0:
+                yield Sample(record=record, chain_id=index)
+            else:
+                yield Sample(record=record, interface_id=index)

boltz/data/sample/distillation.py

@@ -0,0 +1,57 @@
+from typing import Iterator, List
+
+from numpy.random import RandomState
+
+from boltz.data.types import Record
+from boltz.data.sample.sampler import Sample, Sampler
+
+
+class DistillationSampler(Sampler):
+    """A sampler for monomer distillation data."""
+
+    def __init__(self, small_size: int = 200, small_prob: float = 0.01) -> None:
+        """Initialize the sampler.
+
+        Parameters
+        ----------
+        small_size : int, optional
+            The maximum size to be considered small.
+        small_prob : float, optional
+            The probability of sampling a small item.
+
+        """
+        self._size = small_size
+        self._prob = small_prob
+
+    def sample(self, records: List[Record], random: RandomState) -> Iterator[Sample]:
+        """Sample a structure from the dataset infinitely.
+
+        Parameters
+        ----------
+        records : List[Record]
+            The records to sample from.
+        random : RandomState
+            The random state for reproducibility.
+
+        Yields
+        ------
+        Sample
+            A data sample.
+
+        """
+        # Remove records with invalid chains
+        records = [r for r in records if r.chains[0].valid]
+
+        # Split in small and large proteins. We assume that there is only
+        # one chain per record, as is the case for monomer distillation
+        small = [r for r in records if r.chains[0].num_residues <= self._size]
+        large = [r for r in records if r.chains[0].num_residues > self._size]
+
+        # Sample infinitely
+        while True:
+            # Sample small or large
+            samples = small if random.rand() < self._prob else large
+
+            # Sample item from the list
+            index = random.randint(0, len(samples))
+            yield Sample(record=samples[index])

boltz/data/sample/random.py

@@ -0,0 +1,39 @@
+from dataclasses import replace
+from typing import Iterator, List
+
+from numpy.random import RandomState
+
+from boltz.data.types import Record
+from boltz.data.sample.sampler import Sample, Sampler
+
+
+class RandomSampler(Sampler):
+    """A simple random sampler with replacement."""
+
+    def sample(self, records: List[Record], random: RandomState) -> Iterator[Sample]:
+        """Sample a structure from the dataset infinitely.
+
+        Parameters
+        ----------
+        records : List[Record]
+            The records to sample from.
+        random : RandomState
+            The random state for reproducibility.
+
+        Yields
+        ------
+        Sample
+            A data sample.
+
+        """
+        while True:
+            # Sample item from the list
+            index = random.randint(0, len(records))
+            record = records[index]
+
+            # Remove invalid chains and interfaces
+            chains = [c for c in record.chains if c.valid]
+            interfaces = [i for i in record.interfaces if i.valid]
+            record = replace(record, chains=chains, interfaces=interfaces)
+
+            yield Sample(record=record)

boltz/data/sample/sampler.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Iterator, List, Optional
+
+from numpy.random import RandomState
+
+from boltz.data.types import Record
+
+
+@dataclass
+class Sample:
+    """A sample with optional chain and interface IDs.
+
+    Attributes
+    ----------
+    record : Record
+        The record.
+    chain_id : Optional[int]
+        The chain ID.
+    interface_id : Optional[int]
+        The interface ID.
+    """
+
+    record: Record
+    chain_id: Optional[int] = None
+    interface_id: Optional[int] = None
+
+
+class Sampler(ABC):
+    """Abstract base class for samplers."""
+
+    @abstractmethod
+    def sample(self, records: List[Record], random: RandomState) -> Iterator[Sample]:
+        """Sample a structure from the dataset infinitely.
+
+        Parameters
+        ----------
+        records : List[Record]
+            The records to sample from.
+        random : RandomState
+            The random state for reproducibility.
+
+        Yields
+        ------
+        Sample
+            A data sample.
+
+        """
+        raise NotImplementedError

boltz/data/tokenize/boltz.py

@@ -0,0 +1,195 @@
+from dataclasses import astuple, dataclass
+
+import numpy as np
+
+from boltz.data import const
+from boltz.data.tokenize.tokenizer import Tokenizer
+from boltz.data.types import Input, Token, TokenBond, Tokenized
+
+
+@dataclass
+class TokenData:
+    """TokenData datatype."""
+
+    token_idx: int
+    atom_idx: int
+    atom_num: int
+    res_idx: int
+    res_type: int
+    sym_id: int
+    asym_id: int
+    entity_id: int
+    mol_type: int
+    center_idx: int
+    disto_idx: int
+    center_coords: np.ndarray
+    disto_coords: np.ndarray
+    resolved_mask: bool
+    disto_mask: bool
+    cyclic_period: int
+
+
+class BoltzTokenizer(Tokenizer):
+    """Tokenize an input structure for training."""
+
+    def tokenize(self, data: Input) -> Tokenized:
+        """Tokenize the input data.
+
+        Parameters
+        ----------
+        data : Input
+            The input data.
+
+        Returns
+        -------
+        Tokenized
+            The tokenized data.
+
+        """
+        # Get structure data
+        struct = data.structure
+
+        # Create token data
+        token_data = []
+
+        # Keep track of atom_idx to token_idx
+        token_idx = 0
+        atom_to_token = {}
+
+        # Filter to valid chains only
+        chains = struct.chains[struct.mask]
+
+        for chain in chains:
+            # Get residue indices
+            res_start = chain["res_idx"]
+            res_end = chain["res_idx"] + chain["res_num"]
+
+            for res in struct.residues[res_start:res_end]:
+                # Get atom indices
+                atom_start = res["atom_idx"]
+                atom_end = res["atom_idx"] + res["atom_num"]
+
+                # Standard residues are tokens
+                if res["is_standard"]:
+                    # Get center and disto atoms
+                    center = struct.atoms[res["atom_center"]]
+                    disto = struct.atoms[res["atom_disto"]]
+
+                    # Token is present if centers are
+                    is_present = res["is_present"] & center["is_present"]
+                    is_disto_present = res["is_present"] & disto["is_present"]
+
+                    # Apply chain transformation
+                    c_coords = center["coords"]
+                    d_coords = disto["coords"]
+
+                    # Create token
+                    token = TokenData(
+                        token_idx=token_idx,
+                        atom_idx=res["atom_idx"],
+                        atom_num=res["atom_num"],
+                        res_idx=res["res_idx"],
+                        res_type=res["res_type"],
+                        sym_id=chain["sym_id"],
+                        asym_id=chain["asym_id"],
+                        entity_id=chain["entity_id"],
+                        mol_type=chain["mol_type"],
+                        center_idx=res["atom_center"],
+                        disto_idx=res["atom_disto"],
+                        center_coords=c_coords,
+                        disto_coords=d_coords,
+                        resolved_mask=is_present,
+                        disto_mask=is_disto_present,
+                        cyclic_period=chain["cyclic_period"],
+                    )
+                    token_data.append(astuple(token))
+
+                    # Update atom_idx to token_idx
+                    for atom_idx in range(atom_start, atom_end):
+                        atom_to_token[atom_idx] = token_idx
+
+                    token_idx += 1
+
+                # Non-standard are tokenized per atom
+                else:
+                    # We use the unk protein token as res_type
+                    unk_token = const.unk_token["PROTEIN"]
+                    unk_id = const.token_ids[unk_token]
+
+                    # Get atom coordinates
+                    atom_data = struct.atoms[atom_start:atom_end]
+                    atom_coords = atom_data["coords"]
+
+                    # Tokenize each atom
+                    for i, atom in enumerate(atom_data):
+                        # Token is present if atom is
+                        is_present = res["is_present"] & atom["is_present"]
+                        index = atom_start + i
+
+                        # Create token
+                        token = TokenData(
+                            token_idx=token_idx,
+                            atom_idx=index,
+                            atom_num=1,
+                            res_idx=res["res_idx"],
+                            res_type=unk_id,
+                            sym_id=chain["sym_id"],
+                            asym_id=chain["asym_id"],
+                            entity_id=chain["entity_id"],
+                            mol_type=chain["mol_type"],
+                            center_idx=index,
+                            disto_idx=index,
+                            center_coords=atom_coords[i],
+                            disto_coords=atom_coords[i],
+                            resolved_mask=is_present,
+                            disto_mask=is_present,
+                            cyclic_period=chain[
+                                "cyclic_period"
+                            ],  # Enforced to be False in chain parser
+                        )
+                        token_data.append(astuple(token))
+
+                        # Update atom_idx to token_idx
+                        atom_to_token[index] = token_idx
+                        token_idx += 1
+
+        # Create token bonds
+        token_bonds = []
+
+        # Add atom-atom bonds from ligands
+        for bond in struct.bonds:
+            if (
+                bond["atom_1"] not in atom_to_token
+                or bond["atom_2"] not in atom_to_token
+            ):
+                continue
+            token_bond = (
+                atom_to_token[bond["atom_1"]],
+                atom_to_token[bond["atom_2"]],
+            )
+            token_bonds.append(token_bond)
+
+        # Add connection bonds (covalent)
+        for conn in struct.connections:
+            if (
+                conn["atom_1"] not in atom_to_token
+                or conn["atom_2"] not in atom_to_token
+            ):
+                continue
+            token_bond = (
+                atom_to_token[conn["atom_1"]],
+                atom_to_token[conn["atom_2"]],
+            )
+            token_bonds.append(token_bond)
+
+        token_data = np.array(token_data, dtype=Token)
+        token_bonds = np.array(token_bonds, dtype=TokenBond)
+        tokenized = Tokenized(
+            tokens=token_data,
+            bonds=token_bonds,
+            structure=data.structure,
+            msa=data.msa,
+            record=data.record,
+            residue_constraints=data.residue_constraints,
+        )
+        return tokenized