boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/data/msa/mmseqs2.py

@@ -0,0 +1,235 @@
+# From https://github.com/sokrypton/ColabFold/blob/main/colabfold/colabfold.py
+
+import logging
+import os
+import random
+import tarfile
+import time
+from typing import Union
+
+import requests
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+TQDM_BAR_FORMAT = (
+    "{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]"
+)
+
+
+def run_mmseqs2(  # noqa: PLR0912, D103, C901, PLR0915
+    x: Union[str, list[str]],
+    prefix: str = "tmp",
+    use_env: bool = True,
+    use_filter: bool = True,
+    use_pairing: bool = False,
+    pairing_strategy: str = "greedy",
+    host_url: str = "https://api.colabfold.com",
+) -> tuple[list[str], list[str]]:
+    submission_endpoint = "ticket/pair" if use_pairing else "ticket/msa"
+
+    # Set header agent as boltz
+    headers = {}
+    headers["User-Agent"] = "boltz"
+
+    def submit(seqs, mode, N=101):
+        n, query = N, ""
+        for seq in seqs:
+            query += f">{n}\n{seq}\n"
+            n += 1
+
+        error_count = 0
+        while True:
+            try:
+                # https://requests.readthedocs.io/en/latest/user/advanced/#advanced
+                # "good practice to set connect timeouts to slightly larger than a multiple of 3"
+                res = requests.post(
+                    f"{host_url}/{submission_endpoint}",
+                    data={"q": query, "mode": mode},
+                    timeout=6.02,
+                    headers=headers,
+                )
+            except Exception as e:
+                error_count += 1
+                logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
+                logger.warning(f"Error: {e}")
+                if error_count > 5:
+                    raise Exception("Too many failed attempts for the MSA generation request.")
+                time.sleep(5)
+            else:
+                break
+
+        try:
+            out = res.json()
+        except ValueError:
+            logger.error(f"Server didn't reply with json: {res.text}")
+            out = {"status": "ERROR"}
+        return out
+
+    def status(ID):
+        error_count = 0
+        while True:
+            try:
+                res = requests.get(
+                    f"{host_url}/ticket/{ID}", timeout=6.02, headers=headers
+                )
+            except Exception as e:
+                error_count += 1
+                logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
+                logger.warning(f"Error: {e}")
+                if error_count > 5:
+                    raise Exception("Too many failed attempts for the MSA generation request.")
+                time.sleep(5)
+            else:
+                break
+        try:
+            out = res.json()
+        except ValueError:
+            logger.error(f"Server didn't reply with json: {res.text}")
+            out = {"status": "ERROR"}
+        return out
+
+    def download(ID, path):
+        error_count = 0
+        while True:
+            try:
+                res = requests.get(
+                    f"{host_url}/result/download/{ID}", timeout=6.02, headers=headers
+                )
+            except Exception as e:
+                error_count += 1
+                logger.warning(f"Error while fetching result from MSA server. Retrying... ({error_count}/5)")
+                logger.warning(f"Error: {e}")
+                if error_count > 5:
+                    raise Exception("Too many failed attempts for the MSA generation request.")
+                time.sleep(5)
+            else:
+                break
+        with open(path, "wb") as out:
+            out.write(res.content)
+
+    # process input x
+    seqs = [x] if isinstance(x, str) else x
+
+    # setup mode
+    if use_filter:
+        mode = "env" if use_env else "all"
+    else:
+        mode = "env-nofilter" if use_env else "nofilter"
+
+    if use_pairing:
+        mode = ""
+        # greedy is default, complete was the previous behavior
+        if pairing_strategy == "greedy":
+            mode = "pairgreedy"
+        elif pairing_strategy == "complete":
+            mode = "paircomplete"
+        if use_env:
+            mode = mode + "-env"
+
+    # define path
+    path = f"{prefix}_{mode}"
+    if not os.path.isdir(path):
+        os.mkdir(path)
+
+    # call mmseqs2 api
+    tar_gz_file = f"{path}/out.tar.gz"
+    N, REDO = 101, True
+
+    # deduplicate and keep track of order
+    seqs_unique = []
+    # TODO this might be slow for large sets
+    [seqs_unique.append(x) for x in seqs if x not in seqs_unique]
+    Ms = [N + seqs_unique.index(seq) for seq in seqs]
+    # lets do it!
+    if not os.path.isfile(tar_gz_file):
+        TIME_ESTIMATE = 150 * len(seqs_unique)
+        with tqdm(total=TIME_ESTIMATE, bar_format=TQDM_BAR_FORMAT) as pbar:
+            while REDO:
+                pbar.set_description("SUBMIT")
+
+                # Resubmit job until it goes through
+                out = submit(seqs_unique, mode, N)
+                while out["status"] in ["UNKNOWN", "RATELIMIT"]:
+                    sleep_time = 5 + random.randint(0, 5)
+                    logger.error(f"Sleeping for {sleep_time}s. Reason: {out['status']}")
+                    # resubmit
+                    time.sleep(sleep_time)
+                    out = submit(seqs_unique, mode, N)
+
+                if out["status"] == "ERROR":
+                    msg = (
+                        "MMseqs2 API is giving errors. Please confirm your "
+                        " input is a valid protein sequence. If error persists, "
+                        "please try again an hour later."
+                    )
+                    raise Exception(msg)
+
+                if out["status"] == "MAINTENANCE":
+                    msg = (
+                        "MMseqs2 API is undergoing maintenance. "
+                        "Please try again in a few minutes."
+                    )
+                    raise Exception(msg)
+
+                # wait for job to finish
+                ID, TIME = out["id"], 0
+                pbar.set_description(out["status"])
+                while out["status"] in ["UNKNOWN", "RUNNING", "PENDING"]:
+                    t = 5 + random.randint(0, 5)
+                    logger.error(f"Sleeping for {t}s. Reason: {out['status']}")
+                    time.sleep(t)
+                    out = status(ID)
+                    pbar.set_description(out["status"])
+                    if out["status"] == "RUNNING":
+                        TIME += t
+                        pbar.update(n=t)
+
+                if out["status"] == "COMPLETE":
+                    if TIME < TIME_ESTIMATE:
+                        pbar.update(n=(TIME_ESTIMATE - TIME))
+                    REDO = False
+
+                if out["status"] == "ERROR":
+                    REDO = False
+                    msg = (
+                        "MMseqs2 API is giving errors. Please confirm your "
+                        " input is a valid protein sequence. If error persists, "
+                        "please try again an hour later."
+                    )
+                    raise Exception(msg)
+
+            # Download results
+            download(ID, tar_gz_file)
+
+    # prep list of a3m files
+    if use_pairing:
+        a3m_files = [f"{path}/pair.a3m"]
+    else:
+        a3m_files = [f"{path}/uniref.a3m"]
+        if use_env:
+            a3m_files.append(f"{path}/bfd.mgnify30.metaeuk30.smag30.a3m")
+
+    # extract a3m files
+    if any(not os.path.isfile(a3m_file) for a3m_file in a3m_files):
+        with tarfile.open(tar_gz_file) as tar_gz:
+            tar_gz.extractall(path)
+
+    # gather a3m lines
+    a3m_lines = {}
+    for a3m_file in a3m_files:
+        update_M, M = True, None
+        for line in open(a3m_file, "r"):
+            if len(line) > 0:
+                if "\x00" in line:
+                    line = line.replace("\x00", "")
+                    update_M = True
+                if line.startswith(">") and update_M:
+                    M = int(line[1:].rstrip())
+                    update_M = False
+                    if M not in a3m_lines:
+                        a3m_lines[M] = []
+                a3m_lines[M].append(line)
+
+    a3m_lines = ["".join(a3m_lines[n]) for n in Ms]
+    return a3m_lines

boltz/data/pad.py

@@ -0,0 +1,84 @@
+import torch
+from torch import Tensor
+from torch.nn.functional import pad
+
+
+def pad_dim(data: Tensor, dim: int, pad_len: float, value: float = 0) -> Tensor:
+    """Pad a tensor along a given dimension.
+
+    Parameters
+    ----------
+    data : Tensor
+        The input tensor.
+    dim : int
+        The dimension to pad.
+    pad_len : float
+        The padding length.
+    value : int, optional
+        The value to pad with.
+
+    Returns
+    -------
+    Tensor
+        The padded tensor.
+
+    """
+    if pad_len == 0:
+        return data
+
+    total_dims = len(data.shape)
+    padding = [0] * (2 * (total_dims - dim))
+    padding[2 * (total_dims - 1 - dim) + 1] = pad_len
+    return pad(data, tuple(padding), value=value)
+
+
+def pad_to_max(data: list[Tensor], value: float = 0) -> tuple[Tensor, Tensor]:
+    """Pad the data in all dimensions to the maximum found.
+
+    Parameters
+    ----------
+    data : list[Tensor]
+        list of tensors to pad.
+    value : float
+        The value to use for padding.
+
+    Returns
+    -------
+    Tensor
+        The padded tensor.
+    Tensor
+        The padding mask.
+
+    """
+    if isinstance(data[0], str):
+        return data, 0
+
+    # Check if all have the same shape
+    if all(d.shape == data[0].shape for d in data):
+        return torch.stack(data, dim=0), 0
+
+    # Get the maximum in each dimension
+    num_dims = len(data[0].shape)
+    max_dims = [max(d.shape[i] for d in data) for i in range(num_dims)]
+
+    # Get the padding lengths
+    pad_lengths = []
+    for d in data:
+        dims = []
+        for i in range(num_dims):
+            dims.append(0)
+            dims.append(max_dims[num_dims - i - 1] - d.shape[num_dims - i - 1])
+        pad_lengths.append(dims)
+
+    # Pad the data
+    padding = [
+        pad(torch.ones_like(d), pad_len, value=0)
+        for d, pad_len in zip(data, pad_lengths)
+    ]
+    data = [pad(d, pad_len, value=value) for d, pad_len in zip(data, pad_lengths)]
+
+    # Stack the data
+    padding = torch.stack(padding, dim=0)
+    data = torch.stack(data, dim=0)
+
+    return data, padding

boltz/data/parse/a3m.py

@@ -0,0 +1,134 @@
+import gzip
+from pathlib import Path
+from typing import Optional, TextIO
+
+import numpy as np
+
+from boltz.data import const
+from boltz.data.types import MSA, MSADeletion, MSAResidue, MSASequence
+
+
+def _parse_a3m(  # noqa: C901
+    lines: TextIO,
+    taxonomy: Optional[dict[str, str]],
+    max_seqs: Optional[int] = None,
+) -> MSA:
+    """Process an MSA file.
+
+    Parameters
+    ----------
+    lines : TextIO
+        The lines of the MSA file.
+    taxonomy : dict[str, str]
+        The taxonomy database, if available.
+    max_seqs : int, optional
+        The maximum number of sequences.
+
+    Returns
+    -------
+    MSA
+        The MSA object.
+
+    """
+    visited = set()
+    sequences = []
+    deletions = []
+    residues = []
+
+    seq_idx = 0
+    for line in lines:
+        line: str
+        line = line.strip()  # noqa: PLW2901
+        if not line or line.startswith("#"):
+            continue
+
+        # Get taxonomy, if annotated
+        if line.startswith(">"):
+            header = line.split()[0]
+            if taxonomy and header.startswith(">UniRef100"):
+                uniref_id = header.split("_")[1]
+                taxonomy_id = taxonomy.get(uniref_id)
+                if taxonomy_id is None:
+                    taxonomy_id = -1
+            else:
+                taxonomy_id = -1
+            continue
+
+        # Skip if duplicate sequence
+        str_seq = line.replace("-", "").upper()
+        if str_seq not in visited:
+            visited.add(str_seq)
+        else:
+            continue
+
+        # Process sequence
+        residue = []
+        deletion = []
+        count = 0
+        res_idx = 0
+        for c in line:
+            if c != "-" and c.islower():
+                count += 1
+                continue
+            token = const.prot_letter_to_token[c]
+            token = const.token_ids[token]
+            residue.append(token)
+            if count > 0:
+                deletion.append((res_idx, count))
+                count = 0
+            res_idx += 1
+
+        res_start = len(residues)
+        res_end = res_start + len(residue)
+
+        del_start = len(deletions)
+        del_end = del_start + len(deletion)
+
+        sequences.append((seq_idx, taxonomy_id, res_start, res_end, del_start, del_end))
+        residues.extend(residue)
+        deletions.extend(deletion)
+
+        seq_idx += 1
+        if (max_seqs is not None) and (seq_idx >= max_seqs):
+            break
+
+    # Create MSA object
+    msa = MSA(
+        residues=np.array(residues, dtype=MSAResidue),
+        deletions=np.array(deletions, dtype=MSADeletion),
+        sequences=np.array(sequences, dtype=MSASequence),
+    )
+    return msa
+
+
+def parse_a3m(
+    path: Path,
+    taxonomy: Optional[dict[str, str]],
+    max_seqs: Optional[int] = None,
+) -> MSA:
+    """Process an A3M file.
+
+    Parameters
+    ----------
+    path : Path
+        The path to the a3m(.gz) file.
+    taxonomy : Redis
+        The taxonomy database.
+    max_seqs : int, optional
+        The maximum number of sequences.
+
+    Returns
+    -------
+    MSA
+        The MSA object.
+
+    """
+    # Read the file
+    if path.suffix == ".gz":
+        with gzip.open(str(path), "rt") as f:
+            msa = _parse_a3m(f, taxonomy, max_seqs)
+    else:
+        with path.open("r") as f:
+            msa = _parse_a3m(f, taxonomy, max_seqs)
+
+    return msa

boltz/data/parse/csv.py

@@ -0,0 +1,100 @@
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+from boltz.data import const
+from boltz.data.types import MSA, MSADeletion, MSAResidue, MSASequence
+
+
+def parse_csv(
+    path: Path,
+    max_seqs: Optional[int] = None,
+) -> MSA:
+    """Process an A3M file.
+
+    Parameters
+    ----------
+    path : Path
+        The path to the a3m(.gz) file.
+    max_seqs : int, optional
+        The maximum number of sequences.
+
+    Returns
+    -------
+    MSA
+        The MSA object.
+
+    """
+    # Read file
+    data = pd.read_csv(path)
+
+    # Check columns
+    if tuple(sorted(data.columns)) != ("key", "sequence"):
+        msg = "Invalid CSV format, expected columns: ['sequence', 'key']"
+        raise ValueError(msg)
+
+    # Create taxonomy mapping
+    visited = set()
+    sequences = []
+    deletions = []
+    residues = []
+
+    seq_idx = 0
+    for line, key in zip(data["sequence"], data["key"]):
+        line: str
+        line = line.strip()  # noqa: PLW2901
+        if not line:
+            continue
+
+        # Get taxonomy, if annotated
+        taxonomy_id = -1
+        if (str(key) != "nan") and (key is not None) and (key != ""):
+            taxonomy_id = key
+
+        # Skip if duplicate sequence
+        str_seq = line.replace("-", "").upper()
+        if str_seq not in visited:
+            visited.add(str_seq)
+        else:
+            continue
+
+        # Process sequence
+        residue = []
+        deletion = []
+        count = 0
+        res_idx = 0
+        for c in line:
+            if c != "-" and c.islower():
+                count += 1
+                continue
+            token = const.prot_letter_to_token[c]
+            token = const.token_ids[token]
+            residue.append(token)
+            if count > 0:
+                deletion.append((res_idx, count))
+                count = 0
+            res_idx += 1
+
+        res_start = len(residues)
+        res_end = res_start + len(residue)
+
+        del_start = len(deletions)
+        del_end = del_start + len(deletion)
+
+        sequences.append((seq_idx, taxonomy_id, res_start, res_end, del_start, del_end))
+        residues.extend(residue)
+        deletions.extend(deletion)
+
+        seq_idx += 1
+        if (max_seqs is not None) and (seq_idx >= max_seqs):
+            break
+
+    # Create MSA object
+    msa = MSA(
+        residues=np.array(residues, dtype=MSAResidue),
+        deletions=np.array(deletions, dtype=MSADeletion),
+        sequences=np.array(sequences, dtype=MSASequence),
+    )
+    return msa

boltz/data/parse/fasta.py

@@ -0,0 +1,138 @@
+from collections.abc import Mapping
+from pathlib import Path
+
+from Bio import SeqIO
+from rdkit.Chem.rdchem import Mol
+
+from boltz.data.parse.yaml import parse_boltz_schema
+from boltz.data.types import Target
+
+
+def parse_fasta(  # noqa: C901, PLR0912
+    path: Path,
+    ccd: Mapping[str, Mol],
+    mol_dir: Path,
+    boltz2: bool = False,
+) -> Target:
+    """Parse a fasta file.
+
+    The name of the fasta file is used as the name of this job.
+    We rely on the fasta record id to determine the entity type.
+
+    > CHAIN_ID|ENTITY_TYPE|MSA_ID
+    SEQUENCE
+    > CHAIN_ID|ENTITY_TYPE|MSA_ID
+    ...
+
+    Where ENTITY_TYPE is either protein, rna, dna, ccd or smiles,
+    and CHAIN_ID is the chain identifier, which should be unique.
+    The MSA_ID is optional and should only be used on proteins.
+
+    Parameters
+    ----------
+    fasta_file : Path
+        Path to the fasta file.
+    ccd : Dict
+        Dictionary of CCD components.
+    mol_dir : Path
+        Path to the directory containing the molecules.
+    boltz2 : bool
+        Whether to parse the input for Boltz2.
+
+    Returns
+    -------
+    Target
+        The parsed target.
+
+    """
+    # Read fasta file
+    with path.open("r") as f:
+        records = list(SeqIO.parse(f, "fasta"))
+
+    # Make sure all records have a chain id and entity
+    for seq_record in records:
+        if "|" not in seq_record.id:
+            msg = f"Invalid record id: {seq_record.id}"
+            raise ValueError(msg)
+
+        header = seq_record.id.split("|")
+        assert len(header) >= 2, f"Invalid record id: {seq_record.id}"
+
+        chain_id, entity_type = header[:2]
+        if entity_type.lower() not in {"protein", "dna", "rna", "ccd", "smiles"}:
+            msg = f"Invalid entity type: {entity_type}"
+            raise ValueError(msg)
+        if chain_id == "":
+            msg = "Empty chain id in input fasta!"
+            raise ValueError(msg)
+        if entity_type == "":
+            msg = "Empty entity type in input fasta!"
+            raise ValueError(msg)
+
+    # Convert to yaml format
+    sequences = []
+    for seq_record in records:
+        # Get chain id, entity type and sequence
+        header = seq_record.id.split("|")
+        chain_id, entity_type = header[:2]
+        if len(header) == 3 and header[2] != "":
+            assert (
+                entity_type.lower() == "protein"
+            ), "MSA_ID is only allowed for proteins"
+            msa_id = header[2]
+        else:
+            msa_id = None
+
+        entity_type = entity_type.upper()
+        seq = str(seq_record.seq)
+
+        if entity_type == "PROTEIN":
+            molecule = {
+                "protein": {
+                    "id": chain_id,
+                    "sequence": seq,
+                    "modifications": [],
+                    "msa": msa_id,
+                },
+            }
+        elif entity_type == "RNA":
+            molecule = {
+                "rna": {
+                    "id": chain_id,
+                    "sequence": seq,
+                    "modifications": [],
+                },
+            }
+        elif entity_type == "DNA":
+            molecule = {
+                "dna": {
+                    "id": chain_id,
+                    "sequence": seq,
+                    "modifications": [],
+                }
+            }
+        elif entity_type.upper() == "CCD":
+            molecule = {
+                "ligand": {
+                    "id": chain_id,
+                    "ccd": seq,
+                }
+            }
+        elif entity_type.upper() == "SMILES":
+            molecule = {
+                "ligand": {
+                    "id": chain_id,
+                    "smiles": seq,
+                }
+            }
+
+        sequences.append(molecule)
+
+    data = {
+        "sequences": sequences,
+        "bonds": [],
+        "version": 1,
+    }
+
+    name = path.stem
+    return parse_boltz_schema(name, data, ccd, mol_dir, boltz2)