Refactor into modular molecular workloads

2026-07-13 22:50:56 +03:00
parent df4bda6acb
commit 34beb8b0aa
19 changed files with 798 additions and 290 deletions
@@ -1,11 +1,12 @@
 .venv/
 __pycache__/
 *.py[cod]
 *.egg-info/
 # Local ChEMBL download (about 738 MB); obtain it separately before running.
 chembl_*.txt
-# Files generated by scimesh.py
+# Files generated by scimesh workloads
 *_similarities.csv
 test_results.csv
 test_structures/
@@ -1,56 +1,77 @@
 # SciMesh
-Minimal local search for ChEMBL molecules similar to gefitinib (`CHEMBL939`).
+SciMesh is a small local framework for scientific workloads on molecular datasets. It currently provides exact molecular similarity search and exact sparse similarity-graph construction. It runs in one local Python process: there is no network service, multiprocessing, coordinator, database, or dense similarity matrix.
-The script finds `CHEMBL939` in the TSV file and uses its `canonical_smiles` as the reference. It then makes a second streaming pass through the file, generates Morgan fingerprints (`radius=2`, `fpSize=2048`), and ranks the remaining valid SMILES by Tanimoto similarity. Invalid SMILES and `CHEMBL939` itself are skipped. Only the best 20 results (or the value passed to `--top`) are kept in memory.
+The ChEMBL TSV database is intentionally not included in this repository. Download it separately and pass its path to the commands below. The expected columns are `chembl_id` and `canonical_smiles`.
 ## Installation
 SciMesh requires Python 3.10+ and RDKit.
 ```bash
 python -m venv .venv
 source .venv/bin/activate
-pip install -r requirements.txt
+pip install -e .
 ```
-RDKit can also be installed through conda-forge:
+RDKit can alternatively be installed from conda-forge:
 ```bash
 conda install -c conda-forge rdkit
 pip install -e .
 ```
-## Usage
+## Similarity search
 `similarity-search` finds the top-k molecules most similar to a query. The query is supplied either by ChEMBL ID or by SMILES. It uses Morgan fingerprints with `radius=2` and `fpSize=2048`, Tanimoto similarity, streaming TSV reads, and a bounded heap. Invalid SMILES and the query molecule are skipped.
 ```bash
-python scimesh.py chembl_37_chemreps.txt -o gefitinib_similarities.csv
+scimesh similarity-search chembl_37_chemreps.txt \
  --query-id CHEMBL939 \
  --top-k 20 \
  --output results.csv
 ```
-By default, the script writes a CSV with `rank,chembl_id,canonical_smiles,similarity` columns and prints the same top 20 results to the terminal. To choose a different number of results:
+Use a SMILES query when it is not identified by ChEMBL ID:
 ```bash
-python scimesh.py chembl_37_chemreps.txt --top 50 -o top_50.csv
+scimesh similarity-search chembl_37_chemreps.txt \
  --query-smiles 'COc1cc2ncnc(Nc3ccc(F)c(Cl)c3)c2cc1OCCCN1CCOCC1' \
  --top-k 20 \
  --output results.csv
 ```
-During the search, status is written to `stderr` every 100,000 rows: number of processed rows, current and average rates, elapsed time, and the number of invalid SMILES skipped. The interval can be changed or disabled:
+The output CSV contains `rank,chembl_id,canonical_smiles,similarity`. Search progress and valid/invalid-SMILES statistics are written to the terminal. `--max-rows` limits the candidate scan for small tests, while `--progress-every 0` disables progress reports.
 To render the query and retained candidates:
 ```bash
-python scimesh.py chembl_37_chemreps.txt --progress-every 500000
+scimesh similarity-search chembl_37_chemreps.txt \
-python scimesh.py chembl_37_chemreps.txt --progress-every 0
+  --query-id CHEMBL939 \
  --images-dir structures
 ```
-## Quick test on part of the database
+This writes `query.png` and `top_candidates.png` into `structures`.
-The `--max-rows` option limits the second pass to the first `N` TSV rows. `CHEMBL939` is still found in its own streaming pass first, so the reference stays the same. The resulting CSV is the top 20 only within the processed subset, not the full database.
+## Similarity graph
 `similarity-graph` constructs an exact sparse undirected graph. Every valid molecule is a vertex; an edge is emitted only when Tanimoto similarity is at least `--threshold`. Each fingerprint is calculated once. Comparisons are processed block by block, each pair is tested once (`i < j`), and no dense N×N matrix is created or stored.
 ```bash
-python scimesh.py chembl_37_chemreps.txt --max-rows 10000 -o test_results.csv
+scimesh similarity-graph chembl_37_chemreps.txt \
  --max-rows 10000 \
  --threshold 0.7 \
  --block-size 1000 \
  --output similarity_graph.csv
 ```
-## Structure images
+The deterministic edge-list CSV has `source_id,target_id,similarity` columns. The command reports valid molecules, checked pairs, emitted edges, rate, and elapsed time. `--block-size` changes only how comparisons are grouped, not the result.
-Pass a directory to `--images-dir` to create `CHEMBL939_gefitinib.png` for gefitinib and `top_candidates.png` with a grid of top candidates. Candidate images show rank, ChEMBL ID, and Tanimoto similarity.
+## Development
 ```bash
-python scimesh.py chembl_37_chemreps.txt --images-dir structures
+pip install -e '.[dev]'
 pytest
 ```
-The `--image-columns` option controls the number of structures per grid row (default: `4`).
+The package separates common dataset parsing and fingerprints from independent workloads. Add future workloads through the workload registry without changing the main CLI.
@@ -0,0 +1,23 @@
 [build-system]
 requires = ["setuptools>=68"]
 build-backend = "setuptools.build_meta"
 [project]
 name = "scimesh"
 version = "0.1.0"
 description = "Local scientific workloads for molecular similarity analysis"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = ["rdkit>=2024.3"]
 [project.optional-dependencies]
 dev = ["pytest>=8"]
 [project.scripts]
 scimesh = "scimesh.cli:main"
 [tool.setuptools.packages.find]
 include = ["scimesh*"]
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -1 +1 @@
-rdkit>=2024.3
+-e .
@@ -1,270 +0,0 @@
 #!/usr/bin/env python3
 """SciMesh: streaming search for ChEMBL molecules similar to gefitinib."""
 from __future__ import annotations
 import argparse
 import csv
 import heapq
 import sys
 import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator
 from rdkit import Chem, DataStructs, RDLogger
 from rdkit.Chem import Draw, rdFingerprintGenerator
 REFERENCE_CHEMBL_ID = "CHEMBL939"
 FP_RADIUS = 2
 FP_SIZE = 2048
 # Invalid records are intentionally skipped, so do not emit one RDKit error per row.
 RDLogger.DisableLog("rdApp.error")
@dataclass
 class SearchStats:
    """Counters collected while scanning the candidate records."""
    scanned: int = 0
    valid: int = 0
    invalid: int = 0
    stopped_early: bool = False
 def rows(tsv_path: Path) -> Iterator[dict[str, str]]:
    """Read ChEMBL records one at a time without loading the file into memory."""
    with tsv_path.open("r", encoding="utf-8", newline="") as source:
        yield from csv.DictReader(source, delimiter="\t")
 def molecule(smiles: str | None) -> Chem.Mol | None:
    """Return a parsed molecule, or None when a SMILES is empty or invalid."""
    if not smiles:
        return None
    return Chem.MolFromSmiles(smiles)
 def find_reference_smiles(tsv_path: Path) -> str:
    """Locate CHEMBL939 and return its canonical_smiles."""
    for row in rows(tsv_path):
        if row.get("chembl_id") == REFERENCE_CHEMBL_ID:
            smiles = row.get("canonical_smiles", "")
            if molecule(smiles) is None:
                raise ValueError(f"{REFERENCE_CHEMBL_ID} has an invalid canonical_smiles")
            return smiles
    raise ValueError(f"{REFERENCE_CHEMBL_ID} was not found in {tsv_path}")
 def find_similar(
    tsv_path: Path,
    reference_smiles: str,
    limit: int,
    progress_every: int,
    max_rows: int | None,
 ) -> tuple[list[tuple[float, str, str]], SearchStats]:
    """Calculate top similarities, retaining only ``limit`` rows in memory."""
    generator = rdFingerprintGenerator.GetMorganGenerator(
        radius=FP_RADIUS, fpSize=FP_SIZE
    )
    reference_fp = generator.GetFingerprint(molecule(reference_smiles))
    best: list[tuple[float, str, str]] = []
    stats = SearchStats()
    started_at = time.perf_counter()
    last_report_at = started_at
    last_report_rows = 0
    print("Searching candidates...", file=sys.stderr)
    for row in rows(tsv_path):
        if max_rows is not None and stats.scanned >= max_rows:
            stats.stopped_early = True
            break
        stats.scanned += 1
        chembl_id = row.get("chembl_id", "")
        smiles = row.get("canonical_smiles", "")
        if chembl_id != REFERENCE_CHEMBL_ID:
            mol = molecule(smiles)
            if mol is None:
                stats.invalid += 1
            else:
                stats.valid += 1
                similarity = DataStructs.TanimotoSimilarity(
                    reference_fp, generator.GetFingerprint(mol)
                )
                candidate = (similarity, chembl_id, smiles)
                if len(best) < limit:
                    heapq.heappush(best, candidate)
                elif candidate > best[0]:
                    heapq.heapreplace(best, candidate)
        if progress_every and stats.scanned % progress_every == 0:
            now = time.perf_counter()
            interval_seconds = now - last_report_at
            total_seconds = now - started_at
            interval_rows = stats.scanned - last_report_rows
            current_rate = interval_rows / interval_seconds if interval_seconds else 0.0
            average_rate = stats.scanned / total_seconds if total_seconds else 0.0
            print(
                "Processed "
                f"{stats.scanned:,} rows | {current_rate:,.0f} rows/s current | "
                f"{average_rate:,.0f} rows/s average | {total_seconds:.1f}s elapsed | "
                f"{stats.valid:,} valid | {stats.invalid:,} invalid | top {len(best)}",
                file=sys.stderr,
            )
            last_report_at = now
            last_report_rows = stats.scanned
    return sorted(best, reverse=True), stats
 def write_results(output_path: Path, matches: list[tuple[float, str, str]]) -> None:
    """Write ranked matching records as CSV."""
    with output_path.open("w", encoding="utf-8", newline="") as destination:
        writer = csv.DictWriter(
            destination,
            fieldnames=["rank", "chembl_id", "canonical_smiles", "similarity"],
        )
        writer.writeheader()
        for rank, (similarity, chembl_id, smiles) in enumerate(matches, start=1):
            writer.writerow(
                {
                    "rank": rank,
                    "chembl_id": chembl_id,
                    "canonical_smiles": smiles,
                    "similarity": f"{similarity:.6f}",
                }
            )
 def write_images(
    output_dir: Path,
    reference_smiles: str,
    matches: list[tuple[float, str, str]],
    columns: int,
 ) -> tuple[Path, Path]:
    """Create PNG depictions for the reference molecule and ranked matches."""
    output_dir.mkdir(parents=True, exist_ok=True)
    reference_path = output_dir / "CHEMBL939_gefitinib.png"
    reference_image = Draw.MolToImage(
        molecule(reference_smiles), size=(600, 400), legend="CHEMBL939 (gefitinib)"
    )
    reference_image.save(reference_path)
    candidate_path = output_dir / "top_candidates.png"
    candidate_molecules = [molecule(smiles) for _, _, smiles in matches]
    legends = [
        f"#{rank} {chembl_id}\nTanimoto: {similarity:.4f}"
        for rank, (similarity, chembl_id, _) in enumerate(matches, start=1)
    ]
    candidate_image = Draw.MolsToGridImage(
        candidate_molecules,
        molsPerRow=columns,
        subImgSize=(350, 250),
        legends=legends,
    )
    candidate_image.save(candidate_path)
    return reference_path, candidate_path
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Find ChEMBL molecules similar to gefitinib (CHEMBL939)."
    )
    parser.add_argument("input", type=Path, help="Path to ChEMBL TSV file")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("gefitinib_similarities.csv"),
        help="Output CSV path (default: gefitinib_similarities.csv)",
    )
    parser.add_argument(
        "--top",
        type=int,
        default=20,
        help="Number of matches to retain (default: 20)",
    )
    parser.add_argument(
        "--progress-every",
        type=int,
        default=100_000,
        help="Print progress after this many rows; 0 disables it (default: 100000)",
    )
    parser.add_argument(
        "--max-rows",
        type=int,
        help="For testing, scan only the first N TSV rows after locating CHEMBL939",
    )
    parser.add_argument(
        "--images-dir",
        type=Path,
        help="Create reference and top-candidate PNG images in this directory",
    )
    parser.add_argument(
        "--image-columns",
        type=int,
        default=4,
        help="Number of molecules per row in the candidate image (default: 4)",
    )
    return parser.parse_args()
 def main() -> int:
    args = parse_args()
    if args.top < 1:
        print("--top must be a positive integer", file=sys.stderr)
        return 2
    if args.progress_every < 0:
        print("--progress-every cannot be negative", file=sys.stderr)
        return 2
    if args.max_rows is not None and args.max_rows < 1:
        print("--max-rows must be a positive integer", file=sys.stderr)
        return 2
    if args.image_columns < 1:
        print("--image-columns must be a positive integer", file=sys.stderr)
        return 2
    if not args.input.is_file():
        print(f"Input file not found: {args.input}", file=sys.stderr)
        return 2
    try:
        reference_smiles = find_reference_smiles(args.input)
        matches, stats = find_similar(
            args.input,
            reference_smiles,
            args.top,
            args.progress_every,
            args.max_rows,
        )
    except ValueError as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
    write_results(args.output, matches)
    image_paths: tuple[Path, Path] | None = None
    if args.images_dir:
        image_paths = write_images(
            args.images_dir, reference_smiles, matches, args.image_columns
        )
    print(f"Reference {REFERENCE_CHEMBL_ID}: {reference_smiles}")
    print(
        f"Scanned {stats.scanned:,} rows: {stats.valid:,} valid, "
        f"{stats.invalid:,} invalid SMILES."
    )
    if stats.stopped_early:
        print("Stopped early because of --max-rows; results cover only that subset.")
    print(f"Saved {len(matches)} matches to {args.output}")
    if image_paths:
        print(f"Saved reference image to {image_paths[0]}")
        print(f"Saved candidate image to {image_paths[1]}")
    for rank, (similarity, chembl_id, smiles) in enumerate(matches, start=1):
        print(f"{rank:>2}. {chembl_id}\t{similarity:.6f}\t{smiles}")
    return 0
 if __name__ == "__main__":
    raise SystemExit(main())
@@ -0,0 +1 @@
 """SciMesh local scientific workloads."""
@@ -0,0 +1 @@
 """Shared ChEMBL parsing and fingerprint helpers."""
@@ -0,0 +1,85 @@
 """Streaming readers for ChEMBL-style TSV datasets."""
 from __future__ import annotations
 import csv
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator
 from rdkit import Chem, RDLogger
 ID_COLUMN = "chembl_id"
 SMILES_COLUMN = "canonical_smiles"
 # Invalid records are expected in large datasets; suppress one RDKit error per row.
 RDLogger.DisableLog("rdApp.error")
@dataclass
 class DatasetStats:
    """Counters collected while streaming a dataset."""
    scanned: int = 0
    valid: int = 0
    invalid: int = 0
    stopped_early: bool = False
@dataclass(frozen=True)
 class MoleculeRecord:
    """A valid molecule parsed from a ChEMBL TSV row."""
    molecule_id: str
    smiles: str
    molecule: Chem.Mol
 def parse_smiles(smiles: str | None) -> Chem.Mol | None:
    """Return a molecule for a non-empty valid SMILES, otherwise None."""
    if not smiles:
        return None
    return Chem.MolFromSmiles(smiles)
 def iter_rows(tsv_path: Path) -> Iterator[dict[str, str]]:
    """Yield TSV rows without loading the full file into memory."""
    with tsv_path.open("r", encoding="utf-8", newline="") as source:
        reader = csv.DictReader(source, delimiter="\t")
        fieldnames = set(reader.fieldnames or [])
        missing = {ID_COLUMN, SMILES_COLUMN} - fieldnames
        if missing:
            raise ValueError(f"Dataset is missing required columns: {', '.join(sorted(missing))}")
        yield from reader
 def iter_valid_molecules(
    tsv_path: Path, stats: DatasetStats, max_rows: int | None = None
 ) -> Iterator[MoleculeRecord]:
    """Yield valid molecules while updating streaming statistics."""
    for row in iter_rows(tsv_path):
        if max_rows is not None and stats.scanned >= max_rows:
            stats.stopped_early = True
            break
        stats.scanned += 1
        smiles = row.get(SMILES_COLUMN, "")
        molecule = parse_smiles(smiles)
        if molecule is None:
            stats.invalid += 1
            continue
        stats.valid += 1
        yield MoleculeRecord(row.get(ID_COLUMN, ""), smiles, molecule)
 def find_molecule_by_id(tsv_path: Path, molecule_id: str) -> MoleculeRecord:
    """Find and validate a molecule by ChEMBL identifier in a streaming pass."""
    for row in iter_rows(tsv_path):
        if row.get(ID_COLUMN) != molecule_id:
            continue
        smiles = row.get(SMILES_COLUMN, "")
        molecule = parse_smiles(smiles)
        if molecule is None:
            raise ValueError(f"{molecule_id} has an invalid canonical_smiles")
        return MoleculeRecord(molecule_id, smiles, molecule)
    raise ValueError(f"{molecule_id} was not found in {tsv_path}")
@@ -0,0 +1,24 @@
 """Morgan fingerprint utilities shared by molecular workloads."""
 from __future__ import annotations
 from functools import lru_cache
 from typing import Any
 from rdkit import Chem
 from rdkit.Chem import rdFingerprintGenerator
 FP_RADIUS = 2
 FP_SIZE = 2048
@lru_cache(maxsize=1)
 def morgan_generator() -> Any:
    """Return the standard SciMesh Morgan fingerprint generator."""
    return rdFingerprintGenerator.GetMorganGenerator(radius=FP_RADIUS, fpSize=FP_SIZE)
 def fingerprint(molecule: Chem.Mol) -> Any:
    """Build a Morgan radius-2, 2048-bit fingerprint."""
    return morgan_generator().GetFingerprint(molecule)
@@ -0,0 +1,37 @@
 """Command-line entry point for SciMesh."""
 from __future__ import annotations
 import argparse
 import sys
 from scimesh.core.registry import WorkloadRegistry
 from scimesh.workloads import register_workloads
 def build_parser() -> argparse.ArgumentParser:
    """Build the top-level parser from the registered workloads."""
    parser = argparse.ArgumentParser(
        prog="scimesh",
        description="Run local scientific workloads on molecular datasets.",
    )
    subparsers = parser.add_subparsers(dest="workload", required=True)
    registry = WorkloadRegistry()
    register_workloads(registry)
    registry.add_subparsers(subparsers)
    return parser
 def main(argv: list[str] | None = None) -> int:
    """Run a selected workload and return its exit status."""
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return args.handler(args)
    except (OSError, ValueError) as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
 if __name__ == "__main__":
    raise SystemExit(main())
@@ -0,0 +1 @@
 """Core workload abstractions."""
@@ -0,0 +1,27 @@
 """Registry for locally available workloads."""
 from __future__ import annotations
 import argparse
 from scimesh.core.workload import Workload
 class WorkloadRegistry:
    """Collect workloads and expose each one as a CLI subcommand."""
    def __init__(self) -> None:
        self._workloads: dict[str, Workload] = {}
    def register(self, workload: Workload) -> None:
        """Register a workload by its unique command name."""
        if workload.name in self._workloads:
            raise ValueError(f"Workload already registered: {workload.name}")
        self._workloads[workload.name] = workload
    def add_subparsers(self, subparsers: argparse._SubParsersAction) -> None:
        """Add a parser for every registered workload."""
        for workload in self._workloads.values():
            parser = subparsers.add_parser(workload.name, help=workload.help)
            workload.configure_parser(parser)
            parser.set_defaults(handler=workload.run)
@@ -0,0 +1,19 @@
 """Minimal interface implemented by every SciMesh workload."""
 from __future__ import annotations
 import argparse
 from typing import Protocol
 class Workload(Protocol):
    """A workload that can add its CLI and execute from parsed arguments."""
    name: str
    help: str
    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
        """Add workload-specific command-line arguments."""
    def run(self, args: argparse.Namespace) -> int:
        """Execute the workload."""
@@ -0,0 +1,13 @@
 """Built-in SciMesh workloads."""
 from __future__ import annotations
 from scimesh.core.registry import WorkloadRegistry
 from scimesh.workloads.similarity_graph import SimilarityGraphWorkload
 from scimesh.workloads.similarity_search import SimilaritySearchWorkload
 def register_workloads(registry: WorkloadRegistry) -> None:
    """Register built-in workloads in one place, outside the main CLI."""
    registry.register(SimilaritySearchWorkload())
    registry.register(SimilarityGraphWorkload())
@@ -0,0 +1,185 @@
 """Exact sparse molecular similarity graph workload."""
 from __future__ import annotations
 import argparse
 import csv
 import sys
 import time
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 from rdkit import DataStructs
 from scimesh.chemistry.dataset import DatasetStats, MoleculeRecord, iter_valid_molecules
 from scimesh.chemistry.fingerprints import fingerprint
@dataclass(frozen=True)
 class GraphMolecule:
    """A valid record and its fingerprint, built once for graph construction."""
    molecule_id: str
    fingerprint: Any
@dataclass(frozen=True)
 class SimilarityEdge:
    """A thresholded, undirected similarity edge represented once as i < j."""
    source_id: str
    target_id: str
    similarity: float
@dataclass
 class GraphResult:
    """Edges, dataset statistics, and pair-comparison statistics."""
    edges: list[SimilarityEdge]
    stats: DatasetStats
    checked_pairs: int
    elapsed_seconds: float
 def _fingerprinted_molecules(
    tsv_path: Path, max_rows: int | None
 ) -> tuple[list[GraphMolecule], DatasetStats]:
    stats = DatasetStats()
    molecules = [
        GraphMolecule(record.molecule_id, fingerprint(record.molecule))
        for record in iter_valid_molecules(tsv_path, stats, max_rows=max_rows)
    ]
    return molecules, stats
 def build_similarity_graph(
    tsv_path: Path,
    threshold: float,
    block_size: int,
    max_rows: int | None = None,
    progress_every: int = 0,
 ) -> GraphResult:
    """Build an exact sparse graph without creating a dense similarity matrix."""
    if not 0.0 <= threshold <= 1.0:
        raise ValueError("--threshold must be between 0 and 1")
    if block_size < 1:
        raise ValueError("--block-size must be a positive integer")
    molecules, stats = _fingerprinted_molecules(tsv_path, max_rows)
    edges: list[SimilarityEdge] = []
    checked_pairs = 0
    started_at = time.perf_counter()
    next_report = progress_every
    for left_block_start in range(0, len(molecules), block_size):
        left_block_end = min(left_block_start + block_size, len(molecules))
        for right_block_start in range(left_block_start, len(molecules), block_size):
            right_block_end = min(right_block_start + block_size, len(molecules))
            same_block = left_block_start == right_block_start
            for left_index in range(left_block_start, left_block_end):
                right_start = left_index + 1 if same_block else right_block_start
                for right_index in range(right_start, right_block_end):
                    checked_pairs += 1
                    similarity = DataStructs.TanimotoSimilarity(
                        molecules[left_index].fingerprint,
                        molecules[right_index].fingerprint,
                    )
                    if similarity >= threshold:
                        edges.append(
                            SimilarityEdge(
                                molecules[left_index].molecule_id,
                                molecules[right_index].molecule_id,
                                similarity,
                            )
                        )
                    if progress_every and checked_pairs >= next_report:
                        elapsed = time.perf_counter() - started_at
                        rate = checked_pairs / elapsed if elapsed else 0.0
                        print(
                            f"Checked {checked_pairs:,} pairs | {len(edges):,} edges | "
                            f"{rate:,.0f} pairs/s | {elapsed:.1f}s elapsed",
                            file=sys.stderr,
                        )
                        next_report += progress_every
    elapsed_seconds = time.perf_counter() - started_at
    edges.sort(key=lambda edge: (edge.source_id, edge.target_id, -edge.similarity))
    return GraphResult(edges, stats, checked_pairs, elapsed_seconds)
 def write_graph_edges(output_path: Path, edges: list[SimilarityEdge]) -> None:
    """Write a deterministic sparse edge list CSV."""
    with output_path.open("w", encoding="utf-8", newline="") as destination:
        writer = csv.DictWriter(destination, fieldnames=["source_id", "target_id", "similarity"])
        writer.writeheader()
        for edge in edges:
            writer.writerow(
                {
                    "source_id": edge.source_id,
                    "target_id": edge.target_id,
                    "similarity": f"{edge.similarity:.6f}",
                }
            )
 class SimilarityGraphWorkload:
    """CLI adapter for exact block-wise sparse similarity graph construction."""
    name = "similarity-graph"
    help = "Build an exact sparse graph of thresholded molecular similarities."
    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument("input", type=Path, help="Path to ChEMBL TSV file")
        parser.add_argument(
            "--threshold", type=float, required=True,
            help="Create edges at or above this Tanimoto similarity",
        )
        parser.add_argument(
            "--block-size", type=int, default=1_000,
            help="Number of molecules per comparison block (default: 1000)",
        )
        parser.add_argument(
            "--max-rows", type=int,
            help="Read only the first N dataset rows",
        )
        parser.add_argument(
            "--progress-every", type=int, default=100_000,
            help="Print progress after this many pairs; 0 disables it",
        )
        parser.add_argument(
            "-o", "--output", type=Path, default=Path("similarity_graph.csv"),
            help="Output edge-list CSV path",
        )
    def run(self, args: argparse.Namespace) -> int:
        if args.max_rows is not None and args.max_rows < 1:
            raise ValueError("--max-rows must be a positive integer")
        if args.progress_every < 0:
            raise ValueError("--progress-every cannot be negative")
        result = build_similarity_graph(
            args.input,
            args.threshold,
            args.block_size,
            args.max_rows,
            args.progress_every,
        )
        write_graph_edges(args.output, result.edges)
        rate = (
            result.checked_pairs / result.elapsed_seconds
            if result.elapsed_seconds
            else 0.0
        )
        print(
            f"Valid molecules: {result.stats.valid:,} | invalid SMILES: "
            f"{result.stats.invalid:,} | scanned rows: {result.stats.scanned:,}"
        )
        print(
            f"Checked pairs: {result.checked_pairs:,} | edges: {len(result.edges):,} | "
            f"{rate:,.0f} pairs/s | {result.elapsed_seconds:.1f}s elapsed"
        )
        if result.stats.stopped_early:
            print("Stopped early because of --max-rows; graph covers only that subset.")
        print(f"Saved {len(result.edges)} edges to {args.output}")
        return 0
@@ -0,0 +1,225 @@
 """Top-k molecular similarity search workload."""
 from __future__ import annotations
 import argparse
 import csv
 import heapq
 import sys
 import time
 from dataclasses import dataclass
 from pathlib import Path
 from rdkit import Chem, DataStructs
 from rdkit.Chem import Draw
 from scimesh.chemistry.dataset import (
    DatasetStats,
    MoleculeRecord,
    find_molecule_by_id,
    iter_valid_molecules,
    parse_smiles,
 )
 from scimesh.chemistry.fingerprints import fingerprint
@dataclass(frozen=True)
 class SimilarityMatch:
    """A candidate ranked by descending similarity and stable tie-breakers."""
    similarity: float
    molecule_id: str
    smiles: str
    def sort_key(self) -> tuple[float, str, str]:
        return (-self.similarity, self.molecule_id, self.smiles)
@dataclass
 class _HeapEntry:
    """Heap item whose minimum is the worst retained match."""
    match: SimilarityMatch
    def __lt__(self, other: object) -> bool:
        if not isinstance(other, _HeapEntry):
            return NotImplemented
        return self.match.sort_key() > other.match.sort_key()
@dataclass
 class SearchResult:
    """Results and scan statistics for a similarity search."""
    matches: list[SimilarityMatch]
    stats: DatasetStats
 def search_similar(
    tsv_path: Path,
    query: MoleculeRecord,
    top_k: int,
    max_rows: int | None = None,
    progress_every: int = 0,
 ) -> SearchResult:
    """Stream top-k matches, retaining only a bounded heap in memory."""
    if top_k < 1:
        raise ValueError("--top-k must be a positive integer")
    query_fingerprint = fingerprint(query.molecule)
    query_canonical_smiles = Chem.MolToSmiles(query.molecule, canonical=True)
    stats = DatasetStats()
    heap: list[_HeapEntry] = []
    started_at = time.perf_counter()
    last_report_at = started_at
    last_report_rows = 0
    for record in iter_valid_molecules(tsv_path, stats, max_rows=max_rows):
        candidate_canonical_smiles = Chem.MolToSmiles(record.molecule, canonical=True)
        if (
            record.molecule_id == query.molecule_id
            or candidate_canonical_smiles == query_canonical_smiles
        ):
            continue
        match = SimilarityMatch(
            DataStructs.TanimotoSimilarity(query_fingerprint, fingerprint(record.molecule)),
            record.molecule_id,
            record.smiles,
        )
        entry = _HeapEntry(match)
        if len(heap) < top_k:
            heapq.heappush(heap, entry)
        elif match.sort_key() < heap[0].match.sort_key():
            heapq.heapreplace(heap, entry)
        if progress_every and stats.scanned % progress_every == 0:
            now = time.perf_counter()
            interval = now - last_report_at
            total = now - started_at
            current_rate = (stats.scanned - last_report_rows) / interval if interval else 0.0
            average_rate = stats.scanned / total if total else 0.0
            print(
                f"Processed {stats.scanned:,} rows | {current_rate:,.0f} rows/s current | "
                f"{average_rate:,.0f} rows/s average | {total:.1f}s elapsed | "
                f"{stats.valid:,} valid | {stats.invalid:,} invalid | top {len(heap)}",
                file=sys.stderr,
            )
            last_report_at = now
            last_report_rows = stats.scanned
    return SearchResult(sorted((entry.match for entry in heap), key=SimilarityMatch.sort_key), stats)
 def write_search_results(output_path: Path, matches: list[SimilarityMatch]) -> None:
    """Write ranked matches to a deterministic CSV file."""
    with output_path.open("w", encoding="utf-8", newline="") as destination:
        writer = csv.DictWriter(
            destination,
            fieldnames=["rank", "chembl_id", "canonical_smiles", "similarity"],
        )
        writer.writeheader()
        for rank, match in enumerate(matches, start=1):
            writer.writerow(
                {
                    "rank": rank,
                    "chembl_id": match.molecule_id,
                    "canonical_smiles": match.smiles,
                    "similarity": f"{match.similarity:.6f}",
                }
            )
 def write_search_images(
    output_dir: Path, query: MoleculeRecord, matches: list[SimilarityMatch], columns: int
 ) -> tuple[Path, Path]:
    """Create PNG depictions for the query molecule and retained matches."""
    output_dir.mkdir(parents=True, exist_ok=True)
    query_path = output_dir / "query.png"
    Draw.MolToImage(
        query.molecule, size=(600, 400), legend=f"Query: {query.molecule_id}"
    ).save(query_path)
    candidates_path = output_dir / "top_candidates.png"
    molecules = [parse_smiles(match.smiles) for match in matches]
    legends = [
        f"#{rank} {match.molecule_id}\nTanimoto: {match.similarity:.4f}"
        for rank, match in enumerate(matches, start=1)
    ]
    Draw.MolsToGridImage(
        molecules, molsPerRow=columns, subImgSize=(350, 250), legends=legends
    ).save(candidates_path)
    return query_path, candidates_path
 class SimilaritySearchWorkload:
    """CLI adapter for streaming top-k molecular similarity search."""
    name = "similarity-search"
    help = "Find top-k molecules similar to a query molecule."
    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument("input", type=Path, help="Path to ChEMBL TSV file")
        query_group = parser.add_mutually_exclusive_group(required=True)
        query_group.add_argument("--query-id", help="ChEMBL ID of the query molecule")
        query_group.add_argument("--query-smiles", help="SMILES of the query molecule")
        parser.add_argument(
            "--top-k", "--top", dest="top_k", type=int, default=20,
            help="Number of matches to retain (default: 20)",
        )
        parser.add_argument(
            "-o", "--output", type=Path, default=Path("similarity_results.csv"),
            help="Output CSV path",
        )
        parser.add_argument(
            "--progress-every", type=int, default=100_000,
            help="Print progress after this many rows; 0 disables it",
        )
        parser.add_argument(
            "--max-rows", type=int,
            help="Scan only the first N rows after resolving the query",
        )
        parser.add_argument(
            "--images-dir", type=Path,
            help="Directory for query and top-candidate PNG images",
        )
        parser.add_argument(
            "--image-columns", type=int, default=4,
            help="Number of molecules per row in the candidate image",
        )
    def run(self, args: argparse.Namespace) -> int:
        if args.progress_every < 0:
            raise ValueError("--progress-every cannot be negative")
        if args.max_rows is not None and args.max_rows < 1:
            raise ValueError("--max-rows must be a positive integer")
        if args.image_columns < 1:
            raise ValueError("--image-columns must be a positive integer")
        if args.query_id:
            query = find_molecule_by_id(args.input, args.query_id)
        else:
            molecule = parse_smiles(args.query_smiles)
            if molecule is None:
                raise ValueError("--query-smiles is invalid")
            query = MoleculeRecord("query", args.query_smiles, molecule)
        result = search_similar(
            args.input, query, args.top_k, args.max_rows, args.progress_every
        )
        write_search_results(args.output, result.matches)
        image_paths: tuple[Path, Path] | None = None
        if args.images_dir:
            image_paths = write_search_images(
                args.images_dir, query, result.matches, args.image_columns
            )
        print(f"Query {query.molecule_id}: {query.smiles}")
        print(
            f"Scanned {result.stats.scanned:,} rows: {result.stats.valid:,} valid, "
            f"{result.stats.invalid:,} invalid SMILES."
        )
        if result.stats.stopped_early:
            print("Stopped early because of --max-rows; results cover only that subset.")
        print(f"Saved {len(result.matches)} matches to {args.output}")
        if image_paths:
            print(f"Saved query image to {image_paths[0]}")
            print(f"Saved candidate image to {image_paths[1]}")
        return 0
@@ -0,0 +1,21 @@
 from __future__ import annotations
 from pathlib import Path
 import pytest
@pytest.fixture
 def small_dataset(tmp_path: Path) -> Path:
    path = tmp_path / "molecules.tsv"
    path.write_text(
        "chembl_id\tcanonical_smiles\n"
        "QUERY\tCCO\n"
        "ALCOHOL\tCCCO\n"
        "AMINE\tCCN\n"
        "BENZENE\tc1ccccc1\n"
        "BROKEN\tnot-a-smiles\n"
        "DUPLICATE\tCCO\n",
        encoding="utf-8",
    )
    return path
@@ -0,0 +1,55 @@
 from __future__ import annotations
 from pathlib import Path
 from rdkit import DataStructs
 from scimesh.chemistry.dataset import DatasetStats, iter_valid_molecules
 from scimesh.chemistry.fingerprints import fingerprint
 from scimesh.workloads.similarity_graph import (
    SimilarityEdge,
    build_similarity_graph,
    write_graph_edges,
 )
 def _brute_force_edges(dataset: Path, threshold: float) -> list[SimilarityEdge]:
    records = list(iter_valid_molecules(dataset, DatasetStats()))
    edges = []
    for left_index, left in enumerate(records):
        for right in records[left_index + 1 :]:
            similarity = DataStructs.TanimotoSimilarity(
                fingerprint(left.molecule), fingerprint(right.molecule)
            )
            if similarity >= threshold:
                edges.append(SimilarityEdge(left.molecule_id, right.molecule_id, similarity))
    return sorted(edges, key=lambda edge: (edge.source_id, edge.target_id, -edge.similarity))
 def test_graph_matches_brute_force_and_has_unique_non_self_edges(
    small_dataset: Path,
 ) -> None:
    threshold = 0.15
    result = build_similarity_graph(small_dataset, threshold, block_size=2)
    assert result.edges == _brute_force_edges(small_dataset, threshold)
    assert result.checked_pairs == 10
    edge_pairs = [(edge.source_id, edge.target_id) for edge in result.edges]
    assert all(source != target for source, target in edge_pairs)
    assert len(edge_pairs) == len(set(edge_pairs))
    assert result.stats.invalid == 1
 def test_graph_is_block_size_independent_and_deterministic(
    small_dataset: Path, tmp_path: Path
 ) -> None:
    first = build_similarity_graph(small_dataset, threshold=0.15, block_size=1)
    second = build_similarity_graph(small_dataset, threshold=0.15, block_size=3)
    repeated = build_similarity_graph(small_dataset, threshold=0.15, block_size=3)
    assert first.edges == second.edges == repeated.edges
    first_path = tmp_path / "first.csv"
    second_path = tmp_path / "second.csv"
    write_graph_edges(first_path, first.edges)
    write_graph_edges(second_path, repeated.edges)
    assert first_path.read_bytes() == second_path.read_bytes()
@@ -0,0 +1,39 @@
 from __future__ import annotations
 from pathlib import Path
 from rdkit import Chem, DataStructs
 from scimesh.chemistry.dataset import DatasetStats, find_molecule_by_id, iter_valid_molecules
 from scimesh.chemistry.fingerprints import fingerprint
 from scimesh.workloads.similarity_search import SimilarityMatch, search_similar
 def test_search_matches_full_sorting_and_skips_query_and_invalid(
    small_dataset: Path,
 ) -> None:
    query = find_molecule_by_id(small_dataset, "QUERY")
    result = search_similar(small_dataset, query, top_k=2)
    query_smiles = Chem.MolToSmiles(query.molecule, canonical=True)
    expected = []
    for record in iter_valid_molecules(small_dataset, DatasetStats()):
        if record.molecule_id == query.molecule_id:
            continue
        if Chem.MolToSmiles(record.molecule, canonical=True) == query_smiles:
            continue
        expected.append(
            SimilarityMatch(
                DataStructs.TanimotoSimilarity(
                    fingerprint(query.molecule), fingerprint(record.molecule)
                ),
                record.molecule_id,
                record.smiles,
            )
        )
    assert result.matches == sorted(expected, key=SimilarityMatch.sort_key)[:2]
    assert "QUERY" not in {match.molecule_id for match in result.matches}
    assert "DUPLICATE" not in {match.molecule_id for match in result.matches}
    assert result.stats.invalid == 1
    assert result.stats.valid == 5
		`@@ -0,0 +1 @@`
							`"""Shared ChEMBL parsing and fingerprint helpers."""`