diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2ad9c75..fef7c29 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,31 @@
# CHANGELOG
+## Version 4.3.0 2024-05-22
+
+### Features
+
+- Focussed bonds feature has been created to allow users to choose bonds to freeze and break.
+- Functionality has been added to run the broken bonds scorer with MCTS.
+- A disconnection-aware Chemformer expansion strategy has been introduced in the plugins.
+- Weights for single expansion policies can now be provided as input through the config file.
+- The priors returned from the multi-expansion strategy have been rescaled.
+- Functionality to mask reaction templates is now supported in the `TemplateBasedExpansionStrategy`.
+- The `TemplateBasedDirectExpansionStrategy` has been implemented to directly apply the template in the search process.
+- Added support for the C++ version of RDChiral.
+- The multi-objective MCTS core algorithm is now implemented within the MCTS search functionality.
+- A separate GUI component has been introduced for doing MO tree analysis. Additional functionalities have been added to GUI widgets to set two rewards/objectives to MCTS search. Pareto front is automatically plotted if MO-MCTS is run. Route re-ordering is automatically disabled for MO-MCTS.
+- Preprocessing of the tree search can now be done using `aizynthcli`.
+- The `StockAvailablityScorer` has been updated such that it takes and additional `other_source_score` parameter.
+- A `cutoff_number` parameter can be provided to the multi-expansion strategy to obtain only the top predictions.
+- A `BrokenBondsScorer` has been created for scoring nodes and reaction trees based on the breaking of atom bonds.
+- A `RouteSimilarityScorer` has been created for scoring based on an LSTM model for computing Tree Edit Distance to a set of reference routes.
+- A `DeltaSyntheticComplexityScorer` has been created for scoring nodes based on the delta-synthetic-complexity of the node and its parent 'horizon' steps up in the tree.
+
+### Bug-fixes
+
+- Fixed an issue of sending multiple fingerprints to the GPCR Tensorflow serving model.
+- `aizynthcli` has been fixed after updating with multi-objective analysis such that the tool accurately informs the user if a target is solved or not.
+
## Version 4.0.0 2023-11-30
### Features
@@ -283,7 +309,7 @@
- Add tools to train filter policy
- Add logic to prevent cycle forming in MCTS by rejecting creation of parent molecule when expanding
- Introduce new `context` subpackage that contains the `config`, `stock`, `policy` and `scoring` modules
-- The `Stock`, `ExpansionPolicy`, `FilterPolicy` and `ScorerCollection` classes now has a common interface for selection and loading
+- The `Stock`, `ExpansionPolicy`, `FilterPolicy` and `ScorerCollection` classes now has a common interface for selection and loading
- Introduce possibility to remove unsantizable reactions from template library when training
- Catch exceptions from RDChiral more gracefully
diff --git a/README.md b/README.md
index d539e08..c55180c 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
[](https://github.com/MolecularAI/aizynthfinder/blob/master/LICENSE)
[](https://github.com/MolecularAI/aizynthfinder/actions?workflow=tests)
[](https://codecov.io/gh/MolecularAI/aizynthfinder)
-[](https://github.com/python/black)
+[](https://github.com/python/black)
[](https://github.com/MolecularAI/aizynthfinder/releases)
[](https://colab.research.google.com/github/MolecularAI/aizynthfinder/blob/master/contrib/notebook.ipynb)
@@ -28,7 +28,7 @@ The tool has been developed on a Linux platform, but the software has been teste
First time, execute the following command in a console or an Anaconda prompt
- conda create "python>=3.8,<3.10" -n aizynth-env
+ conda create "python>=3.9,<3.11" -n aizynth-env
To install, activate the environment and install the package using pypi
@@ -43,12 +43,12 @@ for a smaller package, without all the functionality, you can also type
First clone the repository using Git.
-Then execute the following commands in the root of the repository
+Then execute the following commands in the root of the repository
conda env create -f env-dev.yml
conda activate aizynth-dev
poetry install --all-extras
-
+
the `aizynthfinder` package is now installed in editable mode.
@@ -68,7 +68,7 @@ To use the tool you need
1. A stock file
2. A trained expansion policy network
3. A trained filter policy network (optional)
-
+
Such files can be downloaded from [figshare](https://figshare.com/articles/AiZynthFinder_a_fast_robust_and_flexible_open-source_software_for_retrosynthetic_planning/12334577) and [here](https://figshare.com/articles/dataset/A_quick_policy_to_filter_reactions_based_on_feasibility_in_AI-guided_retrosynthetic_planning/13280507) or they can be downloaded automatically using
```
@@ -91,7 +91,7 @@ Run the tests using:
The full command run on the CI server is available through an `invoke` command
invoke full-tests
-
+
### Documentation generation
The documentation is generated by Sphinx from hand-written tutorials and docstrings
diff --git a/aizynthfinder/aizynthfinder.py b/aizynthfinder/aizynthfinder.py
index 958af9c..007a856 100644
--- a/aizynthfinder/aizynthfinder.py
+++ b/aizynthfinder/aizynthfinder.py
@@ -15,7 +15,8 @@
)
from aizynthfinder.chem import FixedRetroReaction, Molecule, TreeMolecule
from aizynthfinder.context.config import Configuration
-from aizynthfinder.context.scoring import CombinedScorer
+from aizynthfinder.context.policy import BondFilter
+from aizynthfinder.context.scoring import BrokenBondsScorer, CombinedScorer
from aizynthfinder.reactiontree import ReactionTreeFromExpansion
from aizynthfinder.search.andor_trees import AndOrSearchTreeBase
from aizynthfinder.search.mcts import MctsSearchTree
@@ -82,6 +83,9 @@ def __init__(
self.search_stats: StrDict = dict()
self.routes = RouteCollection([])
self.analysis: Optional[TreeAnalysis] = None
+ self._num_objectives = len(
+ self.config.search.algorithm_config.get("search_rewards", [])
+ )
@property
def target_smiles(self) -> str:
@@ -107,7 +111,7 @@ def target_mol(self, mol: Molecule) -> None:
def build_routes(
self,
selection: Optional[RouteSelectionArguments] = None,
- scorer: Optional[str] = None,
+ scorer: Optional[Union[str, List[str]]] = None,
) -> None:
"""
Build reaction routes
@@ -116,18 +120,10 @@ def build_routes(
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
+ :param scorer: a reference to the object used to score the nodes, can be a list
:raises ValueError: if the search tree not initialized
"""
-
- scorer = scorer or self.config.post_processing.route_scorer
-
- if not self.tree:
- raise ValueError("Search tree not initialized")
-
- _scorer = self.scorers[scorer]
-
- self.analysis = TreeAnalysis(self.tree, scorer=_scorer)
+ self.analysis = self._setup_analysis(scorer=scorer)
config_selection = RouteSelectionArguments(
nmin=self.config.post_processing.min_routes,
nmax=self.config.post_processing.max_routes,
@@ -174,9 +170,13 @@ def prepare_tree(self) -> None:
self.stock.exclude(self.target_mol)
self._logger.debug("Excluding the target compound from the stock")
+ if self.config.search.break_bonds or self.config.search.freeze_bonds:
+ self._setup_focussed_bonds(self.target_mol)
+
self._setup_search_tree()
self.analysis = None
self.routes = RouteCollection([])
+ self.filter_policy.reset_cache()
self.expansion_policy.reset_cache()
def stock_info(self) -> StrDict:
@@ -249,8 +249,40 @@ def tree_search(self, show_progress: bool = False) -> float:
self.search_stats["time"] = time_past
return time_past
+ def _setup_focussed_bonds(self, target_mol: Molecule) -> None:
+ """
+ Setup multi-objective scoring function with 'broken bonds'-scorer and
+ add 'frozen bonds'-filter to filter policy.
+
+ :param target_mol: the target molecule.
+ """
+ target_mol = TreeMolecule(smiles=target_mol.smiles, parent=None)
+
+ bond_filter_key = "__finder_bond_filter"
+ if self.config.search.freeze_bonds:
+ if not target_mol.has_all_focussed_bonds(self.config.search.freeze_bonds):
+ raise ValueError("Bonds in 'freeze_bond' must exist in target molecule")
+ bond_filter = BondFilter(bond_filter_key, self.config)
+ self.filter_policy.load(bond_filter)
+ self.filter_policy.select(bond_filter_key, append=True)
+ elif (
+ self.filter_policy.selection
+ and bond_filter_key in self.filter_policy.selection
+ ):
+ self.filter_policy.deselect(bond_filter_key)
+
+ search_rewards = self.config.search.algorithm_config.get("search_rewards")
+ if not search_rewards:
+ return
+
+ if self.config.search.break_bonds and "broken bonds" in search_rewards:
+ if not target_mol.has_all_focussed_bonds(self.config.search.break_bonds):
+ raise ValueError("Bonds in 'break_bonds' must exist in target molecule")
+ self.scorers.load(BrokenBondsScorer(self.config))
+ self._num_objectives = len(search_rewards)
+
def _setup_search_tree(self) -> None:
- self._logger.debug("Defining tree root: %s" % self.target_smiles)
+ self._logger.debug(f"Defining tree root: {self.target_smiles}")
if self.config.search.algorithm.lower() == "mcts":
self.tree = MctsSearchTree(
root_smiles=self.target_smiles, config=self.config
@@ -259,6 +291,50 @@ def _setup_search_tree(self) -> None:
cls = load_dynamic_class(self.config.search.algorithm)
self.tree = cls(root_smiles=self.target_smiles, config=self.config)
+ def _setup_analysis(
+ self,
+ scorer: Optional[Union[str, List[str]]],
+ ) -> TreeAnalysis:
+ """Configure TreeAnalysis
+
+ :param scorer: a reference to the object used to score the nodes, can be a list
+ :returns: the configured TreeAnalysis
+ :raises ValueError: if the search tree not initialized
+ """
+ if not self.tree:
+ raise ValueError("Search tree not initialized")
+
+ if scorer is None:
+ scorer_names = self.config.post_processing.route_scorers
+ # If not defined, use the same scorer as the search rewards
+ if not scorer_names:
+ search_rewards = self.config.search.algorithm_config.get(
+ "search_rewards"
+ )
+ scorer_names = search_rewards if search_rewards else ["state score"]
+
+ elif isinstance(scorer, str):
+ scorer_names = [scorer]
+ else:
+ scorer_names = list(scorer)
+
+ if "broken bonds" in scorer_names:
+ # Add broken bonds scorer if required
+ self.scorers.load(BrokenBondsScorer(self.config))
+
+ scorers = [self.scorers[name] for name in scorer_names]
+
+ if self.config.post_processing.scorer_weights:
+ scorers = [
+ CombinedScorer(
+ self.config,
+ scorer_names,
+ self.config.post_processing.scorer_weights,
+ )
+ ]
+
+ return TreeAnalysis(self.tree, scorers)
+
class AiZynthExpander:
"""
diff --git a/aizynthfinder/analysis/routes.py b/aizynthfinder/analysis/routes.py
index 1f66f75..611f504 100644
--- a/aizynthfinder/analysis/routes.py
+++ b/aizynthfinder/analysis/routes.py
@@ -2,6 +2,7 @@
"""
from __future__ import annotations
+import copy
from typing import TYPE_CHECKING
import numpy as np
@@ -12,12 +13,12 @@
except ImportError:
pass
-from aizynthfinder.analysis import TreeAnalysis
from aizynthfinder.analysis.utils import CombinedReactionTrees, RouteSelectionArguments
from aizynthfinder.reactiontree import SUPPORT_DISTANCES, ReactionTree
from aizynthfinder.search.mcts import MctsNode, MctsSearchTree
if TYPE_CHECKING:
+ from aizynthfinder.analysis import TreeAnalysis
from aizynthfinder.context.scoring import Scorer
from aizynthfinder.utils.type_utils import (
Any,
@@ -26,6 +27,7 @@
PilImage,
Sequence,
StrDict,
+ Union,
)
@@ -64,11 +66,11 @@ def __init__(self, reaction_trees: Sequence[ReactionTree], **kwargs) -> None:
self._update_route_dict(self.route_metadata, "route_metadata")
self.nodes = self._unpack_kwarg_with_default("nodes", None, **kwargs)
- self.scores = self._unpack_kwarg_with_default("scores", np.nan, **kwargs)
+ self.scores = self._unpack_kwarg_with_default("scores", dict, **kwargs)
self.all_scores = self._unpack_kwarg_with_default("all_scores", dict, **kwargs)
self._dicts: Optional[Sequence[StrDict]] = self._unpack_kwarg("dicts", **kwargs)
- self._images: Optional[Sequence[PilImage]] = self._unpack_kwarg(
+ self._images: Optional[Sequence[Optional[PilImage]]] = self._unpack_kwarg(
"images", **kwargs
)
self._jsons: Optional[Sequence[str]] = self._unpack_kwarg("jsons", **kwargs)
@@ -80,7 +82,9 @@ def __init__(self, reaction_trees: Sequence[ReactionTree], **kwargs) -> None:
@classmethod
def from_analysis(
- cls, analysis: TreeAnalysis, selection: RouteSelectionArguments = None
+ cls,
+ analysis: TreeAnalysis,
+ selection: Optional[RouteSelectionArguments] = None,
) -> "RouteCollection":
"""
Create a collection from a tree analysis.
@@ -90,8 +94,8 @@ def from_analysis(
: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}
+ all_scores = copy.deepcopy(scores)
+ kwargs: Dict[str, Any] = {"scores": scores, "all_scores": all_scores}
if isinstance(analysis.search_tree, MctsSearchTree):
kwargs["nodes"] = items
reaction_trees = [
@@ -120,10 +124,11 @@ def dicts(self) -> Sequence[StrDict]:
return self._dicts
@property
- def images(self) -> Sequence[PilImage]:
+ def images(self) -> Sequence[Optional[PilImage]]:
"""Returns a list of pictoral representation of the routes"""
if self._images is None:
self._images = self.make_images()
+ assert self._images is not None
return self._images
@property
diff --git a/aizynthfinder/analysis/tree_analysis.py b/aizynthfinder/analysis/tree_analysis.py
index 95282b9..ece4352 100644
--- a/aizynthfinder/analysis/tree_analysis.py
+++ b/aizynthfinder/analysis/tree_analysis.py
@@ -1,10 +1,15 @@
""" Module containing classes to perform analysis of the tree search results.
"""
+
from __future__ import annotations
+import operator
from collections import defaultdict
from typing import TYPE_CHECKING
+import numpy as np
+from paretoset import paretorank, paretoset
+
from aizynthfinder.analysis.utils import RouteSelectionArguments
from aizynthfinder.chem import FixedRetroReaction, hash_reactions
from aizynthfinder.context.scoring import StateScorer
@@ -17,6 +22,7 @@
from aizynthfinder.context.scoring import Scorer
from aizynthfinder.utils.type_utils import (
Any,
+ Dict,
Iterable,
List,
Optional,
@@ -26,13 +32,19 @@
Union,
)
+ _Solution = Union[MctsNode, ReactionTree]
+ _AnyListOfSolutions = Union[Sequence[MctsNode], Sequence[ReactionTree]]
+ _AnyListOfReactions = Sequence[
+ Union[Iterable[RetroReaction], Iterable[FixedRetroReaction]]
+ ]
+
class TreeAnalysis:
"""
Class that encapsulate various analysis that can be
performed on a search tree.
- :ivar scorer: the object used to score the nodes
+ :ivar scorers: the objects used to score the nodes
:ivar search_tree: the search tree
:param search_tree: the search tree to do the analysis on
@@ -42,51 +54,94 @@ class TreeAnalysis:
def __init__(
self,
search_tree: Union[MctsSearchTree, AndOrSearchTreeBase],
- scorer: Optional[Scorer] = None,
+ scorer: Optional[Union[Scorer, List[Scorer]]] = None,
) -> None:
self.search_tree = search_tree
- if scorer is None:
- self.scorer: Scorer = StateScorer(search_tree.config)
+ scorer = scorer or StateScorer(search_tree.config)
+ if isinstance(scorer, list):
+ self.scorers: List[Scorer] = scorer
else:
- self.scorer = scorer
+ self.scorers = [scorer]
+ self._direction = "max" # current implementation assumes maximisation
+ self._single_objective = len(self.scorers) == 1
- def best(self) -> Union[MctsNode, ReactionTree]:
+ def best(self) -> _Solution:
"""
Returns the route or MCTS-like node with the highest score.
If several routes have the same score, it will return the first
:return: the top scoring node or route
+ :raises ValueError: if this is a multi-objective analysis, a single best solution cannot be return
"""
+ if not self._single_objective:
+ raise ValueError("Cannot return best item for multi-objective analysis")
+
if isinstance(self.search_tree, MctsSearchTree):
nodes = self._all_nodes()
- sorted_nodes, _, _ = self.scorer.sort(nodes)
+ sorted_nodes, _, _ = self.scorers[0].sort(nodes)
return sorted_nodes[0]
- sorted_routes, _, _ = self.scorer.sort(self.search_tree.routes())
+ sorted_routes, _, _ = self.scorers[0].sort(self.search_tree.routes())
return sorted_routes[0]
+ def pareto_front(self) -> Tuple[_Solution, ...]:
+ """
+ Returns the Routes or MCTS-like Nodes that are the Pareto Front in the Multi-Objective search.
+
+ :returns: the pareto front solutions.
+ :raises ValueError: if this is a single-objective analysis, there is no Pareto-front
+ """
+ if self._single_objective:
+ raise ValueError("Cannot return Pareto front for single-objective analysis")
+
+ if isinstance(self.search_tree, MctsSearchTree):
+ solutions = self.search_tree.nodes()
+ else:
+ solutions = self.search_tree.routes() # type: ignore
+
+ scores_arr = np.array(
+ [[scorer(solution) for scorer in self.scorers] for solution in solutions]
+ )
+ direction_arr = np.repeat(self._direction, len(self.scorers))
+ pareto_mask = paretoset(scores_arr, sense=direction_arr, distinct=False)
+ pareto_idxs = np.arange(len(solutions))[pareto_mask]
+ pareto_front: Sequence[_Solution] = [solutions[idx] for idx in pareto_idxs]
+
+ return tuple(pareto_front)
+
def sort(
self, selection: Optional[RouteSelectionArguments] = None
- ) -> Tuple[Union[Sequence[MctsNode], Sequence[ReactionTree]], Sequence[float]]:
+ ) -> Tuple[_AnyListOfSolutions, Sequence[Dict[str, float]]]:
"""
Sort and select the nodes or routes in the search tree.
+ The score for each solution is a dictionary of scores,
+ one for each of the objectives.
+
:param selection: selection criteria for the routes
- :return: the items
- :return: the score
+ :return: the sorted and selected items
+ :return: the scores of the sorted items
"""
selection = selection or RouteSelectionArguments()
+ if not self._single_objective:
+ return self._pareto_rank_sort(selection)
+
+ scorer = self.scorers[0]
if isinstance(self.search_tree, MctsSearchTree):
nodes = self._all_nodes()
- sorted_items, sorted_scores, _ = self.scorer.sort(nodes)
+ sorted_items, sorted_scores, _ = scorer.sort(nodes)
actions = [node.actions_to() for node in sorted_items]
else:
- sorted_items, sorted_scores, _ = self.scorer.sort(self.search_tree.routes())
+ sorted_items, sorted_scores, _ = scorer.sort(self.search_tree.routes())
actions = [route.reactions() for route in sorted_items]
- return self._collect_top_items(sorted_items, sorted_scores, actions, selection)
+ scores = [{repr(scorer): score} for score in sorted_scores]
+
+ return self._collect_top_items(
+ sorted_items, scores, actions, sorted_scores, selection, True
+ )
def tree_statistics(self) -> StrDict:
"""
@@ -105,29 +160,83 @@ def tree_statistics(self) -> StrDict:
def _all_nodes(self) -> Sequence[MctsNode]:
assert isinstance(self.search_tree, MctsSearchTree)
# This is to keep backwards compatibility, this should be investigate further
- if repr(self.scorer) == "state score":
+ if repr(self.scorers[0]) == "state score":
return list(self.search_tree.graph())
return [node for node in self.search_tree.graph() if not node.children]
+ def _pareto_rank_sort(
+ self,
+ selection: RouteSelectionArguments,
+ ) -> Tuple[_AnyListOfSolutions, Sequence[Dict[str, float]]]:
+ if isinstance(self.search_tree, MctsSearchTree):
+ solutions = self._all_nodes()
+ else:
+ solutions = self.search_tree.routes() # type: ignore
+
+ scores_arr = np.array(
+ [[scorer(solution) for scorer in self.scorers] for solution in solutions]
+ )
+ direction_arr = np.repeat(self._direction, len(self.scorers))
+ pareto_ranks = paretorank(scores_arr, sense=direction_arr, distinct=False)
+
+ sortidx = sorted(range(len(pareto_ranks)), key=pareto_ranks.__getitem__)
+ sorted_pareto_ranks = sorted(pareto_ranks)
+ sorted_scores = [
+ {
+ repr(scorer): scores_arr[idx, scorer_idx]
+ for scorer_idx, scorer in enumerate(self.scorers)
+ }
+ for idx in sortidx
+ ]
+ sorted_items = [solutions[idx] for idx in sortidx]
+
+ if isinstance(self.search_tree, MctsSearchTree):
+ actions = [node.actions_to() for node in sorted_items]
+ else:
+ actions = [route.reactions() for route in sorted_items] # type: ignore
+
+ return self._collect_top_items(
+ sorted_items, sorted_scores, actions, sorted_pareto_ranks, selection, False
+ )
+
+ def _top_nodes(self) -> Tuple[_Solution, ...]:
+ if self._single_objective:
+ return (self.best(),)
+ return self.pareto_front()
+
def _tree_statistics_andor(self) -> StrDict:
assert isinstance(self.search_tree, AndOrSearchTreeBase)
- top_route = self.best()
- assert isinstance(top_route, ReactionTree)
- mols_in_stock = ", ".join(
- mol.smiles for mol in top_route.leafs() if top_route.in_stock(mol)
+ top_routes = self._top_nodes()
+ mols_in_stock = self._top_ranked_join(
+ ", ".join(mol.smiles for mol in route.leafs() if route.in_stock(mol)) # type: ignore
+ for route in top_routes
)
- mols_not_in_stock = ", ".join(
- mol.smiles for mol in top_route.leafs() if not top_route.in_stock(mol)
+ mols_not_in_stock = self._top_ranked_join(
+ ", ".join(mol.smiles for mol in route.leafs() if not route.in_stock(mol)) # type: ignore
+ for route in top_routes
)
all_routes = self.search_tree.routes()
policy_used_counts = self._policy_used_statistics(
[reaction for route in all_routes for reaction in route.reactions()]
)
- availability = ";".join(
- self.search_tree.config.stock.availability_string(mol)
- for mol in top_route.leafs()
+ availability = self._top_ranked_join(
+ ";".join(
+ self.search_tree.config.stock.availability_string(mol)
+ for mol in route.leafs() # type: ignore
+ )
+ for route in top_routes
+ )
+ number_of_precursors_in_stock = self._top_ranked_join(
+ sum(route.in_stock(leaf) for leaf in route.leafs()) for route in top_routes # type: ignore
)
+ # For multi-objective it becomes to messy to have all the scores in this info,
+ # they are nevertheless included with the routes
+ if self._single_objective:
+ top_score = self.scorers[0](top_routes[0])
+ else:
+ top_score = None
+
return {
"number_of_nodes": len(self.search_tree.mol_nodes),
"max_transforms": max(
@@ -138,13 +247,15 @@ def _tree_statistics_andor(self) -> StrDict:
),
"number_of_routes": len(all_routes),
"number_of_solved_routes": sum(route.is_solved for route in all_routes),
- "top_score": self.scorer(top_route),
- "is_solved": top_route.is_solved,
- "number_of_steps": len(list(top_route.reactions())),
- "number_of_precursors": len(list(top_route.leafs())),
- "number_of_precursors_in_stock": sum(
- top_route.in_stock(leaf) for leaf in top_route.leafs()
+ "top_score": top_score,
+ "is_solved": self._top_ranked_join(route.is_solved for route in top_routes),
+ "number_of_steps": self._top_ranked_join(
+ len(list(route.reactions())) for route in top_routes # type: ignore
+ ),
+ "number_of_precursors": self._top_ranked_join(
+ len(list(route.leafs())) for route in top_routes # type: ignore
),
+ "number_of_precursors_in_stock": number_of_precursors_in_stock,
"precursors_in_stock": mols_in_stock,
"precursors_not_in_stock": mols_not_in_stock,
"precursors_availability": availability,
@@ -154,25 +265,38 @@ def _tree_statistics_andor(self) -> StrDict:
def _tree_statistics_mcts(self) -> StrDict:
assert isinstance(self.search_tree, MctsSearchTree)
- top_node = self.best()
- assert isinstance(top_node, MctsNode)
- top_state = top_node.state
+ top_nodes = self._top_nodes()
+ assert isinstance(top_nodes[0], MctsNode)
+ top_states = [node.state for node in top_nodes] # type: ignore
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)
- if instock
+ mols_in_stock = self._top_ranked_join(
+ ", ".join(
+ mol.smiles
+ for mol, instock in zip(state.mols, state.in_stock_list)
+ if instock
+ )
+ for state in top_states
)
- mols_not_in_stock = ", ".join(
- mol.smiles
- for mol, instock in zip(top_state.mols, top_state.in_stock_list)
- if not instock
+ mols_not_in_stock = self._top_ranked_join(
+ ", ".join(
+ mol.smiles
+ for mol, instock in zip(state.mols, state.in_stock_list)
+ if not instock
+ )
+ for state in top_states
)
policy_used_counts = self._policy_used_statistics(
[node[child]["action"] for node in nodes for child in node.children]
)
+ # For multi-objective it becomes to messy to have all the scores in this info,
+ # they are nevertheless included with the routes
+ if self._single_objective:
+ top_score = self.scorers[0](top_nodes[0])
+ else:
+ top_score = None
+
return {
"number_of_nodes": len(nodes),
"max_transforms": max(node.state.max_transforms for node in nodes),
@@ -181,27 +305,59 @@ def _tree_statistics_mcts(self) -> StrDict:
"number_of_solved_routes": sum(
1 for node in nodes if not node.children and node.state.is_solved
),
- "top_score": self.scorer(top_node),
- "is_solved": top_state.is_solved,
- "number_of_steps": top_state.max_transforms,
- "number_of_precursors": len(top_state.mols),
- "number_of_precursors_in_stock": sum(top_state.in_stock_list),
+ "top_score": top_score,
+ "is_solved": self._top_ranked_join(state.is_solved for state in top_states),
+ "number_of_steps": self._top_ranked_join(
+ state.max_transforms for state in top_states
+ ),
+ "number_of_precursors": self._top_ranked_join(
+ len(state.mols) for state in top_states
+ ),
+ "number_of_precursors_in_stock": self._top_ranked_join(
+ sum(state.in_stock_list) for state in top_states
+ ),
"precursors_in_stock": mols_in_stock,
"precursors_not_in_stock": mols_not_in_stock,
- "precursors_availability": ";".join(top_state.stock_availability),
+ "precursors_availability": self._top_ranked_join(
+ ";".join(state.stock_availability) for state in top_states
+ ),
"policy_used_counts": policy_used_counts,
"profiling": getattr(self.search_tree, "profiling", {}),
}
@staticmethod
def _collect_top_items(
- items: Union[Sequence[MctsNode], Sequence[ReactionTree]],
- scores: Sequence[float],
- reactions: Sequence[
- Union[Iterable[RetroReaction], Iterable[FixedRetroReaction]]
- ],
- selection,
- ) -> Tuple[Union[Sequence[MctsNode], Sequence[ReactionTree]], Sequence[float]]:
+ items: _AnyListOfSolutions,
+ scores: Sequence[Dict[str, float]],
+ reactions: _AnyListOfReactions,
+ ranks: Sequence[Union[int, float]],
+ selection: RouteSelectionArguments,
+ min_comp: bool,
+ ) -> Tuple[_AnyListOfSolutions, Sequence[Dict[str, float]]]:
+ """
+ Finds and returns the top-ranked super-nodes or reaction trees and
+ their scores.
+
+ The `selection` argument is used to control the number of items that are returned.
+ If at least one of the items is solved, and `selection` specifies that all items
+ should be returned, all solved items are returned.
+ Otherwise, at least a minimum number of items are returned, but there could be more
+ than this if they have the same rank. However not more than a maximum number of
+ items are returned.
+
+ Duplicated routes are ignored.
+
+ The items are compared based on rank. For single-ojective analysis, this is simply
+ the score and higher is better. For multi-objective analysis, the rank is the Pareto
+ rank and smaller i better.
+
+ :param items: the super-nodes or reaction trees to select from, sorted by rank
+ :param scores: the computed scores of the nodes
+ :param reaction: the reactions for each item, used for checking duplicity
+ :param ranks: the rank of each item, used for determining selection
+ :param selection: min and max solutions, or indication to return all
+ :param min_comp: True if higher rank is better, else False
+ """
if len(items) <= selection.nmin:
return items, scores
@@ -215,9 +371,14 @@ def _collect_top_items(
seen_hashes = set()
best_items: List[Any] = []
best_scores = []
- last_score = 1e16
- for score, item, actions in zip(scores, items, reactions):
- if len(best_items) >= min_return and score < last_score:
+ if min_comp:
+ last_rank = 1e16
+ comp_op = operator.lt
+ else:
+ last_rank = 0
+ comp_op = operator.gt
+ for rank, score, item, actions in zip(ranks, scores, items, reactions):
+ if len(best_items) >= min_return and comp_op(rank, last_rank):
break
route_hash = hash_reactions(actions)
@@ -226,7 +387,7 @@ def _collect_top_items(
seen_hashes.add(route_hash)
best_items.append(item)
best_scores.append(score)
- last_score = score
+ last_rank = rank
if max_return and len(best_items) == max_return:
break
@@ -243,3 +404,10 @@ def _policy_used_statistics(
if policy_used:
policy_used_counts[policy_used] += 1
return dict(policy_used_counts)
+
+ @staticmethod
+ def _top_ranked_join(items: Iterable[Any]) -> Union[str, Any]:
+ items = list(items)
+ if len(items) == 1:
+ return items[0]
+ return "|".join(f"{item}" for item in items)
diff --git a/aizynthfinder/chem/mol.py b/aizynthfinder/chem/mol.py
index 8864151..128937b 100644
--- a/aizynthfinder/chem/mol.py
+++ b/aizynthfinder/chem/mol.py
@@ -8,6 +8,7 @@
from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem, Descriptors
+from aizynthfinder.utils.bonds import sort_bonds
from aizynthfinder.utils.exceptions import MoleculeException
if TYPE_CHECKING:
@@ -320,6 +321,30 @@ def mapped_atom_bonds(self) -> List[Tuple[int, int]]:
]
return self._atom_bonds
+ def get_bonds_in_molecule(
+ self, query_bonds: Sequence[Sequence[int]]
+ ) -> Sequence[Sequence[int]]:
+ """
+ Get bonds (from a list of bonds) that are present in the molecule.
+ :param bonds: List of bond (atom pairs)
+ :return: A list of bonds
+ """
+ molecule_bonds = sort_bonds(self.mapped_atom_bonds)
+ query_bonds = sort_bonds(query_bonds)
+ bonds_in_mol = [bond for bond in query_bonds if bond in molecule_bonds]
+ return bonds_in_mol
+
+ def has_all_focussed_bonds(self, bonds: Sequence[Sequence[int]]) -> bool:
+ """Checks that the focussed bonds exist in the target molecule's atom bonds.
+
+ :param bonds: Focussed bonds.
+ :param target_mol: The target molecule.
+
+ :return: A boolean indicating if the input bonds exist in the target molecule.
+ """
+ bonds_in_mol = self.get_bonds_in_molecule(bonds)
+ return len(bonds_in_mol) == len(bonds)
+
def _set_atom_mappings(self) -> None:
atom_mappings = [
atom.GetAtomMapNum()
diff --git a/aizynthfinder/chem/reaction.py b/aizynthfinder/chem/reaction.py
index 56db3e3..7b34a97 100644
--- a/aizynthfinder/chem/reaction.py
+++ b/aizynthfinder/chem/reaction.py
@@ -9,8 +9,14 @@
import numpy as np
from rdchiral import main as rdc
-from rdchiral.bonds import get_atoms_across_double_bonds
-from rdchiral.initialization import BondDirOpposite
+
+try:
+ from rdchiral.bonds import get_atoms_across_double_bonds
+ from rdchiral.initialization import BondDirOpposite
+except ImportError:
+ RDCHIRAL_CPP = True
+else:
+ RDCHIRAL_CPP = False
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.rdchem import BondDir, BondStereo, ChiralType
@@ -36,6 +42,11 @@
Union,
)
+if RDCHIRAL_CPP:
+ logger().warning(
+ "WARNING: C++ version of RDChiral is supported, but with limited functionality"
+ )
+
class _ReactionInterfaceMixin:
"""
@@ -102,7 +113,7 @@ def reaction_smiles(self) -> str:
"""
reactants = ".".join(mol.smiles for mol in self._reactants_getter()) # type: ignore
products = ".".join(mol.smiles for mol in self._products_getter()) # type: ignore
- return "%s>>%s" % (reactants, products)
+ return f"{reactants}>>{products}"
class RetroReaction(abc.ABC, _ReactionInterfaceMixin):
@@ -631,3 +642,11 @@ def __init__(self, product: TreeMolecule) -> None:
# Get atoms across double bonds defined by mapnum
self.atoms_across_double_bonds = get_atoms_across_double_bonds(self.reactants)
+
+
+if RDCHIRAL_CPP:
+
+ def _wrapper(mol):
+ return rdc.rdchiralReactants(mol.mapped_smiles)
+
+ _RdChiralProductWrapper = _wrapper # type: ignore
diff --git a/aizynthfinder/context/config.py b/aizynthfinder/context/config.py
index 8b8e39a..c55d194 100644
--- a/aizynthfinder/context/config.py
+++ b/aizynthfinder/context/config.py
@@ -8,6 +8,7 @@
from typing import TYPE_CHECKING
import yaml
+
from aizynthfinder.context.policy import ExpansionPolicy, FilterPolicy
from aizynthfinder.context.scoring import ScorerCollection
from aizynthfinder.context.stock import Stock
@@ -23,7 +24,8 @@ class _PostprocessingConfiguration:
max_routes: int = 25
all_routes: bool = False
route_distance_model: Optional[str] = None
- route_scorer: str = "state score"
+ route_scorers: List[str] = field(default_factory=lambda: [])
+ scorer_weights: Optional[List[float]] = field(default_factory=lambda: None)
@dataclass
@@ -35,9 +37,10 @@ class _SearchConfiguration:
"default_prior": 0.5,
"use_prior": True,
"prune_cycles_in_search": True,
- "search_reward": "state score",
+ "search_rewards": ["state score"],
"immediate_instantiation": (),
"mcts_grouping": None,
+ "search_rewards_weights": [],
}
)
max_transforms: int = 6
diff --git a/aizynthfinder/context/policy/__init__.py b/aizynthfinder/context/policy/__init__.py
index 2951826..b9cfd90 100644
--- a/aizynthfinder/context/policy/__init__.py
+++ b/aizynthfinder/context/policy/__init__.py
@@ -1,11 +1,14 @@
""" Sub-package containing policy routines
"""
+
from aizynthfinder.context.policy.expansion_strategies import (
ExpansionStrategy,
MultiExpansionStrategy,
+ TemplateBasedDirectExpansionStrategy,
TemplateBasedExpansionStrategy,
)
from aizynthfinder.context.policy.filter_strategies import (
+ BondFilter,
FilterStrategy,
QuickKerasFilter,
ReactantsCountFilter,
diff --git a/aizynthfinder/context/policy/expansion_strategies.py b/aizynthfinder/context/policy/expansion_strategies.py
index 21c3a79..83e10f6 100644
--- a/aizynthfinder/context/policy/expansion_strategies.py
+++ b/aizynthfinder/context/policy/expansion_strategies.py
@@ -1,5 +1,6 @@
""" Module containing classes that implements different expansion policy strategies
"""
+
from __future__ import annotations
import abc
@@ -8,18 +9,25 @@
import numpy as np
import pandas as pd
-from aizynthfinder.chem import TemplatedRetroReaction
+from aizynthfinder.chem import SmilesBasedRetroReaction, TemplatedRetroReaction
from aizynthfinder.context.policy.utils import _make_fingerprint
from aizynthfinder.utils.exceptions import PolicyException
from aizynthfinder.utils.logging import logger
from aizynthfinder.utils.models import load_model
-from aizynthfinder.utils.math import softmax
if TYPE_CHECKING:
from aizynthfinder.chem import TreeMolecule
from aizynthfinder.chem.reaction import RetroReaction
from aizynthfinder.context.config import Configuration
- from aizynthfinder.utils.type_utils import Any, Dict, List, Sequence, Tuple
+ from aizynthfinder.utils.type_utils import (
+ Any,
+ Dict,
+ List,
+ Optional,
+ Sequence,
+ StrDict,
+ Tuple,
+ )
class ExpansionStrategy(abc.ABC):
@@ -54,7 +62,7 @@ def __init__(self, key: str, config: Configuration, **kwargs: str) -> None:
def __call__(
self,
molecules: Sequence[TreeMolecule],
- cache_molecules: Sequence[TreeMolecule] = None,
+ cache_molecules: Optional[Sequence[TreeMolecule]] = None,
) -> Tuple[List[RetroReaction], List[float]]:
return self.get_actions(molecules, cache_molecules)
@@ -62,7 +70,7 @@ def __call__(
def get_actions(
self,
molecules: Sequence[TreeMolecule],
- cache_molecules: Sequence[TreeMolecule] = None,
+ cache_molecules: Optional[Sequence[TreeMolecule]] = None,
) -> Tuple[List[RetroReaction], List[float]]:
"""
Get all the probable actions of a set of molecules
@@ -88,6 +96,9 @@ class MultiExpansionStrategy(ExpansionStrategy):
:ivar additive_expansion: a conditional setting to specify whether all the actions
and priors of the selected expansion strategies should be combined or not.
Defaults to False.
+ :ivar expansion_strategy_weights: a list of weights for each expansion strategy.
+ The weights should sum to one. Exception is the default, where unity weight
+ is associated to each strategy.
:param key: the key or label
:param config: the configuration of the tree search
@@ -102,17 +113,28 @@ def __init__(
self,
key: str,
config: Configuration,
- **kwargs: str,
+ **kwargs: Any,
) -> None:
super().__init__(key, config, **kwargs)
self._config = config
+ self._expansion_strategies: List[ExpansionStrategy] = []
self.expansion_strategy_keys = kwargs["expansion_strategies"]
+
+ self.cutoff_number = kwargs.get("cutoff_number")
+ if self.cutoff_number:
+ print(f"Setting multi-expansion cutoff_number: {self.cutoff_number}")
+
+ self.expansion_strategy_weights = self._set_expansion_strategy_weights(kwargs)
self.additive_expansion: bool = bool(kwargs.get("additive_expansion", False))
+ self._logger.info(
+ f"Multi-expansion strategy with policies: {self.expansion_strategy_keys}"
+ f", and corresponding weights: {self.expansion_strategy_weights}"
+ )
def get_actions(
self,
molecules: Sequence[TreeMolecule],
- cache_molecules: Sequence[TreeMolecule] = None,
+ cache_molecules: Optional[Sequence[TreeMolecule]] = None,
) -> Tuple[List[RetroReaction], List[float]]:
"""
Get all the probable actions of a set of molecules, using the selected policies.
@@ -123,6 +145,8 @@ def get_actions(
a sub class to combine different expansion strategies in different ways.
:param molecules: the molecules to consider
+ :param cache_molecules: additional molecules to submit to the expansion
+ policy but that only will be cached for later use
:return: the actions and the priors of those actions
:raises: PolicyException: if the policy isn't selected
"""
@@ -130,17 +154,31 @@ def get_actions(
all_possible_actions = []
all_priors = []
- for expansion_strategy in expansion_strategies:
+ for expansion_strategy, expansion_strategy_weight in zip(
+ expansion_strategies, self.expansion_strategy_weights
+ ):
possible_actions, priors = expansion_strategy.get_actions(
molecules, cache_molecules
)
+
all_possible_actions.extend(possible_actions)
- all_priors.extend(priors)
if not self.additive_expansion and all_possible_actions:
+ all_priors.extend(priors)
break
+
+ weighted_prior = [expansion_strategy_weight * p for p in priors]
+
+ all_priors.extend(weighted_prior)
+
+ all_possible_actions, all_priors = self._prune_actions(
+ all_possible_actions, all_priors
+ )
return all_possible_actions, all_priors
def _get_expansion_strategies_from_config(self) -> List[ExpansionStrategy]:
+ if self._expansion_strategies:
+ return self._expansion_strategies
+
if not all(
key in self._config.expansion_policy.items
for key in self.expansion_strategy_keys
@@ -149,10 +187,64 @@ def _get_expansion_strategies_from_config(self) -> List[ExpansionStrategy]:
"The input expansion strategy keys must exist in the "
"expansion policies listed in config"
)
- return [
+ self._expansion_strategies = [
self._config.expansion_policy[key] for key in self.expansion_strategy_keys
]
+ for expansion_strategy, weight in zip(
+ self._expansion_strategies, self.expansion_strategy_weights
+ ):
+ if not getattr(expansion_strategy, "rescale_prior", True) and weight < 1:
+ setattr(expansion_strategy, "rescale_prior", True)
+ self._logger.info(
+ f"Enforcing {expansion_strategy.key}.rescale_prior=True"
+ )
+ return self._expansion_strategies
+
+ def _prune_actions(
+ self, actions: List[RetroReaction], priors: List[float]
+ ) -> Tuple[List[RetroReaction], List[float]]:
+ """
+ Prune the actions if a maximum number of actions is specified.
+
+ :param actions: list of predicted actions
+ :param priors: list of prediction probabilities
+ :return: the top 'self.cutoff_number' actions and corresponding priors.
+ """
+ if not self.cutoff_number:
+ return actions, priors
+
+ sortidx = np.argsort(np.array(priors))[::-1].astype(int)
+ priors = [priors[idx] for idx in sortidx[0 : self.cutoff_number]]
+ actions = [actions[idx] for idx in sortidx[0 : self.cutoff_number]]
+ return actions, priors
+
+ def _set_expansion_strategy_weights(self, kwargs: StrDict) -> List[float]:
+ """
+ Set the weights of each expansion strategy using the input kwargs from config.
+ The weights in the config should sum to one.
+ If not set in the config file, the weights default to one for each strategy
+ (for backwards compatibility).
+
+ :param kwargs: input arguments to the MultiExpansionStrategy
+ :raises: ValueError if weights from the config file do not sum to one.
+ :return: a list of expansion strategy weights
+ """
+ if not "expansion_strategy_weights" in kwargs:
+ return [1.0 for _ in self.expansion_strategy_keys]
+
+ expansion_strategy_weights = kwargs["expansion_strategy_weights"]
+ sum_weights = sum(expansion_strategy_weights)
+
+ if sum_weights != 1:
+ raise ValueError(
+ "The expansion strategy weights in MultiExpansion should "
+ "sum to one. -> "
+ f"sum({expansion_strategy_weights})={sum_weights}."
+ )
+
+ return expansion_strategy_weights
+
class TemplateBasedExpansionStrategy(ExpansionStrategy):
"""
@@ -165,6 +257,8 @@ class TemplateBasedExpansionStrategy(ExpansionStrategy):
:ivar use_remote_models: a boolean to connect to remote TensorFlow servers
:ivar rescale_prior: a boolean to apply softmax to the priors
:ivar chiral_fingerprints: if True will base expansion on chiral fingerprint
+ :ivar mask: a boolean vector of masks for the reaction templates. The length of the vector should be equal to the
+ number of templates. It is set to None if no mask file is provided as input.
:param key: the key or label
:param config: the configuration of the tree search
@@ -184,6 +278,7 @@ def __init__(self, key: str, config: Configuration, **kwargs: str) -> None:
source = kwargs["model"]
templatefile = kwargs["template"]
+ maskfile: str = kwargs.get("mask", "")
self.template_column: str = kwargs.get("template_column", "retro_template")
self.cutoff_cumulative: float = float(kwargs.get("cutoff_cumulative", 0.995))
self.cutoff_number: int = int(kwargs.get("cutoff_number", 50))
@@ -205,6 +300,10 @@ def __init__(self, key: str, config: Configuration, **kwargs: str) -> None:
else:
self.templates = pd.read_hdf(templatefile, "table")
+ self.mask: Optional[np.ndarray] = (
+ self._load_mask_file(maskfile) if maskfile else None
+ )
+
if hasattr(self.model, "output_size") and len(self.templates) != self.model.output_size: # type: ignore
raise PolicyException(
f"The number of templates ({len(self.templates)}) does not agree with the " # type: ignore
@@ -212,11 +311,10 @@ def __init__(self, key: str, config: Configuration, **kwargs: str) -> None:
)
self._cache: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
- # pylint: disable=R0914
def get_actions(
self,
molecules: Sequence[TreeMolecule],
- cache_molecules: Sequence[TreeMolecule] = None,
+ cache_molecules: Optional[Sequence[TreeMolecule]] = None,
) -> Tuple[List[RetroReaction], List[float]]:
"""
Get all the probable actions of a set of molecules, using the selected policies and given cutoffs
@@ -236,7 +334,7 @@ def get_actions(
probable_transforms_idx, probs = self._cache[mol.inchi_key]
possible_moves = self.templates.iloc[probable_transforms_idx]
if self.rescale_prior:
- probs = softmax(probs)
+ probs /= probs.sum()
priors.extend(probs)
for idx, (move_index, move) in enumerate(possible_moves.iterrows()):
metadata = dict(move)
@@ -266,6 +364,8 @@ def _cutoff_predictions(self, predictions: np.ndarray) -> np.ndarray:
* cumulative probability less than a threshold (cutoff_cumulative)
* or at most N (cutoff_number)
"""
+ if self.mask is not None:
+ predictions[~self.mask] = 0
sortidx = np.argsort(predictions)[::-1]
cumsum: np.ndarray = np.cumsum(predictions[sortidx])
if any(cumsum >= self.cutoff_cumulative):
@@ -275,6 +375,15 @@ def _cutoff_predictions(self, predictions: np.ndarray) -> np.ndarray:
maxidx = min(maxidx, self.cutoff_number) or 1
return sortidx[:maxidx]
+ def _load_mask_file(self, maskfile: str) -> np.ndarray:
+ self._logger.info(f"Loading masking of templates from {maskfile} to {self.key}")
+ mask = np.load(maskfile)["arr_0"]
+ if len(mask) != len(self.templates):
+ raise PolicyException(
+ f"The number of masks {len(mask)} does not match the number of templates {len(self.templates)}"
+ )
+ return mask
+
def _update_cache(self, molecules: Sequence[TreeMolecule]) -> None:
pred_inchis = []
fp_list = []
@@ -296,3 +405,46 @@ def _update_cache(self, molecules: Sequence[TreeMolecule]) -> None:
probable_transforms_idx,
pred[probable_transforms_idx],
)
+
+
+class TemplateBasedDirectExpansionStrategy(TemplateBasedExpansionStrategy):
+ """
+ A template-based expansion strategy that will return `SmilesBasedRetroReaction` objects upon expansion
+ by directly applying the template
+
+ :param key: the key or label
+ :param config: the configuration of the tree search
+ :param source: the source of the policy model
+ :param templatefile: the path to a HDF5 file with the templates
+ :raises PolicyException: if the length of the model output vector is not same as the number of templates
+ """
+
+ def get_actions(
+ self,
+ molecules: Sequence[TreeMolecule],
+ cache_molecules: Optional[Sequence[TreeMolecule]] = None,
+ ) -> Tuple[List[RetroReaction], List[float]]:
+ """
+ Get all the probable actions of a set of molecules, using the selected policies and given cutoffs
+
+ :param molecules: the molecules to consider
+ :param cache_molecules: additional molecules to submit to the expansion
+ policy but that only will be cached for later use
+ :return: the actions and the priors of those actions
+ """
+ possible_actions = []
+ priors = []
+
+ super_actions, super_priors = super().get_actions(molecules, cache_molecules)
+ for templated_action, prior in zip(super_actions, super_priors):
+ for reactants in templated_action.reactants:
+ reactants_str = ".".join(mol.smiles for mol in reactants)
+ new_action = SmilesBasedRetroReaction(
+ templated_action.mol,
+ metadata=templated_action.metadata,
+ reactants_str=reactants_str,
+ )
+ possible_actions.append(new_action)
+ priors.append(prior)
+
+ return possible_actions, priors # type: ignore
diff --git a/aizynthfinder/context/policy/filter_strategies.py b/aizynthfinder/context/policy/filter_strategies.py
index f76ef8c..dd13c67 100644
--- a/aizynthfinder/context/policy/filter_strategies.py
+++ b/aizynthfinder/context/policy/filter_strategies.py
@@ -9,14 +9,18 @@
from aizynthfinder.chem import TemplatedRetroReaction
from aizynthfinder.context.policy.utils import _make_fingerprint
-from aizynthfinder.utils.exceptions import PolicyException, RejectionException
+from aizynthfinder.utils.bonds import BrokenBonds
+from aizynthfinder.utils.exceptions import (
+ PolicyException,
+ RejectionException,
+)
from aizynthfinder.utils.logging import logger
from aizynthfinder.utils.models import load_model
if TYPE_CHECKING:
from aizynthfinder.chem.reaction import RetroReaction
from aizynthfinder.context.config import Configuration
- from aizynthfinder.utils.type_utils import Any, List, Tuple
+ from aizynthfinder.utils.type_utils import Any, Dict, List, Optional, Tuple
class FilterStrategy(abc.ABC):
@@ -62,6 +66,33 @@ def apply(self, reaction: RetroReaction) -> None:
"""
+class BondFilter(FilterStrategy):
+ """
+ Check if focussed bonds to freeze stay frozen in a reaction.
+
+ :param key: the key or label
+ :param config: the configuration of the tree search
+ """
+
+ def __init__(self, key: str, config: Configuration, **kwargs: Any) -> None:
+ super().__init__(key, config, **kwargs)
+
+ self._freeze_bonds = config.search.freeze_bonds
+ self._broken_bonds = BrokenBonds(self._freeze_bonds)
+ self._logger.info(
+ f"Loading bond filter to {key} with {len(self._freeze_bonds)} "
+ "bonds to freeze"
+ )
+
+ def apply(self, reaction: RetroReaction) -> None:
+ broken_frozen_bonds = self._broken_bonds(reaction)
+ if len(broken_frozen_bonds) > 0:
+ raise RejectionException(
+ f"{reaction} was filtered out as the focussed bonds "
+ f"'{broken_frozen_bonds}' were found to be broken in the reaction"
+ )
+
+
class QuickKerasFilter(FilterStrategy):
"""
Filter quick-filter trained on artificial negative data
diff --git a/aizynthfinder/context/policy/policies.py b/aizynthfinder/context/policy/policies.py
index cb13152..a72c088 100644
--- a/aizynthfinder/context/policy/policies.py
+++ b/aizynthfinder/context/policy/policies.py
@@ -222,3 +222,12 @@ def load_from_config(self, **config: Any) -> None:
del kwargs["type"]
obj = cls(key, self._config, **kwargs)
self.load(obj)
+
+ def reset_cache(self) -> None:
+ """Reset filtering cache."""
+ if not self.selection:
+ return
+
+ for name in self.selection:
+ if hasattr(self[name], "reset_cache"):
+ self[name].reset_cache()
diff --git a/aizynthfinder/context/scoring/__init__.py b/aizynthfinder/context/scoring/__init__.py
index 92b9e17..c6cb73b 100644
--- a/aizynthfinder/context/scoring/__init__.py
+++ b/aizynthfinder/context/scoring/__init__.py
@@ -1,9 +1,12 @@
""" Sub-package containing scoring routines
"""
+
from aizynthfinder.context.scoring.collection import ScorerCollection
from aizynthfinder.context.scoring.scorers import (
AverageTemplateOccurrenceScorer,
+ BrokenBondsScorer,
CombinedScorer,
+ DeltaSyntheticComplexityScorer,
FractionInStockScorer,
MaxTransformScorerer,
NumberOfPrecursorsInStockScorer,
@@ -12,8 +15,10 @@
PriceSumScorer,
ReactionClassMembershipScorer,
RouteCostScorer,
+ RouteSimilarityScorer,
Scorer,
StateScorer,
StockAvailabilityScorer,
+ SUPPORT_DISTANCES,
)
from aizynthfinder.utils.exceptions import ScorerException
diff --git a/aizynthfinder/context/scoring/collection.py b/aizynthfinder/context/scoring/collection.py
index 3fc84fe..36d5c40 100644
--- a/aizynthfinder/context/scoring/collection.py
+++ b/aizynthfinder/context/scoring/collection.py
@@ -61,24 +61,33 @@ def __init__(self, config: Configuration) -> None:
self._config = config
self.create_default_scorers()
+ def __repr__(self) -> str:
+ if self.selection:
+ return f"{self._collection_name} ({', '.join(self.selection)})"
+
+ return f"{self._collection_name} ({', '.join(self.items)})"
+
def create_default_scorers(self) -> None:
"""
Setup the scores that only need the config as their input.
"""
for cls in _SIMPLE_SCORERS:
- self.load(cls(self._config))
+ self.load(cls(self._config), silent=True)
- def load(self, scorer: Scorer) -> None: # type: ignore
+ def load(self, scorer: Scorer, silent: bool = False) -> None: # type: ignore
"""
Add a pre-initialized scorer object to the collection
:param scorer: the item to add
+ :param silent: if True will not write out the name of the loaded scorer
"""
if not isinstance(scorer, Scorer):
raise ScorerException(
"Only objects of classes inherited from Scorer can be added"
)
self._items[repr(scorer)] = scorer
+ if not silent:
+ self._logger.info(f"Loaded scorer: {repr(scorer)}")
def load_from_config(self, **scorers_config: Any) -> None:
"""
@@ -119,7 +128,8 @@ def make_subset(self, subset_names: List[str]) -> "ScorerCollection":
scorer = self._items[name]
except KeyError:
raise ScorerException(f"Unable to find '{name}' in parent collection")
- new_collection.load(scorer)
+ new_collection.load(scorer, silent=True)
+ new_collection.select(name, append=True)
return new_collection
def names(self) -> List[str]:
diff --git a/aizynthfinder/context/scoring/scorers.py b/aizynthfinder/context/scoring/scorers.py
index b9ac45d..886e4b4 100644
--- a/aizynthfinder/context/scoring/scorers.py
+++ b/aizynthfinder/context/scoring/scorers.py
@@ -1,8 +1,10 @@
""" Module containing classes used to score the reaction routes.
"""
+
from __future__ import annotations
import abc
+import json
from collections import defaultdict
from collections.abc import Sequence as SequenceAbc
from dataclasses import dataclass
@@ -10,14 +12,29 @@
import numpy as np
+try:
+ from route_distances.route_distances import route_distances_calculator
+except ImportError:
+ SUPPORT_DISTANCES = False
+else:
+ SUPPORT_DISTANCES = True
+
from aizynthfinder.chem import TreeMolecule
from aizynthfinder.context.stock import StockException
from aizynthfinder.reactiontree import ReactionTree
from aizynthfinder.search.mcts import MctsNode
+from aizynthfinder.utils.bonds import BrokenBonds
from aizynthfinder.utils.exceptions import ScorerException
+from aizynthfinder.utils.logging import logger
+from aizynthfinder.utils.sc_score import SCScore
if TYPE_CHECKING:
- from aizynthfinder.chem import FixedRetroReaction, Molecule, RetroReaction
+ from aizynthfinder.chem import (
+ FixedRetroReaction,
+ Molecule,
+ RetroReaction,
+ UniqueMolecule,
+ )
from aizynthfinder.context.config import Configuration
from aizynthfinder.utils.type_utils import (
Iterable,
@@ -29,6 +46,7 @@
Union,
)
+ _Molecules = Sequence[Molecule]
_Scoreable = TypeVar("_Scoreable", MctsNode, ReactionTree)
_Scoreables = Sequence[_Scoreable]
_ScorerItemType = Union[_Scoreables, _Scoreable]
@@ -515,7 +533,8 @@ class StockAvailabilityScorer(Scorer):
The score is calculated as a product of a stock score per starting material. The stock
score for each molecule is based on the source of the stock, or a default value if the
- molecule is not in stock.
+ molecule is not in stock. The `other_source_score` parameter can be used to
+ distinguish between "not in stock" and "not in the specificed sources" cases.
"""
def __init__(
@@ -523,11 +542,13 @@ def __init__(
config: Configuration,
source_score: StrDict,
default_score: float = 0.1,
+ other_source_score: Optional[float] = None,
) -> None:
super().__init__(config)
assert self._config is not None
self.source_score = source_score
self.default_score = default_score
+ self.other_source_score = other_source_score
def __repr__(self) -> str:
return "stock availability"
@@ -544,7 +565,12 @@ def _calculate_leaf_costs(
for source in availability
if source in self.source_score
]
- prod *= max(scores) if scores else self.default_score
+ if scores:
+ prod *= max(scores)
+ elif self.other_source_score and mol in self._config.stock:
+ prod *= self.other_source_score
+ else:
+ prod *= self.default_score
return prod
def _score_node(self, node: MctsNode) -> float:
@@ -554,6 +580,269 @@ def _score_reaction_tree(self, tree: ReactionTree) -> float:
return self._calculate_leaf_costs(tree.leafs())
+class BrokenBondsScorer(Scorer):
+ """Class for scoring nodes and reaction trees based on the breaking of atom bonds
+
+ The score is a summation of the depths in the tree where the focussed bonds
+ are found to break in the reaction. If a focussed bond is found to be unbroken
+ in the entire tree, the total length of the tree will be added to the score.
+ """
+
+ scorer_name = "broken bonds"
+
+ def __init__(self, config: Configuration) -> None:
+ super().__init__(config)
+ self._break_bonds = [tuple(bond) for bond in config.search.break_bonds]
+ self._break_bonds_operator = config.search.break_bonds_operator.lower()
+ self._reverse_order = False
+
+ def __repr__(self) -> str:
+ return self.scorer_name
+
+ def _calculate_broken_bonds_score(
+ self, reactions: Sequence[RetroReaction], depths: Sequence[int]
+ ) -> float:
+
+ if not reactions:
+ return 0
+
+ max_score = len(set(depths)) * len(self._break_bonds)
+ broken_focussed_bonds = []
+ scores = []
+
+ # The score should be 0 when no transformations were reported
+ if not depths:
+ return 0
+
+ for reaction, depth in zip(reactions, depths):
+ broken_bonds = BrokenBonds(self._break_bonds)(reaction)
+ broken_untracked_bonds = [
+ bond for bond in broken_bonds if bond not in broken_focussed_bonds
+ ]
+ if len(broken_untracked_bonds) > 0:
+ broken_focussed_bonds += broken_untracked_bonds
+ scores.append(1 - (depth / max_score))
+ if self._break_bonds_operator == "or":
+ break
+
+ if self._break_bonds_operator != "or" or (
+ self._break_bonds_operator == "or" and len(broken_focussed_bonds) == 0
+ ):
+ # type: ignore
+ if unbroken_bonds := set(self._break_bonds) - set(broken_focussed_bonds):
+ scores.append(1 - (len(set(depths)) * len(unbroken_bonds)) / max_score)
+ return sum(scores) / len(scores)
+
+ def _score_node(self, node: MctsNode) -> float:
+ reactions = node.actions_to()
+ depths = []
+ for reaction in reactions:
+ depth = (
+ max(
+ reactant.transform
+ for reactant in reaction.reactants[reaction.index]
+ )
+ - 1
+ )
+ depths.append(depth)
+ return self._calculate_broken_bonds_score(reactions, depths)
+
+ def _score_reaction_tree(self, tree: ReactionTree) -> float:
+ reactions = [
+ reaction.to_smiles_based_retroreaction() for reaction in tree.reactions()
+ ]
+ depths = [tree.depth(reaction) // 2 for reaction in tree.reactions()]
+ return self._calculate_broken_bonds_score(reactions, depths)
+
+
+class RouteSimilarityScorer(Scorer):
+ """
+ Class for scoring based on an LSTM model for computing Tree Edit Distance to
+ a set of reference routes.
+
+ :param config: the configuration of the tree search
+ :param routes_path: the filename of a JSON file with reference routes
+ :param model_path: the filename of a checkpoint file with the LSTM model
+ :param scaler_params: the parameter settings of the scaler
+ :param agg_func: the name of numpy function used to aggregate the distances
+ to the reference routes
+ :param similarity: if True, will compute similarity score else distance scores
+ """
+
+ scorer_name = "route similarity"
+
+ def __init__(
+ self,
+ config: Configuration,
+ routes_path: str,
+ model_path: str,
+ scaler_params: Optional[StrDict] = None,
+ agg_func: str = "min",
+ similarity: bool = False,
+ ) -> None:
+ if not SUPPORT_DISTANCES:
+ raise ValueError(
+ "Distance calculations are not supported by this installation."
+ " Please install aizynthfinder with extras dependencies."
+ )
+
+ # Default scaler from benchmarking
+ if scaler_params is None:
+ scaler_params = {
+ "name": "squash",
+ "slope": 0.5,
+ "xoffset": 10,
+ "yoffset": 0,
+ }
+
+ super().__init__(config, scaler_params)
+ self._reverse_order = similarity
+ self.calculator = route_distances_calculator("lstm", model_path=model_path)
+
+ # Scalar needs to be applied before taking into account of `similarity` parameter,
+ # hence we are reassigning the global scaler and then disabling it.
+ self._local_scaler = self._scaler
+ self._scaler = None
+
+ if not hasattr(np, agg_func):
+ raise ValueError(f"Cannot identify aggregate function {agg_func} in numpy")
+ self.agg_func = getattr(np, agg_func)
+ self.similarity = similarity
+
+ try:
+ with open(routes_path or "") as file:
+ self.routes = json.load(file)
+ except FileNotFoundError:
+ logger().info(
+ f"Could not load reference routes from {routes_path}. Assuming they will be set later"
+ )
+ self.routes = []
+ self.n_routes = len(self.routes)
+
+ def _score_node(self, node: MctsNode) -> float:
+ # We don't have any short-cut to score a node,
+ # so we need to convert it to a reaction tree
+ return self._score_reaction_tree(node.to_reaction_tree())
+
+ def _score_reaction_tree(self, tree: ReactionTree) -> float:
+ if not self.routes:
+ return 0.0 if self.similarity else 1.0
+
+ dist_matrix = self.calculator(self.routes + [tree.to_dict()])
+ distances = dist_matrix[self.n_routes, : self.n_routes]
+ score = self.agg_func(distances)
+ norm_score = self._local_scaler(score)
+ if self.similarity:
+ return 1.0 - float(norm_score)
+ return float(norm_score)
+
+
+class DeltaSyntheticComplexityScorer(Scorer):
+ """
+ Class for scoring nodes based on the delta-synthetic-complexity of the node
+ and its parent 'horizon' steps up in the tree.
+
+ :param config: the configuration the tree search
+ :param sc_score_model: the path to the SCScore model
+ :param scaler_params: the parameter settings of the scaler, defaults to max-min between -1.5 and 4
+ :param horizon: the number of steps backwards to look for parent molecule
+ """
+
+ scorer_name: str = "delta-SC score"
+
+ def __init__(
+ self,
+ config: Configuration,
+ sc_score_model: str,
+ scaler_params: Optional[StrDict] = None,
+ horizon: int = 3,
+ ) -> None:
+ if scaler_params is None:
+ scaler_params = {
+ "name": "min_max",
+ "min_val": -1.5,
+ "max_val": 4,
+ "reverse": False,
+ }
+ super().__init__(config, scaler_params)
+ ## This is necessary because config should not be optional for this scorer
+ self._config: Configuration = config
+
+ self.horizon = horizon
+ self._model = SCScore(sc_score_model)
+
+ def sc_deltas(self, mols: _Molecules, parents: _Molecules) -> Sequence[float]:
+ """
+ Calculate delta SC-score among the list of mol-parent pairs.
+
+ :params mols: the leaves of the tree
+ :param parents: the parent of the leaves
+ :returns: the pair-wise difference in SCScore
+ """
+ delta_sc_scores = []
+ for mol, parent in zip(mols, parents):
+ mol.sanitize()
+ parent.sanitize()
+
+ sc_score_mol = self._model(mol.rd_mol)
+ sc_score_parent = self._model(parent.rd_mol)
+
+ delta_sc_score = sc_score_parent - sc_score_mol
+ delta_sc_scores.append(delta_sc_score)
+ return delta_sc_scores
+
+ def _get_parent_from_reaction_tree(
+ self, tree: ReactionTree, mol: UniqueMolecule, horizon: int
+ ) -> UniqueMolecule:
+ horizon = min(horizon, tree.depth(mol) // 2)
+
+ parent = mol
+ # Step up in the tree, first a reaction, then a molecule
+ for _ in range(horizon):
+ parent = tree.parent_molecule(parent)
+
+ return parent
+
+ def _get_parent_from_tree_molecule(
+ self, mol: TreeMolecule, horizon: int
+ ) -> TreeMolecule:
+ horizon = min(horizon, mol.transform)
+
+ parent = mol
+ # Step up in the tree via parent nodes
+ for _ in range(horizon):
+ grandparent = parent.parent
+ if not grandparent:
+ break
+ parent = grandparent
+
+ return parent
+
+ def _score_node(self, node: MctsNode) -> float:
+ expandable_mols = node.state.expandable_mols
+ if not expandable_mols:
+ expandable_mols = list(node.state.mols)
+
+ parent_mols = [
+ self._get_parent_from_tree_molecule(mol, self.horizon)
+ for mol in expandable_mols
+ ]
+ delta_sc_scores = self.sc_deltas(expandable_mols, parent_mols)
+ return min(delta_sc_scores)
+
+ def _score_reaction_tree(self, tree: ReactionTree) -> float:
+ expandable_mols = [node for node in tree.leafs() if not tree.in_stock(node)]
+ if not expandable_mols:
+ expandable_mols = list(tree.leafs())
+
+ parent_mols = [
+ self._get_parent_from_reaction_tree(tree, node, self.horizon)
+ for node in expandable_mols
+ ]
+ delta_sc_scores = self.sc_deltas(expandable_mols, parent_mols)
+ return min(delta_sc_scores)
+
+
class CombinedScorer(Scorer):
"""Class for scoring nodes and reaction trees by combining weighted scores from a list of scorers
diff --git a/aizynthfinder/context/stock/queries.py b/aizynthfinder/context/stock/queries.py
index 9938ed5..e9d70e4 100644
--- a/aizynthfinder/context/stock/queries.py
+++ b/aizynthfinder/context/stock/queries.py
@@ -1,5 +1,6 @@
""" Module containing classes that interfaces different stocks query classes
"""
+
from __future__ import annotations
import os
@@ -22,11 +23,9 @@
from pymongo.collection import Collection as MongoCollection
from pymongo.database import Database as MongoDatabase
- # pylint: disable=ungrouped-imports
from aizynthfinder.utils.type_utils import Optional, Set, StrDict
-# pylint: disable=no-self-use
class StockQueryMixin:
"""
Mixin class for all query classes, providing a default interface
@@ -120,7 +119,7 @@ def __init__(
if ext in [".h5", ".hdf5"]:
stock_df: pd.DataFrame = pd.read_hdf(path, key="table")
- elif ext == ".csv":
+ else:
stock_df = pd.read_csv(
path,
usecols=[inchi_key_col, price_col] if price_col else [inchi_key_col],
diff --git a/aizynthfinder/context/stock/stock.py b/aizynthfinder/context/stock/stock.py
index fc322c2..e0eb893 100644
--- a/aizynthfinder/context/stock/stock.py
+++ b/aizynthfinder/context/stock/stock.py
@@ -110,10 +110,10 @@ def availability_list(self, mol: Molecule) -> List[str]:
if mol not in self[key]:
continue
try:
- availability.append(self[key].availability_string(mol))
+ availability.extend(self[key].availability_string(mol).split(","))
except (StockException, AttributeError):
availability.append(key)
- return availability
+ return sorted(list(set(availability)))
def availability_string(self, mol: Molecule) -> str:
"""
diff --git a/aizynthfinder/interfaces/aizynthapp.py b/aizynthfinder/interfaces/aizynthapp.py
index 5baca55..c455537 100644
--- a/aizynthfinder/interfaces/aizynthapp.py
+++ b/aizynthfinder/interfaces/aizynthapp.py
@@ -27,6 +27,7 @@
from rdkit import Chem
from aizynthfinder.aizynthfinder import AiZynthFinder
+from aizynthfinder.interfaces.gui.utils import pareto_fronts_plot, route_display
from aizynthfinder.utils.logging import setup_logger
if TYPE_CHECKING:
@@ -53,7 +54,6 @@ class AiZynthApp:
"""
def __init__(self, configfile: str, setup: bool = True) -> None:
- # pylint: disable=used-before-assignment
setup_logger(logging.INFO)
self.finder = AiZynthFinder(configfile=configfile)
self._input: StrDict = dict()
@@ -152,10 +152,24 @@ def _create_input_widgets(self) -> None:
value=self.finder.config.search.max_transforms,
style={"description_width": "initial"},
)
- self._input["exclude_target_from_stock"] = widgets.Checkbox(
- value=self.finder.config.search.exclude_target_from_stock,
- description="Exclude target from stock",
+
+ rewards = self.finder.config.search.algorithm_config.get("search_rewards", [])
+ objective1 = rewards[0] if len(rewards) > 0 else "state score"
+ objective2 = rewards[1] if len(rewards) > 1 else "None"
+ self._input["reward1"] = widgets.Dropdown(
+ options=self.finder.scorers.names(),
+ description="MCTS reward 1:",
+ value=objective1,
+ style={"description_width": "initial"},
+ )
+
+ self._input["reward2"] = widgets.Dropdown(
+ options=["None"] + self.finder.scorers.names(),
+ description="MCTS reward 2:",
+ value=objective2,
+ style={"description_width": "initial"},
)
+
vbox = VBox(
[
self._input["policy"],
@@ -164,7 +178,8 @@ def _create_input_widgets(self) -> None:
max_iter_box,
self._input["return_first"],
self._input["max_transforms"],
- self._input["exclude_target_from_stock"],
+ self._input["reward1"],
+ self._input["reward2"],
]
)
box_options = HBox([box_stocks, vbox])
@@ -176,13 +191,15 @@ def _create_route_widgets(self) -> None:
description="Reorder by:",
style={"description_width": "initial"},
)
+ self._input["scorer"].disabled = True
self._input["scorer"].observe(self._on_change_scorer)
- self._buttons["show_routes"] = Button(description="Show Reactions")
+ self._buttons["show_routes"] = Button(description="Show Routes")
self._buttons["show_routes"].on_click(self._on_display_button_clicked)
self._input["route"] = Dropdown(
options=[],
description="Routes: ",
)
+ self._input["route"].disabled = True
self._input["route"].observe(self._on_change_route_option)
display(
HBox(
@@ -194,6 +211,10 @@ def _create_route_widgets(self) -> None:
)
)
+ self._output["pareto_fronts"] = widgets.Output(
+ layout={"border": "1px solid silver", "width": "99%", "overflow": "auto"}
+ )
+ display(self._output["pareto_fronts"])
self._output["routes"] = widgets.Output(
layout={"border": "1px solid silver", "width": "99%"}
)
@@ -251,13 +272,22 @@ def _on_extend_button_clicked(self, _) -> None:
def _on_display_button_clicked(self, _) -> None:
self._toggle_button(False)
- self.finder.build_routes()
+ rewards = self.finder.config.search.algorithm_config["search_rewards"]
+ self.finder.build_routes(scorer=rewards)
self.finder.routes.make_images()
self.finder.routes.compute_scores(*self.finder.scorers.objects())
self._input["route"].options = [
f"Option {i}" for i, _ in enumerate(self.finder.routes, 1) # type: ignore
]
+
+ self._output["pareto_fronts"].clear_output()
+ if len(rewards) > 1:
+ with self._output["pareto_fronts"]:
+ pareto_fronts_plot(self.finder.routes)
+
self._show_route(0)
+ self._input["scorer"].disabled = len(rewards) > 1
+ self._input["route"].disabled = False
self._toggle_button(True)
def _prepare_search(self) -> None:
@@ -288,12 +318,19 @@ def _prepare_search(self) -> None:
self.finder.config.search.iteration_limit = self._input[
"iteration_limit"
].value
- self.finder.config.search.exclude_target_from_stock = self._input[
- "exclude_target_from_stock"
- ].value
+
+ rewards = [self._input["reward1"].value]
+ if self._input["reward2"].value != "None":
+ rewards.append(self._input["reward2"].value)
+ self.finder.config.search.algorithm_config["search_rewards"] = rewards
+ if self.finder.config.search.algorithm != "mcts":
+ print(
+ "Only MCTS algorithm is supported in GUI interface. Resetting to MCTS"
+ )
+ self.finder.config.search.algorithm = "mcts"
smiles = self._input["smiles"].value
- print("Setting target molecule with smiles: %s" % smiles)
+ print(f"Setting target molecule with smiles: {smiles}")
self.finder.target_smiles = smiles
self.finder.prepare_tree()
@@ -304,29 +341,7 @@ def _show_mol(self, change) -> None:
display(mol)
def _show_route(self, index) -> None:
- if (
- index is None
- or self.finder.routes is None
- or index >= len(self.finder.routes)
- ):
- return
-
- route = self.finder.routes[index]
- state = route["node"].state
- status = "Solved" if state.is_solved else "Not Solved"
-
- self._output["routes"].clear_output()
- with self._output["routes"]:
- display(HTML("
%s" % status))
- table_content = "".join(
- f"
{name}
{score:.4f}
"
- for name, score in route["all_score"].items()
- )
- display(HTML(f"
{table_content}
"))
- display(HTML("
Compounds to Procure"))
- display(state.to_image())
- display(HTML("
Steps"))
- display(self.finder.routes[index]["image"])
+ route_display(index, self.finder.routes, self._output["routes"])
def _toggle_button(self, on_) -> None:
for button in self._buttons.values():
diff --git a/aizynthfinder/interfaces/aizynthcli.py b/aizynthfinder/interfaces/aizynthcli.py
index 221414e..81f7dfb 100644
--- a/aizynthfinder/interfaces/aizynthcli.py
+++ b/aizynthfinder/interfaces/aizynthcli.py
@@ -1,5 +1,6 @@
""" Module containing classes and routines for the CLI
"""
+
from __future__ import annotations
import argparse
@@ -36,6 +37,7 @@
)
_PostProcessingJob = Callable[[AiZynthFinder], StrDict]
+ _PreProcessingJob = Callable[[AiZynthFinder, int], None]
def _do_clustering(
@@ -115,6 +117,9 @@ def _get_arguments() -> argparse.Namespace:
nargs="+",
help="a number of modules that performs post-processing tasks",
)
+ parser.add_argument(
+ "--pre_processing", help="a module that perform pre-processing tasks"
+ )
parser.add_argument(
"--checkpoint",
required=False,
@@ -137,6 +142,21 @@ def _load_postprocessing_jobs(modules: Optional[List[str]]) -> List[_PostProcess
return jobs
+def _load_preprocessing_job(module_name: Optional[str]) -> Optional[_PreProcessingJob]:
+ if not module_name:
+ return None
+
+ try:
+ module = importlib.import_module(module_name)
+ except ModuleNotFoundError:
+ pass
+ else:
+ if hasattr(module, "pre_processing"):
+ print(f"Adding pre-processing job from {module_name}")
+ return getattr(module, "pre_processing")
+ return None
+
+
def _select_stocks(finder: AiZynthFinder, args: argparse.Namespace) -> None:
stocks = list(args.stocks)
try:
@@ -176,9 +196,12 @@ def _process_single_smiles(
do_clustering: bool,
route_distance_model: Optional[str],
post_processing: List[_PostProcessingJob],
+ pre_processing: Optional[_PreProcessingJob],
) -> None:
output_name = output_name or "trees.json"
finder.target_smiles = smiles
+ if pre_processing:
+ pre_processing(finder, -1)
try:
finder.prepare_tree()
except ValueError as err:
@@ -196,9 +219,6 @@ def _process_single_smiles(
)
logger().info(f"Trees saved to {output_name}")
- scores = ", ".join("%.4f" % score for score in finder.routes.scores)
- logger().info(f"Scores for best routes: {scores}")
-
stats = finder.extract_statistics()
if do_clustering:
_do_clustering(
@@ -218,6 +238,7 @@ def _process_multi_smiles(
do_clustering: bool,
route_distance_model: Optional[str],
post_processing: List[_PostProcessingJob],
+ pre_processing: Optional[_PreProcessingJob],
checkpoint: Optional[str],
) -> None:
output_name = output_name or "output.json.gz"
@@ -237,7 +258,9 @@ def _process_multi_smiles(
for key, value in checkpoint_data.items()
if key != "processed_smiles"
}
- for smi in smiles:
+ for idx, smi in enumerate(smiles):
+ if pre_processing:
+ pre_processing(finder, idx)
processed_results = {}
finder.target_smiles = smi
try:
@@ -261,9 +284,6 @@ def _process_multi_smiles(
for key, value in stats.items():
processed_results[key] = value
processed_results["stock_info"] = finder.stock_info()
- processed_results["top_scores"] = ", ".join(
- "%.4f" % score for score in finder.routes.scores
- )
processed_results["trees"] = finder.routes.dict_with_extra(
include_metadata=True, include_scores=True
)
@@ -355,6 +375,7 @@ def main() -> None:
finder = AiZynthFinder(configfile=args.config)
_select_stocks(finder, args)
post_processing = _load_postprocessing_jobs(args.post_processing)
+ pre_processing = _load_preprocessing_job(args.pre_processing)
finder.expansion_policy.select(args.policy or finder.expansion_policy.items[0])
if args.filter:
finder.filter_policy.select(args.filter)
@@ -368,6 +389,7 @@ def main() -> None:
args.cluster,
args.route_distance_model,
post_processing,
+ pre_processing,
args.checkpoint,
]
if multi_smiles:
diff --git a/aizynthfinder/interfaces/gui/pareto_fronts.py b/aizynthfinder/interfaces/gui/pareto_fronts.py
new file mode 100644
index 0000000..75ff36c
--- /dev/null
+++ b/aizynthfinder/interfaces/gui/pareto_fronts.py
@@ -0,0 +1,173 @@
+"""Module containing a GUI extension for plotting pareto fronts
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from IPython.display import display
+from ipywidgets import Button, Checkbox, Dropdown, HBox, IntRangeSlider, Output, VBox
+from aizynthfinder.analysis import (
+ RouteCollection,
+ RouteSelectionArguments,
+ TreeAnalysis,
+)
+from aizynthfinder.interfaces.gui.utils import pareto_fronts_plot, route_display
+
+if TYPE_CHECKING:
+ from aizynthfinder.context.scoring.collection import ScorerCollection
+ from aizynthfinder.interfaces.aizynthapp import AiZynthApp
+ from aizynthfinder.search.andor_trees import AndOrSearchTreeBase
+ from aizynthfinder.search.mcts import MctsSearchTree
+ from aizynthfinder.utils.type_utils import Optional, StrDict, Union
+
+
+class ParetoFrontsGUI:
+ """GUI extension for plotting pareto fronts of routes.
+
+ :param routes: the routes to plot
+ :param scorers: the available scorers
+ """
+
+ def __init__(
+ self,
+ tree: Union[MctsSearchTree, AndOrSearchTreeBase],
+ scorers: ScorerCollection,
+ ) -> None:
+ self._tree = tree
+ self._scorers = scorers
+ self._routes: Optional[RouteCollection] = None
+ self._input: StrDict = dict()
+ self._output: StrDict = dict()
+ self._buttons: StrDict = dict()
+ self._create_input()
+ self._create_route_widgets()
+
+ @classmethod
+ def from_app(cls, app: AiZynthApp):
+ """,
+ Helper function to create a GUI from a GUI app interface
+
+ :param app: the app to extract the routes from
+ :return: the GUI object
+ """
+ if app.finder.tree is None:
+ raise ValueError("Cannot initialize GUI from no tree")
+ return ParetoFrontsGUI(app.finder.tree, app.finder.scorers)
+
+ def _toggle_button(self, on_) -> None:
+ for button in self._buttons.values():
+ button.disabled = not on_
+
+ def _create_input(self) -> None:
+ self._input["scorer_1"] = Dropdown(
+ options=self._scorers.names(),
+ description="Scorer 1:",
+ style={"description_width": "initial"},
+ )
+ self._input["scorer_2"] = Dropdown(
+ options=self._scorers.names(),
+ description="Scorer 2:",
+ style={"description_width": "initial"},
+ )
+ self._input["select_routes"] = IntRangeSlider(
+ value=[5, 25],
+ min=1,
+ max=50,
+ step=1,
+ description="Select #Routes:",
+ disabled=False,
+ continuous_update=True,
+ orientation="horizontal",
+ readout=True,
+ readout_format="d",
+ style={"description_width": "initial"},
+ )
+ self._input["select_all_routes"] = Checkbox(
+ value=False,
+ description="Select All Routes",
+ style={"description_width": "initial"},
+ layout={"justify": "left"},
+ )
+ selection_box = HBox(
+ [
+ self._input["select_routes"],
+ self._input["select_all_routes"],
+ ]
+ )
+ scorers_box = HBox(
+ [
+ self._input["scorer_1"],
+ self._input["scorer_2"],
+ ]
+ )
+ display(
+ VBox([selection_box, scorers_box], layout={"border": "1px solid silver"})
+ )
+
+ def _create_route_widgets(self) -> None:
+ self._buttons["show_routes"] = Button(description="Show Routes")
+ self._buttons["show_routes"].on_click(self._on_display_routes_button_clicked)
+ self._input["route"] = Dropdown(
+ options=[],
+ description="Routes: ",
+ )
+ self._input["route"].observe(self._on_change_route_option)
+ display(
+ HBox(
+ [
+ self._buttons["show_routes"],
+ self._input["route"],
+ ]
+ )
+ )
+
+ self._output["pareto_fronts"] = Output(
+ layout={"border": "1px solid silver", "width": "99%", "overflow": "auto"}
+ )
+ display(self._output["pareto_fronts"])
+
+ self._output["routes"] = Output(
+ layout={"border": "1px solid silver", "width": "99%", "overflow": "auto"}
+ )
+ display(self._output["routes"])
+
+ def _on_display_routes_button_clicked(self, _) -> None:
+ self._toggle_button(False)
+
+ objectives = [self._input["scorer_1"].value, self._input["scorer_2"].value]
+ mo_scorers = [self._scorers[name] for name in objectives]
+ analysis = TreeAnalysis(self._tree, scorer=mo_scorers)
+
+ config_selection = RouteSelectionArguments(
+ nmin=self._input["select_routes"].value[0],
+ nmax=self._input["select_routes"].value[1],
+ return_all=self._input["select_all_routes"].value,
+ )
+ self._routes = RouteCollection.from_analysis(analysis, config_selection)
+ self._routes.compute_scores(*self._scorers.objects())
+ self._routes.make_images()
+
+ self._input["route"].options = [
+ f"Option {i}" for i, _ in enumerate(self._routes, 1) # type: ignore
+ ]
+
+ with self._output["pareto_fronts"]:
+ pareto_fronts_plot(self._routes)
+
+ self._show_route(0)
+ self._toggle_button(True)
+
+ def _on_change_route_option(self, change) -> None:
+ if change["name"] != "index":
+ return
+ self._show_route(self._input["route"].index)
+
+ def _show_route(self, index) -> None:
+ if self._routes is None:
+ return
+ route_display(index, self._routes, self._output["routes"])
+
+ def _reset_widgets(self):
+ self._output["pareto_fronts"].clear_output()
+ self._output["routes"].clear_output()
+ self._input["route"].options = []
diff --git a/aizynthfinder/interfaces/gui/utils.py b/aizynthfinder/interfaces/gui/utils.py
new file mode 100644
index 0000000..7e62e7a
--- /dev/null
+++ b/aizynthfinder/interfaces/gui/utils.py
@@ -0,0 +1,129 @@
+"""Module containing utility functions for GUI.
+"""
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import ipywidgets as widgets
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from IPython.display import HTML, display
+from paretoset import paretorank
+
+if TYPE_CHECKING:
+ from aizynthfinder.analysis.routes import RouteCollection
+ from aizynthfinder.utils.type_utils import List, Optional
+
+
+def pareto_fronts_plot(
+ routes: RouteCollection,
+) -> None:
+ """Plot the pareto front(s).
+
+ :param routes: the route collection to plot as Pareto fronts
+ """
+ scorer_names = list(routes.scores[0].keys())
+ scores = np.array(
+ [[score_dict[name] for name in scorer_names] for score_dict in routes.scores]
+ )
+ direction_arr = np.repeat("max", len(scorer_names))
+ pareto_ranks = paretorank(scores, sense=direction_arr, distinct=False)
+
+ pareto_fronts = pd.DataFrame(scores, columns=scorer_names)
+ pareto_fronts.loc[:, "pareto_rank"] = pareto_ranks
+ pareto_fronts.loc[:, "route"] = np.arange(1, scores.shape[0] + 1)
+ pareto_fronts_unique = pareto_fronts.drop_duplicates(scorer_names)
+ # Apply the default theme
+ sns.set_theme()
+ fig = sns.relplot(
+ data=pareto_fronts_unique,
+ x=scorer_names[0],
+ y=scorer_names[1],
+ hue="pareto_rank",
+ kind="line",
+ markers=True,
+ )
+ fig.set(
+ xlabel=f"{scorer_names[0]}",
+ ylabel=f"{scorer_names[1]}",
+ title="Pareto Fronts",
+ )
+
+ objectives2solutions = defaultdict(list)
+ for _, row in pareto_fronts.iterrows():
+ x_val = row[scorer_names[0]]
+ y_val = row[scorer_names[1]]
+ objectives2solutions[(x_val, y_val)].append(int(row["route"]))
+
+ # Add route label on pareto line(s)
+ for _, row in pareto_fronts_unique.iterrows():
+ x_val = row[scorer_names[0]]
+ y_val = row[scorer_names[1]]
+ point_val = f"Option {_values_to_string(objectives2solutions[(x_val, y_val)])}"
+ plt.text(x=x_val, y=y_val, s=point_val, size=10)
+ plt.show()
+
+
+def route_display(
+ index: Optional[int], routes: RouteCollection, output_widget: widgets.Output
+) -> None:
+ """
+ Display a route with auxillary information in a widget
+
+ :param index: the index of the route to display
+ :param routes: the route collection
+ :param output_widget: the widget to display the route on
+ """
+ if index is None or routes is None or index >= len(routes):
+ return
+
+ route = routes[index]
+ state = route["node"].state
+ status = "Solved" if state.is_solved else "Not Solved"
+
+ output_widget.clear_output()
+ with output_widget:
+ display(HTML(f"
{status}"))
+ table_content = "".join(
+ f"
{name}
{score:.4f}
"
+ if isinstance(score, (float, int))
+ else f"
{name}
({', '.join(f'{val:.4f}' for val in score)})
"
+ for name, score in route["all_score"].items()
+ )
+ display(HTML(f"
{table_content}
"))
+ display(HTML("
Compounds to Procure"))
+ display(state.to_image())
+ display(HTML("