def sample_subgraphs(smiles, num_samples=10, frac=0.5, vis_dir=None):
mol = Chem.MolFromSmiles(smiles)
clusters, atom_cls = find_clusters(mol)
cluster_sizes = [len(cls) for cls in clusters]
p = np.array(cluster_sizes).astype('float')
p /= p.sum()
selected_atoms_list = []
for n in range(num_samples):
selected_clusters = np.random.choice(len(clusters), int(frac * len(clusters)), p=p, replace=False)
selected_atoms = set()
for i in selected_clusters:
for j in clusters[i]:
selected_atoms.add(j)
minimum_smiles, _ = extract_subgraph(smiles, selected_atoms)
selected_atoms_list.append(selected_atoms)
if vis_dir is not None:
png_f = f'subgraph_{n}.png'
Draw.MolToFile(mol, filename=os.path.join(vis_dir, png_f), highlightAtoms=selected_atoms)
png_f = f'subgraph_{n}_extracted.png'
print(minimum_smiles)
Draw.MolToFile(Chem.MolFromSmiles(minimum_smiles), filename=os.path.join(vis_dir, png_f))
return selected_atoms_list
def mcts_rollout(node, state_map, orig_smiles, clusters, atom_cls, nei_cls,
scoring_function):
#print('cur_node', node.smiles, node.P, node.N, node.W)
cur_atoms = node.atoms
if len(cur_atoms) <= MIN_ATOMS:
return node.P
# Expand if this node has never been visited
if len(node.children) == 0:
cur_cls = set([i for i, x in enumerate(clusters) if x <= cur_atoms])
for i in cur_cls:
leaf_atoms = [
a for a in clusters[i] if len(atom_cls[a] & cur_cls) == 1
]
if len(nei_cls[i] & cur_cls) == 1 or len(
clusters[i]) == 2 and len(leaf_atoms) == 1:
new_atoms = cur_atoms - set(leaf_atoms)
new_smiles, _ = extract_subgraph(orig_smiles, new_atoms)
#print('new_smiles', node.smiles, '->', new_smiles)
if new_smiles in state_map:
new_node = state_map[new_smiles] # merge identical states
else:
new_node = MCTSNode(new_smiles, new_atoms)
if new_smiles:
node.children.append(new_node)
state_map[node.smiles] = node
if len(node.children) == 0: return node.P # cannot find leaves
scores = scoring_function([x.smiles for x in node.children])
for child, score in zip(node.children, scores):
child.P = score
sum_count = sum([c.N for c in node.children])
selected_node = max(node.children, key=lambda x: x.Q() + x.U(sum_count))
v = mcts_rollout(selected_node, state_map, orig_smiles, clusters, atom_cls,
nei_cls, scoring_function)
selected_node.W += v
selected_node.N += 1
return v
def extract_selected_subgraph(smiles, selected_atoms):
mol = Chem.MolFromSmiles(smiles)
clusters, atom_cls = find_clusters(mol)
selected_clusters = []
for cls in clusters:
if len(cls) > 2:
num_selected = 0
for atom in cls:
num_selected += atom in selected_atoms
if num_selected >= 2:
# print('select the whole aromatic ring since 2 or more atoms are selected')
selected_clusters.append(cls)
for cls in selected_clusters:
for atom in cls:
selected_atoms.add(atom)
minimum_smiles, _ = extract_subgraph(smiles, selected_atoms)
# print(selected_atoms)
# print(f'{smiles} --> {minimum_smiles}')
return minimum_smiles
def extract_selected_subgraph_for_gcn(smiles, selected_atoms, vis_dir=None):
mol = Chem.MolFromSmiles(smiles)
clusters, atom_cls = find_clusters(mol)
selected_clusters = set()
for atom in selected_atoms:
for cls in atom_cls[atom]:
selected_clusters.add(clusters[cls])
# print(selected_clusters)
for cls in selected_clusters:
for atom in cls:
selected_atoms.add(atom)
minimum_smiles, _ = extract_subgraph(smiles, selected_atoms)
# print(selected_atoms)
# print(f'{smiles} --> {minimum_smiles}')
if vis_dir is not None:
png_f = f'atoms_selected{len(selected_atoms)}.png'
Draw.MolToFile(mol, filename=os.path.join(vis_dir, png_f), highlightAtoms=selected_atoms)
# png_f = f'atoms_minimum_extracted{len(selected_atoms)}.png'
# Draw.MolToFile(Chem.MolFromSmiles(minimum_smiles), filename=os.path.join(vis_dir, png_f))
return minimum_smiles
def find_minimum_subgraph(smiles, selected_atoms, vis_dir=None):
mol = Chem.MolFromSmiles(smiles)
clusters, atom_cls = find_clusters(mol)
selected_clusters = set()
cluster_votes = {}
# First iteration: select a cluster when,
# 1. An atom uniquely belongs to this cluster,
# 2. Two atoms belong to this cluster.
for atom in selected_atoms:
assert len(atom_cls[atom]) > 0
if len(atom_cls[atom]) == 1:
selected_clusters.add(atom_cls[atom][0])
else:
for cls in atom_cls[atom]:
if cls not in cluster_votes:
cluster_votes[cls] = 0
cluster_votes[cls] += 1
if cluster_votes[cls] >= 2:
selected_clusters.add(cls)
# Second iteration: randomly select a cluster for the remaining atoms.
for atom in selected_atoms:
selected = False
for cls in atom_cls[atom]:
if cls in selected_clusters:
selected = True
break
if not selected:
selected_clusters.add(atom_cls[atom][0])
cluster_neighbor = {}
for i in range(len(clusters)):
cluster_neighbor[i] = set()
for atom in clusters[i]:
cluster_neighbor[i].update(atom_cls[atom])
cluster_neighbor[i].remove(i)
# remove degree-1 unselected clusters iteratively
leaf_clusters = set()
while True:
updated = False
for i in range(len(clusters)):
if i in selected_clusters or i in leaf_clusters:
continue
if len(cluster_neighbor[i]) > 1:
removable = True
neighbor_pairs = [(j, k) for j in cluster_neighbor[i] for k in cluster_neighbor[i] if j < k]
for j, k in neighbor_pairs:
if j not in cluster_neighbor[k] or k not in cluster_neighbor[j]:
removable = False
break
if not removable:
continue
leaf_clusters.add(i)
for j in cluster_neighbor[i]:
cluster_neighbor[j].remove(i)
updated = True
if not updated:
break
minimum_atoms = set()
for i in range(len(clusters)):
if i not in leaf_clusters:
minimum_atoms.update(clusters[i])
minimum_smiles, _ = extract_subgraph(smiles, minimum_atoms)
# print(f'{smiles} --> {minimum_smiles}')
if vis_dir is not None:
png_f = f'atoms_selected{len(selected_atoms)}.png'
Draw.MolToFile(mol, filename=os.path.join(vis_dir, png_f), highlightAtoms=selected_atoms)
png_f = f'atoms_minimum{len(selected_atoms)}.png'
Draw.MolToFile(mol, filename=os.path.join(vis_dir, png_f), highlightAtoms=minimum_atoms)
png_f = f'atoms_minimum_extracted{len(selected_atoms)}.png'
Draw.MolToFile(Chem.MolFromSmiles(minimum_smiles), filename=os.path.join(vis_dir, png_f))
return minimum_smiles
import sys
from rdkit import Chem
from multiobj_rationale.fuseprop import enum_subgraph, extract_subgraph
ratio_list = [0.3, 0.4, 0.5, 0.6, 0.7]
next(sys.stdin)
for line in sys.stdin:
smiles = line.strip("\r\n ").split(',')[0]
mol = Chem.MolFromSmiles(smiles)
selections = enum_subgraph(mol, ratio_list)
res = []
for selected_atoms in selections:
subgraph, _ = extract_subgraph(smiles, selected_atoms)
if subgraph is not None:
res.append(subgraph)
for subgraph in set(res):
print(smiles, subgraph)