250 lines
10 KiB
Python
250 lines
10 KiB
Python
"""Scientific core for the SDK-built ``similarity-graph`` workload.
|
|
|
|
Molecules are parsed once into deterministic row-ordered blocks; every block
|
|
pair ``(i, j)`` with ``i <= j`` becomes one map task (diagonal tasks compare
|
|
pairs ``a < b`` inside a block, off-diagonal tasks compare every molecule
|
|
across two blocks). The reducer enforces the CTX-10 pair-coverage invariant:
|
|
the union of task pair sets must equal all unordered molecule pairs exactly
|
|
once, and the merged edge set must contain no duplicate unordered pair.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
from pathlib import Path
|
|
from typing import Iterable, Sequence
|
|
|
|
from rdkit import Chem, DataStructs
|
|
|
|
from scimesh.chemistry.dataset import iter_valid_molecules
|
|
from scimesh.chemistry.fingerprints import fingerprint
|
|
|
|
EDGE_COLUMNS = ("source_id", "target_id", "similarity")
|
|
|
|
MoleculeBlock = list[tuple[str, str]] # (chembl_id, smiles), row-ordered
|
|
|
|
|
|
def parse_molecule_blocks(
|
|
input_path: Path,
|
|
block_size: int,
|
|
max_rows: int | None = None,
|
|
) -> tuple[list[MoleculeBlock], dict[str, int]]:
|
|
"""Parse valid molecules into deterministic row-ordered blocks.
|
|
|
|
Mirrors the local reference's strictness: an empty or duplicate
|
|
``chembl_id`` fails the run, because the edge identity is the molecule id.
|
|
Invalid SMILES rows are skipped and counted.
|
|
"""
|
|
if (
|
|
isinstance(block_size, bool)
|
|
or not isinstance(block_size, int)
|
|
or block_size < 1
|
|
):
|
|
raise ValueError("block_size must be a positive integer")
|
|
from scimesh.chemistry.dataset import DatasetStats
|
|
|
|
stats = DatasetStats()
|
|
blocks: list[MoleculeBlock] = []
|
|
current: MoleculeBlock = []
|
|
seen_ids: set[str] = set()
|
|
for record in iter_valid_molecules(input_path, stats, max_rows=max_rows):
|
|
if not record.molecule_id:
|
|
raise ValueError("dataset contains an empty chembl_id")
|
|
if record.molecule_id in seen_ids:
|
|
raise ValueError(
|
|
f"dataset contains a duplicate chembl_id: {record.molecule_id}"
|
|
)
|
|
seen_ids.add(record.molecule_id)
|
|
current.append((record.molecule_id, record.smiles))
|
|
if len(current) == block_size:
|
|
blocks.append(current)
|
|
current = []
|
|
if current:
|
|
blocks.append(current)
|
|
if not blocks:
|
|
raise ValueError("dataset has no valid molecules")
|
|
return blocks, {
|
|
"rows_scanned": stats.scanned,
|
|
"valid_molecules": stats.valid,
|
|
"invalid_smiles": stats.invalid,
|
|
"block_count": len(blocks),
|
|
}
|
|
|
|
|
|
def write_block_tsv(rows: MoleculeBlock, path: Path) -> None:
|
|
"""Write one molecule block as a header TSV with the input column names."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8", newline="") as destination:
|
|
writer = csv.DictWriter(
|
|
destination,
|
|
fieldnames=["chembl_id", "canonical_smiles"],
|
|
delimiter="\t",
|
|
lineterminator="\n",
|
|
)
|
|
writer.writeheader()
|
|
for molecule_id, smiles in rows:
|
|
writer.writerow({"chembl_id": molecule_id, "canonical_smiles": smiles})
|
|
|
|
|
|
def read_block_rows(path: Path) -> MoleculeBlock:
|
|
rows: MoleculeBlock = []
|
|
with path.open("r", encoding="utf-8", newline="") as source:
|
|
reader = csv.DictReader(source, delimiter="\t")
|
|
for row in reader:
|
|
molecule_id = row.get("chembl_id", "")
|
|
smiles = row.get("canonical_smiles", "")
|
|
if not molecule_id or not smiles:
|
|
raise ValueError("block artifact contains an invalid row")
|
|
rows.append((molecule_id, smiles))
|
|
return rows
|
|
|
|
|
|
def compute_block_edges(
|
|
left: MoleculeBlock,
|
|
right: MoleculeBlock,
|
|
threshold: float,
|
|
threshold_direction: str,
|
|
) -> list[tuple[str, str, float]]:
|
|
"""Compare one planned block pair and emit only thresholded edges."""
|
|
if not 0.0 <= threshold <= 1.0:
|
|
raise ValueError("threshold must be between 0 and 1")
|
|
if threshold_direction not in {"greater", "less"}:
|
|
raise ValueError("threshold_direction must be 'greater' or 'less'")
|
|
left_fingerprints = [
|
|
(molecule_id, fingerprint(Chem.MolFromSmiles(smiles)))
|
|
for molecule_id, smiles in left
|
|
]
|
|
right_fingerprints = [
|
|
(molecule_id, fingerprint(Chem.MolFromSmiles(smiles)))
|
|
for molecule_id, smiles in right
|
|
]
|
|
diagonal = left is right or left == right
|
|
edges: list[tuple[str, str, float]] = []
|
|
for left_index, (left_id, left_fp) in enumerate(left_fingerprints):
|
|
right_start = left_index + 1 if diagonal else 0
|
|
for right_index in range(right_start, len(right_fingerprints)):
|
|
right_id, right_fp = right_fingerprints[right_index]
|
|
similarity = DataStructs.TanimotoSimilarity(left_fp, right_fp)
|
|
matches_threshold = (
|
|
similarity >= threshold
|
|
if threshold_direction == "greater"
|
|
else similarity <= threshold
|
|
)
|
|
if matches_threshold:
|
|
edges.append((left_id, right_id, similarity))
|
|
return edges
|
|
|
|
|
|
def write_edge_csv(output_path: Path, edges: Iterable[tuple[str, str, float]]) -> None:
|
|
"""Write an edge table CSV with six-decimal similarity values.
|
|
|
|
Uses the CSV module's default ``\\r\\n`` line terminator so the bytes match
|
|
the local ``write_graph_edges`` reference exactly.
|
|
"""
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with output_path.open("w", encoding="utf-8", newline="") as destination:
|
|
writer = csv.DictWriter(destination, fieldnames=list(EDGE_COLUMNS))
|
|
writer.writeheader()
|
|
for source_id, target_id, similarity in edges:
|
|
writer.writerow(
|
|
{
|
|
"source_id": source_id,
|
|
"target_id": target_id,
|
|
"similarity": f"{similarity:.6f}",
|
|
}
|
|
)
|
|
|
|
|
|
def read_edge_csv(path: Path) -> list[tuple[str, str, float]]:
|
|
"""Read a materialized edge CSV with strict row validation."""
|
|
if not path.is_file():
|
|
raise ValueError("materialized edge partial is missing")
|
|
edges: list[tuple[str, str, float]] = []
|
|
with path.open("r", encoding="utf-8", newline="") as source:
|
|
reader = csv.DictReader(source)
|
|
if tuple(reader.fieldnames or ()) != EDGE_COLUMNS:
|
|
raise ValueError("edge partial has an invalid CSV header")
|
|
for row in reader:
|
|
if set(row) != set(EDGE_COLUMNS):
|
|
raise ValueError("edge partial has an invalid row")
|
|
source_id = row["source_id"]
|
|
target_id = row["target_id"]
|
|
if not source_id or not target_id:
|
|
raise ValueError("edge partial contains an empty molecule id")
|
|
try:
|
|
similarity = float(row["similarity"])
|
|
except (TypeError, ValueError) as error:
|
|
raise ValueError("edge partial has an invalid similarity") from error
|
|
if not 0 <= similarity <= 1:
|
|
raise ValueError("edge partial has an invalid similarity")
|
|
edges.append((source_id, target_id, similarity))
|
|
return edges
|
|
|
|
|
|
def block_pair_from_key(key: str) -> tuple[int, int]:
|
|
"""Parse ``map.<i>x<j>`` partial keys into block indices."""
|
|
prefix = "map."
|
|
if not key.startswith(prefix):
|
|
raise ValueError("graph partial key must use map.<left>x<right>")
|
|
raw = key[len(prefix) :]
|
|
left_raw, separator, right_raw = raw.partition("x")
|
|
if not separator or not left_raw.isdigit() or not right_raw.isdigit():
|
|
raise ValueError("graph partial key must use map.<left>x<right>")
|
|
return int(left_raw), int(right_raw)
|
|
|
|
|
|
def check_pair_coverage(pairs: Sequence[tuple[int, int]]) -> None:
|
|
"""Enforce the pair-coverage invariant: every block pair exactly once.
|
|
|
|
``pairs`` must contain every ``(i, j)`` with ``0 <= i <= j < n`` exactly
|
|
once, where ``n`` is derived from the largest referenced block index.
|
|
"""
|
|
if not pairs:
|
|
raise ValueError("graph partial keys cover no block pairs")
|
|
unique = set(pairs)
|
|
if len(unique) != len(pairs):
|
|
raise ValueError("graph partial keys contain a duplicate block pair")
|
|
if any(left > right or left < 0 or right < 0 for left, right in unique):
|
|
raise ValueError("graph partial keys reference an invalid block pair")
|
|
n = max(right for _, right in unique) + 1
|
|
expected = {(left, right) for left in range(n) for right in range(left, n)}
|
|
missing = sorted(expected - unique)
|
|
if missing:
|
|
raise ValueError(
|
|
"graph partial keys do not cover the full block pair set: "
|
|
+ ", ".join(f"{left}x{right}" for left, right in missing)
|
|
)
|
|
unexpected = sorted(unique - expected)
|
|
if unexpected:
|
|
raise ValueError(
|
|
"graph partial keys cover pairs outside the block pair set: "
|
|
+ ", ".join(f"{left}x{right}" for left, right in unexpected)
|
|
)
|
|
|
|
|
|
def merge_edge_partials(
|
|
partial_paths: Sequence[Path],
|
|
output_path: Path,
|
|
) -> dict[str, int]:
|
|
"""Merge edge partials with duplicate detection and deterministic sort.
|
|
|
|
The merged edge list is sorted by ``(source_id, target_id, -similarity)``,
|
|
matching the local brute-force reference exactly.
|
|
"""
|
|
if not partial_paths:
|
|
raise ValueError("graph reducer requires at least one edge partial")
|
|
edges: list[tuple[str, str, float]] = []
|
|
seen_pairs: set[tuple[str, str]] = set()
|
|
for path in partial_paths:
|
|
for source_id, target_id, similarity in read_edge_csv(path):
|
|
unordered = (min(source_id, target_id), max(source_id, target_id))
|
|
if unordered in seen_pairs:
|
|
raise ValueError(
|
|
f"graph partials contain a duplicate unordered pair: {unordered[0]}, {unordered[1]}"
|
|
)
|
|
seen_pairs.add(unordered)
|
|
edges.append((source_id, target_id, similarity))
|
|
edges.sort(key=lambda edge: (edge[0], edge[1], -edge[2]))
|
|
write_edge_csv(output_path, edges)
|
|
return {"partial_count": len(partial_paths), "edges_emitted": len(edges)}
|