Fix max_transforms in mcts

MolecularAI · Dec 5, 2023 · 6e9e809 · 6e9e809
1 parent 801c55d
commit 6e9e809
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/aizynthfinder/search/mcts/state.py b/aizynthfinder/search/mcts/state.py
@@ -51,7 +51,7 @@ def __init__(self, mols: Sequence[TreeMolecule], config: Configuration) -> None:
         self.is_solved = all(self.in_stock_list)
         self.max_transforms = max(mol.transform for mol in self.mols)
         self.is_terminal = (
-            self.max_transforms > config.search.max_transforms
+            self.max_transforms >= config.search.max_transforms
         ) or self.is_solved
 
         inchis = [mol.inchi_key for mol in self.mols]

diff --git a/tests/test_finder.py b/tests/test_finder.py
@@ -227,7 +227,7 @@ def test_three_expansions_not_solved(setup_aizynthfinder):
     }
     finder = setup_aizynthfinder(lookup, [child1_smi[0], child1_smi[2], child2_smi[0]])
     finder.config.search.return_first = True
-    finder.config.search.max_transforms = 2
+    finder.config.search.max_transforms = 3
     finder.config.search.iteration_limit = 15
 
     finder.tree_search()
@@ -569,7 +569,7 @@ def test_two_expansions_prune_cyclic(setup_aizynthfinder):
     finder.tree_search()
 
     nodes = list(finder.tree.graph())
-    assert len(nodes) == 8
+    assert len(nodes) == 7
     assert state_smiles(nodes[0].state) == [root_smi]
     assert state_smiles(nodes[1].state) == child1_smi
     assert state_smiles(nodes[2].state) == [root_smi]