protein-quest 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

protein_quest/parallel.py

@@ -2,8 +2,10 @@
 
 import logging
 import os
+from collections.abc import Callable, Collection
+from typing import Concatenate, ParamSpec, cast
 
-from dask.distributed import LocalCluster
+from dask.distributed import Client, LocalCluster, progress
 from distributed.deploy.cluster import Cluster
 from psutil import cpu_count
 
@@ -66,3 +68,37 @@ def _configure_cpu_dask_scheduler(nproc: int, name: str) -> LocalCluster:
     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/fetch.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Mapping
 from pathlib import Path
 
-from protein_quest.utils import retrieve_files
+from protein_quest.utils import retrieve_files, run_async
 
 
 def _map_id_mmcif(pdb_id: str) -> tuple[str, str]:
@@ -49,3 +49,17 @@ async def fetch(ids: Iterable[str], save_dir: Path, max_parallel_downloads: int
 
     await retrieve_files(urls, save_dir, max_parallel_downloads, desc="Downloading PDBe mmCIF files")
     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/pdbe/io.py

@@ -2,15 +2,22 @@
 
 import gzip
 import logging
-from collections.abc import Generator
+from collections.abc import Generator, Iterable
+from datetime import UTC, datetime
 from pathlib import Path
 
 import gemmi
 
-from protein_quest import __version__
+from protein_quest.__version__ import __version__
+from protein_quest.utils import CopyMethod, copyfile
 
 logger = logging.getLogger(__name__)
 
+# TODO remove once v0.7.4 of gemmi is released,
+# as uv pip install git+https://github.com/project-gemmi/gemmi.git installs 0.7.4.dev0 which does not print leaks
+# Swallow gemmi leaked function warnings
+gemmi.set_leak_warnings(False)
+
 
 def nr_residues_in_chain(file: Path | str, chain: str = "A") -> int:
     """Returns the number of residues in a specific chain from a mmCIF/pdb file.
@@ -23,14 +30,21 @@ def nr_residues_in_chain(file: Path | str, chain: str = "A") -> int:
         The number of residues in the specified chain.
     """
     structure = gemmi.read_structure(str(file))
-    model = structure[0]
-    gchain = find_chain_in_model(model, chain)
+    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 find_chain_in_structure(structure: gemmi.Structure, wanted_chain: str) -> gemmi.Chain | None:
+    for model in structure:
+        chain = find_chain_in_model(model, wanted_chain)
+        if chain is not None:
+            return chain
+    return None
+
+
 def find_chain_in_model(model: gemmi.Model, wanted_chain: str) -> gemmi.Chain | None:
     chain = model.find_chain(wanted_chain)
     if chain is None:
@@ -63,10 +77,12 @@ def write_structure(structure: gemmi.Structure, path: Path):
         with gzip.open(path, "wt") as f:
             f.write(body)
     elif path.name.endswith(".cif"):
-        doc = structure.make_mmcif_document()
+        # do not write chem_comp so it is viewable by molstar
+        # see https://github.com/project-gemmi/gemmi/discussions/362
+        doc = structure.make_mmcif_document(gemmi.MmcifOutputGroups(True, chem_comp=False))
         doc.write_file(str(path))
     elif path.name.endswith(".cif.gz"):
-        doc = structure.make_mmcif_document()
+        doc = structure.make_mmcif_document(gemmi.MmcifOutputGroups(True, chem_comp=False))
         cif_str = doc.as_string()
         with gzip.open(path, "wt") as f:
             f.write(cif_str)
@@ -106,14 +122,17 @@ def locate_structure_file(root: Path, pdb_id: str) -> Path:
     Raises:
         FileNotFoundError: If no structure file is found for the given PDB ID.
     """
-    exts = [".cif.gz", ".cif", ".pdb.gz", ".pdb"]
-    # files downloaded from https://www.ebi.ac.uk/pdbe/ website
-    # have file names like pdb6t5y.ent or pdb6t5y.ent.gz for a PDB formatted file.
-    # TODO support pdb6t5y.ent or pdb6t5y.ent.gz file names
+    exts = [".cif.gz", ".cif", ".pdb.gz", ".pdb", ".ent", ".ent.gz"]
     for ext in exts:
-        candidate = root / f"{pdb_id.lower()}{ext}"
-        if candidate.exists():
-            return candidate
+        candidates = (
+            root / f"{pdb_id}{ext}",
+            root / f"{pdb_id.lower()}{ext}",
+            root / f"{pdb_id.upper()}{ext}",
+            root / f"pdb{pdb_id.lower()}{ext}",
+        )
+        for candidate in candidates:
+            if candidate.exists():
+                return candidate
     msg = f"No structure file found for {pdb_id} in {root}"
     raise FileNotFoundError(msg)
 
@@ -131,55 +150,132 @@ def glob_structure_files(input_dir: Path) -> Generator[Path]:
         yield from input_dir.glob(f"*{ext}")
 
 
+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 _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."""
+    return {c for model in structure for c in model}
+
+
 def write_single_chain_pdb_file(
-    input_file: Path, chain2keep: str, output_dir: Path, out_chain: str = "A"
-) -> Path | None:
+    input_file: Path,
+    chain2keep: str,
+    output_dir: Path,
+    out_chain: str = "A",
+    copy_method: CopyMethod = "copy",
+) -> Path:
     """Write a single chain from a mmCIF/pdb file to a new mmCIF/pdb 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 mmCIF/pdb 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 mmCIF/pdb file or None if not created.
+        Path to the output mmCIF/pdb 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 = gemmi.read_structure(str(input_file))
-    model = structure[0]
+    structure.setup_entities()
 
-    # Only count residues of polymer
-    model.remove_ligands_and_waters()
-
-    chain = find_chain_in_model(model, chain2keep)
+    chain = find_chain_in_structure(structure, chain2keep)
+    chainnames_in_structure = {c.name for c in chains_in_structure(structure)}
     if chain is None:
-        logger.warning(
-            "Chain %s not found in %s. Skipping.",
-            chain2keep,
-            input_file,
-        )
-        return 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}"
+    output_file = output_dir / f"{name}_{chain_name}2{out_chain}{extension}"
 
-    new_structure = gemmi.Structure()
-    new_structure.resolution = structure.resolution
-    new_id = structure.name + f"{chain2keep}2{out_chain}"
-    new_structure.name = new_id
-    new_structure.info["_entry.id"] = new_id
-    new_title = f"From {structure.info['_entry.id']} chain {chain2keep} to {out_chain}"
-    new_structure.info["_struct.title"] = new_title
-    new_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"
-    new_si.version = str(__version__)
-    new_structure.meta.software.append(new_si)
-    new_model = gemmi.Model(1)
-    chain.name = out_chain
-    new_model.add_chain(chain)
-    new_structure.add_model(new_model)
-    write_structure(new_structure, output_file)
+    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

protein_quest/ss.py

@@ -0,0 +1,264 @@
+"""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, read_structure, set_leak_warnings
+
+from protein_quest.converter import PositiveInt, Ratio, converter
+
+logger = logging.getLogger(__name__)
+
+# TODO remove once v0.7.4 of gemmi is released,
+# as uv pip install git+https://github.com/project-gemmi/gemmi.git installs 0.7.4.dev0 which does not print leaks
+# Swallow gemmi leaked function warnings
+set_leak_warnings(False)
+
+# 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 _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(str(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/taxonomy.py

@@ -9,9 +9,9 @@ from typing import Literal, get_args
 from aiohttp.client import ClientResponse
 from aiohttp_retry import RetryClient
 from cattrs.gen import make_dict_structure_fn, override
-from cattrs.preconf.orjson import make_converter
 from yarl import URL
 
+from protein_quest.converter import converter
 from protein_quest.go import TextIOWrapper
 from protein_quest.utils import friendly_session
 
@@ -20,6 +20,16 @@ logger = logging.getLogger(__name__)
 
 @dataclass(frozen=True, slots=True)
 class Taxon:
+    """Dataclass representing a taxon.
+
+    Arguments:
+        taxon_id: The unique identifier for the taxon.
+        scientific_name: The scientific name of the taxon.
+        rank: The taxonomic rank of the taxon (e.g., species, genus).
+        common_name: The common name of the taxon (if available).
+        other_names: A set of other names for the taxon (if available).
+    """
+
     taxon_id: str
     scientific_name: str
     rank: str
@@ -32,8 +42,6 @@ class SearchTaxonResponse:
     results: list[Taxon]
 
 
-converter = make_converter()
-
 converter.register_structure_hook(
     Taxon,
     make_dict_structure_fn(
@@ -47,7 +55,9 @@ converter.register_structure_hook(
 )
 
 SearchField = Literal["tax_id", "scientific", "common", "parent"]
+"""Type of search field"""
 search_fields: set[SearchField | None] = set(get_args(SearchField)) | {None}
+"""Set of valid search fields"""
 
 
 def _get_next_page(response: ClientResponse) -> URL | str | None: