fix typing errors and add ignore stock to retro even if in stock

Author: bruceunx

aizynthfinder/aizynthfinder.py

@@ -1,4 +1,4 @@
-from typing import Optional, Dict
+from typing import Optional, Dict, Any
 import time
 
 from .analysis.routes import RouteCollection
@@ -22,6 +22,8 @@ def __init__(self, configfile: str) -> None:
         self.scorers = self.config.scorers
 
         self._target_mol: Optional[Molecule] = None
+        self.analysis: TreeAnalysis | None = None
+        self.tree: MctsSearchTree | None = None
 
     @property
     def target_smiles(self) -> str:
@@ -43,6 +45,8 @@ def target_mol(self, mol: Molecule) -> None:
         self._target_mol = mol
 
     def build_routes(self, scorer: str = "state score") -> None:
+        if self.tree is None:
+            return None
         self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
         config_selection = RouteSelectionArguments(
             nmin=self.config.post_processing.min_routes,
@@ -53,7 +57,7 @@ def build_routes(self, scorer: str = "state score") -> None:
                                                     config_selection)
 
     def extract_statistics(self) -> Dict:
-        if not self.analysis:
+        if self.analysis is None:
             return {}
         stats = {
             "target":
@@ -81,19 +85,21 @@ def prepare_tree(self) -> None:
 
     def tree_search(self) -> None:
         self.prepare_tree()
-        self.search_stats = {"returned_first": False, "iterations": 0}
+        assert self.tree is not None
+        self.search_stats: Dict[str, Any] = {
+            "returned_first": False,
+            "iterations": 0
+        }
 
         time0 = time.time()
         time_past = .0
         i = 1
         while time_past < self.config.time_limit and i <= self.config.iteration_limit:
             self.search_stats["iterations"] += 1
             is_solved = self.tree.one_iteration()
-
             if is_solved and "first_solution_time" not in self.search_stats:
                 self.search_stats["first_solution_time"] = time.time() - time0
                 self.search_stats["first_solution_iteration"] = i
-
             if self.config.return_first and is_solved:
                 self.search_stats["returned_first"] = True
                 break

aizynthfinder/analysis/tree_analysis.py

@@ -28,7 +28,8 @@ def sort(
         sorted_items, sorted_scores, _ = self.scorer.sort(nodes)
         actions = [node.actions_to() for node in sorted_items]
 
-        return self._collect_top_items(sorted_items, sorted_scores, actions, selection)
+        return self._collect_top_items(sorted_items, sorted_scores, actions,
+                                       selection)  # type: ignore
 
     def tree_statistics(self) -> Dict:
         return self._tree_statistics_mcts()
@@ -47,14 +48,17 @@ def _tree_statistics_mcts(self) -> Dict:
         top_state = top_node.state
         nodes = list(self.search_tree.graph())
         mols_in_stock = ", ".join(
-            mol.smiles for mol, instock in zip(top_state.mols, top_state.in_stock_list)
+            mol.smiles
+            for mol, instock in zip(top_state.mols, top_state.in_stock_list)
             if instock)
         mols_not_in_stock = ", ".join(
-            mol.smiles for mol, instock in zip(top_state.mols, top_state.in_stock_list)
+            mol.smiles
+            for mol, instock in zip(top_state.mols, top_state.in_stock_list)
             if not instock)
 
-        policy_used_counts = self._policy_used_statistics(
-            [node[child]["action"] for node in nodes for child in node.children])
+        policy_used_counts = self._policy_used_statistics([
+            node[child]["action"] for node in nodes for child in node.children
+        ])
 
         return {
             "number_of_nodes":
@@ -66,7 +70,8 @@ def _tree_statistics_mcts(self) -> Dict:
             "number_of_routes":
             sum(1 for node in nodes if not node.children),
             "number_of_solved_routes":
-            sum(1 for node in nodes if not node.children and node.state.is_solved),
+            sum(1 for node in nodes
+                if not node.children and node.state.is_solved),
             "top_score":
             self.scorer(top_node),
             "is_solved":
@@ -96,7 +101,8 @@ def _collect_top_items(
         reactions: Sequence[Union[Iterable[RetroReaction],
                                   Iterable[FixedRetroReaction]]],
         selection,
-    ) -> Tuple[Union[Sequence[MctsNode], Sequence[ReactionTree]], Sequence[float]]:
+    ) -> Tuple[Union[Sequence[MctsNode], Sequence[ReactionTree]],
+               Sequence[float]]:
         if len(items) <= selection.nmin:
             return items, scores
 
@@ -130,7 +136,8 @@ def _collect_top_items(
 
     @staticmethod
     def _policy_used_statistics(
-            reactions: Iterable[Union[RetroReaction, FixedRetroReaction]]) -> Dict:
+            reactions: Iterable[Union[RetroReaction,
+                                      FixedRetroReaction]]) -> Dict:
         policy_used_counts: Dict = defaultdict(int)
         for reaction in reactions:
             policy_used = reaction.metadata.get("policy_name")

aizynthfinder/context/config.py

@@ -10,32 +10,33 @@
 
 @dataclass
 class _PostprocessingConfiguration:
-    min_routes: int = 1
-    max_routes: int = 8
-    all_routes: bool = False
+    min_routes = 1
+    max_routes = 8
+    all_routes = False
     route_distance_model: Optional[str] = None
 
 
 @dataclass
 class Configuration:
 
-    C: float = 1.4
-    cutoff_cumulative: float = 0.995
-    cutoff_number: int = 50
-    additive_expansion: bool = False
-    use_rdchiral: bool = True
-    max_transforms: int = 6
-    default_prior: float = 0.5
-    use_prior: bool = True
-    iteration_limit: int = 100
-    return_first: bool = False
-    time_limit: int = 20
-    filter_cutoff: float = 0.05
-    exclude_target_from_stock: bool = True
-    template_column: str = "retro_template"
-    prune_cycles_in_search: bool = True
-    search_algorithm: str = "mcts"
+    C = 1.4
+    cutoff_cumulative = 0.995
+    cutoff_number = 50
+    additive_expansion = False
+    use_rdchiral = True
+    max_transforms = 6
+    default_prior = 0.5
+    use_prior = True
+    iteration_limit = 100
+    return_first = False
+    time_limit = 20
+    filter_cutoff = 0.05
+    exclude_target_from_stock = True
+    template_column = "retro_template"
+    prune_cycles_in_search = True
+    search_algorithm = "mcts"
     post_processing = _PostprocessingConfiguration()
+    ignore_stock = False
 
     def __post_init__(self) -> None:
         self._properties: Dict = {}
@@ -54,7 +55,8 @@ def from_dict(cls, source: Dict) -> "Configuration":
         config_obj = Configuration()
         config_obj._update_from_config(source)
 
-        config_obj.expansion_policy.load_from_config(**source.get("policy", {}))
+        config_obj.expansion_policy.load_from_config(
+            **source.get("policy", {}))
         config_obj.filter_policy.load_from_config(**source.get("filter", {}))
         config_obj.stock.load_from_config(**source.get("stock", {}))
 

aizynthfinder/mcts/state.py

@@ -12,9 +12,13 @@ class MctsState:
     def __init__(self, mols: Sequence[TreeMolecule], config: Any) -> None:
         self.mols = mols
         self.stock = config.stock
-        self.in_stock_list = [mol in self.stock for mol in self.mols]
+        if config.ignore_stock:
+            self.in_stock_list = [False for mol in self.mols]
+        else:
+            self.in_stock_list = [mol in self.stock for mol in self.mols]
         self.expandable_mols = [
-            mol for mol, in_stock in zip(self.mols, self.in_stock_list) if not in_stock
+            mol for mol, in_stock in zip(self.mols, self.in_stock_list)
+            if not in_stock
         ]
         self._stock_availability: Optional[List[str]] = None
         self.is_solved = all(self.in_stock_list)
@@ -39,7 +43,9 @@ def __str__(self) -> str:
         string = "%s\n%s\n%s\n%s\nScore: %0.3F Solved: %s" % (
             str([mol.smiles for mol in self.mols]),
             str([mol.transform for mol in self.mols]),
-            str([mol.parent.smiles if mol.parent else "-" for mol in self.mols]),
+            str([
+                mol.parent.smiles if mol.parent else "-" for mol in self.mols
+            ]),
             str(self.in_stock_list),
             self.score,
             self.is_solved,

usage.py

@@ -8,18 +8,20 @@
 finder.expansion_policy.select("uspto")
 finder.filter_policy.select("uspto")
 
-finder.target_smiles = "Cc1cccc(c1N(CC(=O)Nc2ccc(cc2)c3ncon3)C(=O)C4CCS(=O)(=O)CC4)C"
+# Cc1cccc(c1N(CC(=O)Nc2ccc(cc2)c3ncon3)C(=O)C4CCS(=O)(=O)CC4)C
+finder.target_smiles = "COc1cc(C=CC(=O)O)cc(OC)c1OC"
 # print(psutil.Process().memory_info().rss / (1024 * 1024), "MB")
 finder.tree_search()
 finder.build_routes()
 
-stats = finder.extract_statistics()
+# stats = finder.extract_statistics()
 # routes details
 # for d in finder.routes.dicts:
 #     print("---" * 20)
 #     print(d)
+breakpoint()
 plt.imshow(finder.routes.images[0])
 plt.show()
 
 # stats
-print(stats)
+# print(stats)