def test_find_repetetive_patterns_created_tree(default_config, mock_stock,
shared_datadir):
mock_stock(default_config, Molecule(smiles="CC"), Molecule(smiles="C"))
# Try one with 2 repetetive units
search_tree = SearchTree.from_json(
shared_datadir / "tree_with_repetition.json", default_config)
analysis = TreeAnalysis(search_tree)
rt = ReactionTree.from_analysis(analysis)
assert rt.has_repeating_patterns
hidden_nodes = [
node for node in rt.graph if rt.graph.nodes[node].get("hide", False)
]
assert len(hidden_nodes) == 5
# Try one with 3 repetetive units
search_tree = SearchTree.from_json(
shared_datadir / "tree_with_3_repetitions.json", default_config)
analysis = TreeAnalysis(search_tree)
rt = ReactionTree.from_analysis(analysis)
assert rt.has_repeating_patterns
hidden_nodes = [
node for node in rt.graph if rt.graph.nodes[node].get("hide", False)
]
assert len(hidden_nodes) == 10
def test_find_repetetive_patterns_created_tree_no_patterns(
default_config, mock_stock, shared_datadir):
mock_stock(default_config, Molecule(smiles="CC"), Molecule(smiles="CCCO"))
# Try with a short tree (3 nodes, 1 reaction)
search_tree = SearchTree.from_json(
shared_datadir / "tree_without_repetition.json", default_config)
analysis = TreeAnalysis(search_tree)
rt = ReactionTree.from_analysis(analysis)
assert not rt.has_repeating_patterns
hidden_nodes = [
node for node in rt.graph if rt.graph.nodes[node].get("hide", False)
]
assert len(hidden_nodes) == 0
# Try with something longer
search_tree = SearchTree.from_json(
shared_datadir / "tree_without_repetition_longer.json", default_config)
analysis = TreeAnalysis(search_tree)
rt = ReactionTree.from_analysis(analysis)
assert not rt.has_repeating_patterns
def build_routes(self, min_nodes=5):
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param min_nodes: the minimum number of top-ranked nodes to consider, defaults to 5
:type min_nodes: int, optional
"""
self.analysis = TreeAnalysis(self.tree)
self.routes = RouteCollection.from_analysis(self.analysis, min_nodes)
def test_sort_nodes(setup_complete_tree):
tree, nodes = setup_complete_tree
analysis = TreeAnalysis(tree)
best_nodes = analysis.sort_nodes()
assert len(best_nodes) == 3
assert best_nodes[0].state.score == 0.99
assert best_nodes[0] is nodes[2]
best_nodes = analysis.sort_nodes(min_return=0)
assert len(best_nodes) == 0
def build_routes(self, min_nodes=5, scorer="state score"):
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param min_nodes: the minimum number of top-ranked nodes to consider, defaults to 5
:type min_nodes: int, optional
:param scorer: the object used to score the nodes
:type scorer: str, optional
"""
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, min_nodes)
def test_create_combine_tree_dict_from_tree(mock_stock, default_config,
load_reaction_tree,
shared_datadir):
mock_stock(
default_config,
"Nc1ccc(NC(=S)Nc2ccccc2)cc1",
"Cc1ccc2nc3ccccc3c(Cl)c2c1",
"Nc1ccc(N)cc1",
"S=C=Nc1ccccc1",
"Cc1ccc2nc3ccccc3c(N)c2c1",
"Nc1ccc(Br)cc1",
)
search_tree = SearchTree.from_json(
shared_datadir / "tree_for_clustering.json", default_config)
analysis = TreeAnalysis(search_tree)
collection = RouteCollection.from_analysis(analysis, 3)
expected = load_reaction_tree("combined_example_tree.json")
combined_dict = collection.combined_reaction_trees().to_dict()
assert len(combined_dict["children"]) == 2
assert combined_dict["children"][0]["is_reaction"]
assert len(combined_dict["children"][0]["children"]) == 2
assert len(combined_dict["children"][1]["children"]) == 2
assert len(combined_dict["children"][1]["children"][0]["children"]) == 2
assert combined_dict["children"][1]["children"][0]["children"][0][
"is_reaction"]
assert combined_dict == expected
def build_routes(self, min_nodes: int = 5, scorer: str = "state score") -> None:
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param min_nodes: the minimum number of top-ranked nodes to consider, defaults to 5
:param scorer: a reference to the object used to score the nodes
:raises ValueError: if the search tree not initialized
"""
if not self.tree:
raise ValueError("Search tree not initialized")
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, min_nodes)
def test_reactiontree_to_json(setup_complete_tree, shared_datadir):
filename = str(shared_datadir / "sample_reaction.json")
with open(filename, "r") as fileobj:
expected = json.load(fileobj)
tree, nodes = setup_complete_tree
analysis = TreeAnalysis(tree)
resp = ReactionTree.from_analysis(analysis).to_json()
assert json.loads(resp) == expected
def build_routes(self,
selection: RouteSelectionArguments = None,
scorer: str = "state score") -> None:
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param selection: the selection criteria for the routes
:param scorer: a reference to the object used to score the nodes
:raises ValueError: if the search tree not initialized
"""
if not self.tree:
raise ValueError("Search tree not initialized")
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, selection)
def test_route_node_depth_from_analysis(default_config, mock_stock,
shared_datadir):
mock_stock(default_config, Molecule(smiles="CC"), Molecule(smiles="CCCO"))
search_tree = SearchTree.from_json(
shared_datadir / "tree_without_repetition.json", default_config)
analysis = TreeAnalysis(search_tree)
rt = ReactionTree.from_analysis(analysis)
mols = list(rt.molecules())
assert rt.depth(mols[0]) == 0
assert rt.depth(mols[1]) == 2
assert rt.depth(mols[2]) == 2
rxns = list(rt.reactions())
assert rt.depth(rxns[0]) == 1
for mol in rt.molecules():
assert rt.depth(mol) == 2 * rt.graph.nodes[mol]["transform"]
def test_route_to_reactiontree(setup_complete_tree):
tree, nodes = setup_complete_tree
analysis = TreeAnalysis(tree)
reaction_tree = ReactionTree.from_analysis(analysis).graph
reaction_nodes = [
node.inchi_key for node in reaction_tree if isinstance(node, Molecule)
]
tree_nodes = [
mol.inchi_key for mol in nodes[0].state.mols + nodes[1].state.mols +
nodes[2].state.mols
]
assert len(reaction_nodes) == 6
assert reaction_nodes == tree_nodes
reaction_nodes = [
node for node in reaction_tree if isinstance(node, Reaction)
]
assert len(reaction_nodes) == 2
def test_create_route_collection(setup_complete_tree, mocker):
tree, nodes = setup_complete_tree
analysis = TreeAnalysis(tree)
mocker.patch("aizynthfinder.analysis.ReactionTree.to_dict")
mocker.patch("aizynthfinder.analysis.json.dumps")
routes = RouteCollection.from_analysis(analysis, 5)
assert len(routes) == 3
assert routes[0]["score"] == 0.99
assert routes[0]["node"] is nodes[2]
reaction_nodes = [
node for node in routes[0]["reaction_tree"].graph
if isinstance(node, Reaction)
]
assert len(reaction_nodes) == 2
# Just see that the code does not crash, does not verify content
assert len(routes.images) == 3
assert len(routes.dicts) == 3
assert len(routes.jsons) == 3
def from_analysis(
cls,
analysis: TreeAnalysis,
selection: RouteSelectionArguments = None) -> "RouteCollection":
"""
Create a collection from a tree analysis.
:param analysis: the tree analysis to use
:param selection: selection criteria for the routes
:return: the created collection
"""
items, scores = analysis.sort(selection)
all_scores = [{repr(analysis.scorer): score} for score in scores]
kwargs = {"scores": scores, "all_scores": all_scores}
if isinstance(analysis.search_tree, MctsSearchTree):
kwargs["nodes"] = items
reaction_trees = [
from_node.to_reaction_tree() for from_node in items
if isinstance(from_node, MctsNode)
]
else:
reaction_trees = items # type: ignore
return cls(reaction_trees, **kwargs)
class AiZynthFinder:
"""
Public API to the aizynthfinder tool
If intantiated with the path to a yaml file or dictionary of settings
the stocks and policy networks are loaded directly.
Otherwise, the user is responsible for loading them prior to
executing the tree search.
:ivar config: the configuration of the search
:vartype config: Configuration
:ivar policy: the policy model
:vartype policy: Policy
:ivar stock: the stock
:vartype stock: Stock
:ivar tree: the search tree
:vartype tree: SearchTree
:ivar analysis: the tree analysis
:vartype analysis: TreeAnalysis
:ivar routes: the top-ranked routes
:vartype routes: RouteCollection
:ivar search_stats: statistics of the latest search: time, number of iterations and if it returned first solution
:vartype search_stats: dict
:param configfile: the path to yaml file with configuration (has priority over configdict), defaults to None
:type configfile: str, optional
:param configdict: the config as a dictionary source, defaults to None
:type configdict: dict, optional
"""
def __init__(self, configfile=None, configdict=None):
self._logger = logger()
if configfile:
self.config = Configuration.from_file(configfile)
elif configdict:
self.config = Configuration.from_dict(configdict)
else:
self.config = Configuration()
self.expansion_policy = self.config.expansion_policy
self.filter_policy = self.config.filter_policy
self.stock = self.config.stock
self.scorers = self.config.scorers
self.tree = None
self._target_mol = None
self.search_stats = {}
self.routes = None
self.analysis = None
@property
def target_smiles(self):
"""
The SMILES representation of the molecule to predict routes on.
:return: the SMILES
:rvalue: str
"""
return self._target_mol.smiles
@target_smiles.setter
def target_smiles(self, smiles):
self.target_mol = Molecule(smiles=smiles)
@property
def target_mol(self):
"""
The molecule to predict routes on
:return: the molecule
:rvalue: Molecule
"""
return self._target_mol
@target_mol.setter
def target_mol(self, mol):
self.tree = None
self._target_mol = mol
def build_routes(self, min_nodes=5, scorer="state score"):
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param min_nodes: the minimum number of top-ranked nodes to consider, defaults to 5
:type min_nodes: int, optional
:param scorer: the object used to score the nodes
:type scorer: str, optional
"""
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, min_nodes)
def extract_statistics(self):
""" Extracts tree statistics as a dictionary
"""
if not self.analysis:
return {}
stats = {
"target": self.target_smiles,
"search_time": self.search_stats["time"]
}
stats.update(self.analysis.tree_statistics())
return stats
def prepare_tree(self):
""" Setup the tree for searching
"""
self.stock.reset_exclusion_list()
if self.config.exclude_target_from_stock and self.target_mol in self.stock:
self.stock.exclude(self.target_mol)
self._logger.debug("Excluding the target compound from the stock")
self._logger.debug("Defining tree root: %s" % self.target_smiles)
self.tree = SearchTree(root_smiles=self.target_smiles,
config=self.config)
self.analysis = None
self.routes = None
@deprecated(version="2.1.0", reason="Not supported anymore")
def run_from_json(self, params):
"""
Run a search tree by reading settings from a JSON
:param params: the parameters of the tree search
:type params: dict
:return: dictionary with all settings and top scored routes
:rtype: dict
"""
self.stock.select(params["stocks"])
self.expansion_policy.select(
params.get("policy", params.get("policies", "")))
if "filter" in params:
self.filter_policy.select(params["filter"])
else:
self.filter_policy.deselect()
self.config.C = params["C"]
self.config.max_transforms = params["max_transforms"]
self.config.cutoff_cumulative = params["cutoff_cumulative"]
self.config.cutoff_number = params["cutoff_number"]
self.target_smiles = params["smiles"]
self.config.return_first = params["return_first"]
self.config.time_limit = params["time_limit"]
self.config.iteration_limit = params["iteration_limit"]
self.config.exclude_target_from_stock = params[
"exclude_target_from_stock"]
self.config.filter_cutoff = params["filter_cutoff"]
self.prepare_tree()
self.tree_search()
self.build_routes()
if not params.get("score_trees", False):
return {
"request": self._get_settings(),
"trees": self.routes.dicts,
}
self.routes.compute_scores(*self.scorers.objects())
return {
"request": self._get_settings(),
"trees": self.routes.dict_with_scores(),
}
def tree_search(self, show_progress=False):
"""
Perform the actual tree search
:param show_progress: if True, shows a progress bar
:type show_progress: bool
:return: the time past in seconds
:rtype: float
"""
if not self.tree:
self.prepare_tree()
self.search_stats = {"returned_first": False, "iterations": 0}
time0 = time.time()
i = 1
self._logger.debug("Starting search")
time_past = time.time() - time0
if show_progress:
pbar = tqdm(total=self.config.iteration_limit)
while time_past < self.config.time_limit and i <= self.config.iteration_limit:
if show_progress:
pbar.update(1)
self.search_stats["iterations"] += 1
leaf = self.tree.select_leaf()
leaf.expand()
while not leaf.is_terminal():
child = leaf.promising_child()
if child:
child.expand()
leaf = child
self.tree.backpropagate(leaf, leaf.state.score)
if self.config.return_first and leaf.state.is_solved:
self._logger.debug("Found first solved route")
self.search_stats["returned_first"] = True
break
i = i + 1
time_past = time.time() - time0
if show_progress:
pbar.close()
self._logger.debug("Search completed")
self.search_stats["time"] = time_past
return time_past
def _get_settings(self):
"""Get the current settings as a dictionary
"""
# To be backward-compatible
if len(self.expansion_policy.selection) == 1:
policy_value = self.expansion_policy.selection[0]
policy_key = "policy"
else:
policy_value = self.expansion_policy.selection
policy_key = "policies"
dict_ = {
"stocks": self.stock.selection,
policy_key: policy_value,
"C": self.config.C,
"max_transforms": self.config.max_transforms,
"cutoff_cumulative": self.config.cutoff_cumulative,
"cutoff_number": self.config.cutoff_number,
"smiles": self.target_smiles,
"return_first": self.config.return_first,
"time_limit": self.config.time_limit,
"iteration_limit": self.config.iteration_limit,
"exclude_target_from_stock": self.config.exclude_target_from_stock,
"filter_cutoff": self.config.filter_cutoff,
}
if self.filter_policy.selection:
dict_["filter"] = self.filter_policy.selection
return dict_
def wrapper(scorer=None):
return TreeAnalysis(tree, scorer=scorer)
class AiZynthFinder:
"""
Public API to the aizynthfinder tool
If instantiated with the path to a yaml file or dictionary of settings
the stocks and policy networks are loaded directly.
Otherwise, the user is responsible for loading them prior to
executing the tree search.
:ivar config: the configuration of the search
:ivar expansion_policy: the expansion policy model
:ivar filter_policy: the filter policy model
:ivar stock: the stock
:ivar scorers: the loaded scores
:ivar tree: the search tree
:ivar analysis: the tree analysis
:ivar routes: the top-ranked routes
:ivar search_stats: statistics of the latest search
:param configfile: the path to yaml file with configuration (has priority over configdict), defaults to None
:param configdict: the config as a dictionary source, defaults to None
"""
def __init__(self,
configfile: str = None,
configdict: StrDict = None) -> None:
self._logger = logger()
if configfile:
self.config = Configuration.from_file(configfile)
elif configdict:
self.config = Configuration.from_dict(configdict)
else:
self.config = Configuration()
self.expansion_policy = self.config.expansion_policy
self.filter_policy = self.config.filter_policy
self.stock = self.config.stock
self.scorers = self.config.scorers
self.tree: Optional[Union[MctsSearchTree, AndOrSearchTreeBase]] = None
self._target_mol: Optional[Molecule] = None
self.search_stats: StrDict = dict()
self.routes = RouteCollection([])
self.analysis: Optional[TreeAnalysis] = None
@property
def target_smiles(self) -> str:
"""The SMILES representation of the molecule to predict routes on."""
if not self._target_mol:
return ""
return self._target_mol.smiles
@target_smiles.setter
def target_smiles(self, smiles: str) -> None:
self.target_mol = Molecule(smiles=smiles)
@property
def target_mol(self) -> Optional[Molecule]:
"""The molecule to predict routes on"""
return self._target_mol
@target_mol.setter
def target_mol(self, mol: Molecule) -> None:
self.tree = None
self._target_mol = mol
def build_routes(self,
selection: RouteSelectionArguments = None,
scorer: str = "state score") -> None:
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param selection: the selection criteria for the routes
:param scorer: a reference to the object used to score the nodes
:raises ValueError: if the search tree not initialized
"""
if not self.tree:
raise ValueError("Search tree not initialized")
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, selection)
def extract_statistics(self) -> StrDict:
"""Extracts tree statistics as a dictionary"""
if not self.analysis:
return {}
stats = {
"target":
self.target_smiles,
"search_time":
self.search_stats["time"],
"first_solution_time":
self.search_stats.get("first_solution_time", 0),
"first_solution_iteration":
self.search_stats.get("first_solution_iteration", 0),
}
stats.update(self.analysis.tree_statistics())
return stats
def prepare_tree(self) -> None:
"""
Setup the tree for searching
:raises ValueError: if the target molecule was not set
"""
if not self.target_mol:
raise ValueError("No target molecule set")
self.stock.reset_exclusion_list()
if self.config.exclude_target_from_stock and self.target_mol in self.stock:
self.stock.exclude(self.target_mol)
self._logger.debug("Excluding the target compound from the stock")
self._setup_search_tree()
self.analysis = None
self.routes = RouteCollection([])
def tree_search(self, show_progress: bool = False) -> float:
"""
Perform the actual tree search
:param show_progress: if True, shows a progress bar
:return: the time past in seconds
"""
if not self.tree:
self.prepare_tree()
# This is for type checking, prepare_tree is creating it.
assert self.tree is not None
self.search_stats = {"returned_first": False, "iterations": 0}
time0 = time.time()
i = 1
self._logger.debug("Starting search")
time_past = time.time() - time0
if show_progress:
pbar = tqdm(total=self.config.iteration_limit, leave=False)
while time_past < self.config.time_limit and i <= self.config.iteration_limit:
if show_progress:
pbar.update(1)
self.search_stats["iterations"] += 1
try:
is_solved = self.tree.one_iteration()
except StopIteration:
break
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._logger.debug("Found first solved route")
self.search_stats["returned_first"] = True
break
i = i + 1
time_past = time.time() - time0
if show_progress:
pbar.close()
time_past = time.time() - time0
self._logger.debug("Search completed")
self.search_stats["time"] = time_past
return time_past
def _setup_search_tree(self):
self._logger.debug("Defining tree root: %s" % self.target_smiles)
if self.config.search_algorithm.lower() == "mcts":
self.tree = MctsSearchTree(root_smiles=self.target_smiles,
config=self.config)
else:
cls = load_dynamic_class(self.config.search_algorithm)
self.tree: AndOrSearchTreeBase = cls(
root_smiles=self.target_smiles, config=self.config)
class AiZynthFinder:
"""
Public API to the aizynthfinder tool
If instantiated with the path to a yaml file or dictionary of settings
the stocks and policy networks are loaded directly.
Otherwise, the user is responsible for loading them prior to
executing the tree search.
:ivar config: the configuration of the search
:ivar expansion_policy: the expansion policy model
:ivar filter_policy: the filter policy model
:ivar stock: the stock
:ivar scorers: the loaded scores
:ivar tree: the search tree
:ivar analysis: the tree analysis
:ivar routes: the top-ranked routes
:ivar search_stats: statistics of the latest search
:param configfile: the path to yaml file with configuration (has priority over configdict), defaults to None
:param configdict: the config as a dictionary source, defaults to None
"""
def __init__(self,
configfile: str = None,
configdict: StrDict = None) -> None:
self._logger = logger()
if configfile:
self.config = Configuration.from_file(configfile)
elif configdict:
self.config = Configuration.from_dict(configdict)
else:
self.config = Configuration()
self.expansion_policy = self.config.expansion_policy
self.filter_policy = self.config.filter_policy
self.stock = self.config.stock
self.scorers = self.config.scorers
self.tree: Optional[Union[MctsSearchTree, AndOrSearchTreeBase]] = None
self._target_mol: Optional[Molecule] = None
self.search_stats: StrDict = dict()
self.routes = RouteCollection([])
self.analysis: Optional[TreeAnalysis] = None
@property
def target_smiles(self) -> str:
"""The SMILES representation of the molecule to predict routes on."""
if not self._target_mol:
return ""
return self._target_mol.smiles
@target_smiles.setter
def target_smiles(self, smiles: str) -> None:
self.target_mol = Molecule(smiles=smiles)
@property
def target_mol(self) -> Optional[Molecule]:
"""The molecule to predict routes on"""
return self._target_mol
@target_mol.setter
def target_mol(self, mol: Molecule) -> None:
self.tree = None
self._target_mol = mol
def build_routes(self,
min_nodes: int = 5,
scorer: str = "state score") -> None:
"""
Build reaction routes
This is necessary to call after the tree search has completed in order
to extract results from the tree search.
:param min_nodes: the minimum number of top-ranked nodes to consider, defaults to 5
:param scorer: a reference to the object used to score the nodes
:raises ValueError: if the search tree not initialized
"""
if not self.tree:
raise ValueError("Search tree not initialized")
self.analysis = TreeAnalysis(self.tree, scorer=self.scorers[scorer])
self.routes = RouteCollection.from_analysis(self.analysis, min_nodes)
def extract_statistics(self) -> StrDict:
"""Extracts tree statistics as a dictionary"""
if not self.analysis:
return {}
stats = {
"target":
self.target_smiles,
"search_time":
self.search_stats["time"],
"first_solution_time":
self.search_stats.get("first_solution_time", 0),
"first_solution_iteration":
self.search_stats.get("first_solution_iteration", 0),
}
stats.update(self.analysis.tree_statistics())
return stats
def prepare_tree(self) -> None:
"""
Setup the tree for searching
:raises ValueError: if the target molecule was not set
"""
if not self.target_mol:
raise ValueError("No target molecule set")
self.stock.reset_exclusion_list()
if self.config.exclude_target_from_stock and self.target_mol in self.stock:
self.stock.exclude(self.target_mol)
self._logger.debug("Excluding the target compound from the stock")
self._setup_search_tree()
self.analysis = None
self.routes = RouteCollection([])
@deprecated(version="2.1.0", reason="Not supported anymore")
def run_from_json(self, params: StrDict) -> StrDict:
"""
Run a search tree by reading settings from a JSON
:param params: the parameters of the tree search
:return: dictionary with all settings and top scored routes
"""
self.stock.select(params["stocks"])
self.expansion_policy.select(
params.get("policy", params.get("policies", "")))
if "filter" in params:
self.filter_policy.select(params["filter"])
else:
self.filter_policy.deselect()
self.config.C = params["C"]
self.config.max_transforms = params["max_transforms"]
self.config.cutoff_cumulative = params["cutoff_cumulative"]
self.config.cutoff_number = params["cutoff_number"]
self.target_smiles = params["smiles"]
self.config.return_first = params["return_first"]
self.config.time_limit = params["time_limit"]
self.config.iteration_limit = params["iteration_limit"]
self.config.exclude_target_from_stock = params[
"exclude_target_from_stock"]
self.config.filter_cutoff = params["filter_cutoff"]
self.prepare_tree()
self.tree_search()
self.build_routes()
if not params.get("score_trees", False):
return {
"request": self._get_settings(),
"trees": self.routes.dicts,
}
self.routes.compute_scores(*self.scorers.objects())
return {
"request": self._get_settings(),
"trees": self.routes.dict_with_scores(),
}
def tree_search(self, show_progress: bool = False) -> float:
"""
Perform the actual tree search
:param show_progress: if True, shows a progress bar
:return: the time past in seconds
"""
if not self.tree:
self.prepare_tree()
assert (self.tree is not None
) # This is for type checking, prepare_tree is creating it.
self.search_stats = {"returned_first": False, "iterations": 0}
time0 = time.time()
i = 1
self._logger.debug("Starting search")
time_past = time.time() - time0
if show_progress:
pbar = tqdm(total=self.config.iteration_limit, leave=False)
while time_past < self.config.time_limit and i <= self.config.iteration_limit:
if show_progress:
pbar.update(1)
self.search_stats["iterations"] += 1
try:
is_solved = self.tree.one_iteration()
except StopIteration:
break
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._logger.debug("Found first solved route")
self.search_stats["returned_first"] = True
break
i = i + 1
time_past = time.time() - time0
if show_progress:
pbar.close()
time_past = time.time() - time0
self._logger.debug("Search completed")
self.search_stats["time"] = time_past
return time_past
def _get_settings(self) -> StrDict:
"""Get the current settings as a dictionary"""
# To be backward-compatible
if (self.expansion_policy.selection
and len(self.expansion_policy.selection) == 1):
policy_value = self.expansion_policy.selection[0]
policy_key = "policy"
else:
policy_value = self.expansion_policy.selection # type: ignore
policy_key = "policies"
dict_ = {
"stocks": self.stock.selection,
policy_key: policy_value,
"C": self.config.C,
"max_transforms": self.config.max_transforms,
"cutoff_cumulative": self.config.cutoff_cumulative,
"cutoff_number": self.config.cutoff_number,
"smiles": self.target_smiles,
"return_first": self.config.return_first,
"time_limit": self.config.time_limit,
"iteration_limit": self.config.iteration_limit,
"exclude_target_from_stock": self.config.exclude_target_from_stock,
"filter_cutoff": self.config.filter_cutoff,
}
if self.filter_policy.selection:
dict_["filter"] = self.filter_policy.selection
return dict_
def _setup_search_tree(self):
self._logger.debug("Defining tree root: %s" % self.target_smiles)
if self.config.search_algorithm.lower() == "mcts":
self.tree = MctsSearchTree(root_smiles=self.target_smiles,
config=self.config)
else:
module_name, cls_name = self.config.search_algorithm.rsplit(
".", maxsplit=1)
try:
module_obj = importlib.import_module(module_name)
except ImportError:
raise ValueError(f"Could not import module {module_name}")
if not hasattr(module_obj, cls_name):
raise ValueError(
f"Could not identify class {cls_name} in module")
self.tree: AndOrSearchTreeBase = getattr(module_obj, cls_name)(
root_smiles=self.target_smiles, config=self.config)
