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

torch_molecule/generator/gdss/toremove/ScoreNetworkA.py

@@ -0,0 +1,167 @@
+import torch
+from torch.nn import Parameter
+import torch.nn.functional as F
+
+from .utils import mask_adjs, pow_tensor, mask_x
+from .layers import DenseGCNConv, MLP, AttentionLayer
+
+# def node_feature_to_matrix(x):
+#     """
+#     :param x:  BS x N x F
+#     :return:
+#     x_pair: BS x N x N x 2F
+#     """
+#     x_b = x.unsqueeze(-2).expand(x.size(0), x.size(1), x.size(1), -1)  # BS x N x N x F
+#     x_pair = torch.cat([x_b, x_b.transpose(1, 2)], dim=-1)  # BS x N x N x 2F
+#     return x_pair
+
+# class BaselineNetworkLayer(torch.nn.Module):
+
+#     def __init__(self, num_linears, conv_input_dim, conv_output_dim, input_dim, output_dim, batch_norm=False):
+
+#         super(BaselineNetworkLayer, self).__init__()
+
+#         self.convs = torch.nn.ModuleList()
+#         for _ in range(input_dim):
+#             self.convs.append(DenseGCNConv(conv_input_dim, conv_output_dim))
+#         self.hidden_dim = max(input_dim, output_dim)
+#         self.mlp_in_dim = input_dim + 2*conv_output_dim
+#         self.mlp = MLP(num_linears, self.mlp_in_dim, self.hidden_dim, output_dim, 
+#                             use_bn=False, activate_func=F.elu)
+#         self.multi_channel = MLP(2, input_dim*conv_output_dim, self.hidden_dim, conv_output_dim, 
+#                                     use_bn=False, activate_func=F.elu)
+        
+#     def forward(self, x, adj, flags):
+    
+#         x_list = []
+#         for _ in range(len(self.convs)):
+#             _x = self.convs[_](x, adj[:,_,:,:])
+#             x_list.append(_x)
+#         x_out = mask_x(self.multi_channel(torch.cat(x_list, dim=-1)) , flags)
+#         x_out = torch.tanh(x_out)
+
+#         x_matrix = node_feature_to_matrix(x_out)
+#         mlp_in = torch.cat([x_matrix, adj.permute(0,2,3,1)], dim=-1)
+#         shape = mlp_in.shape
+#         mlp_out = self.mlp(mlp_in.view(-1, shape[-1]))
+#         _adj = mlp_out.view(shape[0], shape[1], shape[2], -1).permute(0,3,1,2)
+#         _adj = _adj + _adj.transpose(-1,-2)
+#         adj_out = mask_adjs(_adj, flags)
+
+#         return x_out, adj_out
+
+
+# class BaselineNetwork(torch.nn.Module):
+
+#     def __init__(self, max_feat_num, max_node_num, nhid, num_layers, num_linears, 
+#                     c_init, c_hid, c_final, adim, num_heads=4, conv='GCN'):
+
+#         super(BaselineNetwork, self).__init__()
+
+#         self.nfeat = max_feat_num
+#         self.max_node_num = max_node_num
+#         self.nhid  = nhid
+#         self.num_layers = num_layers
+#         self.num_linears = num_linears
+#         self.c_init = c_init
+#         self.c_hid = c_hid
+#         self.c_final = c_final
+
+#         self.layers = torch.nn.ModuleList()
+#         for _ in range(self.num_layers):
+#             if _==0:
+#                 self.layers.append(BaselineNetworkLayer(self.num_linears, self.nfeat, self.nhid, self.c_init, self.c_hid))
+
+#             elif _==self.num_layers-1:
+#                 self.layers.append(BaselineNetworkLayer(self.num_linears, self.nhid, self.nhid, self.c_hid, self.c_final))
+
+#             else:
+#                 self.layers.append(BaselineNetworkLayer(self.num_linears, self.nhid, self.nhid, self.c_hid, self.c_hid)) 
+
+#         self.fdim = self.c_hid*(self.num_layers-1) + self.c_final + self.c_init
+#         self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=1, 
+#                             use_bn=False, activate_func=F.elu)
+#         self.mask = torch.ones([self.max_node_num, self.max_node_num]) - torch.eye(self.max_node_num)
+#         self.mask.unsqueeze_(0)   
+
+#     def forward(self, x, adj, flags=None):
+
+#         adjc = pow_tensor(adj, self.c_init)
+
+#         adj_list = [adjc]
+#         for _ in range(self.num_layers):
+
+#             x, adjc = self.layers[_](x, adjc, flags)
+#             adj_list.append(adjc)
+        
+#         adjs = torch.cat(adj_list, dim=1).permute(0,2,3,1)
+#         out_shape = adjs.shape[:-1] # B x N x N
+#         score = self.final(adjs).view(*out_shape)
+
+#         self.mask = self.mask.to(score.device)
+#         score = score * self.mask
+
+#         score = mask_adjs(score, flags)
+
+#         return score
+
+
+# class ScoreNetworkA(BaselineNetwork):
+    # def __init__(self, max_feat_num, max_node_num, nhid, num_layers, num_linears, 
+    #                 c_init, c_hid, c_final, adim, num_heads=4, conv='GCN'):
+    #     super(ScoreNetworkA, self).__init__(max_feat_num, max_node_num, nhid, num_layers, num_linears, 
+    #                                         c_init, c_hid, c_final, adim, num_heads=4, conv='GCN')
+class ScoreNetworkA(torch.nn.Module):
+    def __init__(self, max_feat_num, max_node_num, nhid, num_layers, num_linears, 
+                    c_init, c_hid, c_final, adim, num_heads=4, conv='GCN'):
+        super(ScoreNetworkA, self).__init__()
+        
+        self.adim = adim
+        self.num_heads = num_heads
+        self.conv = conv
+        self.num_layers = num_layers
+        self.num_linears = num_linears
+        self.nfeat = max_feat_num
+        self.nhid = nhid
+        self.c_init = c_init
+        self.c_hid = c_hid
+        self.c_final = c_final
+        self.max_node_num = max_node_num
+
+        self.layers = torch.nn.ModuleList()
+        for i in range(self.num_layers):
+            if i==0:
+                self.layers.append(AttentionLayer(self.num_linears, self.nfeat, self.nhid, self.nhid, self.c_init, 
+                                                    self.c_hid, self.num_heads, self.conv))
+            elif i==self.num_layers-1:
+                self.layers.append(AttentionLayer(self.num_linears, self.nhid, self.adim, self.nhid, self.c_hid, 
+                                                    self.c_final, self.num_heads, self.conv))
+            else:
+                self.layers.append(AttentionLayer(self.num_linears, self.nhid, self.adim, self.nhid, self.c_hid, 
+                                                    self.c_hid, self.num_heads, self.conv))
+
+        self.fdim = self.c_hid*(self.num_layers-1) + self.c_final + self.c_init
+        self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=1, 
+                            use_bn=False, activate_func=F.elu)
+        self.mask = torch.ones([self.max_node_num, self.max_node_num]) - torch.eye(self.max_node_num)
+        self.mask.unsqueeze_(0)  
+
+    def forward(self, x, adj, flags):
+
+        adjc = pow_tensor(adj, self.c_init)
+
+        adj_list = [adjc]
+        for i in range(self.num_layers):
+            x, adjc = self.layers[i](x, adjc, flags)
+            adj_list.append(adjc)
+        
+        adjs = torch.cat(adj_list, dim=1).permute(0,2,3,1)
+        out_shape = adjs.shape[:-1] # B x N x N
+        score = self.final(adjs).view(*out_shape)
+        
+        self.mask = self.mask.to(score.device)
+        score = score * self.mask
+
+        score = mask_adjs(score, flags)
+
+        return score

torch_molecule/generator/gdss/toremove/ScoreNetworkX.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn.functional as F
+
+from .layers import DenseGCNConv, MLP, AttentionLayer
+from .utils import mask_x, pow_tensor
+
+class ScoreNetworkX(torch.nn.Module):
+    def __init__(self, max_feat_num, depth, nhid):
+        super(ScoreNetworkX, self).__init__()
+
+        self.nfeat = max_feat_num
+        self.depth = depth
+        self.nhid = nhid
+
+        self.layers = torch.nn.ModuleList()
+        for _ in range(self.depth):
+            if _ == 0:
+                self.layers.append(DenseGCNConv(self.nfeat, self.nhid))
+            else:
+                self.layers.append(DenseGCNConv(self.nhid, self.nhid))
+
+        self.fdim = self.nfeat + self.depth * self.nhid
+        self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=self.nfeat, 
+                            use_bn=False, activate_func=F.elu)
+
+        self.activation = torch.tanh
+
+    def forward(self, x, adj, flags):
+
+        x_list = [x]
+        for _ in range(self.depth):
+            x = self.layers[_](x, adj)
+            x = self.activation(x)
+            x_list.append(x)
+
+        xs = torch.cat(x_list, dim=-1) # B x N x (F + num_layers x H)
+        out_shape = (adj.shape[0], adj.shape[1], -1)
+        x = self.final(xs).view(*out_shape)
+
+        x = mask_x(x, flags)
+
+        return x
+
+
+# class ScoreNetworkX_GMH(torch.nn.Module):
+#     def __init__(self, max_feat_num, depth, nhid, num_linears,
+#                  c_init, c_hid, c_final, adim, num_heads=4, conv='GCN'):
+#         super().__init__()
+
+#         self.depth = depth
+#         self.c_init = c_init
+
+#         self.layers = torch.nn.ModuleList()
+#         for _ in range(self.depth):
+#             if _ == 0:
+#                 self.layers.append(AttentionLayer(num_linears, max_feat_num, nhid, nhid, c_init, 
+#                                                   c_hid, num_heads, conv))
+#             elif _ == self.depth - 1:
+#                 self.layers.append(AttentionLayer(num_linears, nhid, adim, nhid, c_hid, 
+#                                                   c_final, num_heads, conv))
+#             else:
+#                 self.layers.append(AttentionLayer(num_linears, nhid, adim, nhid, c_hid, 
+#                                                   c_hid, num_heads, conv))
+
+#         fdim = max_feat_num + depth * nhid
+#         self.final = MLP(num_layers=3, input_dim=fdim, hidden_dim=2*fdim, output_dim=max_feat_num, 
+#                          use_bn=False, activate_func=F.elu)
+
+#         self.activation = torch.tanh
+
+#     def forward(self, x, adj, flags):
+#         adjc = pow_tensor(adj, self.c_init)
+
+#         x_list = [x]
+#         for _ in range(self.depth):
+#             x, adjc = self.layers[_](x, adjc, flags)
+#             x = self.activation(x)
+#             x_list.append(x)
+
+#         xs = torch.cat(x_list, dim=-1) # B x N x (F + num_layers x H)
+#         out_shape = (adj.shape[0], adj.shape[1], -1)
+#         x = self.final(xs).view(*out_shape)
+#         x = mask_x(x, flags)
+
+#         return x

torch_molecule/generator/gdss/utils.py

@@ -0,0 +1,240 @@
+import json
+import numpy as np
+import torch
+from torch_geometric.utils import to_dense_adj, to_dense_batch, remove_self_loops
+
+from rdkit import Chem
+
+def mask_x(x, flags):
+    if flags is None:
+        flags = torch.ones((x.shape[0], x.shape[1]), device=x.device)
+    return x * flags[:,:,None]
+
+# -------- Mask batch of adjacency matrices with 0-1 flags tensor --------
+def mask_adjs(adjs, flags):
+    """
+    :param adjs:  B x N x N or B x C x N x N
+    :param flags: B x N
+    :return:
+    """
+    if flags is None:
+        flags = torch.ones((adjs.shape[0], adjs.shape[-1]), device=adjs.device)
+
+    if len(adjs.shape) == 4:
+        flags = flags.unsqueeze(1)  # B x 1 x N
+    adjs = adjs * flags.unsqueeze(-1)
+    adjs = adjs * flags.unsqueeze(-2)
+    return adjs
+
+def quantize_mol(adjs):
+    adjs[adjs >= 3.5] = 4
+    adjs[torch.bitwise_and(adjs >= 2.5, adjs < 3.5)] = 3
+    adjs[torch.bitwise_and(adjs >= 1.5, adjs < 2.5)] = 2
+    adjs[torch.bitwise_and(adjs >= 0.5, adjs < 1.5)] = 1
+    adjs[adjs < 0.5] = 0
+    adjs[torch.isnan(adjs)] = 0
+    return adjs.long()
+
+# -------- Generate noise --------
+def gen_noise(x, flags, sym=True):
+    z = torch.randn_like(x)
+    if sym:
+        z = z.triu(1)
+        z = z + z.transpose(-1, -2)
+        z = mask_adjs(z, flags)
+    else:
+        z = mask_x(z, flags)
+    return z
+
+# -------- Create higher order adjacency matrices --------
+def pow_tensor(x, cnum):
+    # x : B x N x N
+    x_ = x.clone()
+    xc = [x.unsqueeze(1)]
+    for _ in range(cnum-1):
+        x_ = torch.bmm(x_, x)
+        xc.append(x_.unsqueeze(1))
+    xc = torch.cat(xc, dim=1)
+
+    return xc
+
+# -------- Common utils --------
+def to_dense(x, edge_index, edge_attr, batch, max_num_nodes):
+    X, node_mask = to_dense_batch(x=x, batch=batch, max_num_nodes=max_num_nodes)
+    edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
+    E = to_dense_adj(
+        edge_index=edge_index,
+        batch=batch,
+        edge_attr=edge_attr,
+        max_num_nodes=max_num_nodes,
+    )
+    return X, E, node_mask
+
+# class PlaceHolder:
+#     def __init__(self, X, E, y):
+#         self.X = X
+#         self.E = E
+#         self.y = y
+
+#     def type_as(self, x: torch.Tensor, categorical: bool = False):
+#         """ Changes the device and dtype of X, E, y. """
+#         self.X = self.X.type_as(x)
+#         self.E = self.E.type_as(x)
+#         if categorical:
+#             self.y = self.y.type_as(x)
+#         return self
+
+#     def mask(self, node_mask, collapse=False):
+#         x_mask = node_mask.unsqueeze(-1)          # bs, n, 1
+#         e_mask1 = x_mask.unsqueeze(2)             # bs, n, 1, 1
+#         e_mask2 = x_mask.unsqueeze(1)             # bs, 1, n, 1
+
+#         if collapse:
+#             self.X = torch.argmax(self.X, dim=-1)
+#             self.E = torch.argmax(self.E, dim=-1)
+
+#             self.X[node_mask == 0] = - 1
+#             self.E[(e_mask1 * e_mask2).squeeze(-1) == 0] = - 1
+#         else:
+#             self.X = self.X * x_mask
+#             self.E = self.E * e_mask1 * e_mask2
+#             assert torch.allclose(self.E, torch.transpose(self.E, 1, 2))
+#         return self
+
+def compute_dataset_info(smiles_or_mol_list, cache_path=None):
+    pt = Chem.GetPeriodicTable()
+    atom_name_list = []
+    atom_count_list = []
+    for i in range(2, 119):
+        atom_name_list.append(pt.GetElementSymbol(i))
+        atom_count_list.append(0)
+    atom_name_list.append('*')
+    atom_count_list.append(0)
+    
+    bond_count_list = [0, 0, 0, 0, 0]
+    bond_type_to_index = {Chem.BondType.SINGLE: 1, Chem.BondType.DOUBLE: 2, Chem.BondType.TRIPLE: 3, Chem.BondType.AROMATIC: 4}
+
+    n_atom_list = []
+    n_bond_list = []
+    n_atoms_per_mol_count = {}
+    max_node = 0
+    for i, sms_or_mol in enumerate(smiles_or_mol_list):
+        if isinstance(sms_or_mol, str):
+            mol = Chem.MolFromSmiles(sms_or_mol)
+        else:
+            mol = sms_or_mol
+
+        n_atom = mol.GetNumHeavyAtoms()
+        n_bond = mol.GetNumBonds()
+        max_node = max(max_node, n_atom)
+        
+        n_atom_list.append(n_atom)
+        n_bond_list.append(n_bond)
+
+        # Count atoms per molecule size
+        if n_atom in n_atoms_per_mol_count:
+            n_atoms_per_mol_count[n_atom] += 1
+        else:
+            n_atoms_per_mol_count[n_atom] = 1
+            
+        cur_atom_count_arr = np.zeros(118)
+        for atom in mol.GetAtoms():
+            symbol = atom.GetSymbol()
+            if symbol == 'H':
+                continue
+            elif symbol == '*':
+                atom_count_list[-1] += 1
+                cur_atom_count_arr[-1] += 1
+            else:
+                atom_count_list[atom.GetAtomicNum()-2] += 1
+                cur_atom_count_arr[atom.GetAtomicNum()-2] += 1
+        
+        tansition_E_temp = np.zeros((118, 118, 5))
+        for bond in mol.GetBonds():
+            start_atom, end_atom = bond.GetBeginAtom(), bond.GetEndAtom()
+            if start_atom.GetSymbol() == 'H' or end_atom.GetSymbol() == 'H':
+                continue
+            
+            if start_atom.GetSymbol() == '*':
+                start_index = 117
+            else:
+                start_index = start_atom.GetAtomicNum() - 2
+            if end_atom.GetSymbol() == '*':
+                end_index = 117
+            else:
+                end_index = end_atom.GetAtomicNum() - 2
+
+            bond_type = bond.GetBondType()
+            bond_index = bond_type_to_index[bond_type]
+            bond_count_list[bond_index] += 2
+    
+    # Create n_atoms_per_mol array with proper size
+    n_atoms_per_mol = [0] * (max_node + 1)
+    for n_atom, count in n_atoms_per_mol_count.items():
+        n_atoms_per_mol[n_atom] = count
+    
+    n_atoms_per_mol = np.array(n_atoms_per_mol) / np.sum(n_atoms_per_mol)
+    n_atoms_per_mol = n_atoms_per_mol.tolist()
+
+    atom_count_list = np.array(atom_count_list) / np.sum(atom_count_list)
+    active_atoms = np.array(atom_name_list)[atom_count_list > 0]
+    active_atoms = active_atoms.tolist()
+    atom_count_list = atom_count_list.tolist()
+
+    bond_count_list = np.array(bond_count_list) / np.sum(bond_count_list)
+    bond_count_list = bond_count_list.tolist()
+
+
+    node_types = torch.Tensor(atom_count_list)
+    edge_types = torch.Tensor(bond_count_list)
+    active_index = (node_types > 0).nonzero().squeeze()
+
+    x_margins = node_types.float() / torch.sum(node_types)
+    x_margins = x_margins[active_index]
+    e_margins = edge_types.float() / torch.sum(edge_types)
+
+    num_nodes_dist = DistributionNodes(torch.Tensor(n_atoms_per_mol))
+
+    meta_dict = {
+        'active_index': active_index,
+        'x_margins': x_margins,
+        'e_margins': e_margins,
+        'atom_decoder': active_atoms,
+        'num_nodes_dist': num_nodes_dist,
+        'max_node': max_node, 
+        }
+
+    if cache_path is not None:
+        with open(cache_path, "w") as f:
+            json.dump(meta_dict, f)
+
+    return meta_dict
+
+class DistributionNodes:
+    def __init__(self, histogram):
+        """ Compute the distribution of the number of nodes in the dataset, and sample from this distribution.
+            historgram: dict. The keys are num_nodes, the values are counts
+        """
+
+        if type(histogram) == dict:
+            max_n_nodes = max(histogram.keys())
+            prob = torch.zeros(max_n_nodes + 1)
+            for num_nodes, count in histogram.items():
+                prob[num_nodes] = count
+        else:
+            prob = histogram
+
+        self.prob = prob / prob.sum()
+        self.m = torch.distributions.Categorical(prob)
+
+    def sample_n(self, n_samples, device):
+        idx = self.m.sample((n_samples,))
+        return idx.to(device)
+
+    def log_prob(self, batch_n_nodes):
+        assert len(batch_n_nodes.size()) == 1
+        p = self.prob.to(batch_n_nodes.device)
+
+        probas = p[batch_n_nodes]
+        log_p = torch.log(probas + 1e-30)
+        return log_p

torch_molecule/generator/graph_dit/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_graph_dit import GraphDITMolecularGenerator
+
+__all__ = ['GraphDITMolecularGenerator']

torch_molecule/generator/graph_dit/diffusion.py

@@ -0,0 +1,188 @@
+import torch
+import numpy as np
+from torch.nn import functional as F
+from .utils import PlaceHolder
+
+def cosine_beta_schedule_discrete(timesteps, s=0.008):
+    """ Cosine schedule as proposed in https://openreview.net/forum?id=-NEXDKk8gZ. """
+    steps = timesteps + 2
+    x = np.linspace(0, steps, steps)
+
+    alphas_cumprod = np.cos(0.5 * np.pi * ((x / steps) + s) / (1 + s)) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    betas = 1 - alphas
+    return betas.squeeze()
+
+class NoiseScheduleDiscrete(torch.nn.Module):
+    def __init__(self, timesteps):
+        super(NoiseScheduleDiscrete, self).__init__()
+        self.timesteps = timesteps
+
+        betas = cosine_beta_schedule_discrete(timesteps)
+        self.register_buffer('betas', torch.from_numpy(betas).float())
+
+        # 0.9999
+        self.alphas = 1 - torch.clamp(self.betas, min=0, max=1)
+
+        log_alpha = torch.log(self.alphas)
+        log_alpha_bar = torch.cumsum(log_alpha, dim=0)
+        self.alphas_bar = torch.exp(log_alpha_bar)
+
+    def forward(self, t_normalized=None, t_int=None):
+        assert int(t_normalized is None) + int(t_int is None) == 1
+        if t_int is None:
+            t_int = torch.round(t_normalized * self.timesteps)
+        return self.betas[t_int.long()]
+
+    def get_alpha_bar(self, t_normalized=None, t_int=None):
+        assert int(t_normalized is None) + int(t_int is None) == 1
+        if t_int is None:
+            t_int = torch.round(t_normalized * self.timesteps)
+        self.alphas_bar = self.alphas_bar.to(t_int.device)
+        return self.alphas_bar[t_int.long()]
+    
+class MarginalTransition:
+    def __init__(self, x_marginals, e_marginals, xe_conditions, ex_conditions, n_nodes):
+        self.X_classes = len(x_marginals)
+        self.E_classes = len(e_marginals)
+        self.x_marginals = x_marginals # Dx
+        self.e_marginals = e_marginals # Dx, De
+        self.xe_conditions = xe_conditions
+
+        self.u_x = x_marginals.unsqueeze(0).expand(self.X_classes, -1).unsqueeze(0) # 1, Dx, Dx
+        self.u_e = e_marginals.unsqueeze(0).expand(self.E_classes, -1).unsqueeze(0) # 1, De, De
+        self.u_xe = xe_conditions.unsqueeze(0) # 1, Dx, De
+        self.u_ex = ex_conditions.unsqueeze(0) # 1, De, Dx
+        self.u = self.get_union_transition(self.u_x, self.u_e, self.u_xe, self.u_ex, n_nodes) # 1, Dx + n*De, Dx + n*De
+
+    def get_union_transition(self, u_x, u_e, u_xe, u_ex, n_nodes):
+        u_e = u_e.repeat(1, n_nodes, n_nodes) # (1, n*de, n*de)
+        u_xe = u_xe.repeat(1, 1, n_nodes) # (1, dx, n*de)
+        u_ex = u_ex.repeat(1, n_nodes, 1) # (1, n*de, dx)
+        u0 = torch.cat([u_x, u_xe], dim=2) # (1, dx, dx + n*de)
+        u1 = torch.cat([u_ex, u_e], dim=2) # (1, n*de, dx + n*de)
+        u = torch.cat([u0, u1], dim=1) # (1, dx + n*de, dx + n*de)
+        return u
+
+    def index_edge_margin(self, X, q_e, n_bond=5):
+        # q_e: (bs, dx, de) --> (bs, n, de)
+        bs, n, n_atom = X.shape
+        node_indices = X.argmax(-1)  # (bs, n)
+        ind = node_indices[ :, :, None].expand(bs, n, n_bond)
+        q_e = torch.gather(q_e, 1, ind)
+        return q_e
+    
+    def get_Qt(self, beta_t, device):
+        """ Returns one-step transition matrices for X and E, from step t - 1 to step t.
+        Qt = (1 - beta_t) * I + beta_t / K
+        beta_t: (bs)
+        returns: q (bs, d0, d0)
+        """
+        bs = beta_t.size(0)
+        d0 = self.u.size(-1)
+        self.u = self.u.to(device)
+        u = self.u.expand(bs, d0, d0)
+
+        beta_t = beta_t.to(device)
+        beta_t = beta_t.view(bs, 1, 1)
+        q = beta_t * u + (1 - beta_t) * torch.eye(d0, device=device).unsqueeze(0)
+
+        return PlaceHolder(X=q, E=None, y=None)
+    
+    def get_Qt_bar(self, alpha_bar_t, device):
+        """ Returns t-step transition matrices for X and E, from step 0 to step t.
+        Qt = prod(1 - beta_t) * I + (1 - prod(1 - beta_t)) * K
+        alpha_bar_t: (bs, 1) roduct of the (1 - beta_t) for each time step from 0 to t.
+        returns: q (bs, d0, d0)
+        """
+        bs = alpha_bar_t.size(0)
+        d0 = self.u.size(-1)
+        alpha_bar_t = alpha_bar_t.to(device)
+        alpha_bar_t = alpha_bar_t.view(bs, 1, 1)
+        self.u = self.u.to(device)
+        q = alpha_bar_t * torch.eye(d0, device=device).unsqueeze(0) + (1 - alpha_bar_t) * self.u
+
+        return PlaceHolder(X=q, E=None, y=None)
+
+def sample_discrete_features(probX, probE, node_mask):
+    """ Sample features from multinomial distribution with given probabilities (probX, probE, proby)
+        :param probX: bs, n, dx_out        node features
+        :param probE: bs, n, n, de_out     edge features
+        :param proby: bs, dy_out           global features.
+    """
+    bs, n, _ = probX.shape
+
+    # Noise X
+    # The masked rows should define probability distributions as well
+    probX[~node_mask] = 1 / probX.shape[-1]
+
+    # Flatten the probability tensor to sample with multinomial
+    probX = probX.reshape(bs * n, -1)  # (bs * n, dx_out)
+
+    # Sample X
+    probX = probX + 1e-12
+    probX = probX / probX.sum(dim=-1, keepdim=True)
+    X_t = probX.multinomial(1)  # (bs * n, 1)
+    X_t = X_t.reshape(bs, n)  # (bs, n)
+
+    # Noise E
+    # The masked rows should define probability distributions as well
+    inverse_edge_mask = ~(node_mask.unsqueeze(1) * node_mask.unsqueeze(2))
+    diag_mask = torch.eye(n).unsqueeze(0).expand(bs, -1, -1)
+
+    probE[inverse_edge_mask] = 1 / probE.shape[-1]
+    probE[diag_mask.bool()] = 1 / probE.shape[-1]
+    probE = probE.reshape(bs * n * n, -1)  # (bs * n * n, de_out)
+    probE = probE + 1e-12
+    probE = probE / probE.sum(dim=-1, keepdim=True)
+
+    # Sample E
+    E_t = probE.multinomial(1).reshape(bs, n, n)  # (bs, n, n)
+    E_t = torch.triu(E_t, diagonal=1)
+    E_t = (E_t + torch.transpose(E_t, 1, 2))
+
+    return PlaceHolder(X=X_t, E=E_t, y=None)
+
+def sample_discrete_feature_noise(limit_dist, node_mask):
+    """Sample from the limit distribution of the diffusion process"""
+    bs, n_max = node_mask.shape
+    x_limit = limit_dist.X[None, None, :].expand(bs, n_max, -1)
+    U_X = x_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max)
+    U_X = F.one_hot(U_X.long(), num_classes=x_limit.shape[-1]).type_as(x_limit)
+
+    e_limit = limit_dist.E[None, None, None, :].expand(bs, n_max, n_max, -1)
+    U_E = e_limit.flatten(end_dim=-2).multinomial(1).reshape(bs, n_max, n_max)
+    U_E = F.one_hot(U_E.long(), num_classes=e_limit.shape[-1]).type_as(x_limit)
+
+    U_X = U_X.to(node_mask.device)
+    U_E = U_E.to(node_mask.device)
+
+    # Get upper triangular part of edge noise, without main diagonal
+    upper_triangular_mask = torch.zeros_like(U_E)
+    indices = torch.triu_indices(row=U_E.size(1), col=U_E.size(2), offset=1)
+    upper_triangular_mask[:, indices[0], indices[1], :] = 1
+
+    U_E = U_E * upper_triangular_mask
+    U_E = U_E + torch.transpose(U_E, 1, 2)
+
+    assert (U_E == torch.transpose(U_E, 1, 2)).all()
+    return PlaceHolder(X=U_X, E=U_E, y=None).mask(node_mask)
+
+def reverse_diffusion(predX_0, X_t, Qt, Qsb, Qtb):
+    """M: X or E
+    Compute xt @ Qt.T * x0 @ Qsb / x0 @ Qtb @ xt.T for each possible value of x0
+    X_t: bs, n, dt          or bs, n, n, dt
+    Qt: bs, d_t-1, dt
+    Qsb: bs, d0, d_t-1
+    Qtb: bs, d0, dt.
+    """
+    Qt_T = Qt.transpose(-1, -2)  # bs, N, dt
+    assert Qt.dim() == 3
+    left_term = X_t @ Qt_T  # bs, N, d_t-1
+    right_term = predX_0 @ Qsb
+    numerator = left_term * right_term  # bs, N, d_t-1
+
+    denominator = Qtb @ X_t.transpose(-1, -2)  # bs, d0, N
+    denominator = denominator.transpose(-1, -2)  # bs, N, d0
+    return numerator / denominator.clamp_min(1e-5)