boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/data/filter/dynamic/max_residues.py

@@ -0,0 +1,37 @@
+from boltz.data.types import Record
+from boltz.data.filter.dynamic.filter import DynamicFilter
+
+
+class MaxResiduesFilter(DynamicFilter):
+    """A filter that filters structures based on their size."""
+
+    def __init__(self, min_residues: int = 1, max_residues: int = 500) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        min_chains : int
+            The minimum number of chains allowed.
+        max_chains : int
+            The maximum number of chains allowed.
+
+        """
+        self.min_residues = min_residues
+        self.max_residues = max_residues
+
+    def filter(self, record: Record) -> bool:
+        """Filter structures based on their resolution.
+
+        Parameters
+        ----------
+        record : Record
+            The record to filter.
+
+        Returns
+        -------
+        bool
+            Whether the record should be filtered.
+
+        """
+        num_residues = sum(chain.num_residues for chain in record.chains)
+        return num_residues <= self.max_residues and num_residues >= self.min_residues

boltz/data/filter/dynamic/resolution.py

@@ -0,0 +1,34 @@
+from boltz.data.types import Record
+from boltz.data.filter.dynamic.filter import DynamicFilter
+
+
+class ResolutionFilter(DynamicFilter):
+    """A filter that filters complexes based on their resolution."""
+
+    def __init__(self, resolution: float = 9.0) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        resolution : float, optional
+            The maximum allowed resolution.
+
+        """
+        self.resolution = resolution
+
+    def filter(self, record: Record) -> bool:
+        """Filter complexes based on their resolution.
+
+        Parameters
+        ----------
+        record : Record
+            The record to filter.
+
+        Returns
+        -------
+        bool
+            Whether the record should be filtered.
+
+        """
+        structure = record.structure
+        return structure.resolution <= self.resolution

boltz/data/filter/dynamic/size.py

@@ -0,0 +1,38 @@
+from boltz.data.types import Record
+from boltz.data.filter.dynamic.filter import DynamicFilter
+
+
+class SizeFilter(DynamicFilter):
+    """A filter that filters structures based on their size."""
+
+    def __init__(self, min_chains: int = 1, max_chains: int = 300) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        min_chains : int
+            The minimum number of chains allowed.
+        max_chains : int
+            The maximum number of chains allowed.
+
+        """
+        self.min_chains = min_chains
+        self.max_chains = max_chains
+
+    def filter(self, record: Record) -> bool:
+        """Filter structures based on their resolution.
+
+        Parameters
+        ----------
+        record : Record
+            The record to filter.
+
+        Returns
+        -------
+        bool
+            Whether the record should be filtered.
+
+        """
+        num_chains = record.structure.num_chains
+        num_valid = sum(1 for chain in record.chains if chain.valid)
+        return num_chains <= self.max_chains and num_valid >= self.min_chains

boltz/data/filter/dynamic/subset.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+
+from boltz.data.types import Record
+from boltz.data.filter.dynamic.filter import DynamicFilter
+
+
+class SubsetFilter(DynamicFilter):
+    """Filter a data record based on a subset of the data."""
+
+    def __init__(self, subset: str, reverse: bool = False) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        subset : str
+            The subset of data to consider, one per line.
+
+        """
+        with Path(subset).open("r") as f:
+            subset = f.read().splitlines()
+
+        self.subset = {s.lower() for s in subset}
+        self.reverse = reverse
+
+    def filter(self, record: Record) -> bool:
+        """Filter a data record.
+
+        Parameters
+        ----------
+        record : Record
+            The object to consider filtering in / out.
+
+        Returns
+        -------
+        bool
+            True if the data passes the filter, False otherwise.
+
+        """
+        if self.reverse:
+            return record.id.lower() not in self.subset
+        else:  # noqa: RET505
+            return record.id.lower() in self.subset

boltz/data/filter/static/filter.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+from boltz.data.types import Structure
+
+
+class StaticFilter(ABC):
+    """Base class for structure filters."""
+
+    @abstractmethod
+    def filter(self, structure: Structure) -> np.ndarray:
+        """Filter chains in a structure.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        raise NotImplementedError

boltz/data/filter/static/ligand.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from boltz.data import const
+from boltz.data.filter.static.filter import StaticFilter
+from boltz.data.types import Structure
+
+
+class ExcludedLigands(StaticFilter):
+    """Filter excluded ligands."""
+
+    def filter(self, structure: Structure) -> np.ndarray:
+        """Filter excluded ligands.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        valid = np.ones(len(structure.chains), dtype=bool)
+
+        for i, chain in enumerate(structure.chains):
+            if chain["mol_type"] != const.chain_type_ids["NONPOLYMER"]:
+                continue
+
+            res_start = chain["res_idx"]
+            res_end = res_start + chain["res_num"]
+            residues = structure.residues[res_start:res_end]
+            if any(res["name"] in const.ligand_exclusion for res in residues):
+                valid[i] = 0
+
+        return valid

boltz/data/filter/static/polymer.py

@@ -0,0 +1,299 @@
+import itertools
+from dataclasses import dataclass
+
+import numpy as np
+from sklearn.neighbors import KDTree
+
+from boltz.data import const
+from boltz.data.filter.static.filter import StaticFilter
+from boltz.data.types import Structure
+
+
+class MinimumLengthFilter(StaticFilter):
+    """Filter polymers based on their length.
+
+    We use the number of resolved residues when considering
+    the minimum, and the sequence length for the maximum.
+
+    """
+
+    def __init__(self, min_len: int = 4, max_len: int = 5000) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        min_len : float, optional
+            The minimum allowed length.
+        max_len : float, optional
+            The maximum allowed length.
+
+        """
+        self._min = min_len
+        self._max = max_len
+
+    def filter(self, structure: Structure) -> np.ndarray:
+        """Filter a chains based on their length.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        valid = np.ones(len(structure.chains), dtype=bool)
+
+        for i, chain in enumerate(structure.chains):
+            if chain["mol_type"] == const.chain_type_ids["NONPOLYMER"]:
+                continue
+
+            res_start = chain["res_idx"]
+            res_end = res_start + chain["res_num"]
+            residues = structure.residues[res_start:res_end]
+            resolved = residues["is_present"].sum()
+
+            if (resolved < self._min) or (resolved > self._max):
+                valid[i] = 0
+
+        return valid
+
+
+class UnknownFilter(StaticFilter):
+    """Filter proteins with all unknown residues."""
+
+    def filter(self, structure: Structure) -> np.ndarray:
+        """Filter proteins with all unknown residues.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        valid = np.ones(len(structure.chains), dtype=bool)
+        unk_toks = {
+            const.chain_type_ids["PROTEIN"]: const.unk_token_ids["PROTEIN"],
+            const.chain_type_ids["DNA"]: const.unk_token_ids["DNA"],
+            const.chain_type_ids["RNA"]: const.unk_token_ids["RNA"],
+        }
+
+        for i, chain in enumerate(structure.chains):
+            if chain["mol_type"] == const.chain_type_ids["NONPOLYMER"]:
+                continue
+
+            res_start = chain["res_idx"]
+            res_end = res_start + chain["res_num"]
+            residues = structure.residues[res_start:res_end]
+
+            unk_id = unk_toks[chain["mol_type"]]
+            if np.all(residues["res_type"] == unk_id):
+                valid[i] = 0
+
+        return valid
+
+
+class ConsecutiveCA(StaticFilter):
+    """Filter proteins with consecutive CA atoms above a threshold."""
+
+    def __init__(self, max_dist: int = 10.0) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        max_dist : float, optional
+            The maximum allowed distance.
+
+        """
+        self._max_dist = max_dist
+
+    def filter(self, structure: Structure) -> np.ndarray:
+        """Filter protein if consecutive CA atoms above a threshold.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        valid = np.ones(len(structure.chains), dtype=bool)
+
+        # Remove chain if consecutive CA atoms are above threshold
+        for i, chain in enumerate(structure.chains):
+            # Skip non-protein chains
+            if chain["mol_type"] != const.chain_type_ids["PROTEIN"]:
+                continue
+
+            # Get residues
+            res_start = chain["res_idx"]
+            res_end = res_start + chain["res_num"]
+            residues = structure.residues[res_start:res_end]
+
+            # Get c-alphas
+            ca_ids = residues["atom_center"]
+            ca_atoms = structure.atoms[ca_ids]
+
+            res_valid = residues["is_present"]
+            ca_valid = ca_atoms["is_present"] & res_valid
+            ca_coords = ca_atoms["coords"]
+
+            # Compute distances between consecutive atoms
+            dist = np.linalg.norm(ca_coords[1:] - ca_coords[:-1], axis=1)
+            dist = dist > self._max_dist
+            dist = dist[ca_valid[1:] & ca_valid[:-1]]
+
+            # Remove the chain if any valid pair is above threshold
+            if np.any(dist):
+                valid[i] = 0
+
+        return valid
+
+
+@dataclass(frozen=True)
+class Clash:
+    """A clash between two chains."""
+
+    chain: int
+    other: int
+    num_atoms: int
+    num_clashes: int
+
+
+class ClashingChainsFilter(StaticFilter):
+    """A filter that filters clashing chains.
+
+    Clashing chains are defined as those with >30% of atoms
+    within 1.7 Å of an atom in another chain. If two chains
+    are clashing with each other, the chain with the greater
+    percentage of clashing atoms will be removed. If the same
+    fraction of atoms are clashing, the chain with fewer total
+    atoms is removed. If the chains have the same number of
+    atoms, then the chain with the larger chain id is removed.
+
+    """
+
+    def __init__(self, dist: float = 1.7, freq: float = 0.3) -> None:
+        """Initialize the filter.
+
+        Parameters
+        ----------
+        dist : float, optional
+            The maximum distance for a clash.
+        freq : float, optional
+            The maximum allowed frequency of clashes.
+
+        """
+        self._dist = dist
+        self._freq = freq
+
+    def filter(self, structure: Structure) -> np.ndarray:  # noqa: PLR0912, C901
+        """Filter out clashing chains.
+
+        Parameters
+        ----------
+        structure : Structure
+            The structure to filter chains from.
+
+        Returns
+        -------
+        np.ndarray
+            The chains to keep, as a boolean mask.
+
+        """
+        num_chains = len(structure.chains)
+        if num_chains < 2:  # noqa: PLR2004
+            return np.ones(num_chains, dtype=bool)
+
+        # Get unique chain pairs
+        pairs = itertools.combinations(range(num_chains), 2)
+
+        # Compute clashes
+        clashes: list[Clash] = []
+        for i, j in pairs:
+            # Get the chains
+            c1 = structure.chains[i]
+            c2 = structure.chains[j]
+
+            # Get the atoms from each chain
+            c1_start = c1["atom_idx"]
+            c2_start = c2["atom_idx"]
+            c1_end = c1_start + c1["atom_num"]
+            c2_end = c2_start + c2["atom_num"]
+
+            atoms1 = structure.atoms[c1_start:c1_end]
+            atoms2 = structure.atoms[c2_start:c2_end]
+            atoms1 = atoms1[atoms1["is_present"]]
+            atoms2 = atoms2[atoms2["is_present"]]
+
+            # Skip if either chain has no atoms
+            if len(atoms1) == 0 or len(atoms2) == 0:
+                continue
+
+            # Compute the number of clashes
+            # Compute the distance matrix
+            tree = KDTree(atoms1["coords"], metric="euclidean")
+            query = tree.query_radius(atoms2["coords"], self._dist)
+
+            c2_clashes = sum(len(neighbors) > 0 for neighbors in query)
+            c1_clashes = len(set(itertools.chain.from_iterable(query)))
+
+            # Save results
+            if (c1_clashes / len(atoms1)) > self._freq:
+                clashes.append(Clash(i, j, len(atoms1), c1_clashes))
+            if (c2_clashes / len(atoms2)) > self._freq:
+                clashes.append(Clash(j, i, len(atoms2), c2_clashes))
+
+        # Compute indices to clash map
+        removed = set()
+        ids_to_clash = {(c.chain, c.other): c for c in clashes}
+
+        # Filter out chains according to ruleset
+        for clash in clashes:
+            # If either is already removed, skip
+            if clash.chain in removed or clash.other in removed:
+                continue
+
+            # Check if the two chains clash with each other
+            other_clash = ids_to_clash.get((clash.other, clash.chain))
+            if other_clash is not None:
+                # Remove the chain with the most clashes
+                clash1_freq = clash.num_clashes / clash.num_atoms
+                clash2_freq = other_clash.num_clashes / other_clash.num_atoms
+                if clash1_freq > clash2_freq:
+                    removed.add(clash.chain)
+                elif clash1_freq < clash2_freq:
+                    removed.add(clash.other)
+
+                # If same, remove the chain with fewer atoms
+                elif clash.num_atoms < other_clash.num_atoms:
+                    removed.add(clash.chain)
+                elif clash.num_atoms > other_clash.num_atoms:
+                    removed.add(clash.other)
+
+                # If same, remove the chain with the larger chain id
+                else:
+                    removed.add(max(clash.chain, clash.other))
+
+            # Otherwise, just remove the chain directly
+            else:
+                removed.add(clash.chain)
+
+        # Remove the chains
+        valid = np.ones(len(structure.chains), dtype=bool)
+        for i in removed:
+            valid[i] = 0
+
+        return valid