Refactor into modular molecular workloads

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}")