Refactor into modular molecular workloads
This commit is contained in:
Emil
+2
-1
@@ -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
|
||||
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.
|
||||
|
||||
@@ -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
@@ -1 +1 @@
|
||||
rdkit>=2024.3
|
||||
-e .
|
||||
|
||||
-270
@@ -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
|
||||
Reference in New Issue
Block a user