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/

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 'Radii' 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.
@@ -515,22 +551,6 @@ def _match_conditions(self, smiles: str):
                 return False
         return True
 
-    def _get_ec_numbers(self, reaction_rule_id: str):
-        """
-        Get the EC numbers associated with a reaction rule.
-
-        Parameters
-        ----------
-        reaction_rule_id: str
-            The reaction rule id.
-
-        Returns
-        -------
-        str
-            The EC numbers associated with the reaction rule.
-        """
-        return self._reaction_rules[self._reaction_rules.InternalID == reaction_rule_id].EC_Numbers.values[0]
-
     def process_results(self, save: bool = True, overwrite: bool = True):
         """
         Process the results of the reactor.
@@ -588,13 +608,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 +627,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 'Radii' 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,
@@ -102,7 +112,8 @@ def biocatalyzer_cli(compounds,
 
         output_path: Path to the output directory.
     """
-    logging.basicConfig(filename=f'{output_path}logging.log', level=logging.DEBUG)
+    Path(output_path).mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(filename=Path(output_path) / 'logging.log', level=logging.DEBUG)
     if reaction_rules is None:
         logging.info(f"Using default reaction rules file.")
         reaction_rules = DATA_FILES / "../data/reactionrules/reaction_rules_biocatalyzer.tsv.bz2"
@@ -112,6 +123,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 'Radii' 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,
@@ -81,6 +91,9 @@ def bioreactor_cli(compounds,
 
         output_path: Path to the output directory.
     """
+    # Configure logging before anything logs: a prior logging call makes basicConfig a no-op.
+    Path(output_path).mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(filename=Path(output_path) / 'logging.log', level=logging.DEBUG)
     if reaction_rules is None:
         reaction_rules = DATA_FILES / "../data/reactionrules/reaction_rules_biocatalyzer.tsv.bz2"
         reaction_rules = reaction_rules.resolve()
@@ -89,11 +102,11 @@ 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,
                     n_jobs=n_jobs)
-    logging.basicConfig(filename=f'{output_path}_logging.log', level=logging.DEBUG)
     br.react()
 
 

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
+            are left untouched.
 
         Returns
         -------
@@ -70,9 +79,9 @@ def load_reaction_rules(path: str, orgs: Union[str, List[str]] = 'ALL') -> pd.Da
             raise FileNotFoundError(f"File {path} not found.")
         path = Path(path)
         if path.suffix == '.bz2':
-            rules = pd.read_csv(path, header=0, sep='\t', compression='bz2')
+            rules = pd.read_csv(path, header=0, sep='\t', compression='bz2', dtype={'Radii': str})
         else:
-            rules = pd.read_csv(path, header=0, sep='\t')
+            rules = pd.read_csv(path, header=0, sep='\t', dtype={'Radii': str})
         if 'InternalID' not in rules.columns:
             raise ValueError('The reaction rules file must contain a column named "InternalID".')
         if 'Reactants' not in rules.columns:
@@ -95,8 +104,59 @@ 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.strip().upper() == '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))]
+                if len(rules) == 0:
+                    logging.warning(f'No reaction rules are modelled at radius in {sorted(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]
+        spec = str(radius).strip()
+        try:
+            if ':' in spec:
+                start, end = (int(r) for r in spec.split(':'))
+                radii = list(range(start, end + 1))
+            else:
+                radii = [int(r) for r in spec.split(';') if r.strip() != '']
+        except ValueError:
+            radii = []
+        if not radii:
+            raise ValueError(f"Invalid radius '{radius}': expected 'ALL', an integer (6), "
+                             f"a ;-separated list ('4;6;8') or an ascending range ('4:8').")
+        return radii
+
     @staticmethod
     def load_organisms(path: str) -> Union[str, List[str]]:
         """