Add sort_reactants option to ReactionEnumerator

AUTHORS.md

@@ -9,6 +9,7 @@
 
 We welcome contributions from the community. Contributors will be listed here.
 
+- **Matt Ball** — minor QoL changes
 <!--
 Example entry:
 - **Name** — contribution area (e.g., SMARTS templates, performance work, documentation)

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 This follows the guideline on [keep a changelog](https://keepachangelog.com/en/1.1.0/).
 
+## [Unreleased]
+
+### Added
+
+- `sort_reactants` argument (default `True`) on `ReactionEnumerator.enumerate_pair`, `products_for_pair`, `enumerate_pairs` and `enumerate_pairs_lazy`. With `sort_reactants=False`, `reactant_a` and `reactant_b` of each result keep the order of the input pair, so a caller can supply reactants in a deterministic order. The default behaviour is unchanged.
+
 ## [2.0.1] 2026-09-23
 
 ### Fixed

README.md

@@ -60,6 +60,9 @@ results = enumerator.enumerate_pair("c1ccc(Br)cc1", "c1ccc(B(O)O)cc1")
 
 # Pass clean_smiles=True to standardize SMILES before enumeration if desired
 results = enumerator.enumerate_pair("c1ccc(Br)cc1", "c1ccc(B(O)O)cc1", clean_smiles=True)
+
+# By default reactant_a and reactant_b are sorted. Pass sort_reactants=False to keep the input order.
+results = enumerator.enumerate_pair("c1ccc(B(O)O)cc1", "c1ccc(Br)cc1", sort_reactants=False)
 ```
 
 #### Batch Processing
@@ -159,7 +162,7 @@ Any SMILES not present in `keys_map` are classified on-the-fly as a fallback.
 - **Preprocessing**: `preprocess_smiles()` classifies each molecule once and returns a `dict[str, set[str]]` that can be passed to `enumerate_pairs(precomputed_keys=...)`. This is most useful when the same molecules appear in many pairs, e.g. when enumerating all pairwise combinations of a library.
 - **SMILES standardization**: All enumeration methods and `preprocess_smiles` accept `clean_smiles=True` to normalize SMILES before processing (removes isotopes, strips salts/solvents, neutralizes charges, keeps the largest fragment).
 - **Reaction Selection**: All 36 available reaction types are used by default. Pass a custom list to restrict enumeration to specific reactions, e.g. `reaction_list=["suzuki", "amide_coupling"]`.
-- **Result Format**: `ReactionResult` objects with `reactant_a`, `reactant_b` (sorted alphabetically), `reaction_name`, and `products` (canonical SMILES).
+- **Result Format**: `ReactionResult` objects with `reactant_a`, `reactant_b` (sorted alphabetically by default, or in input order with `sort_reactants=False`), `reaction_name`, and `products` (canonical SMILES).
 
 > **Note on RDKit logging.** Importing `smartreact` silences RDKit's C++ log channel (`rdApp.*`) process-wide via `RDLogger.DisableLog`. This keeps batch enumeration from flooding stderr when parsing unreliable inputs, but it also mutes RDKit warnings produced by *other* code in the same process. If you need RDKit's warnings, re-enable them after import with `RDLogger.EnableLog("rdApp.*")`.
 

src/smartreact/enumerator.py

@@ -109,6 +109,7 @@ def _worker_run_pair(
     m2: Chem.Mol | None,
     keys1: set[str],
     keys2: set[str],
+    sort_reactants: bool = True,
 ) -> list[ReactionResult]:
     """Run all applicable templates for one pair in a worker process.
 
@@ -126,7 +127,7 @@ def _worker_run_pair(
     if m1 is None or m2 is None:
         return []
 
-    ra, rb = tuple(sorted([s1, s2]))
+    ra, rb = tuple(sorted([s1, s2])) if sort_reactants else (s1, s2)
     results: list[ReactionResult] = []
     candidates = candidate_templates(keys1, keys2, _WORKER_TEMPLATE_INDEX)
     for idx in sorted(candidates):
@@ -159,6 +160,7 @@ def _worker_run_pairs_batch(
     pairs: list[tuple[str, str]],
     smiles_mols: dict[str, Chem.Mol | None],
     smiles_keys: dict[str, set[str]],
+    sort_reactants: bool = True,
 ) -> list[ReactionResult]:
     """Process a batch of pairs in a single worker call.
 
@@ -179,7 +181,13 @@ def _worker_run_pairs_batch(
     for s1, s2 in pairs:
         results.extend(
             _worker_run_pair(
-                s1, s2, smiles_mols[s1], smiles_mols[s2], smiles_keys[s1], smiles_keys[s2]
+                s1,
+                s2,
+                smiles_mols[s1],
+                smiles_mols[s2],
+                smiles_keys[s1],
+                smiles_keys[s2],
+                sort_reactants,
             )
         )
     return results
@@ -268,9 +276,14 @@ def _enumerate_pair_from_mols(
         m2: Chem.Mol | None,
         keys1: set[str],
         keys2: set[str],
+        sort_reactants: bool = True,
     ) -> list[ReactionResult]:
         """Core enumeration logic given already-parsed molecules.
 
+        If ``sort_reactants`` is ``True`` (default), ``reactant_a`` and
+        ``reactant_b`` of each result hold the two SMILES in sorted order.
+        Otherwise they hold ``s1`` and ``s2`` in the order passed in.
+
         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
@@ -304,7 +317,7 @@ def _enumerate_pair_from_mols(
                 )
                 continue
 
-            ra, rb = tuple(sorted([s1, s2]))
+            ra, rb = tuple(sorted([s1, s2])) if sort_reactants else (s1, s2)
             results.append(
                 ReactionResult(
                     reactant_a=ra,
@@ -321,16 +334,19 @@ def _enumerate_pair_from_keys(
         s2: str,
         keys1: set[str],
         keys2: set[str],
+        sort_reactants: bool = True,
     ) -> list[ReactionResult]:
         m1 = Chem.MolFromSmiles(s1)
         m2 = Chem.MolFromSmiles(s2)
-        return self._enumerate_pair_from_mols(s1, s2, m1, m2, keys1, keys2)
+        return self._enumerate_pair_from_mols(s1, s2, m1, m2, keys1, keys2, sort_reactants)
 
     # ------------------------------------------------------------------
     # Public API
     # ------------------------------------------------------------------
 
-    def enumerate_pair(self, s1: str, s2: str, clean_smiles: bool = False) -> list[ReactionResult]:
+    def enumerate_pair(
+        self, s1: str, s2: str, clean_smiles: bool = False, sort_reactants: bool = True
+    ) -> list[ReactionResult]:
         """
         Enumerate all products for a single pair of SMILES across all
         configured reaction templates.
@@ -342,6 +358,10 @@ def enumerate_pair(self, s1: str, s2: str, clean_smiles: bool = False) -> list[R
         clean_smiles:
             If ``True``, standardize each SMILES before enumeration.
             Returns ``[]`` if either molecule cannot be parsed after cleaning.
+        sort_reactants:
+            If ``True`` (default), ``reactant_a`` and ``reactant_b`` of each result
+            hold the two SMILES in sorted order. If ``False``, they hold ``s1``
+            and ``s2`` in the order given, so the caller controls the reactant order.
         """
         if clean_smiles:
             c1 = _clean_smiles(s1)
@@ -358,17 +378,21 @@ def enumerate_pair(self, s1: str, s2: str, clean_smiles: bool = False) -> list[R
 
         k1 = self._keys_for_smiles(s1)
         k2 = self._keys_for_smiles(s2)
-        return self._enumerate_pair_from_keys(s1, s2, k1, k2)
+        return self._enumerate_pair_from_keys(s1, s2, k1, k2, sort_reactants)
 
-    def products_for_pair(self, s1: str, s2: str, clean_smiles: bool = False) -> list[str]:
+    def products_for_pair(
+        self, s1: str, s2: str, clean_smiles: bool = False, sort_reactants: bool = True
+    ) -> list[str]:
         """Return only unique product SMILES for a pair.
 
         Parameters
         ----------
         clean_smiles:
             If ``True``, standardize each SMILES before enumeration.
+        sort_reactants:
+            Passed to :meth:`enumerate_pair`.
         """
-        res = self.enumerate_pair(s1, s2, clean_smiles=clean_smiles)
+        res = self.enumerate_pair(s1, s2, clean_smiles=clean_smiles, sort_reactants=sort_reactants)
         return sorted({p for r in res for p in r.products})
 
     def _process_chunk(
@@ -377,6 +401,7 @@ def _process_chunk(
         parallel: bool,
         clean_smiles: bool,
         precomputed_keys: dict[str, set[str]] | None,
+        sort_reactants: bool = True,
     ) -> Iterator[ReactionResult]:
         """Enumerate reactions for a materialised list of pairs and yield results."""
         if not pair_list:
@@ -430,7 +455,13 @@ def _process_chunk(
         if not parallel or self.n_cores == 1:
             for a, b in pair_list:
                 yield from self._enumerate_pair_from_mols(
-                    a, b, smiles_mols[a], smiles_mols[b], smiles_keys[a], smiles_keys[b]
+                    a,
+                    b,
+                    smiles_mols[a],
+                    smiles_mols[b],
+                    smiles_keys[a],
+                    smiles_keys[b],
+                    sort_reactants,
                 )
             return
 
@@ -447,6 +478,7 @@ def _process_chunk(
                 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},
+                sort_reactants,
             )
             for batch in batches
         ]
@@ -471,6 +503,7 @@ def enumerate_pairs_lazy(
         clean_smiles: bool = False,
         precomputed_keys: dict[str, set[str]] | None = None,
         chunk_size: int = 50_000,
+        sort_reactants: bool = True,
     ) -> Iterator[ReactionResult]:
         """
         Enumerate reactions for many reactant pairs, yielding results as each
@@ -492,17 +525,23 @@ def enumerate_pairs_lazy(
         chunk_size:
             Number of pairs processed per iteration.  Larger values amortise
             key-classification overhead at the cost of higher peak memory.
+        sort_reactants:
+            If ``True`` (default), each result lists its reactants in sorted
+            order. If ``False``, it lists them in the order of the input pair.
         """
         pairs_iter = iter(pairs)
         while chunk := [(str(a), str(b)) for a, b in itertools.islice(pairs_iter, chunk_size)]:
-            yield from self._process_chunk(chunk, parallel, clean_smiles, precomputed_keys)
+            yield from self._process_chunk(
+                chunk, parallel, clean_smiles, precomputed_keys, sort_reactants
+            )
 
     def enumerate_pairs(
         self,
         pairs: Iterable[tuple[str, str]],
         parallel: bool = False,
         clean_smiles: bool = False,
         precomputed_keys: dict[str, set[str]] | None = None,
+        sort_reactants: bool = True,
     ) -> list[ReactionResult]:
         """
         Enumerate reactions for many reactant pairs.
@@ -521,9 +560,14 @@ def enumerate_pairs(
             by :func:`~smartreact.preprocessing.preprocess_smiles`. Avoids
             re-classifying molecules that have already been processed. Keys for
             any SMILES not present in the mapping are computed on the fly.
+        sort_reactants:
+            If ``True`` (default), each result lists its reactants in sorted
+            order. If ``False``, it lists them in the order of the input pair.
         """
         pair_list = [(str(a), str(b)) for a, b in pairs]
-        return list(self._process_chunk(pair_list, parallel, clean_smiles, precomputed_keys))
+        return list(
+            self._process_chunk(pair_list, parallel, clean_smiles, precomputed_keys, sort_reactants)
+        )
 
 
 __all__ = ["ReactionEnumerator"]

tests/test_enumerator.py

@@ -79,6 +79,23 @@ def test_reactants_sorted(self, enumerator: ReactionEnumerator):
         for r in results:
             assert r.reactant_a <= r.reactant_b
 
+    def test_sort_reactants_false_keeps_input_order(self, enumerator: ReactionEnumerator):
+        # Test both argument orders. Only one of them matches the sorted order.
+        for s1, s2 in [(BROMOBENZENE, PHENYLBORONIC_ACID), (PHENYLBORONIC_ACID, BROMOBENZENE)]:
+            results = enumerator.enumerate_pair(s1, s2, sort_reactants=False)
+            assert results
+            for r in results:
+                assert (r.reactant_a, r.reactant_b) == (s1, s2)
+
+    def test_sort_reactants_false_pairs(self, enumerator: ReactionEnumerator):
+        pair = (PHENYLBORONIC_ACID, BROMOBENZENE)
+        for parallel in (False, True):
+            results = enumerator.enumerate_pairs([pair], parallel=parallel, sort_reactants=False)
+            assert results
+            assert all((r.reactant_a, r.reactant_b) == pair for r in results)
+        lazy = list(enumerator.enumerate_pairs_lazy([pair], sort_reactants=False))
+        assert lazy and all((r.reactant_a, r.reactant_b) == pair for r in lazy)
+
     def test_invalid_smiles(self, enumerator: ReactionEnumerator):
         results = enumerator.enumerate_pair("invalid", BROMOBENZENE)
         assert results == []
@@ -305,9 +322,9 @@ def test_workers_receive_only_their_own_batch_molecules(self):
         real_submit = ex.submit
         captured: list[tuple[list[tuple[str, str]], set[str], set[str]]] = []
 
-        def _spy(fn, batch, mols, keys):
+        def _spy(fn, batch, mols, keys, *extra):
             captured.append((batch, set(mols), set(keys)))
-            return real_submit(fn, batch, mols, keys)
+            return real_submit(fn, batch, mols, keys, *extra)
 
         with patch.object(ex, "submit", _spy):
             enum_par.enumerate_pairs(pairs, parallel=True)