| | |
| |
|
| | import os |
| | import numpy as np |
| | import logging |
| | import time |
| | from .mol_tree import MolTree |
| |
|
| | def molstar(target_mol, target_mol_id, starting_mols, expand_fn, value_fn, |
| | iterations, viz=False, viz_dir=None, max_time=300): |
| | |
| | mol_tree = MolTree( |
| | target_mol=target_mol, |
| | known_mols=starting_mols, |
| | value_fn=value_fn |
| | ) |
| |
|
| | i = -1 |
| | start_time = time.time() |
| |
|
| | if not mol_tree.succ: |
| | for i in range(iterations): |
| | if time.time() - start_time > max_time: |
| | break |
| | |
| | scores = [] |
| | for m in mol_tree.mol_nodes: |
| | if m.open: |
| | scores.append(m.v_target()) |
| | else: |
| | scores.append(np.inf) |
| | scores = np.array(scores) |
| |
|
| | if np.min(scores) == np.inf: |
| | break |
| |
|
| | metric = scores |
| |
|
| | mol_tree.search_status = np.min(metric) |
| | m_next = mol_tree.mol_nodes[np.argmin(metric)] |
| | assert m_next.open |
| |
|
| | result = expand_fn(m_next.mol) |
| |
|
| | if result is not None and (len(result['scores']) > 0): |
| | reactants = result['reactants'] |
| | scores = result['scores'] |
| | analysis_tokens = result['analysis'] |
| | costs = 0.0 - np.log(np.clip(np.array(scores), 1e-3, 1.0)) |
| | templates = result['templates'] |
| | |
| | reactant_lists = [] |
| | for j in range(len(scores)): |
| | reactant_list = list(set(reactants[j].split('.'))) |
| | reactant_lists.append(reactant_list) |
| |
|
| | assert m_next.open |
| | succ = mol_tree.expand(m_next, reactant_lists, costs, templates, analysis_tokens) |
| |
|
| | if succ: |
| | break |
| |
|
| | |
| | if mol_tree.root.succ_value <= mol_tree.search_status: |
| | break |
| |
|
| | else: |
| | mol_tree.expand(m_next, None, None, None, None) |
| |
|
| | search_time = time.time() - start_time |
| |
|
| | best_route = None |
| | if mol_tree.succ: |
| | best_route = mol_tree.get_best_route() |
| | assert best_route is not None |
| |
|
| | return mol_tree.succ, best_route, i+1 |
