protein-quest 0.9.0__py3-none-any.whl

protein_quest/converter.py

@@ -0,0 +1,46 @@
+"""Convert json or dict to Python objects."""
+
+from cattrs.preconf.orjson import make_converter
+from yarl import URL
+
+type Percentage = float
+"""Type alias for percentage values (0.0-100.0)."""
+type Ratio = float
+"""Type alias for ratio values (0.0-1.0)."""
+type PositiveInt = int
+"""Type alias for positive integer values (>= 0)."""
+
+converter = make_converter()
+"""cattrs converter to read JSON document or dict to Python objects."""
+converter.register_structure_hook(URL, lambda v, _: URL(v))
+converter.register_unstructure_hook(URL, lambda u: str(u))
+
+
+@converter.register_structure_hook
+def percentage_hook(val, _) -> Percentage:
+    value = float(val)
+    """Cattrs hook to validate percentage values."""
+    if not 0.0 <= value <= 100.0:
+        msg = f"Value {value} is not a valid percentage (0.0-100.0)"
+        raise ValueError(msg)
+    return value
+
+
+@converter.register_structure_hook
+def ratio_hook(val, _) -> Ratio:
+    """Cattrs hook to validate ratio values."""
+    value = float(val)
+    if not 0.0 <= value <= 1.0:
+        msg = f"Value {value} is not a valid ratio (0.0-1.0)"
+        raise ValueError(msg)
+    return value
+
+
+@converter.register_structure_hook
+def positive_int_hook(val, _) -> PositiveInt:
+    """Cattrs hook to validate positive integer values."""
+    value = int(val)
+    if value < 0:
+        msg = f"Value {value} is not a valid positive integer (>= 0)"
+        raise ValueError(msg)
+    return value

protein_quest/emdb.py

@@ -0,0 +1,37 @@
+"""Module dealing with Electron Microscopy Data Bank (EMDB)."""
+
+from collections.abc import Iterable, Mapping
+from pathlib import Path
+
+from protein_quest.utils import Cacher, retrieve_files
+
+
+def _map_id2volume_url(emdb_id: str) -> tuple[str, str]:
+    # https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-19583/map/emd_19583.map.gz
+    fn = emdb_id.lower().replace("emd-", "emd_") + ".map.gz"
+    url = f"https://ftp.ebi.ac.uk/pub/databases/emdb/structures/{emdb_id}/map/{fn}"
+    return url, fn
+
+
+async def fetch(
+    emdb_ids: Iterable[str], save_dir: Path, max_parallel_downloads: int = 1, cacher: Cacher | None = None
+) -> Mapping[str, Path]:
+    """Fetches volume files from the EMDB database.
+
+    Args:
+        emdb_ids: A list of EMDB IDs to fetch.
+        save_dir: The directory to save the downloaded files.
+        max_parallel_downloads: The maximum number of parallel downloads.
+        cacher: An optional cacher to use for caching downloaded files.
+
+    Returns:
+        A mapping of EMDB IDs to their downloaded files.
+    """
+    id2urls = {emdb_id: _map_id2volume_url(emdb_id) for emdb_id in emdb_ids}
+    urls = list(id2urls.values())
+    id2paths = {emdb_id: save_dir / fn for emdb_id, (_, fn) in id2urls.items()}
+
+    # TODO show progress of each item
+    # TODO handle failed downloads, by skipping them instead of raising an error
+    await retrieve_files(urls, save_dir, max_parallel_downloads, desc="Downloading EMDB volume files", cacher=cacher)
+    return id2paths

protein_quest/filters.py

@@ -0,0 +1,163 @@
+"""Module for filtering structure files and their contents."""
+
+import logging
+from collections.abc import Collection, Generator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from dask.distributed import Client
+from distributed.deploy.cluster import Cluster
+from tqdm.auto import tqdm
+
+from protein_quest.parallel import configure_dask_scheduler, dask_map_with_progress
+from protein_quest.structure import nr_residues_in_chain, write_single_chain_structure_file
+from protein_quest.utils import CopyMethod, copyfile
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChainFilterStatistics:
+    input_file: Path
+    chain_id: str
+    passed: bool = False
+    output_file: Path | None = None
+    discard_reason: Exception | None = None
+
+
+def filter_file_on_chain(
+    file_and_chain: tuple[Path, str],
+    output_dir: Path,
+    out_chain: str = "A",
+    copy_method: CopyMethod = "copy",
+) -> ChainFilterStatistics:
+    input_file, chain_id = file_and_chain
+    logger.debug("Filtering %s on chain %s", input_file, chain_id)
+    try:
+        output_file = write_single_chain_structure_file(
+            input_file, chain_id, output_dir, out_chain=out_chain, copy_method=copy_method
+        )
+        return ChainFilterStatistics(
+            input_file=input_file,
+            chain_id=chain_id,
+            output_file=output_file,
+            passed=True,
+        )
+    except Exception as e:  # noqa: BLE001 - error is handled downstream
+        return ChainFilterStatistics(input_file=input_file, chain_id=chain_id, discard_reason=e)
+
+
+def _filter_files_on_chain_sequentially(
+    file2chains: Collection[tuple[Path, str]],
+    output_dir: Path,
+    out_chain: str = "A",
+    copy_method: CopyMethod = "copy",
+) -> list[ChainFilterStatistics]:
+    results = []
+    for file_and_chain in tqdm(file2chains, unit="file"):
+        result = filter_file_on_chain(
+            file_and_chain,
+            output_dir=output_dir,
+            out_chain=out_chain,
+            copy_method=copy_method,
+        )
+        results.append(result)
+    return results
+
+
+def filter_files_on_chain(
+    file2chains: Collection[tuple[Path, str]],
+    output_dir: Path,
+    out_chain: str = "A",
+    scheduler_address: str | Cluster | Literal["sequential"] | None = None,
+    copy_method: CopyMethod = "copy",
+) -> list[ChainFilterStatistics]:
+    """Filter mmcif/PDB files by chain.
+
+    Args:
+        file2chains: Which chain to keep for each PDB file.
+            First item is the PDB file path, second item is the chain ID.
+        output_dir: The directory where the filtered files will be written.
+        out_chain: Under what name to write the kept chain.
+        scheduler_address: The address of the Dask scheduler.
+            If not provided, will create a local cluster.
+            If set to `sequential` will run tasks sequentially.
+        copy_method: How to copy when a direct copy is possible.
+
+    Returns:
+        Result of the filtering process.
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+    if scheduler_address == "sequential":
+        return _filter_files_on_chain_sequentially(
+            file2chains, output_dir, out_chain=out_chain, copy_method=copy_method
+        )
+
+    # TODO make logger.debug in filter_file_on_chain show to user when --log
+    # GPT-5 generated a fairly difficult setup with a WorkerPlugin, need to find a simpler approach
+    scheduler_address = configure_dask_scheduler(
+        scheduler_address,
+        name="filter-chain",
+    )
+
+    with Client(scheduler_address) as client:
+        client.forward_logging()
+        return dask_map_with_progress(
+            client,
+            filter_file_on_chain,
+            file2chains,
+            output_dir=output_dir,
+            out_chain=out_chain,
+            copy_method=copy_method,
+        )
+
+
+@dataclass
+class ResidueFilterStatistics:
+    """Statistics for filtering files based on residue count in a specific chain.
+
+    Parameters:
+        input_file: The path to the input file.
+        residue_count: The number of residues.
+        passed: Whether the file passed the filtering criteria.
+        output_file: The path to the output file, if passed.
+    """
+
+    input_file: Path
+    residue_count: int
+    passed: bool
+    output_file: Path | None
+
+
+def filter_files_on_residues(
+    input_files: list[Path],
+    output_dir: Path,
+    min_residues: int,
+    max_residues: int,
+    chain: str = "A",
+    copy_method: CopyMethod = "copy",
+) -> Generator[ResidueFilterStatistics]:
+    """Filter PDB/mmCIF files by number of residues in given chain.
+
+    Args:
+        input_files: The list of input PDB/mmCIF files.
+        output_dir: The directory where the filtered files will be written.
+        min_residues: The minimum number of residues in chain.
+        max_residues: The maximum number of residues in chain.
+        chain: The chain to count residues of.
+        copy_method: How to copy passed files to output directory:
+
+    Yields:
+        Objects containing information about the filtering process for each input file.
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for input_file in tqdm(input_files, unit="file"):
+        residue_count = nr_residues_in_chain(input_file, chain=chain)
+        passed = min_residues <= residue_count <= max_residues
+        if passed:
+            output_file = output_dir / input_file.name
+            copyfile(input_file, output_file, copy_method)
+            yield ResidueFilterStatistics(input_file, residue_count, True, output_file)
+        else:
+            yield ResidueFilterStatistics(input_file, residue_count, False, None)

protein_quest/go.py

@@ -0,0 +1,165 @@
+"""Module for Gene Ontology (GO) functions."""
+
+import csv
+import logging
+from collections.abc import Generator
+from dataclasses import dataclass
+from io import TextIOWrapper
+from typing import Literal, get_args
+
+from cattrs.gen import make_dict_structure_fn, override
+
+from protein_quest.converter import converter
+from protein_quest.utils import friendly_session
+
+logger = logging.getLogger(__name__)
+
+Aspect = Literal["cellular_component", "biological_process", "molecular_function"]
+"""The aspect of the GO term."""
+allowed_aspects = set(get_args(Aspect))
+"""Allowed aspects for GO terms."""
+
+
+@dataclass(frozen=True, slots=True)
+class GoTerm:
+    """A Gene Ontology (GO) term.
+
+    Parameters:
+        id: The unique identifier for the GO term, e.g., 'GO:0043293'.
+        is_obsolete: Whether the GO term is obsolete.
+        name: The name of the GO term.
+        definition: The definition of the GO term.
+        aspect: The aspect of the GO term.
+    """
+
+    id: str
+    is_obsolete: bool
+    name: str
+    definition: str
+    aspect: Aspect
+
+
+@dataclass(frozen=True, slots=True)
+class PageInfo:
+    current: int
+    total: int
+
+
+@dataclass(frozen=True, slots=True)
+class SearchResponse:
+    results: list[GoTerm]
+    number_of_hits: int
+    page_info: PageInfo
+
+
+def flatten_definition(definition, _context) -> str:
+    return definition["text"]
+
+
+# Use hook to convert incoming camelCase to snake_case
+# and to flatten definition {text} to text
+# see https://catt.rs/en/stable/customizing.html#rename
+converter.register_structure_hook(
+    GoTerm,
+    make_dict_structure_fn(
+        GoTerm,
+        converter,
+        is_obsolete=override(rename="isObsolete"),
+        definition=override(struct_hook=flatten_definition),
+    ),
+)
+converter.register_structure_hook(
+    SearchResponse,
+    make_dict_structure_fn(
+        SearchResponse, converter, number_of_hits=override(rename="numberOfHits"), page_info=override(rename="pageInfo")
+    ),
+)
+
+
+async def search_gene_ontology_term(
+    term: str, aspect: Aspect | None = None, include_obsolete: bool = False, limit: int = 100
+) -> list[GoTerm]:
+    """Search for a Gene Ontology (GO) term by its name or ID.
+
+    Calls the EBI QuickGO API at https://www.ebi.ac.uk/QuickGO/api/index.html .
+
+    Examples:
+        To search for `apoptosome` terms do.
+
+        >>> from protein_quest.go import search_go_term
+        >>> r = await search_go_term('apoptosome')
+        >>> len(r)
+        5
+        >>> r[0]
+        GoTerm(id='GO:0043293', is_obsolete=False, name='apoptosome', definition='A multisubunit protein ...')
+
+    Args:
+        term: The GO term to search for. For example `nucleus` or `GO:0006816`.
+        aspect: The aspect to filter by. If not given, all aspects are included.
+        include_obsolete: Whether to include obsolete terms. By default, obsolete terms are excluded.
+        limit: The maximum number of results to return.
+
+    Returns:
+        List of GO terms
+
+    Raises:
+        ValueError: If the aspect is invalid.
+    """
+    url = "https://www.ebi.ac.uk/QuickGO/services/ontology/go/search"
+    page_limit = 100
+    params = {"query": term, "limit": str(page_limit), "page": "1"}
+    if aspect is not None and aspect not in allowed_aspects:
+        msg = f"Invalid aspect: {aspect}. Allowed aspects are: {allowed_aspects} or None."
+        raise ValueError(msg)
+    logger.debug("Fetching GO terms from %s with params %s", url, params)
+    async with friendly_session() as session:
+        # Fetch first page to learn how many pages there are
+        async with session.get(url, params=params) as response:
+            response.raise_for_status()
+            raw_data = await response.read()
+            data = converter.loads(raw_data, SearchResponse)
+
+        terms = list(_filter_go_terms(data.results, aspect, include_obsolete))
+        if len(terms) >= limit:
+            # Do not fetch additional pages if we have enough results
+            return terms[:limit]
+        total_pages = data.page_info.total
+        logger.debug("GO search returned %s pages (current=%s)", total_pages, data.page_info.current)
+
+        # Retrieve remaining pages (if any) and extend results
+        if total_pages > 1:
+            for page in range(2, total_pages + 1):
+                params["page"] = str(page)
+                logger.debug("Fetching additional GO terms page %s/%s with params %s", page, total_pages, params)
+                async with session.get(url, params=params) as response:
+                    response.raise_for_status()
+                    raw_data = await response.read()
+                    data = converter.loads(raw_data, SearchResponse)
+                terms.extend(_filter_go_terms(data.results, aspect, include_obsolete))
+                if len(terms) >= limit:
+                    # Do not fetch additional pages if we have enough results
+                    break
+
+        return terms[:limit]
+
+
+def _filter_go_terms(terms: list[GoTerm], aspect: Aspect | None, include_obsolete: bool) -> Generator[GoTerm]:
+    for oboterm in terms:
+        if not include_obsolete and oboterm.is_obsolete:
+            continue
+        if aspect and oboterm.aspect != aspect:
+            continue
+        yield oboterm
+
+
+def write_go_terms_to_csv(terms: list[GoTerm], csv_file: TextIOWrapper) -> None:
+    """Write a list of GO terms to a CSV file.
+
+    Args:
+        terms: The list of GO terms to write.
+        csv_file: The CSV file to write to.
+    """
+    writer = csv.writer(csv_file)
+    writer.writerow(["id", "name", "aspect", "definition"])
+    for term in terms:
+        writer.writerow([term.id, term.name, term.aspect, term.definition])