torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/generator/jtvae/jtnn/mol_tree.py

@@ -0,0 +1,111 @@
+import rdkit
+import rdkit.Chem as Chem
+from .chemutils import get_clique_mol, tree_decomp, get_mol, get_smiles, set_atommap, enum_assemble, decode_stereo
+from .vocab import *
+
+class MolTreeNode(object):
+
+    def __init__(self, smiles, clique=[]):
+        self.smiles = smiles
+        self.mol = get_mol(self.smiles)
+
+        self.clique = [x for x in clique] #copy
+        self.neighbors = []
+        
+    def add_neighbor(self, nei_node):
+        self.neighbors.append(nei_node)
+
+    def recover(self, original_mol):
+        clique = []
+        clique.extend(self.clique)
+        if not self.is_leaf:
+            for cidx in self.clique:
+                original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(self.nid)
+
+        for nei_node in self.neighbors:
+            clique.extend(nei_node.clique)
+            if nei_node.is_leaf: #Leaf node, no need to mark 
+                continue
+            for cidx in nei_node.clique:
+                #allow singleton node override the atom mapping
+                if cidx not in self.clique or len(nei_node.clique) == 1:
+                    atom = original_mol.GetAtomWithIdx(cidx)
+                    atom.SetAtomMapNum(nei_node.nid)
+
+        clique = list(set(clique))
+        label_mol = get_clique_mol(original_mol, clique)
+        self.label = Chem.MolToSmiles(Chem.MolFromSmiles(get_smiles(label_mol)))
+
+        for cidx in clique:
+            original_mol.GetAtomWithIdx(cidx).SetAtomMapNum(0)
+
+        return self.label
+    
+    def assemble(self):
+        neighbors = [nei for nei in self.neighbors if nei.mol.GetNumAtoms() > 1]
+        neighbors = sorted(neighbors, key=lambda x:x.mol.GetNumAtoms(), reverse=True)
+        singletons = [nei for nei in self.neighbors if nei.mol.GetNumAtoms() == 1]
+        neighbors = singletons + neighbors
+
+        cands,aroma = enum_assemble(self, neighbors)
+        new_cands = [cand for i,cand in enumerate(cands) if aroma[i] >= 0]
+        if len(new_cands) > 0: cands = new_cands
+
+        if len(cands) > 0:
+            self.cands, _ = zip(*cands)
+            self.cands = list(self.cands)
+        else:
+            self.cands = []
+
+class MolTree(object):
+
+    def __init__(self, smiles, voc=None):
+        self.smiles = smiles
+        self.mol = get_mol(smiles)
+        self.voc = voc
+
+        #Stereo Generation (currently disabled)
+        #mol = Chem.MolFromSmiles(smiles)
+        #self.smiles3D = Chem.MolToSmiles(mol, isomericSmiles=True)
+        #self.smiles2D = Chem.MolToSmiles(mol)
+        #self.stereo_cands = decode_stereo(self.smiles2D)
+
+        cliques, edges = tree_decomp(self.mol)
+        self.nodes = []
+        root = 0
+        for i,c in enumerate(cliques):
+            cmol = get_clique_mol(self.mol, c)
+            node = MolTreeNode(get_smiles(cmol), c)
+            self.nodes.append(node)
+            if min(c) == 0: root = i
+
+        for x,y in edges:
+            self.nodes[x].add_neighbor(self.nodes[y])
+            self.nodes[y].add_neighbor(self.nodes[x])
+        
+        if root > 0:
+            self.nodes[0],self.nodes[root] = self.nodes[root],self.nodes[0]
+
+        for i,node in enumerate(self.nodes):
+            node.nid = i + 1
+            if len(node.neighbors) > 1: #Leaf node mol is not marked
+                set_atommap(node.mol, node.nid)
+            node.is_leaf = (len(node.neighbors) == 1)
+
+    def size(self):
+        return len(self.nodes)
+
+    def recover(self):
+        for node in self.nodes:
+            node.recover(self.mol)
+
+    def assemble(self):
+        for node in self.nodes:
+            node.assemble()
+
+def dfs(node, fa_idx):
+    max_depth = 0
+    for child in node.neighbors:
+        if child.idx == fa_idx: continue
+        max_depth = max(max_depth, dfs(child, node.idx))
+    return max_depth + 1

torch_molecule/generator/jtvae/jtnn/mpn.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+import rdkit.Chem as Chem
+import torch.nn.functional as F
+
+
+from .nnutils import *
+from .chemutils import get_mol
+
+ELEM_LIST = ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe', 'Al', 'I', 'B', 'K', 'Se', 'Zn', 'H', 'Cu', 'Mn', 'unknown']
+
+ATOM_FDIM = len(ELEM_LIST) + 6 + 5 + 4 + 1
+BOND_FDIM = 5 + 6
+MAX_NB = 6
+
+def onek_encoding_unk(x, allowable_set):
+    if x not in allowable_set:
+        x = allowable_set[-1]
+    return list(map(lambda s: x == s, allowable_set))
+
+def atom_features(atom):
+    return torch.Tensor(onek_encoding_unk(atom.GetSymbol(), ELEM_LIST) 
+            + onek_encoding_unk(atom.GetDegree(), [0,1,2,3,4,5]) 
+            + onek_encoding_unk(atom.GetFormalCharge(), [-1,-2,1,2,0])
+            + onek_encoding_unk(int(atom.GetChiralTag()), [0,1,2,3])
+            + [atom.GetIsAromatic()])
+
+def bond_features(bond):
+    bt = bond.GetBondType()
+    stereo = int(bond.GetStereo())
+    fbond = [bt == Chem.rdchem.BondType.SINGLE, bt == Chem.rdchem.BondType.DOUBLE, bt == Chem.rdchem.BondType.TRIPLE, bt == Chem.rdchem.BondType.AROMATIC, bond.IsInRing()]
+    fstereo = onek_encoding_unk(stereo, [0,1,2,3,4,5])
+    return torch.Tensor(fbond + fstereo)
+
+class MPN(nn.Module):
+
+    def __init__(self, hidden_size, depth):
+        super(MPN, self).__init__()
+        self.hidden_size = hidden_size
+        self.depth = depth
+
+        self.W_i = nn.Linear(ATOM_FDIM + BOND_FDIM, hidden_size, bias=False)
+        self.W_h = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.W_o = nn.Linear(ATOM_FDIM + hidden_size, hidden_size)
+
+    def forward(self, fatoms, fbonds, agraph, bgraph, scope):
+        fatoms = create_var(fatoms)
+        fbonds = create_var(fbonds)
+        agraph = create_var(agraph)
+        bgraph = create_var(bgraph)
+
+        binput = self.W_i(fbonds)
+        message = F.relu(binput)
+
+        for i in range(self.depth - 1):
+            nei_message = index_select_ND(message, 0, bgraph)
+            nei_message = nei_message.sum(dim=1)
+            nei_message = self.W_h(nei_message)
+            message = F.relu(binput + nei_message)
+
+        nei_message = index_select_ND(message, 0, agraph)
+        nei_message = nei_message.sum(dim=1)
+        ainput = torch.cat([fatoms, nei_message], dim=1)
+        atom_hiddens = F.relu(self.W_o(ainput))
+
+        max_len = max([x for _,x in scope])
+        batch_vecs = []
+        for st,le in scope:
+            cur_vecs = atom_hiddens[st : st + le].mean(dim=0)
+            batch_vecs.append( cur_vecs )
+
+        mol_vecs = torch.stack(batch_vecs, dim=0)
+        return mol_vecs 
+
+    @staticmethod
+    def tensorize(mol_batch):
+        padding = torch.zeros(ATOM_FDIM + BOND_FDIM)
+        fatoms,fbonds = [],[padding] #Ensure bond is 1-indexed
+        in_bonds,all_bonds = [],[(-1,-1)] #Ensure bond is 1-indexed
+        scope = []
+        total_atoms = 0
+
+        for smiles in mol_batch:
+            mol = get_mol(smiles)
+            #mol = Chem.MolFromSmiles(smiles)
+            n_atoms = mol.GetNumAtoms()
+            for atom in mol.GetAtoms():
+                fatoms.append( atom_features(atom) )
+                in_bonds.append([])
+
+            for bond in mol.GetBonds():
+                a1 = bond.GetBeginAtom()
+                a2 = bond.GetEndAtom()
+                x = a1.GetIdx() + total_atoms
+                y = a2.GetIdx() + total_atoms
+
+                b = len(all_bonds) 
+                all_bonds.append((x,y))
+                fbonds.append( torch.cat([fatoms[x], bond_features(bond)], 0) )
+                in_bonds[y].append(b)
+
+                b = len(all_bonds)
+                all_bonds.append((y,x))
+                fbonds.append( torch.cat([fatoms[y], bond_features(bond)], 0) )
+                in_bonds[x].append(b)
+            
+            scope.append((total_atoms,n_atoms))
+            total_atoms += n_atoms
+
+        total_bonds = len(all_bonds)
+        fatoms = torch.stack(fatoms, 0)
+        fbonds = torch.stack(fbonds, 0)
+        agraph = torch.zeros(total_atoms,MAX_NB).long()
+        bgraph = torch.zeros(total_bonds,MAX_NB).long()
+
+        for a in range(total_atoms):
+            for i,b in enumerate(in_bonds[a]):
+                if i >= MAX_NB: break
+                agraph[a,i] = b
+
+        for b1 in range(1, total_bonds):
+            x,y = all_bonds[b1]
+            for i,b2 in enumerate(in_bonds[x]):
+                if all_bonds[b2][0] != y:
+                    if i >= MAX_NB: break
+                    bgraph[b1,i] = b2
+
+        return (fatoms, fbonds, agraph, bgraph, scope)
+

torch_molecule/generator/jtvae/jtnn/nnutils.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+def create_var(tensor, requires_grad=None):
+    if requires_grad is None:
+        return Variable(tensor).cuda()
+    else:
+        return Variable(tensor, requires_grad=requires_grad).cuda()
+
+def index_select_ND(source, dim, index):
+    index_size = index.size()
+    suffix_dim = source.size()[1:]
+    final_size = index_size + suffix_dim
+    target = source.index_select(dim, index.view(-1))
+    return target.view(final_size)
+
+def avg_pool(all_vecs, scope, dim):
+    size = create_var(torch.Tensor([le for _,le in scope]))
+    return all_vecs.sum(dim=dim) / size.unsqueeze(-1)
+
+def stack_pad_tensor(tensor_list):
+    max_len = max([t.size(0) for t in tensor_list])
+    for i,tensor in enumerate(tensor_list):
+        pad_len = max_len - tensor.size(0)
+        tensor_list[i] = F.pad( tensor, (0,0,0,pad_len) )
+    return torch.stack(tensor_list, dim=0)
+
+#3D padded tensor to 2D matrix, with padded zeros removed
+def flatten_tensor(tensor, scope):
+    assert tensor.size(0) == len(scope)
+    tlist = []
+    for i,tup in enumerate(scope):
+        le = tup[1]
+        tlist.append( tensor[i, 0:le] )
+    return torch.cat(tlist, dim=0)
+
+#2D matrix to 3D padded tensor
+def inflate_tensor(tensor, scope): 
+    max_len = max([le for _,le in scope])
+    batch_vecs = []
+    for st,le in scope:
+        cur_vecs = tensor[st : st + le]
+        cur_vecs = F.pad( cur_vecs, (0,0,0,max_len-le) )
+        batch_vecs.append( cur_vecs )
+
+    return torch.stack(batch_vecs, dim=0)
+
+def GRU(x, h_nei, W_z, W_r, U_r, W_h):
+    hidden_size = x.size()[-1]
+    sum_h = h_nei.sum(dim=1)
+    z_input = torch.cat([x,sum_h], dim=1)
+    z = F.sigmoid(W_z(z_input))
+
+    r_1 = W_r(x).view(-1,1,hidden_size)
+    r_2 = U_r(h_nei)
+    r = F.sigmoid(r_1 + r_2)
+    
+    gated_h = r * h_nei
+    sum_gated_h = gated_h.sum(dim=1)
+    h_input = torch.cat([x,sum_gated_h], dim=1)
+    pre_h = F.tanh(W_h(h_input))
+    new_h = (1.0 - z) * sum_h + z * pre_h
+    return new_h
+

torch_molecule/generator/jtvae/jtnn/vocab.py

@@ -0,0 +1,31 @@
+import rdkit
+import rdkit.Chem as Chem
+import copy
+
+def get_slots(smiles):
+    mol = Chem.MolFromSmiles(smiles)
+    return [(atom.GetSymbol(), atom.GetFormalCharge(), atom.GetTotalNumHs()) for atom in mol.GetAtoms()]
+
+class Vocab(object):
+    benzynes = ['C1=CC=CC=C1', 'C1=CC=NC=C1', 'C1=CC=NN=C1', 'C1=CN=CC=N1', 'C1=CN=CN=C1', 'C1=CN=NC=N1', 'C1=CN=NN=C1', 'C1=NC=NC=N1', 'C1=NN=CN=N1']
+    penzynes = ['C1=C[NH]C=C1', 'C1=C[NH]C=N1', 'C1=C[NH]N=C1', 'C1=C[NH]N=N1', 'C1=COC=C1', 'C1=COC=N1', 'C1=CON=C1', 'C1=CSC=C1', 'C1=CSC=N1', 'C1=CSN=C1', 'C1=CSN=N1', 'C1=NN=C[NH]1', 'C1=NN=CO1', 'C1=NN=CS1', 'C1=N[NH]C=N1', 'C1=N[NH]N=C1', 'C1=N[NH]N=N1', 'C1=NN=N[NH]1', 'C1=NN=NS1', 'C1=NOC=N1', 'C1=NON=C1', 'C1=NSC=N1', 'C1=NSN=C1']
+
+    def __init__(self, smiles_list):
+        self.vocab = smiles_list
+        self.vmap = {x:i for i,x in enumerate(self.vocab)}
+        self.slots = [get_slots(smiles) for smiles in self.vocab]
+        Vocab.benzynes = [s for s in smiles_list if s.count('=') >= 2 and Chem.MolFromSmiles(s).GetNumAtoms() == 6] + ['C1=CCNCC1']
+        Vocab.penzynes = [s for s in smiles_list if s.count('=') >= 2 and Chem.MolFromSmiles(s).GetNumAtoms() == 5] + ['C1=NCCN1','C1=NNCC1']
+        
+    def get_index(self, smiles):
+        return self.vmap[smiles]
+
+    def get_smiles(self, idx):
+        return self.vocab[idx]
+
+    def get_slots(self, idx):
+        return copy.deepcopy(self.slots[idx])
+
+    def size(self):
+        return len(self.vocab)
+

torch_molecule/generator/jtvae/jtnn_vae.py

@@ -0,0 +1,246 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import copy
+import rdkit
+import rdkit.Chem as Chem
+
+from .jtnn.mol_tree import Vocab, MolTree
+from .jtnn.nnutils import create_var, flatten_tensor, avg_pool
+from .jtnn.jtnn_enc import JTNNEncoder
+from .jtnn.jtnn_dec import JTNNDecoder
+from .jtnn.mpn import MPN
+from .jtnn.jtmpn import JTMPN
+from .jtnn.datautils import tensorize
+from .jtnn.chemutils import enum_assemble, set_atommap, copy_edit_mol, attach_mols
+
+class JTNNVAE(nn.Module):
+    def __init__(self, vocab, hidden_size, latent_size, depthT, depthG):
+        super(JTNNVAE, self).__init__()
+        self.hidden_size = hidden_size
+        assert latent_size % 2 == 0, 'latent_size must be even'
+        self.latent_size = latent_size = int(latent_size / 2) #Tree and Mol has two vectors
+
+        if vocab is None:
+            raise ValueError('JTNNVAE is not properly initialized. vocab is required')
+
+        self.vocab = vocab
+        self.jtnn = JTNNEncoder(hidden_size, depthT, nn.Embedding(vocab.size(), hidden_size))
+        self.decoder = JTNNDecoder(vocab, hidden_size, latent_size, nn.Embedding(vocab.size(), hidden_size))
+
+        self.jtmpn = JTMPN(hidden_size, depthG)
+        self.mpn = MPN(hidden_size, depthG)
+
+        self.A_assm = nn.Linear(latent_size, hidden_size, bias=False)
+        self.assm_loss = nn.CrossEntropyLoss(reduction='sum')
+
+        self.T_mean = nn.Linear(hidden_size, latent_size)
+        self.T_var = nn.Linear(hidden_size, latent_size)
+        self.G_mean = nn.Linear(hidden_size, latent_size)
+        self.G_var = nn.Linear(hidden_size, latent_size)
+
+        self.initialize_parameters()
+    
+    def initialize_parameters(self, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+    def compute_loss(self, x_batch):
+        word_loss, topo_loss, assm_loss, kl_div, _, _, _ = self.forward(x_batch)
+        return word_loss, topo_loss, assm_loss, kl_div
+        
+    def encode(self, jtenc_holder, mpn_holder):
+        tree_vecs, tree_mess = self.jtnn(*jtenc_holder)
+        mol_vecs = self.mpn(*mpn_holder)
+        return tree_vecs, tree_mess, mol_vecs
+    
+    def encode_from_smiles(self, smiles_list):
+        tree_batch = [MolTree(s) for s in smiles_list]
+        _, jtenc_holder, mpn_holder = tensorize(tree_batch, self.vocab, assm=False)
+        tree_vecs, _, mol_vecs = self.encode(jtenc_holder, mpn_holder)
+        return torch.cat([tree_vecs, mol_vecs], dim=-1)
+
+    def encode_latent(self, jtenc_holder, mpn_holder):
+        tree_vecs, _ = self.jtnn(*jtenc_holder)
+        mol_vecs = self.mpn(*mpn_holder)
+        tree_mean = self.T_mean(tree_vecs)
+        mol_mean = self.G_mean(mol_vecs)
+        tree_var = -torch.abs(self.T_var(tree_vecs))
+        mol_var = -torch.abs(self.G_var(mol_vecs))
+        return torch.cat([tree_mean, mol_mean], dim=1), torch.cat([tree_var, mol_var], dim=1)
+
+    def rsample(self, z_vecs, W_mean, W_var):
+        batch_size = z_vecs.size(0)
+        z_mean = W_mean(z_vecs)
+        z_log_var = -torch.abs(W_var(z_vecs)) # Following Mueller et al.
+        kl_loss = -0.5 * torch.sum(1.0 + z_log_var - z_mean * z_mean - torch.exp(z_log_var)) / batch_size
+        epsilon = create_var(torch.randn_like(z_mean))
+        z_vecs = z_mean + torch.exp(z_log_var / 2) * epsilon
+        return z_vecs, kl_loss
+
+    def sample_prior(self, device, prob_decode=False):
+        z_tree = torch.randn(1, self.latent_size).to(device)
+        z_mol = torch.randn(1, self.latent_size).to(device)
+        return self.decode(z_tree, z_mol, prob_decode, device)
+
+    def forward(self, x_batch):
+        x_batch, x_jtenc_holder, x_mpn_holder, x_jtmpn_holder = x_batch
+        x_tree_vecs, x_tree_mess, x_mol_vecs = self.encode(x_jtenc_holder, x_mpn_holder)
+        z_tree_vecs, tree_kl = self.rsample(x_tree_vecs, self.T_mean, self.T_var)
+        z_mol_vecs, mol_kl = self.rsample(x_mol_vecs, self.G_mean, self.G_var)
+
+        kl_div = tree_kl + mol_kl
+        word_loss, topo_loss, word_acc, topo_acc = self.decoder(x_batch, z_tree_vecs)
+        assm_loss, assm_acc = self.assm(x_batch, x_jtmpn_holder, z_mol_vecs, x_tree_mess)
+
+        # return word_loss + topo_loss + assm_loss + beta * kl_div, kl_div.item(), word_acc, topo_acc, assm_acc
+        return word_loss, topo_loss, assm_loss, kl_div, word_acc, topo_acc, assm_acc
+    
+    def assm(self, mol_batch, jtmpn_holder, x_mol_vecs, x_tree_mess):
+        jtmpn_holder,batch_idx = jtmpn_holder
+        fatoms,fbonds,agraph,bgraph,scope = jtmpn_holder
+        batch_idx = create_var(batch_idx)
+
+        cand_vecs = self.jtmpn(fatoms, fbonds, agraph, bgraph, scope, x_tree_mess)
+
+        x_mol_vecs = x_mol_vecs.index_select(0, batch_idx)
+        x_mol_vecs = self.A_assm(x_mol_vecs) #bilinear
+        scores = torch.bmm(
+                x_mol_vecs.unsqueeze(1),
+                cand_vecs.unsqueeze(-1)
+        ).squeeze()
+        
+        cnt,tot,acc = 0,0,0
+        all_loss = []
+        for i,mol_tree in enumerate(mol_batch):
+            comp_nodes = [node for node in mol_tree.nodes if len(node.cands) > 1 and not node.is_leaf]
+            cnt += len(comp_nodes)
+            for node in comp_nodes:
+                label = node.cands.index(node.label)
+                ncand = len(node.cands)
+                cur_score = scores.narrow(0, tot, ncand)
+                tot += ncand
+
+                if cur_score.data[label] >= cur_score.max().item():
+                    acc += 1
+
+                label = create_var(torch.LongTensor([label]))
+                all_loss.append( self.assm_loss(cur_score.view(1,-1), label) )
+        
+        all_loss = sum(all_loss) / len(mol_batch)
+        return all_loss, acc * 1.0 / cnt
+
+    def decode(self, x_tree_vecs, x_mol_vecs, prob_decode, device):
+        #currently do not support batch decoding
+        assert x_tree_vecs.size(0) == 1 and x_mol_vecs.size(0) == 1
+
+        pred_root, pred_nodes = self.decoder.decode(x_tree_vecs, prob_decode)
+        if len(pred_nodes) == 0: return None
+        elif len(pred_nodes) == 1: return pred_root.smiles
+
+        #Mark nid & is_leaf & atommap
+        for i, node in enumerate(pred_nodes):
+            node.nid = i + 1
+            node.is_leaf = (len(node.neighbors) == 1)
+            if len(node.neighbors) > 1:
+                set_atommap(node.mol, node.nid)
+
+        scope = [(0, len(pred_nodes))]
+        jtenc_holder,mess_dict = JTNNEncoder.tensorize_nodes(pred_nodes, scope)
+        _, tree_mess = self.jtnn(*jtenc_holder)
+        tree_mess = (tree_mess, mess_dict) #Important: tree_mess is a matrix, mess_dict is a python dict
+
+        x_mol_vecs = self.A_assm(x_mol_vecs).squeeze() #bilinear
+
+        cur_mol = copy_edit_mol(pred_root.mol)
+        global_amap = [{}] + [{} for node in pred_nodes]
+        global_amap[1] = {atom.GetIdx():atom.GetIdx() for atom in cur_mol.GetAtoms()}
+
+        cur_mol,_ = self.dfs_assemble(tree_mess, x_mol_vecs, pred_nodes, cur_mol, global_amap, [], pred_root, None, prob_decode, check_aroma=True, device=device)
+        if cur_mol is None: 
+            cur_mol = copy_edit_mol(pred_root.mol)
+            global_amap = [{}] + [{} for node in pred_nodes]
+            global_amap[1] = {atom.GetIdx():atom.GetIdx() for atom in cur_mol.GetAtoms()}
+            cur_mol, pre_mol = self.dfs_assemble(tree_mess, x_mol_vecs, pred_nodes, cur_mol, global_amap, [], pred_root, None, prob_decode, check_aroma=False, device=device)
+            if cur_mol is None: cur_mol = pre_mol
+
+        if cur_mol is None: 
+            return None
+
+        cur_mol = cur_mol.GetMol()
+        set_atommap(cur_mol)
+        cur_mol = Chem.MolFromSmiles(Chem.MolToSmiles(cur_mol))
+        return Chem.MolToSmiles(cur_mol) if cur_mol is not None else 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, device):
+        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 = list(zip(*cands))
+        aroma_score = torch.Tensor(aroma_score).to(device)
+        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, device)
+                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