protein-quest 0.8.0__tar.gz → 0.9.0__tar.gz

{protein_quest-0.8.0 → protein_quest-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: protein_quest
-Version: 0.8.0
+Version: 0.9.0
 Summary: Search/retrieve/filter proteins and protein structures
 Project-URL: Homepage, https://github.com/haddocking/protein-quest
 Project-URL: Issues, https://github.com/haddocking/protein-quest/issues

{protein_quest-0.8.0 → protein_quest-0.9.0}/src/protein_quest/__version__.py

@@ -1,2 +1,2 @@
-__version__ = "0.8.0"
+__version__ = "0.9.0"
 """The version of the package."""

{protein_quest-0.8.0 → protein_quest-0.9.0}/src/protein_quest/alphafold/confidence.py

@@ -4,11 +4,16 @@ import logging
 from collections.abc import Generator
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal
 
 import gemmi
+from dask.distributed import Client
+from distributed.deploy.cluster import Cluster
+from tqdm.auto import tqdm
 
 from protein_quest.converter import Percentage, PositiveInt, converter
 from protein_quest.io import read_structure, write_structure
+from protein_quest.parallel import configure_dask_scheduler, dask_map_with_progress
 from protein_quest.ss import nr_of_residues_in_total
 from protein_quest.utils import CopyMethod, copyfile
 
@@ -112,7 +117,7 @@ class ConfidenceFilterResult:
     filtered_file: Path | None = None
 
 
-def filter_file_on_residues(
+def filter_file_on_confidence(
     file: Path, query: ConfidenceFilterQuery, filtered_dir: Path, copy_method: CopyMethod = "copy"
 ) -> ConfidenceFilterResult:
     """Filter a single AlphaFoldDB structure file (*.pdb[.gz], *.cif[.gz]) based on confidence.
@@ -155,12 +160,31 @@ def filter_file_on_residues(
     )
 
 
+def _filter_files_on_confidence_sequentially(
+    alphafold_pdb_files: list[Path],
+    query: ConfidenceFilterQuery,
+    filtered_dir: Path,
+    copy_method: CopyMethod = "copy",
+) -> list[ConfidenceFilterResult]:
+    results = []
+    for file in tqdm(
+        alphafold_pdb_files,
+        total=len(alphafold_pdb_files),
+        desc="Filtering on confidence",
+        unit="file",
+    ):
+        result = filter_file_on_confidence(file, query, filtered_dir, copy_method)
+        results.append(result)
+    return results
+
+
 def filter_files_on_confidence(
     alphafold_pdb_files: list[Path],
     query: ConfidenceFilterQuery,
     filtered_dir: Path,
     copy_method: CopyMethod = "copy",
-) -> Generator[ConfidenceFilterResult]:
+    scheduler_address: str | Cluster | Literal["sequential"] | None = None,
+) -> list[ConfidenceFilterResult]:
     """Filter AlphaFoldDB structures based on confidence.
 
     Args:
@@ -168,13 +192,35 @@ def filter_files_on_confidence(
         query: The confidence filter query containing the confidence thresholds.
         filtered_dir: Directory where the filtered mmcif/PDB files will be saved.
         copy_method: How to copy when a direct copy is possible.
+        scheduler_address: The address of the Dask scheduler.
+            If not provided, will create a local cluster.
+            If set to `sequential` will run tasks sequentially.
 
-    Yields:
-        For each mmcif/PDB files yields whether it was filtered or not,
+    Returns:
+        For each mmcif/PDB files returns whether it was filtered or not,
             and number of residues with pLDDT above the confidence threshold.
     """
-    # Note on why code looks duplicated:
-    # In ../filter.py:filter_files_on_residues() we filter on number of residues on a file level
-    # here we filter on file level and inside file remove low confidence residues
-    for pdb_file in alphafold_pdb_files:
-        yield filter_file_on_residues(pdb_file, query, filtered_dir, copy_method)
+    filtered_dir.mkdir(parents=True, exist_ok=True)
+    if scheduler_address == "sequential":
+        return _filter_files_on_confidence_sequentially(
+            alphafold_pdb_files,
+            query,
+            filtered_dir,
+            copy_method=copy_method,
+        )
+
+    scheduler_address = configure_dask_scheduler(
+        scheduler_address,
+        name="filter-confidence",
+    )
+
+    with Client(scheduler_address) as client:
+        client.forward_logging()
+        return dask_map_with_progress(
+            client,
+            filter_file_on_confidence,
+            alphafold_pdb_files,
+            query=query,
+            filtered_dir=filtered_dir,
+            copy_method=copy_method,
+        )

{protein_quest-0.8.0 → protein_quest-0.9.0}/src/protein_quest/cli.py

@@ -508,6 +508,15 @@ def _add_retrieve_emdb_parser(subparsers: argparse._SubParsersAction):
     _add_cacher_arguments(parser)
 
 
+def _add_scheduler_address_argument(parser):
+    parser.add_argument(
+        "--scheduler-address",
+        help=dedent("""Address of the Dask scheduler to connect to.
+            If not provided, will create a local cluster.
+            If set to `sequential` will run tasks sequentially."""),
+    )
+
+
 def _add_filter_confidence_parser(subparsers: argparse._SubParsersAction):
     """Add filter confidence subcommand parser."""
     parser = subparsers.add_parser(
@@ -542,6 +551,7 @@ def _add_filter_confidence_parser(subparsers: argparse._SubParsersAction):
             In CSV format with `<input_file>,<residue_count>,<passed>,<output_file>` columns.
             Use `-` for stdout."""),
     ).complete = shtab.FILE
+    _add_scheduler_address_argument(parser)
     _add_copy_method_arguments(parser)
 
 
@@ -576,12 +586,7 @@ def _add_filter_chain_parser(subparsers: argparse._SubParsersAction):
         help=dedent("""\
         Directory to write the single-chain PDB/mmCIF files. Output files are in same format as input files."""),
     ).complete = shtab.DIRECTORY
-    parser.add_argument(
-        "--scheduler-address",
-        help=dedent("""Address of the Dask scheduler to connect to.
-            If not provided, will create a local cluster.
-            If set to `sequential` will run tasks sequentially."""),
-    )
+    _add_scheduler_address_argument(parser)
     _add_copy_method_arguments(parser)
 
 
@@ -1060,6 +1065,7 @@ def _handle_filter_confidence(args: argparse.Namespace):
     max_residues = args.max_residues
     stats_file: TextIOWrapper | None = args.write_stats
     copy_method: CopyMethod = structure(args.copy_method, CopyMethod)  # pyright: ignore[reportArgumentType]
+    scheduler_address = structure(args.scheduler_address, str | None)  # pyright: ignore[reportArgumentType]
 
     output_dir.mkdir(parents=True, exist_ok=True)
     input_files = sorted(glob_structure_files(input_dir))
@@ -1078,16 +1084,14 @@ def _handle_filter_confidence(args: argparse.Namespace):
         writer.writerow(["input_file", "residue_count", "passed", "output_file"])
 
     passed_count = 0
-    for r in tqdm(
-        filter_files_on_confidence(input_files, query, output_dir, copy_method=copy_method),
-        total=len(input_files),
-        unit="file",
-    ):
+    results = filter_files_on_confidence(
+        input_files, query, output_dir, copy_method=copy_method, scheduler_address=scheduler_address
+    )
+    for r in results:
         if r.filtered_file:
             passed_count += 1
         if stats_file:
             writer.writerow([r.input_file, r.count, r.filtered_file is not None, r.filtered_file])  # pyright: ignore[reportPossiblyUnboundVariable]
-
     rprint(f"Filtered {passed_count} mmcif/PDB files by confidence, written to {output_dir} directory")
     if stats_file:
         rprint(f"Statistics written to {_name_of(stats_file)}")

{protein_quest-0.8.0 → protein_quest-0.9.0}/src/protein_quest/filters.py

@@ -48,6 +48,24 @@ def filter_file_on_chain(
         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,
@@ -72,11 +90,9 @@ def filter_files_on_chain(
     """
     output_dir.mkdir(parents=True, exist_ok=True)
     if scheduler_address == "sequential":
-
-        def task(file_and_chain: tuple[Path, str]) -> ChainFilterStatistics:
-            return filter_file_on_chain(file_and_chain, output_dir, out_chain=out_chain, copy_method=copy_method)
-
-        return list(map(task, file2chains))
+        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

{protein_quest-0.8.0 → protein_quest-0.9.0}/src/protein_quest/mcp_server.py

@@ -40,7 +40,7 @@ from typing import Annotated
 from fastmcp import FastMCP
 from pydantic import Field
 
-from protein_quest.alphafold.confidence import ConfidenceFilterQuery, ConfidenceFilterResult, filter_file_on_residues
+from protein_quest.alphafold.confidence import ConfidenceFilterQuery, ConfidenceFilterResult, filter_file_on_confidence
 from protein_quest.alphafold.fetch import AlphaFoldEntry, DownloadableFormat
 from protein_quest.alphafold.fetch import fetch_many as alphafold_fetch
 from protein_quest.emdb import fetch as emdb_fetch
@@ -199,7 +199,7 @@ def alphafold_confidence_filter(file: Path, query: ConfidenceFilterQuery, filter
 
     If passes filter writes file to filtered_dir with residues above confidence threshold.
     """
-    return filter_file_on_residues(file, query, filtered_dir)
+    return filter_file_on_confidence(file, query, filtered_dir)
 
 
 mcp.tool(filter_file_on_secondary_structure)

{protein_quest-0.8.0 → protein_quest-0.9.0}/tests/alphafold/test_confidence.py

@@ -49,7 +49,7 @@ def test_filter_files_on_confidence(sample_pdb_file: Path, tmp_path: Path):
         min_residues=10,
     )
 
-    results = list(filter_files_on_confidence(input_files, query, tmp_path))
+    results = filter_files_on_confidence(input_files, query, tmp_path)
 
     expected = [
         ConfidenceFilterResult(