diff --git a/CHANGELOG.md b/CHANGELOG.md index 2359548..56d12fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,18 @@ -# Changelog +# Changelog for smartreact + +This follows the guideline on [keep a changelog](https://keepachangelog.com/en/1.1.0/). + +## [1.1.0] 2026-07-28 + +### Added + +- `build_template_index` and `candidate_templates` for looking up applicable reaction templates from a reactant pair's key sets + +### Changed + +- `ReactionEnumerator` selects templates through the new index instead of scanning the full template library for every reactant pair +- Each unique SMILES is parsed once per call rather than once per pair occurrence, and parallel workers receive pre-parsed molecules restricted to the molecules their own batch references +- Enumeration results and their ordering are unchanged ## [1.0.0] 2026-06-08 diff --git a/src/smartreact/__init__.py b/src/smartreact/__init__.py index 65a2fa2..ac509f7 100644 --- a/src/smartreact/__init__.py +++ b/src/smartreact/__init__.py @@ -4,7 +4,12 @@ from .conversion import detect_format, read_sdf_block, read_sdf_file, to_smiles from .enumerator import ReactionEnumerator from .keygen import KeyGenerator, KeysLibrary, load_rules -from .keys import extract_key_strings, orders_for_template +from .keys import ( + build_template_index, + candidate_templates, + extract_key_strings, + orders_for_template, +) from .preprocessing import load_preprocessed, preprocess_smiles, save_preprocessed from .templates import load_templates from .types import KeyMatch, KeysResult, ReactionResult, ReactionTemplate, Rule @@ -27,6 +32,8 @@ "Rule", "extract_key_strings", "orders_for_template", + "build_template_index", + "candidate_templates", "clean_smiles", "to_smiles", "detect_format", diff --git a/src/smartreact/enumerator.py b/src/smartreact/enumerator.py index efb979a..844ec82 100644 --- a/src/smartreact/enumerator.py +++ b/src/smartreact/enumerator.py @@ -13,7 +13,7 @@ from .cleaning import clean_smiles as _clean_smiles from .keygen import KeyGenerator -from .keys import extract_key_strings, orders_for_template +from .keys import build_template_index, candidate_templates, extract_key_strings from .templates import load_templates from .types import ReactionResult, ReactionTemplate @@ -23,6 +23,7 @@ _RXN_CACHE: dict[str, AllChem.ChemicalReaction] = {} _WORKER_TEMPLATES: list[ReactionTemplate] | None = None +_WORKER_TEMPLATE_INDEX: dict[tuple[str, str], list[int]] | None = None def _get_rxn(smarts: str) -> AllChem.ChemicalReaction | None: @@ -75,12 +76,15 @@ def _collect_products( def _worker_init( templates: list[ReactionTemplate] | None = None, + template_index: dict[tuple[str, str], list[int]] | None = None, log_file: str | None = None, ) -> None: - """Configure templates and logging in worker processes.""" - global _WORKER_TEMPLATES + """Configure templates, the key-pair template index, and logging in worker processes.""" + global _WORKER_TEMPLATES, _WORKER_TEMPLATE_INDEX if templates is not None: _WORKER_TEMPLATES = templates + if template_index is not None: + _WORKER_TEMPLATE_INDEX = template_index if not log_file: return try: @@ -101,27 +105,33 @@ def _worker_init( def _worker_run_pair( s1: str, s2: str, + m1: Chem.Mol | None, + m2: Chem.Mol | None, keys1: set[str], keys2: set[str], ) -> list[ReactionResult]: """Run all applicable templates for one pair in a worker process. - Templates are pre-loaded into ``_WORKER_TEMPLATES`` via the pool initializer, - so they are not serialized per-task. Molecules are parsed once per call. + Templates and the key-pair template index are pre-loaded into + ``_WORKER_TEMPLATES`` / ``_WORKER_TEMPLATE_INDEX`` via the pool initializer, + so they are not serialized per-task. ``m1``/``m2`` are parsed once per + unique molecule by the caller and reused across all of that molecule's + pairs, rather than re-parsed for every pair occurrence. Candidate + templates are found via an index lookup over the pair's own (typically + small) key sets, rather than scanning every template. """ assert _WORKER_TEMPLATES is not None, "_WORKER_TEMPLATES not initialised in worker" + assert _WORKER_TEMPLATE_INDEX is not None, "_WORKER_TEMPLATE_INDEX not initialised in worker" - m1 = Chem.MolFromSmiles(s1) - m2 = Chem.MolFromSmiles(s2) if m1 is None or m2 is None: return [] ra, rb = tuple(sorted([s1, s2])) results: list[ReactionResult] = [] - for templ in _WORKER_TEMPLATES: - orders = orders_for_template(keys1, keys2, templ.reactant_categories) - if not orders: - continue + candidates = candidate_templates(keys1, keys2, _WORKER_TEMPLATE_INDEX) + for idx in sorted(candidates): + templ = _WORKER_TEMPLATES[idx] + orders = candidates[idx] rxn = _get_rxn(templ.smarts) if rxn is None: continue @@ -147,16 +157,31 @@ def _worker_run_pair( def _worker_run_pairs_batch( pairs: list[tuple[str, str]], + smiles_mols: dict[str, Chem.Mol | None], smiles_keys: dict[str, set[str]], ) -> list[ReactionResult]: """Process a batch of pairs in a single worker call. Receiving a whole batch minimises IPC round-trips: the main process - sends exactly ``n_cores`` tasks instead of one per pair. + sends exactly ``n_cores`` tasks instead of one per pair. ``smiles_mols`` + is parsed once per unique molecule by the caller rather than re-parsed + here for every pair; parsing costs roughly 5x what shipping the parsed + ``Mol`` does, so this is a clear win whenever molecules recur across + pairs (as they do in combinatorial enumeration). + + ``smiles_mols`` and ``smiles_keys`` are restricted by the caller to the + molecules this batch actually references. Every ``submit`` pickles its + arguments independently, so passing the chunk-wide dicts would serialise + them once per batch -- ``n_cores`` copies of data each worker mostly + cannot use. """ results: list[ReactionResult] = [] for s1, s2 in pairs: - results.extend(_worker_run_pair(s1, s2, smiles_keys[s1], smiles_keys[s2])) + results.extend( + _worker_run_pair( + s1, s2, smiles_mols[s1], smiles_mols[s2], smiles_keys[s1], smiles_keys[s2] + ) + ) return results @@ -189,6 +214,7 @@ def __init__( self.keygen = KeyGenerator(n_cores=self.n_cores) self.templates: list[ReactionTemplate] = load_templates(reaction_list=reaction_list) + self.template_index = build_template_index(self.templates) self.log_file = log_file self._pool: ProcessPoolExecutor | None = None @@ -196,7 +222,9 @@ def _get_pool(self) -> ProcessPoolExecutor: if self._pool is None: self._pool = ProcessPoolExecutor( max_workers=self.n_cores, - initializer=partial(_worker_init, self.templates, self.log_file), + initializer=partial( + _worker_init, self.templates, self.template_index, self.log_file + ), ) return self._pool @@ -232,23 +260,30 @@ def __del__(self) -> None: def _keys_for_smiles(self, smiles: str) -> set[str]: return extract_key_strings(self.keygen.classify(str(smiles))) - def _enumerate_pair_from_keys( + def _enumerate_pair_from_mols( self, s1: str, s2: str, + m1: Chem.Mol | None, + m2: Chem.Mol | None, keys1: set[str], keys2: set[str], ) -> list[ReactionResult]: - m1 = Chem.MolFromSmiles(s1) - m2 = Chem.MolFromSmiles(s2) + """Core enumeration logic given already-parsed molecules. + + Callers that process many pairs (``_process_chunk``) parse each + unique SMILES once and reuse the ``Mol`` across all of that + molecule's pairs, rather than calling this once per pair with a + freshly parsed molecule every time. + """ if m1 is None or m2 is None: return [] results: list[ReactionResult] = [] - for templ in self.templates: - orders = orders_for_template(keys1, keys2, templ.reactant_categories) - if not orders: - continue + candidates = candidate_templates(keys1, keys2, self.template_index) + for idx in sorted(candidates): + templ = self.templates[idx] + orders = candidates[idx] rxn = _get_rxn(templ.smarts) if rxn is None: @@ -280,6 +315,17 @@ def _enumerate_pair_from_keys( ) return results + def _enumerate_pair_from_keys( + self, + s1: str, + s2: str, + keys1: set[str], + keys2: set[str], + ) -> list[ReactionResult]: + m1 = Chem.MolFromSmiles(s1) + m2 = Chem.MolFromSmiles(s2) + return self._enumerate_pair_from_mols(s1, s2, m1, m2, keys1, keys2) + # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ @@ -363,6 +409,7 @@ def _process_chunk( return unique_smiles = list({s for pair in pair_list for s in pair}) + smiles_mols: dict[str, Chem.Mol | None] = {s: Chem.MolFromSmiles(s) for s in unique_smiles} if precomputed_keys is not None: missing = [s for s in unique_smiles if s not in precomputed_keys] @@ -382,14 +429,27 @@ def _process_chunk( if not parallel or self.n_cores == 1: for a, b in pair_list: - yield from self._enumerate_pair_from_keys(a, b, smiles_keys[a], smiles_keys[b]) + yield from self._enumerate_pair_from_mols( + a, b, smiles_mols[a], smiles_mols[b], smiles_keys[a], smiles_keys[b] + ) return chunk_size = max(1, len(pair_list) // self.n_cores) batches = [pair_list[i : i + chunk_size] for i in range(0, len(pair_list), chunk_size)] ex = self._get_pool() - futures = [ex.submit(_worker_run_pairs_batch, batch, smiles_keys) for batch in batches] + # Ship each worker only the molecules its own batch references: submit + # pickles its arguments per call, so passing the chunk-wide dicts would + # serialise them n_cores times over. + futures = [ + ex.submit( + _worker_run_pairs_batch, + batch, + {s: smiles_mols[s] for pair in batch for s in pair}, + {s: smiles_keys[s] for pair in batch for s in pair}, + ) + for batch in batches + ] try: for fut in futures: yield from fut.result() diff --git a/src/smartreact/keys.py b/src/smartreact/keys.py index fba0901..75fc865 100644 --- a/src/smartreact/keys.py +++ b/src/smartreact/keys.py @@ -1,7 +1,11 @@ from __future__ import annotations +from collections import defaultdict +from collections.abc import Sequence from typing import Any +from .types import ReactionTemplate + def extract_key_strings(res: Any) -> set[str]: """ @@ -61,4 +65,52 @@ def orders_for_template( return orders -__all__ = ["extract_key_strings", "orders_for_template"] +def build_template_index(templates: Sequence[ReactionTemplate]) -> dict[tuple[str, str], list[int]]: + """Map each (left_req, right_req) requirement pair to the template indices that need it. + + Built once from the template library. Lets :func:`candidate_templates` find + applicable templates via lookups over a pair's own key sets, instead of + scanning every template for every reactant pair. + """ + index: dict[tuple[str, str], list[int]] = defaultdict(list) + for i, templ in enumerate(templates): + index[templ.reactant_categories].append(i) + return dict(index) + + +def candidate_templates( + keys1: set[str], + keys2: set[str], + index: dict[tuple[str, str], list[int]], +) -> dict[int, list[int]]: + """Find template indices whose reactant_categories are satisfied by (keys1, keys2). + + Equivalent to calling :func:`orders_for_template` for every template in the + library and keeping the non-empty results, but costs O(|keys1| * |keys2|) + index lookups rather than O(len(templates)). + + Returns + ------- + dict[int, list[int]] + Maps template index to its list of valid orders (``[0]``, ``[1]``, or + ``[0, 1]``), keyed and ordered exactly as :func:`orders_for_template` + would produce for that template. + """ + candidates: dict[int, set[int]] = defaultdict(set) + for ka in keys1: + for kb in keys2: + for t in index.get((ka, kb), ()): + candidates[t].add(0) + for ka in keys2: + for kb in keys1: + for t in index.get((ka, kb), ()): + candidates[t].add(1) + return {t: sorted(orders) for t, orders in candidates.items()} + + +__all__ = [ + "extract_key_strings", + "orders_for_template", + "build_template_index", + "candidate_templates", +] diff --git a/tests/test_enumerator.py b/tests/test_enumerator.py index 51bdf6d..d1aa087 100644 --- a/tests/test_enumerator.py +++ b/tests/test_enumerator.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import logging from unittest.mock import patch @@ -74,6 +75,33 @@ def test_multiple_pairs(self, enumerator: ReactionEnumerator): def test_empty_input(self, enumerator: ReactionEnumerator): assert enumerator.enumerate_pairs([]) == [] + def test_parses_each_unique_molecule_once(self, enumerator: ReactionEnumerator): + """Regression test: enumerate_pairs must parse each unique SMILES once, + not once per pair occurrence, even when a molecule appears in many pairs. + + Keys are precomputed up front so the patched window only covers + enumerator-level Mol parsing, not KeyGenerator's own (separate) + internal MolFromSmiles call during classification. + """ + from smartreact.preprocessing import preprocess_smiles + + pairs = [ + (BROMOBENZENE, PHENYLBORONIC_ACID), + (BROMOBENZENE, ETHYLAMINE), + (BROMOBENZENE, ACETIC_ACID), + (IODOBENZENE, ETHYLAMINE), + ] + unique_smiles = {s for pair in pairs for s in pair} + keys_map = preprocess_smiles(list(unique_smiles), enumerator.keygen) + + with patch( + "smartreact.enumerator.Chem.MolFromSmiles", + wraps=Chem.MolFromSmiles, + ) as mock_parse: + enumerator.enumerate_pairs(pairs, parallel=False, precomputed_keys=keys_map) + + assert mock_parse.call_count == len(unique_smiles) + def test_parallel_matches_sequential(self): pairs = [ (BROMOBENZENE, PHENYLBORONIC_ACID), @@ -109,11 +137,56 @@ def _key(r): assert {_key(r) for r in results_seq} == {_key(r) for r in results_par} + def test_workers_receive_only_their_own_batch_molecules(self): + """Each submitted batch must carry only the molecules it references. + + ``submit`` pickles its arguments per call, so passing the chunk-wide + mol/key dicts to every batch would serialise them n_cores times over. + """ + molecules = [ + BROMOBENZENE, + PHENYLBORONIC_ACID, + ETHYLAMINE, + ACETIC_ACID, + IODOBENZENE, + ETHANOL, + ] + pairs = list(itertools.combinations(molecules, 2)) + + enum_par = ReactionEnumerator(n_cores=3) + ex = enum_par._get_pool() + real_submit = ex.submit + captured: list[tuple[list[tuple[str, str]], set[str], set[str]]] = [] + + def _spy(fn, batch, mols, keys): + captured.append((batch, set(mols), set(keys))) + return real_submit(fn, batch, mols, keys) + + with patch.object(ex, "submit", _spy): + enum_par.enumerate_pairs(pairs, parallel=True) + + assert len(captured) > 1, "expected the pair list to be split across several batches" + for batch, mols, keys in captured: + needed = {s for pair in batch for s in pair} + assert mols == needed + assert keys == needed + + # At least one batch must be a strict subset, or the assertions above + # would also pass for code that broadcasts everything to everyone. + assert any(mols < set(molecules) for _, mols, _ in captured) + assert set().union(*(mols for _, mols, _ in captured)) == set(molecules) + class TestEnumeratorInit: def test_default_reactions(self, enumerator: ReactionEnumerator): assert len(enumerator.templates) > 0 + def test_template_index_covers_all_templates(self, enumerator: ReactionEnumerator): + """The key-pair index used for candidate lookup must reference every + loaded template exactly once.""" + indexed_ids = sorted(idx for idxs in enumerator.template_index.values() for idx in idxs) + assert indexed_ids == list(range(len(enumerator.templates))) + def test_all_reactions(self, enumerator_all: ReactionEnumerator): default_enum = ReactionEnumerator(n_cores=1) assert len(enumerator_all.templates) >= len(default_enum.templates) diff --git a/tests/test_keys.py b/tests/test_keys.py index 3129e11..47a7f0e 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -1,7 +1,23 @@ from __future__ import annotations -from smartreact.keys import extract_key_strings, orders_for_template -from smartreact.types import KeyMatch, KeysResult +from smartreact.keys import ( + build_template_index, + candidate_templates, + extract_key_strings, + orders_for_template, +) +from smartreact.types import KeyMatch, KeysResult, ReactionTemplate + + +def _make_template( + left: str, right: str, template_id: int = 0, name: str = "t" +) -> ReactionTemplate: + return ReactionTemplate( + smarts="[#6:1]>>[#6:1]", + reactant_categories=(left, right), + name=name, + template_id=template_id, + ) class TestOrdersForTemplate: @@ -52,3 +68,62 @@ def test_excludes_higher_levels(self): def test_empty_matches(self): kr = KeysResult(smiles="C", matches=[]) assert extract_key_strings(kr) == set() + + +class TestBuildTemplateIndex: + def test_groups_by_requirement_pair(self): + templates = [ + _make_template("A", "B", template_id=0), + _make_template("A", "B", template_id=1), + _make_template("C", "D", template_id=2), + ] + index = build_template_index(templates) + assert index[("A", "B")] == [0, 1] + assert index[("C", "D")] == [2] + + def test_empty_templates(self): + assert build_template_index([]) == {} + + +class TestCandidateTemplates: + def test_forward_only(self): + index = build_template_index([_make_template("A", "B")]) + assert candidate_templates({"A"}, {"B"}, index) == {0: [0]} + + def test_reverse_only(self): + index = build_template_index([_make_template("A", "B")]) + assert candidate_templates({"B"}, {"A"}, index) == {0: [1]} + + def test_both_orders(self): + index = build_template_index([_make_template("A", "B")]) + assert candidate_templates({"A", "B"}, {"A", "B"}, index) == {0: [0, 1]} + + def test_no_match_returns_empty_dict(self): + index = build_template_index([_make_template("A", "B")]) + assert candidate_templates({"X"}, {"Y"}, index) == {} + + def test_equivalent_to_scanning_orders_for_template(self): + """candidate_templates must find exactly the same (template, orders) as + scanning every template with orders_for_template, for arbitrary key sets. + + This is the correctness guarantee behind replacing the O(len(templates)) + per-pair scan with an O(|keys1| * |keys2|) index lookup. + """ + templates = [ + _make_template("A", "B", template_id=0), + _make_template("B", "A", template_id=1), + _make_template("A", "A", template_id=2), + _make_template("C", "D", template_id=3), + _make_template("A", "B", template_id=4), # duplicate requirement pair + ] + index = build_template_index(templates) + + key_sets = [{"A"}, {"B"}, {"A", "B"}, {"C", "D"}, {"A", "C"}, set(), {"Z"}] + for keys1 in key_sets: + for keys2 in key_sets: + expected = {} + for i, templ in enumerate(templates): + orders = orders_for_template(keys1, keys2, templ.reactant_categories) + if orders: + expected[i] = orders + assert candidate_templates(keys1, keys2, index) == expected