Add option to stop search after search graph exceeds a certain size (#85)

This option is called `limit_graph_nodes`. I also added a basic test for
this.

---------

Co-authored-by: Krzysztof Maziarz <krzysztof.maziarz@microsoft.com>

Author: AustinT

CHANGELOG.md

@@ -21,6 +21,7 @@ and the project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.
 - Integrate the Graph2Edits model ([#65](https://github.com/microsoft/syntheseus/pull/65), [#66](https://github.com/microsoft/syntheseus/pull/66)) ([@kmaziarz])
 - Improve the docs and add tutorials ([#54](https://github.com/microsoft/syntheseus/pull/54), [#77](https://github.com/microsoft/syntheseus/pull/77), [#78](https://github.com/microsoft/syntheseus/pull/78), [#79](https://github.com/microsoft/syntheseus/pull/79), [#82](https://github.com/microsoft/syntheseus/pull/82)) ([@kmaziarz], [@austint])
 - Add random search algorithm as a simple baseline ([#83](https://github.com/microsoft/syntheseus/pull/83)) ([@austint])
+- Add optional argument `limit_graph_nodes` to base search algorithm class to stop search after the search graph exceeds a certain number of nodes ([#85](https://github.com/microsoft/syntheseus/pull/85)) ([@austint])
 
 ### Fixed
 

syntheseus/cli/search.py

@@ -34,6 +34,7 @@
 from syntheseus.reaction_prediction.utils.config import get_config as cli_get_config
 from syntheseus.reaction_prediction.utils.misc import set_random_seed
 from syntheseus.reaction_prediction.utils.model_loading import get_model
+from syntheseus.search import INT_INF
 from syntheseus.search.algorithms.best_first.retro_star import RetroStarSearch
 from syntheseus.search.algorithms.mcts import base as mcts_base
 from syntheseus.search.algorithms.mcts.molset import MolSetMCTS
@@ -114,6 +115,7 @@ class SearchConfig(BackwardModelConfig):
     time_limit_s: float = 600
     limit_reaction_model_calls: int = 1_000_000
     limit_iterations: int = 1_000_000
+    limit_graph_nodes: int = INT_INF
     prevent_repeat_mol_in_trees: bool = True
 
     use_gpu: bool = True  # Whether to use a GPU
@@ -186,6 +188,7 @@ def run_from_config(config: SearchConfig) -> Path:
                 "time_limit_s",
                 "limit_reaction_model_calls",
                 "limit_iterations",
+                "limit_graph_nodes",
                 "prevent_repeat_mol_in_trees",
             ]
         }

syntheseus/search/algorithms/base.py

@@ -70,6 +70,7 @@ def __init__(
         self,
         limit_iterations: int = INT_INF,
         limit_reaction_model_calls: int = INT_INF,
+        limit_graph_nodes: int = INT_INF,
         time_limit_s: float = math.inf,
         max_expansion_depth: int = 50,
         expand_purchasable_mols: bool = False,
@@ -85,6 +86,7 @@ def __init__(
 
         self.limit_iterations = limit_iterations
         self.limit_reaction_model_calls = limit_reaction_model_calls
+        self.limit_graph_nodes = limit_graph_nodes
         self.time_limit_s = time_limit_s
         self.max_expansion_depth = max_expansion_depth
         self.expand_purchasable_mols = expand_purchasable_mols
@@ -176,6 +178,7 @@ def should_stop_search(self, graph) -> bool:
         return (
             (elapsed_time >= self.time_limit_s)
             or (self.reaction_model.num_calls() >= self.limit_reaction_model_calls)
+            or (len(graph) >= self.limit_graph_nodes)
             or (self.stop_on_first_solution and graph.root_node.has_solution)
         )
 

syntheseus/search/algorithms/mcts/base.py

@@ -137,15 +137,18 @@ def _check_infinite_runtime(self) -> None:
         no_limit_iter = self.limit_iterations >= INT_INF
         no_limit_rxn = self.limit_reaction_model_calls >= INT_INF
         no_limit_time = self.time_limit_s >= math.inf
+        no_limit_nodes = self.limit_graph_nodes >= INT_INF
 
-        if no_limit_iter and no_limit_rxn and no_limit_time:
+        if no_limit_iter and no_limit_rxn and no_limit_time and no_limit_nodes:
             warnings.warn(
                 "No kind of run limit set. This algorithm will almost certainty run forever."
             )
         elif no_limit_iter and no_limit_time:
             warnings.warn(
-                "No iteration or time limit set. It is possible (but not certain) that MCTS "
-                "will run forever with these settings. At the very least, it could run for a very long time."
+                "No iteration or time limit set (although a reaction model call and/or graph node limit was set)."
+                " Under these conditions, it is possible (but not certain) that MCTS "
+                "will run forever (for example if there are no leaf nodes eligible for expansion in the graph)."
+                " At the very least, it could run for an unexpectedly long time."
                 " It is recommended to set either an iteration limit or a time limit."
             )
 

syntheseus/tests/search/algorithms/test_base.py

@@ -157,6 +157,45 @@ def test_limit_reaction_model_calls(
             # Default case is to match exactly
             assert alg.reaction_model.num_calls() == limit
 
+    @pytest.mark.parametrize("limit", [0, 1, 2, 25, 100])
+    def test_limit_graph_nodes(
+        self,
+        retrosynthesis_task6: RetrosynthesisTask,
+        limit: int,
+    ) -> None:
+        """
+        Test that limiting the number of nodes in the graph works as intended.
+        The algorithm should run until the node limit is reached, and then stop
+        (without adding *too many* extra nodes).
+
+        `retrosynthesis_task6` is chosen because it can create a very large graph.
+        """
+
+        # Run algorithm
+        alg = self.setup_algorithm(
+            reaction_model=retrosynthesis_task6.reaction_model,
+            mol_inventory=retrosynthesis_task6.inventory,
+            limit_graph_nodes=limit,
+            limit_iterations=int(1e6),  # a very high limit, but avoids MCTS warnings
+        )
+        output_graph, _ = alg.run_from_mol(retrosynthesis_task6.target_mol)
+
+        # The algorithm will stop running when the graph size meets or exceeds the limit.
+        # However, since multiple nodes are added during each expansion, the node count might
+        # not exactly equal the limit. Therefore, we choose a variable tolerance.
+        # "Tolerance" here is len(graph) - limit
+        if limit == 0:
+            tolerance = 1  # will stop search immediately with only the root node
+        elif limit == 1:
+            tolerance = 0  # should not expand root node
+        elif limit == 2:
+            tolerance = 19  # a very high number, since first expansion brings node count to 21 for AND/OR graphs
+        else:
+            tolerance = 20  # a fairly high tolerance (should always be enough for one expansion)
+
+        # The actual test
+        assert limit <= len(output_graph) <= tolerance + limit
+
     @pytest.mark.parametrize("limit", [0, 1, 2, 100])
     def test_limit_iterations(
         self,