protein-quest 0.9.0__py3-none-any.whl

protein_quest/parallel.py

@@ -0,0 +1,104 @@
+"""Dask helper functions."""
+
+import logging
+import os
+from collections.abc import Callable, Collection
+from typing import Concatenate, ParamSpec, cast
+
+from dask.distributed import Client, LocalCluster, progress
+from distributed.deploy.cluster import Cluster
+from psutil import cpu_count
+
+logger = logging.getLogger(__name__)
+
+
+def configure_dask_scheduler(
+    scheduler_address: str | Cluster | None,
+    name: str,
+    nproc: int = 1,
+) -> str | Cluster:
+    """Configure the Dask scheduler by reusing existing or creating a new cluster.
+
+    Args:
+        scheduler_address: Address of the Dask scheduler to connect to, or None for local cluster.
+        name: Name for the Dask cluster.
+        nproc: Number of processes to use per worker for CPU support.
+
+    Returns:
+        A Dask Cluster instance or a string address for the scheduler.
+    """
+    if scheduler_address is None:
+        scheduler_address = _configure_cpu_dask_scheduler(nproc, name)
+        logger.info(f"Using local Dask cluster: {scheduler_address}")
+
+    return scheduler_address
+
+
+def nr_cpus() -> int:
+    """Determine the number of CPU cores to use.
+
+    If the environment variables SLURM_CPUS_PER_TASK or OMP_NUM_THREADS are set,
+    their value is used. Otherwise, the number of physical CPU cores is returned.
+
+    Returns:
+        The number of CPU cores to use.
+
+    Raises:
+        ValueError: If the number of physical CPU cores cannot be determined.
+    """
+    physical_cores = cpu_count(logical=False)
+    if physical_cores is None:
+        msg = "Cannot determine number of logical CPU cores."
+        raise ValueError(msg)
+    for var in ["SLURM_CPUS_PER_TASK", "OMP_NUM_THREADS"]:
+        value = os.environ.get(var)
+        if value is not None:
+            logger.warning(
+                'Not using all CPU cores (%s) of machine, environment variable "%s" is set to %s.',
+                physical_cores,
+                var,
+                value,
+            )
+            return int(value)
+    return physical_cores
+
+
+def _configure_cpu_dask_scheduler(nproc: int, name: str) -> LocalCluster:
+    total_cpus = nr_cpus()
+    n_workers = total_cpus // nproc
+    # Use single thread per worker to prevent GIL slowing down the computations
+    return LocalCluster(name=name, threads_per_worker=1, n_workers=n_workers)
+
+
+# Generic type parameters used across helpers
+P = ParamSpec("P")
+
+
+def dask_map_with_progress[T, R, **P](
+    client: Client,
+    func: Callable[Concatenate[T, P], R],
+    iterable: Collection[T],
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> list[R]:
+    """
+    Wrapper for map, progress, and gather of Dask that returns a correctly typed list.
+
+    Args:
+        client: Dask client.
+        func: Function to map; first parameter comes from ``iterable`` and any
+            additional parameters can be provided positionally via ``*args`` or
+            as keyword arguments via ``**kwargs``.
+        iterable: Collection of arguments to map over.
+        *args: Additional positional arguments to pass to client.map().
+        **kwargs: Additional keyword arguments to pass to client.map().
+
+    Returns:
+        List of results of type returned by `func` function.
+    """
+    if client.dashboard_link:
+        logger.info(f"Follow progress on dask dashboard at: {client.dashboard_link}")
+    futures = client.map(func, iterable, *args, **kwargs)
+    progress(futures)
+    results = client.gather(futures)
+    return cast("list[R]", results)

protein_quest/pdbe/__init__.py

@@ -0,0 +1 @@
+"""Modules related to PDBe (Protein Data Bank in Europe)."""

protein_quest/pdbe/fetch.py

@@ -0,0 +1,68 @@
+"""Module for fetching structures from PDBe."""
+
+from collections.abc import Iterable, Mapping
+from pathlib import Path
+
+from protein_quest.utils import Cacher, retrieve_files, run_async
+
+
+def _map_id_mmcif(pdb_id: str) -> tuple[str, str]:
+    """
+    Map PDB id to a download gzipped mmCIF url and file.
+
+    For example for PDB id "8WAS", the url will be
+    "https://www.ebi.ac.uk/pdbe/entry-files/download/8was.cif.gz" and the file will be "8was.cif.gz".
+
+    Args:
+        pdb_id: The PDB ID to map.
+
+    Returns:
+        A tuple containing the URL to download the mmCIF file and the filename.
+    """
+    fn = f"{pdb_id.lower()}.cif.gz"
+    # On PDBe you can sometimes download an updated mmCIF file,
+    # Current url is for the archive mmCIF file
+    # TODO check if archive is OK, or if we should try to download the updated file
+    # this will cause many more requests, so we should only do this if needed
+    url = f"https://www.ebi.ac.uk/pdbe/entry-files/download/{fn}"
+    return url, fn
+
+
+async def fetch(
+    ids: Iterable[str], save_dir: Path, max_parallel_downloads: int = 5, cacher: Cacher | None = None
+) -> Mapping[str, Path]:
+    """Fetches mmCIF files from the PDBe database.
+
+    Args:
+        ids: A set of PDB IDs to fetch.
+        save_dir: The directory to save the fetched mmCIF files to.
+        max_parallel_downloads: The maximum number of parallel downloads.
+        cacher: An optional cacher to use for caching downloaded files.
+
+    Returns:
+        A dict of id and paths to the downloaded mmCIF files.
+    """
+
+    # The future result, is in a different order than the input ids,
+    # so we need to map the ids to the urls and filenames.
+
+    id2urls = {pdb_id: _map_id_mmcif(pdb_id) for pdb_id in ids}
+    urls = list(id2urls.values())
+    id2paths = {pdb_id: save_dir / fn for pdb_id, (_, fn) in id2urls.items()}
+
+    await retrieve_files(urls, save_dir, max_parallel_downloads, desc="Downloading PDBe mmCIF files", cacher=cacher)
+    return id2paths
+
+
+def sync_fetch(ids: Iterable[str], save_dir: Path, max_parallel_downloads: int = 5) -> Mapping[str, Path]:
+    """Synchronously fetches mmCIF files from the PDBe database.
+
+    Args:
+        ids: A set of PDB IDs to fetch.
+        save_dir: The directory to save the fetched mmCIF files to.
+        max_parallel_downloads: The maximum number of parallel downloads.
+
+    Returns:
+        A dict of id and paths to the downloaded mmCIF files.
+    """
+    return run_async(fetch(ids, save_dir, max_parallel_downloads))

protein_quest/ss.py

@@ -0,0 +1,280 @@
+"""Module for dealing with secondary structure."""
+
+import logging
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from pathlib import Path
+
+from gemmi import Structure
+
+from protein_quest.converter import PositiveInt, Ratio, converter
+from protein_quest.io import read_structure
+
+logger = logging.getLogger(__name__)
+
+# TODO if a structure has no secondary structure information, calculate it with `gemmi ss`.
+# https://github.com/MonomerLibrary/monomers/wiki/Installation as --monomers dir
+# gemmi executable is in https://pypi.org/project/gemmi-program/
+# `gemmi ss` only prints secondary structure to stdout with `-v` flag.
+
+
+def nr_of_residues_in_total(structure: Structure) -> int:
+    """Count the total number of residues in the structure.
+
+    Args:
+        structure: The gemmi Structure object to analyze.
+
+    Returns:
+        The total number of residues in the structure.
+    """
+    count = 0
+    for model in structure:
+        for chain in model:
+            count += len(chain)
+    return count
+
+
+def nr_of_residues_in_helix(structure: Structure) -> int:
+    """Count the number of residues in alpha helices.
+
+    Requires structure to have secondary structure information.
+
+    Args:
+        structure: The gemmi Structure object to analyze.
+
+    Returns:
+        The number of residues in alpha helices.
+    """
+    # For cif files from AlphaFold the helix.length is set to -1
+    # so use resid instead
+    count = 0
+    for helix in structure.helices:
+        end = helix.end.res_id.seqid.num
+        start = helix.start.res_id.seqid.num
+        if end is None or start is None:
+            logger.warning(f"Invalid helix coordinates: {helix.end} or {helix.start}")
+            continue
+        length = end - start + 1
+        count += length
+    return count
+
+
+def nr_of_residues_in_sheet(structure: Structure) -> int:
+    """Count the number of residues in beta sheets.
+
+    Requires structure to have secondary structure information.
+
+    Args:
+        structure: The gemmi Structure object to analyze.
+
+    Returns:
+        The number of residues in beta sheets.
+    """
+    count = 0
+    for sheet in structure.sheets:
+        for strand in sheet.strands:
+            end = strand.end.res_id.seqid.num
+            start = strand.start.res_id.seqid.num
+            if end is None or start is None:
+                logger.warning(f"Invalid strand coordinates: {strand.end} or {strand.start}")
+                continue
+            length = end - start + 1
+            count += length
+    return count
+
+
+@dataclass
+class SecondaryStructureFilterQuery:
+    """Query object to filter on secondary structure.
+
+    Parameters:
+        abs_min_helix_residues: Minimum number of residues in helices (absolute).
+        abs_max_helix_residues: Maximum number of residues in helices (absolute).
+        abs_min_sheet_residues: Minimum number of residues in sheets (absolute).
+        abs_max_sheet_residues: Maximum number of residues in sheets (absolute).
+        ratio_min_helix_residues: Minimum number of residues in helices (relative).
+        ratio_max_helix_residues: Maximum number of residues in helices (relative).
+        ratio_min_sheet_residues: Minimum number of residues in sheets (relative).
+        ratio_max_sheet_residues: Maximum number of residues in sheets (relative).
+    """
+
+    abs_min_helix_residues: PositiveInt | None = None
+    abs_max_helix_residues: PositiveInt | None = None
+    abs_min_sheet_residues: PositiveInt | None = None
+    abs_max_sheet_residues: PositiveInt | None = None
+    ratio_min_helix_residues: Ratio | None = None
+    ratio_max_helix_residues: Ratio | None = None
+    ratio_min_sheet_residues: Ratio | None = None
+    ratio_max_sheet_residues: Ratio | None = None
+
+    def is_actionable(self) -> bool:
+        """Check if the secondary structure query has any actionable filters.
+
+        Returns:
+            True if any of the filters are set, False otherwise.
+        """
+        return any(
+            field is not None
+            for field in [
+                self.abs_min_helix_residues,
+                self.abs_max_helix_residues,
+                self.abs_min_sheet_residues,
+                self.abs_max_sheet_residues,
+                self.ratio_min_helix_residues,
+                self.ratio_max_helix_residues,
+                self.ratio_min_sheet_residues,
+                self.ratio_max_sheet_residues,
+            ]
+        )
+
+
+def _check_range(min_val, max_val, label):
+    if min_val is not None and max_val is not None and min_val >= max_val:
+        msg = f"Invalid {label} range: min {min_val} must be smaller than max {max_val}"
+        raise ValueError(msg)
+
+
+base_query_hook = converter.get_structure_hook(SecondaryStructureFilterQuery)
+
+
+@converter.register_structure_hook
+def secondary_structure_filter_query_hook(value, _type) -> SecondaryStructureFilterQuery:
+    result: SecondaryStructureFilterQuery = base_query_hook(value, _type)
+    _check_range(result.abs_min_helix_residues, result.abs_max_helix_residues, "absolute helix residue")
+    _check_range(result.abs_min_sheet_residues, result.abs_max_sheet_residues, "absolute sheet residue")
+    _check_range(result.ratio_min_helix_residues, result.ratio_max_helix_residues, "ratio helix residue")
+    _check_range(result.ratio_min_sheet_residues, result.ratio_max_sheet_residues, "ratio sheet residue")
+    return result
+
+
+@dataclass
+class SecondaryStructureStats:
+    """Statistics about the secondary structure of a protein.
+
+    Parameters:
+        nr_residues: Total number of residues in the structure.
+        nr_helix_residues: Number of residues in helices.
+        nr_sheet_residues: Number of residues in sheets.
+        helix_ratio: Ratio of residues in helices.
+        sheet_ratio: Ratio of residues in sheets.
+    """
+
+    nr_residues: PositiveInt
+    nr_helix_residues: PositiveInt
+    nr_sheet_residues: PositiveInt
+    helix_ratio: Ratio
+    sheet_ratio: Ratio
+
+
+@dataclass
+class SecondaryStructureFilterResult:
+    """Result of filtering on secondary structure.
+
+    Parameters:
+        stats: The secondary structure statistics.
+        passed: Whether the structure passed the filtering criteria.
+    """
+
+    stats: SecondaryStructureStats
+    passed: bool = False
+
+
+def _gather_stats(structure: Structure) -> SecondaryStructureStats:
+    nr_total_residues = nr_of_residues_in_total(structure)
+    nr_helix_residues = nr_of_residues_in_helix(structure)
+    nr_sheet_residues = nr_of_residues_in_sheet(structure)
+    if nr_total_residues == 0:
+        msg = "Structure has zero residues; cannot compute secondary structure ratios."
+        raise ValueError(msg)
+    helix_ratio = nr_helix_residues / nr_total_residues
+    sheet_ratio = nr_sheet_residues / nr_total_residues
+    return SecondaryStructureStats(
+        nr_residues=nr_total_residues,
+        nr_helix_residues=nr_helix_residues,
+        nr_sheet_residues=nr_sheet_residues,
+        helix_ratio=helix_ratio,
+        sheet_ratio=sheet_ratio,
+    )
+
+
+def filter_on_secondary_structure(
+    structure: Structure,
+    query: SecondaryStructureFilterQuery,
+) -> SecondaryStructureFilterResult:
+    """Filter a structure based on secondary structure criteria.
+
+    Args:
+        structure: The gemmi Structure object to analyze.
+        query: The filtering criteria to apply.
+
+    Returns:
+        Filtering statistics and whether structure passed.
+    """
+    stats = _gather_stats(structure)
+    conditions: list[bool] = []
+
+    # Helix absolute thresholds
+    if query.abs_min_helix_residues is not None:
+        conditions.append(stats.nr_helix_residues >= query.abs_min_helix_residues)
+    if query.abs_max_helix_residues is not None:
+        conditions.append(stats.nr_helix_residues <= query.abs_max_helix_residues)
+
+    # Helix ratio thresholds
+    if query.ratio_min_helix_residues is not None:
+        conditions.append(stats.helix_ratio >= query.ratio_min_helix_residues)
+    if query.ratio_max_helix_residues is not None:
+        conditions.append(stats.helix_ratio <= query.ratio_max_helix_residues)
+
+    # Sheet absolute thresholds
+    if query.abs_min_sheet_residues is not None:
+        conditions.append(stats.nr_sheet_residues >= query.abs_min_sheet_residues)
+    if query.abs_max_sheet_residues is not None:
+        conditions.append(stats.nr_sheet_residues <= query.abs_max_sheet_residues)
+
+    # Sheet ratio thresholds
+    if query.ratio_min_sheet_residues is not None:
+        conditions.append(stats.sheet_ratio >= query.ratio_min_sheet_residues)
+    if query.ratio_max_sheet_residues is not None:
+        conditions.append(stats.sheet_ratio <= query.ratio_max_sheet_residues)
+
+    if not conditions:
+        msg = "No filtering conditions provided. Please specify at least one condition."
+        raise ValueError(msg)
+    passed = all(conditions)
+    return SecondaryStructureFilterResult(stats=stats, passed=passed)
+
+
+def filter_file_on_secondary_structure(
+    file_path: Path,
+    query: SecondaryStructureFilterQuery,
+) -> SecondaryStructureFilterResult:
+    """Filter a structure file based on secondary structure criteria.
+
+    Args:
+        file_path: The path to the structure file to analyze.
+        query: The filtering criteria to apply.
+
+    Returns:
+        Filtering statistics and whether file passed.
+    """
+    structure = read_structure(file_path)
+    return filter_on_secondary_structure(structure, query)
+
+
+def filter_files_on_secondary_structure(
+    file_paths: Iterable[Path],
+    query: SecondaryStructureFilterQuery,
+) -> Generator[tuple[Path, SecondaryStructureFilterResult]]:
+    """Filter multiple structure files based on secondary structure criteria.
+
+    Args:
+        file_paths: A list of paths to the structure files to analyze.
+        query: The filtering criteria to apply.
+
+    Yields:
+        For each file returns the filtering statistics and whether structure passed.
+    """
+    # TODO check if quick enough in serial mode, if not switch to dask map
+    for file_path in file_paths:
+        result = filter_file_on_secondary_structure(file_path, query)
+        yield file_path, result

protein_quest/structure.py

@@ -0,0 +1,232 @@
+"""Module for querying and modifying [gemmi structures][gemmi.Structure]."""
+
+import logging
+from collections.abc import Iterable
+from datetime import UTC, datetime
+from pathlib import Path
+
+import gemmi
+
+from protein_quest.__version__ import __version__
+from protein_quest.io import read_structure, split_name_and_extension, write_structure
+from protein_quest.utils import CopyMethod, copyfile
+
+logger = logging.getLogger(__name__)
+
+
+def find_chain_in_model(model: gemmi.Model, wanted_chain: str) -> gemmi.Chain | None:
+    """Find a chain in a model.
+
+    Args:
+        model: The gemmi model to search in.
+        wanted_chain: The chain identifier to search for.
+
+    Returns:
+        The found chain or None if not found.
+    """
+    chain = model.find_chain(wanted_chain)
+    if chain is None:
+        # For chain A in 4v92 the find_chain method returns None,
+        # however it is prefixed with 'B',
+        # so we try again as last char of chain name
+        mchains = [c for c in model if c.name.endswith(wanted_chain)]
+        if mchains:
+            return mchains[0]
+    return chain
+
+
+def find_chain_in_structure(structure: gemmi.Structure, wanted_chain: str) -> gemmi.Chain | None:
+    """Find a chain in a structure.
+
+    Args:
+        structure: The gemmi structure to search in.
+        wanted_chain: The chain identifier to search for.
+
+    Returns:
+        The found chain or None if not found.
+    """
+    for model in structure:
+        chain = find_chain_in_model(model, wanted_chain)
+        if chain is not None:
+            return chain
+    return None
+
+
+def nr_residues_in_chain(file: Path, chain: str = "A") -> int:
+    """Returns the number of residues in a specific chain from a structure file.
+
+    Args:
+        file: Path to the input structure file.
+        chain: Chain to count residues of.
+
+    Returns:
+        The number of residues in the specified chain.
+    """
+    structure = read_structure(file)
+    gchain = find_chain_in_structure(structure, chain)
+    if gchain is None:
+        logger.warning("Chain %s not found in %s. Returning 0.", chain, file)
+        return 0
+    return len(gchain)
+
+
+def _dedup_helices(structure: gemmi.Structure):
+    helix_starts: set[str] = set()
+    duplicate_helix_indexes: list[int] = []
+    for hindex, helix in enumerate(structure.helices):
+        if str(helix.start) in helix_starts:
+            logger.debug(f"Duplicate start helix found: {hindex} {helix.start}, removing")
+            duplicate_helix_indexes.append(hindex)
+        else:
+            helix_starts.add(str(helix.start))
+    for helix_index in reversed(duplicate_helix_indexes):
+        structure.helices.pop(helix_index)
+
+
+def _dedup_sheets(structure: gemmi.Structure, chain2keep: str):
+    duplicate_sheet_indexes: list[int] = []
+    for sindex, sheet in enumerate(structure.sheets):
+        if sheet.name != chain2keep:
+            duplicate_sheet_indexes.append(sindex)
+    for sheet_index in reversed(duplicate_sheet_indexes):
+        structure.sheets.pop(sheet_index)
+
+
+def _add_provenance_info(structure: gemmi.Structure, chain2keep: str, out_chain: str):
+    old_id = structure.name
+    new_id = structure.name + f"{chain2keep}2{out_chain}"
+    structure.name = new_id
+    structure.info["_entry.id"] = new_id
+    new_title = f"From {old_id} chain {chain2keep} to {out_chain}"
+    structure.info["_struct.title"] = new_title
+    structure.info["_struct_keywords.pdbx_keywords"] = new_title.upper()
+    new_si = gemmi.SoftwareItem()
+    new_si.classification = gemmi.SoftwareItem.Classification.DataExtraction
+    new_si.name = "protein-quest.pdbe.io.write_single_chain_pdb_file"
+    new_si.version = str(__version__)
+    new_si.date = str(datetime.now(tz=UTC).date())
+    structure.meta.software = [*structure.meta.software, new_si]
+
+
+def chains_in_structure(structure: gemmi.Structure) -> set[gemmi.Chain]:
+    """Get a list of chains in a structure.
+
+    Args:
+        structure: The gemmi structure to get chains from.
+
+    Returns:
+        A set of chains in the structure.
+    """
+    return {c for model in structure for c in model}
+
+
+class ChainNotFoundError(IndexError):
+    """Exception raised when a chain is not found in a structure."""
+
+    def __init__(self, chain: str, file: Path | str, available_chains: Iterable[str]):
+        super().__init__(f"Chain {chain} not found in {file}. Available chains are: {available_chains}")
+        self.chain_id = chain
+        self.file = file
+
+
+def write_single_chain_structure_file(
+    input_file: Path,
+    chain2keep: str,
+    output_dir: Path,
+    out_chain: str = "A",
+    copy_method: CopyMethod = "copy",
+) -> Path:
+    """Write a single chain from a structure file to a new structure file.
+
+    Also
+
+    - removes ligands and waters
+    - renumbers atoms ids
+    - removes chem_comp section from cif files
+    - adds provenance information to the header like software and input file+chain
+
+    This function is equivalent to the following gemmi commands:
+
+    ```shell
+    gemmi convert --remove-lig-wat --select=B --to=cif chain-in/3JRS.cif - | \\
+    gemmi convert --from=cif --rename-chain=B:A - chain-out/3JRS_B2A.gemmi.cif
+    ```
+
+    Args:
+        input_file: Path to the input structure file.
+        chain2keep: The chain to keep.
+        output_dir: Directory to save the output file.
+        out_chain: The chain identifier for the output file.
+        copy_method: How to copy when no changes are needed to output file.
+
+    Returns:
+        Path to the output structure file
+
+    Raises:
+        FileNotFoundError: If the input file does not exist.
+        ChainNotFoundError: If the specified chain is not found in the input file.
+    """
+
+    logger.debug(f"chain2keep: {chain2keep}, out_chain: {out_chain}")
+    structure = read_structure(input_file)
+    structure.setup_entities()
+
+    chain = find_chain_in_structure(structure, chain2keep)
+    chainnames_in_structure = {c.name for c in chains_in_structure(structure)}
+    if chain is None:
+        raise ChainNotFoundError(chain2keep, input_file, chainnames_in_structure)
+    chain_name = chain.name
+    name, extension = split_name_and_extension(input_file.name)
+    output_file = output_dir / f"{name}_{chain_name}2{out_chain}{extension}"
+
+    if output_file.exists():
+        logger.info("Output file %s already exists for input file %s. Skipping.", output_file, input_file)
+        return output_file
+
+    if chain_name == out_chain and len(chainnames_in_structure) == 1:
+        logger.info(
+            "%s only has chain %s and out_chain is also %s. Copying file to %s.",
+            input_file,
+            chain_name,
+            out_chain,
+            output_file,
+        )
+        copyfile(input_file, output_file, copy_method)
+        return output_file
+
+    gemmi.Selection(chain_name).remove_not_selected(structure)
+    for m in structure:
+        m.remove_ligands_and_waters()
+    structure.setup_entities()
+    structure.rename_chain(chain_name, out_chain)
+    _dedup_helices(structure)
+    _dedup_sheets(structure, out_chain)
+    _add_provenance_info(structure, chain_name, out_chain)
+
+    write_structure(structure, output_file)
+
+    return output_file
+
+
+def structure2uniprot_accessions(structure: gemmi.Structure) -> set[str]:
+    """Extract UniProt accessions from a gemmi Structure object.
+
+    Logs a warning and returns an empty set if no accessions are found in structure.
+
+    Args:
+        structure: The gemmi Structure object to extract UniProt accessions from.
+
+    Returns:
+        A set of UniProt accessions found in the structure.
+    """
+    block = structure.make_mmcif_block(gemmi.MmcifOutputGroups(False, struct_ref=True))
+    struct_ref = block.get_mmcif_category("_struct_ref.")
+    uniprot_accessions: set[str] = set()
+    for i, db_name in enumerate(struct_ref["db_name"]):
+        if db_name != "UNP":
+            continue
+        pdbx_db_accession = struct_ref["pdbx_db_accession"][i]
+        uniprot_accessions.add(pdbx_db_accession)
+    if not uniprot_accessions:
+        logger.warning("No UniProt accessions found in structure %s", structure.name)
+    return uniprot_accessions