Radius filter for RetroRules v3 rule sets, plus a performance fix in the expansion loop

.gitignore

@@ -136,3 +136,7 @@ src/biocatalyzer/data/organisms/organisms_paper.tsv
 src/biocatalyzer/data/compounds/drugs_paper_subset.csv
 src/biocatalyzer/data/compounds/drugs_paper_all.tsv
 src/biocatalyzer/data/reactionrules/all_reaction_rules_forward_no_smarts_duplicates.tsv
+
+.vscode/
+
+.pytest_cache/

README.md

@@ -159,7 +159,7 @@ For the `matcher_cli` see [readme_matcher_cli.md](readme_matcher_cli.md).
 
 ## Cite
 
-Manuscript under preparation!
+Manuscript under preparation.
 
 ### Credits and License
 

src/biocatalyzer/bioreactor.py

@@ -30,6 +30,7 @@ def __init__(self,
                  neutralize_compounds: bool = False,
                  reaction_rules_path: str = 'default',
                  organisms_path: str = None,
+                 radius: Union[str, int] = 6,
                  molecules_to_remove_path: Union[str, None] = 'default',
                  patterns_to_remove_path: Union[str, None] = 'default',
                  min_atom_count: int = 5,
@@ -49,6 +50,10 @@ def __init__(self,
             The path to the file containing the reaction rules.
         organisms_path: str
             The path to the file containing the organisms to filter the reaction rules by.
+        radius: Union[str, int]
+            The radius (or radii) of the reaction rules to use: an integer (6), a
+            ;-separated list ('4;6;8') or an inclusive range ('4:8'). 'ALL' disables
+            the filter. Only effective for rule sets that provide a 'Radius' column.
         molecules_to_remove_path: str
             The path to the file containing the molecules to remove from the products.
         patterns_to_remove_path: str
@@ -63,14 +68,17 @@ def __init__(self,
         self._compounds_path = compounds_path
         self._neutralize = neutralize_compounds
         self._organisms_path = organisms_path
+        self._radius = radius
         self._reaction_rules_path = reaction_rules_path
         self._molecules_to_remove_path = molecules_to_remove_path
         self._patterns_to_remove_path = patterns_to_remove_path
         self._set_up_files()
         self._orgs = Loaders.load_organisms(self._organisms_path)
-        self._reaction_rules = Loaders.load_reaction_rules(self._reaction_rules_path, orgs=self._orgs)
+        self._reaction_rules = Loaders.load_reaction_rules(self._reaction_rules_path, orgs=self._orgs,
+                                                          radius=self._radius)
         self._set_output_path(output_path)
         self._compounds = Loaders.load_compounds(self._compounds_path, self._neutralize)
+        self._index_reaction_rules()
         self._molecules_to_remove = Loaders.load_byproducts_to_remove(self._molecules_to_remove_path)
         self._patterns_to_remove = Loaders.load_patterns_to_remove(self._patterns_to_remove_path)
         self._min_atom_count = min_atom_count
@@ -106,6 +114,7 @@ def compounds(self, compounds_path: str):
         if compounds_path != self._compounds_path:
             self._compounds_path = compounds_path
             self._compounds = Loaders.load_compounds(self._compounds_path, self._neutralize)
+            self._index_reaction_rules()
         if self._new_compounds is not None:
             logging.warning('Results should be generated again for the new information provided!')
 
@@ -132,7 +141,9 @@ def reaction_rules(self, reaction_rules_path: str):
             The path to the file containing the reaction rules to use.
         """
         if reaction_rules_path != self._reaction_rules_path:
-            self._reaction_rules = Loaders.load_reaction_rules(reaction_rules_path, orgs=self._orgs)
+            self._reaction_rules = Loaders.load_reaction_rules(reaction_rules_path, orgs=self._orgs,
+                                                              radius=self._radius)
+            self._index_reaction_rules()
             self._reaction_rules_path = reaction_rules_path
         if self._new_compounds is not None:
             logging.warning('Results should be generated again for the new information provided!')
@@ -216,6 +227,7 @@ def compounds_path(self, compounds_path: str):
             self._compounds_path = compounds_path
             logging.info('Loading compounds again with the new path information...')
             self._compounds = Loaders.load_compounds(self._compounds_path, self._neutralize)
+            self._index_reaction_rules()
         if self._new_compounds is not None:
             logging.warning('Results should be generated again for the new information provided!')
 
@@ -245,6 +257,7 @@ def neutralize(self, neutralize: bool):
             self._neutralize = neutralize
             logging.info('Loading compounds again with the new neutralize information...')
             self._compounds = Loaders.load_compounds(self._compounds_path, self._neutralize)
+            self._index_reaction_rules()
         if self._new_compounds is not None:
             logging.warning('Results should be generated again for the new information provided!')
 
@@ -275,7 +288,9 @@ def organisms_path(self, organisms_path: str):
             logging.info('Loading organisms again with the new path information...')
             self._orgs = Loaders.load_organisms(self._organisms_path)
             logging.info('Loading reaction rules again with the new organisms information...')
-            self._reaction_rules = Loaders.load_reaction_rules(self._reaction_rules_path, orgs=self._orgs)
+            self._reaction_rules = Loaders.load_reaction_rules(self._reaction_rules_path, orgs=self._orgs,
+                                                               radius=self._radius)
+            self._index_reaction_rules()
         if self._new_compounds is not None:
             logging.warning('Results should be generated again for the new information provided!')
 
@@ -419,6 +434,27 @@ def _set_output_path(self, output_path: str):
                 )
         self._output_path = output_path
 
+    def _index_reaction_rules(self):
+        """
+        Index the reaction rules by their SMARTS string.
+
+        `_react_single` needs the reactants, the identifier and the EC numbers of the
+        rule it is applying. Resolving them with a boolean mask over the reaction rules
+        dataframe costs a full table scan per (compound, rule) pair, which dominates the
+        runtime once the rule set grows past a few thousand entries. The mapping is
+        therefore built once, when the rules are loaded. SMARTS strings are unique in a
+        BioCatalyzer rule set, so the lookup returns exactly what the mask returned.
+        """
+        self._rules_by_smarts = {
+            smarts: (reactants, internal_id, ec_numbers)
+            for smarts, reactants, internal_id, ec_numbers in zip(
+                self._reaction_rules.SMARTS,
+                self._reaction_rules.Reactants,
+                self._reaction_rules.InternalID,
+                self._reaction_rules.EC_Numbers)
+        }
+        self._compound_ids_by_smiles = dict(zip(self._compounds.smiles, self._compounds.compound_id))
+
     def _match_patterns(self, smiles: str):
         """
         Check if mol matches patterns to remove.
@@ -588,13 +624,12 @@ def _react_single(self, smiles: str, smarts: str, result_queue: multiprocessing.
         result_queue: multiprocessing.Queue
             The queue to store the results.
         """
-        reactants = self._reaction_rules[self._reaction_rules.SMARTS == smarts].Reactants.values[0]
+        reactants, smarts_id, ec_numbers = self._rules_by_smarts[smarts]
         reactants = reactants.replace("Any", smiles).split(';')
         results = ChemUtils.react(reactants, smarts)
         if len(results) == 0:
             return
-        smiles_id = self._compounds[self._compounds.smiles == smiles].compound_id.values[0]
-        smarts_id = self._reaction_rules[self._reaction_rules.SMARTS == smarts].InternalID.values[0]
+        smiles_id = self._compound_ids_by_smiles[smiles]
         most_similar_products_set = set()
         # Collect results in a list
         output_rows = []
@@ -608,7 +643,7 @@ def _react_single(self, smiles: str, smarts: str, result_queue: multiprocessing.
                 if self._match_conditions(most_similar_product):
                     if self._neutralize:
                         most_similar_product = ChemUtils.uncharge_smiles(most_similar_product)
-                    ecs = self._get_ec_numbers(smarts_id)
+                    ecs = ec_numbers
                     output_rows.append(f"{smiles_id}\t{smiles}\t{smarts_id}\t{smiles_id}_{uuid.uuid4()}\t"
                                        f"{most_similar_product}\t{result}\t{ecs}\n")
 

src/biocatalyzer/clis/cli.py

@@ -35,6 +35,15 @@
               default=None,
               help="The path to the user defined file containing the organisms to filter the reaction rules.",
               )
+@click.option("--radius",
+              "radius",
+              type=str,
+              default=6,
+              show_default=True,
+              help="Radius of the reaction rules to use: an integer (6), a ;-separated list "
+                   "('4;6;8') or an inclusive range ('4:8'). Only applied to rule sets that "
+                   "provide a 'Radius' column (e.g. RetroRules v3).",
+              )
 @click.option("--patterns_to_remove",
               "patterns_to_remove",
               type=str,
@@ -87,6 +96,7 @@ def biocatalyzer_cli(compounds,
                      neutralize,
                      reaction_rules,
                      organisms,
+                     radius,
                      patterns_to_remove,
                      molecules_to_remove,
                      min_atom_count,
@@ -112,6 +122,7 @@ def biocatalyzer_cli(compounds,
                     reaction_rules_path=reaction_rules,
                     neutralize_compounds=neutralize,
                     organisms_path=organisms,
+                    radius=radius,
                     patterns_to_remove_path=patterns_to_remove,
                     molecules_to_remove_path=molecules_to_remove,
                     min_atom_count=min_atom_count,

src/biocatalyzer/clis/cli_bioreactor.py

@@ -36,6 +36,15 @@
               default=None,
               help="The path to the user defined file containing the organisms to filter the reaction rules.",
               )
+@click.option("--radius",
+              "radius",
+              type=str,
+              default=6,
+              show_default=True,
+              help="Radius of the reaction rules to use: an integer (6), a ;-separated list "
+                   "('4;6;8') or an inclusive range ('4:8'). Only applied to rule sets that "
+                   "provide a 'Radius' column (e.g. RetroRules v3).",
+              )
 @click.option("--patterns_to_remove",
               "patterns_to_remove",
               type=str,
@@ -69,6 +78,7 @@ def bioreactor_cli(compounds,
                    neutralize,
                    reaction_rules,
                    organisms,
+                   radius,
                    patterns_to_remove,
                    molecules_to_remove,
                    min_atom_count,
@@ -89,6 +99,7 @@ def bioreactor_cli(compounds,
                     reaction_rules_path=reaction_rules,
                     neutralize_compounds=neutralize,
                     organisms_path=organisms,
+                    radius=radius,
                     patterns_to_remove_path=patterns_to_remove,
                     molecules_to_remove_path=molecules_to_remove,
                     min_atom_count=min_atom_count,

src/biocatalyzer/clis/cli_matcher.py

@@ -36,7 +36,8 @@ def matcher_cli(ms_data,
                 compounds_to_match,
                 output_path,
                 tolerance,
-                n_jobs):
+                n_jobs,
+                ):
     """Run the MSDataMatcher.
 
     Mandatory arguments:
@@ -51,7 +52,8 @@ def matcher_cli(ms_data,
                        compounds_to_match_path=compounds_to_match,
                        output_path=output_path,
                        tolerance=tolerance,
-                       n_jobs=n_jobs)
+                       n_jobs=n_jobs,
+                      )
     logging.basicConfig(filename=f'{output_path}_logging.log', level=logging.DEBUG)
     ms.generate_ms_results()
 

src/biocatalyzer/io_utils/loaders.py

@@ -50,7 +50,9 @@ def load_compounds(path: str, neutralize: bool = False):
             raise FileNotFoundError(f"File {path} not found.")
 
     @staticmethod
-    def load_reaction_rules(path: str, orgs: Union[str, List[str]] = 'ALL') -> pd.DataFrame:
+    def load_reaction_rules(path: str,
+                            orgs: Union[str, List[str]] = 'ALL',
+                            radius: Union[str, int, List[int]] = 6) -> pd.DataFrame:
         """
         Load the reaction rules to use.
 
@@ -60,6 +62,13 @@ def load_reaction_rules(path: str, orgs: Union[str, List[str]] = 'ALL') -> pd.Da
             Path to the reaction rules.
         orgs: Union[list, str]
             List of organisms to use. If 'ALL', all organisms will be used.
+        radius: Union[str, int, List[int]]
+            Reaction rule radius (or radii) to keep. If 'ALL', no radius filter is applied.
+            Accepts an integer (6), a ;-separated list ('4;6;8') or a range ('4:8').
+            A rule is kept when any of these radii appears in its 'Radii' field, the
+            same membership test already used for 'Organisms'. Only applied when the
+            reaction rules file provides a 'Radii' column; rule sets without it
+            (e.g. the bundled ones) are left untouched.
 
         Returns
         -------
@@ -95,8 +104,49 @@ def match_org(value, orgs_list):
             rules['has_org'] = rules.apply(lambda x: match_org(x['Organisms'], orgs), axis=1)
             rules = rules[rules['has_org']]
             rules.drop('has_org', axis=1, inplace=True)
+
+        if not (isinstance(radius, str) and radius == 'ALL'):
+            if 'Radii' not in rules.columns:
+                logging.warning(f"The radius filter (in effect: {radius}) was not applied: this "
+                                f"reaction rules file declares no 'Radii' column. All "
+                                f"{len(rules)} rules were kept. Pass radius='ALL' to silence this.")
+            else:
+                radii = Loaders._parse_radius(radius)
+
+                def match_radius(value, radii_list):
+                    if isinstance(value, str):
+                        return any(int(r) in radii_list for r in value.split(',') if r != '')
+                    return False
+
+                rules = rules[rules['Radii'].apply(lambda v: match_radius(v, radii))]
+                logging.info(f'Using {len(rules)} reaction rules modelled at radius in {sorted(radii)}.')
         return rules
 
+    @staticmethod
+    def _parse_radius(radius: Union[str, int, List[int]]) -> List[int]:
+        """
+        Parse the radius specification into an explicit list of radii.
+
+        Parameters
+        ----------
+        radius: Union[str, int, List[int]]
+            An integer (6), a ;-separated list ('4;6;8') or an inclusive range ('4:8').
+
+        Returns
+        -------
+        List[int]:
+            The radii to keep.
+        """
+        if isinstance(radius, int):
+            return [radius]
+        if isinstance(radius, (list, tuple)):
+            return [int(r) for r in radius]
+        radius = str(radius).strip()
+        if ':' in radius:
+            start, end = radius.split(':')
+            return list(range(int(start), int(end) + 1))
+        return [int(r) for r in radius.split(';') if r != '']
+
     @staticmethod
     def load_organisms(path: str) -> Union[str, List[str]]:
         """