def dfs_assemble_graph_conv(self, x_mol_vecs, all_nodes, cur_mol, global_amap, fa_amap, cur_node, fa_node): fa_nid = fa_node.nid if fa_node is not None else -1 prev_nodes = [fa_node] if fa_node is not None else [] children = [nei for nei in cur_node.neighbors if nei.nid != fa_nid] neighbors = [nei for nei in children if nei.mol.GetNumAtoms() > 1] neighbors = sorted(neighbors, key=lambda x: x.mol.GetNumAtoms(), reverse=True) singletons = [nei for nei in children if nei.mol.GetNumAtoms() == 1] neighbors = singletons + neighbors cur_amap = [(fa_nid, a2, a1) for nid, a1, a2 in fa_amap if nid == cur_node.nid] cands = enum_assemble(cur_node, neighbors, prev_nodes, cur_amap) if len(cands) == 0: return None cand_smiles, cand_mols, cand_amap = zip(*cands) cands = [(smiles, all_nodes, cur_node) for smiles in cand_smiles] # jtmpn_holder = JTMessPassNet.tensorize(cands, y_tree_mess[1]) # fatoms, fbonds, agraph, bgraph, scope = jtmpn_holder # cand_vecs = self.jtmpn(fatoms, fbonds, agraph, bgraph, scope, y_tree_mess[0]) jt_graph_enc_holder = MolGraphEncoder.tensorize(cand_smiles) cand_vecs = self.graph_enc(*jt_graph_enc_holder) scores = torch.mv(cand_vecs, x_mol_vecs) _, cand_idx = torch.sort(scores, descending=True) backup_mol = Chem.RWMol(cur_mol) for i in range(cand_idx.numel()): cur_mol = Chem.RWMol(backup_mol) pred_amap = cand_amap[cand_idx[i].item()] new_global_amap = copy.deepcopy(global_amap) for nei_id, ctr_atom, nei_atom in pred_amap: if nei_id == fa_nid: continue new_global_amap[nei_id][nei_atom] = new_global_amap[cur_node.nid][ctr_atom] cur_mol = attach_mols(cur_mol, children, [], new_global_amap) # father is already attached new_mol = cur_mol.GetMol() new_mol = Chem.MolFromSmiles(Chem.MolToSmiles(new_mol)) if new_mol is None: continue result = True for nei_node in children: if nei_node.is_leaf: continue cur_mol = self.dfs_assemble_graph_conv(x_mol_vecs, all_nodes, cur_mol, new_global_amap, pred_amap, nei_node, cur_node) if cur_mol is None: result = False break if result: return cur_mol
def dfs_assemble(self, tree_mess, mol_vec, all_nodes, cur_mol, global_amap, fa_amap, cur_node, fa_node, prob_decode): fa_nid = fa_node.nid if fa_node is not None else -1 prev_nodes = [fa_node] if fa_node is not None else [] children = [nei for nei in cur_node.neighbors if nei.nid != fa_nid] neighbors = [nei for nei in children if nei.mol.GetNumAtoms() > 1] neighbors = sorted(neighbors, key=lambda x:x.mol.GetNumAtoms(), reverse=True) singletons = [nei for nei in children if nei.mol.GetNumAtoms() == 1] neighbors = singletons + neighbors cur_amap = [(fa_nid,a2,a1) for nid,a1,a2 in fa_amap if nid == cur_node.nid] cands = enum_assemble(cur_node, neighbors, prev_nodes, cur_amap) if len(cands) == 0: return None cand_smiles,cand_mols,cand_amap = zip(*cands) cands = [(candmol, all_nodes, cur_node) for candmol in cand_mols] cand_vecs = self.jtmpn(cands, tree_mess) cand_vecs = self.G_mean(cand_vecs) mol_vec = mol_vec.squeeze() scores = torch.mv(cand_vecs, mol_vec) * 20 if prob_decode: probs = nn.Softmax()(scores.view(1,-1)).squeeze() + 1e-5 #prevent prob = 0 cand_idx = torch.multinomial(probs, probs.numel()) else: _,cand_idx = torch.sort(scores, descending=True) backup_mol = Chem.RWMol(cur_mol) for i in xrange(cand_idx.numel()): cur_mol = Chem.RWMol(backup_mol) pred_amap = cand_amap[cand_idx[i].data[0]] new_global_amap = copy.deepcopy(global_amap) for nei_id,ctr_atom,nei_atom in pred_amap: if nei_id == fa_nid: continue new_global_amap[nei_id][nei_atom] = new_global_amap[cur_node.nid][ctr_atom] cur_mol = attach_mols(cur_mol, children, [], new_global_amap) #father is already attached new_mol = cur_mol.GetMol() new_mol = Chem.MolFromSmiles(Chem.MolToSmiles(new_mol)) if new_mol is None: continue result = True for nei_node in children: if nei_node.is_leaf: continue cur_mol = self.dfs_assemble(tree_mess, mol_vec, all_nodes, cur_mol, new_global_amap, pred_amap, nei_node, cur_node, prob_decode) if cur_mol is None: result = False break if result: return cur_mol return None
def dfs_assemble(self, y_tree_mess, x_mol_vecs, all_nodes, cur_mol, global_amap, fa_amap, cur_node, fa_node, prob_decode, check_aroma): fa_nid = fa_node.nid if fa_node is not None else -1 prev_nodes = [fa_node] if fa_node is not None else [] children = [nei for nei in cur_node.neighbors if nei.nid != fa_nid] neighbors = [nei for nei in children if nei.mol.GetNumAtoms() > 1] neighbors = sorted(neighbors, key=lambda x:x.mol.GetNumAtoms(), reverse=True) singletons = [nei for nei in children if nei.mol.GetNumAtoms() == 1] neighbors = singletons + neighbors cur_amap = [(fa_nid,a2,a1) for nid,a1,a2 in fa_amap if nid == cur_node.nid] cands,aroma_score = enum_assemble(cur_node, neighbors, prev_nodes, cur_amap) if len(cands) == 0 or (sum(aroma_score) < 0 and check_aroma): return None, cur_mol cand_smiles,cand_amap = zip(*cands) aroma_score = torch.Tensor(aroma_score).cuda() cands = [(smiles, all_nodes, cur_node) for smiles in cand_smiles] if len(cands) > 1: jtmpn_holder = JTMPN.tensorize(cands, y_tree_mess[1]) fatoms,fbonds,agraph,bgraph,scope = jtmpn_holder cand_vecs = self.jtmpn(fatoms, fbonds, agraph, bgraph, scope, y_tree_mess[0]) scores = torch.mv(cand_vecs, x_mol_vecs) + aroma_score else: scores = torch.Tensor([1.0]) if prob_decode: probs = F.softmax(scores.view(1,-1), dim=1).squeeze() + 1e-7 #prevent prob = 0 cand_idx = torch.multinomial(probs, probs.numel()) else: _,cand_idx = torch.sort(scores, descending=True) backup_mol = Chem.RWMol(cur_mol) pre_mol = cur_mol for i in range(cand_idx.numel()): cur_mol = Chem.RWMol(backup_mol) pred_amap = cand_amap[cand_idx[i].item()] new_global_amap = copy.deepcopy(global_amap) for nei_id,ctr_atom,nei_atom in pred_amap: if nei_id == fa_nid: continue new_global_amap[nei_id][nei_atom] = new_global_amap[cur_node.nid][ctr_atom] cur_mol = attach_mols(cur_mol, children, [], new_global_amap) #father is already attached new_mol = cur_mol.GetMol() new_mol = Chem.MolFromSmiles(Chem.MolToSmiles(new_mol)) if new_mol is None: continue has_error = False for nei_node in children: if nei_node.is_leaf: continue tmp_mol, tmp_mol2 = self.dfs_assemble(y_tree_mess, x_mol_vecs, all_nodes, cur_mol, new_global_amap, pred_amap, nei_node, cur_node, prob_decode, check_aroma) if tmp_mol is None: has_error = True if i == 0: pre_mol = tmp_mol2 break cur_mol = tmp_mol if not has_error: return cur_mol, cur_mol return None, pre_mol