From 34beb8b0aa1f2ccb3cabd4856117bd608190b162 Mon Sep 17 00:00:00 2001 From: Emil Date: Mon, 13 Jul 2026 22:50:56 +0300 Subject: [PATCH] Refactor into modular molecular workloads --- .gitignore | 3 +- README.md | 57 ++++-- pyproject.toml | 23 +++ requirements.txt | 2 +- scimesh.py | 270 ------------------------- scimesh/__init__.py | 1 + scimesh/chemistry/__init__.py | 1 + scimesh/chemistry/dataset.py | 85 ++++++++ scimesh/chemistry/fingerprints.py | 24 +++ scimesh/cli.py | 37 ++++ scimesh/core/__init__.py | 1 + scimesh/core/registry.py | 27 +++ scimesh/core/workload.py | 19 ++ scimesh/workloads/__init__.py | 13 ++ scimesh/workloads/similarity_graph.py | 185 +++++++++++++++++ scimesh/workloads/similarity_search.py | 225 +++++++++++++++++++++ tests/conftest.py | 21 ++ tests/test_similarity_graph.py | 55 +++++ tests/test_similarity_search.py | 39 ++++ 19 files changed, 798 insertions(+), 290 deletions(-) create mode 100644 pyproject.toml delete mode 100644 scimesh.py create mode 100644 scimesh/__init__.py create mode 100644 scimesh/chemistry/__init__.py create mode 100644 scimesh/chemistry/dataset.py create mode 100644 scimesh/chemistry/fingerprints.py create mode 100644 scimesh/cli.py create mode 100644 scimesh/core/__init__.py create mode 100644 scimesh/core/registry.py create mode 100644 scimesh/core/workload.py create mode 100644 scimesh/workloads/__init__.py create mode 100644 scimesh/workloads/similarity_graph.py create mode 100644 scimesh/workloads/similarity_search.py create mode 100644 tests/conftest.py create mode 100644 tests/test_similarity_graph.py create mode 100644 tests/test_similarity_search.py diff --git a/.gitignore b/.gitignore index 85ca7bc..e3f7150 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/README.md b/README.md index 2c372bf..8408288 100644 --- a/README.md +++ b/README.md @@ -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 -python scimesh.py chembl_37_chemreps.txt --progress-every 0 +scimesh similarity-search chembl_37_chemreps.txt \ + --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. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7c37a2a --- /dev/null +++ b/pyproject.toml @@ -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"] diff --git a/requirements.txt b/requirements.txt index 5f46f60..d6e1198 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1 @@ -rdkit>=2024.3 +-e . diff --git a/scimesh.py b/scimesh.py deleted file mode 100644 index 4be5ce7..0000000 --- a/scimesh.py +++ /dev/null @@ -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()) diff --git a/scimesh/__init__.py b/scimesh/__init__.py new file mode 100644 index 0000000..0f00222 --- /dev/null +++ b/scimesh/__init__.py @@ -0,0 +1 @@ +"""SciMesh local scientific workloads.""" diff --git a/scimesh/chemistry/__init__.py b/scimesh/chemistry/__init__.py new file mode 100644 index 0000000..e28e02c --- /dev/null +++ b/scimesh/chemistry/__init__.py @@ -0,0 +1 @@ +"""Shared ChEMBL parsing and fingerprint helpers.""" diff --git a/scimesh/chemistry/dataset.py b/scimesh/chemistry/dataset.py new file mode 100644 index 0000000..85e53e1 --- /dev/null +++ b/scimesh/chemistry/dataset.py @@ -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}") diff --git a/scimesh/chemistry/fingerprints.py b/scimesh/chemistry/fingerprints.py new file mode 100644 index 0000000..347e696 --- /dev/null +++ b/scimesh/chemistry/fingerprints.py @@ -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) diff --git a/scimesh/cli.py b/scimesh/cli.py new file mode 100644 index 0000000..c144de6 --- /dev/null +++ b/scimesh/cli.py @@ -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()) diff --git a/scimesh/core/__init__.py b/scimesh/core/__init__.py new file mode 100644 index 0000000..4ec593f --- /dev/null +++ b/scimesh/core/__init__.py @@ -0,0 +1 @@ +"""Core workload abstractions.""" diff --git a/scimesh/core/registry.py b/scimesh/core/registry.py new file mode 100644 index 0000000..e18c44e --- /dev/null +++ b/scimesh/core/registry.py @@ -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) diff --git a/scimesh/core/workload.py b/scimesh/core/workload.py new file mode 100644 index 0000000..c1b6b1a --- /dev/null +++ b/scimesh/core/workload.py @@ -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.""" diff --git a/scimesh/workloads/__init__.py b/scimesh/workloads/__init__.py new file mode 100644 index 0000000..e74bfc1 --- /dev/null +++ b/scimesh/workloads/__init__.py @@ -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()) diff --git a/scimesh/workloads/similarity_graph.py b/scimesh/workloads/similarity_graph.py new file mode 100644 index 0000000..bc84d7c --- /dev/null +++ b/scimesh/workloads/similarity_graph.py @@ -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 diff --git a/scimesh/workloads/similarity_search.py b/scimesh/workloads/similarity_search.py new file mode 100644 index 0000000..7d1167d --- /dev/null +++ b/scimesh/workloads/similarity_search.py @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..46bac69 --- /dev/null +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_similarity_graph.py b/tests/test_similarity_graph.py new file mode 100644 index 0000000..23e685b --- /dev/null +++ b/tests/test_similarity_graph.py @@ -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() diff --git a/tests/test_similarity_search.py b/tests/test_similarity_search.py new file mode 100644 index 0000000..a8b1538 --- /dev/null +++ b/tests/test_similarity_search.py @@ -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