Refactor into modular molecular workloads

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()