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

torch_molecule/generator/digress/transformer.py

@@ -0,0 +1,360 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.dropout import Dropout
+from torch.nn.modules.linear import Linear
+from torch.nn.modules.normalization import LayerNorm
+from torch.nn import functional as F
+from torch import Tensor
+
+from .utils import PlaceHolder
+from .diffusion import assert_correctly_masked
+
+class GraphTransformer(nn.Module):
+    def __init__(self, num_layer: int, input_dim_X: int, input_dim_E: int, input_dim_y: int, hidden_size_X: int, hidden_size_E: int, hidden_size_y: int, n_head: int,dropout: float):
+        super().__init__()
+        self.num_layer = num_layer
+        self.out_dim_X = input_dim_X
+        self.out_dim_E = input_dim_E
+        self.out_dim_y = input_dim_y
+        act_fn_in = nn.ReLU()
+        act_fn_out = nn.ReLU()
+
+        self.mlp_in_X = nn.Sequential(nn.Linear(input_dim_X, hidden_size_X), act_fn_in,
+                                      nn.Linear(hidden_size_X, hidden_size_X), act_fn_in)
+
+        self.mlp_in_E = nn.Sequential(nn.Linear(input_dim_E, hidden_size_E), act_fn_in,
+                                      nn.Linear(hidden_size_E, hidden_size_E), act_fn_in)
+
+        self.mlp_in_y = nn.Sequential(nn.Linear(input_dim_y, hidden_size_y), act_fn_in,
+                                      nn.Linear(hidden_size_y, hidden_size_y), act_fn_in)
+
+        self.tf_layers = nn.ModuleList([XEyTransformerLayer(dx=hidden_size_X,
+                                                            de=hidden_size_E,
+                                                            dy=hidden_size_y,
+                                                            n_head=n_head,
+                                                            dim_ffX=hidden_size_X,
+                                                            dim_ffE=hidden_size_E,
+                                                            dropout=dropout)
+                                        for i in range(num_layer)])
+
+        self.mlp_out_X = nn.Sequential(nn.Linear(hidden_size_X, hidden_size_X), act_fn_out,
+                                       nn.Linear(hidden_size_X, self.out_dim_X))
+
+        self.mlp_out_E = nn.Sequential(nn.Linear(hidden_size_E, hidden_size_E), act_fn_out,
+                                       nn.Linear(hidden_size_E, self.out_dim_E))
+
+        self.mlp_out_y = nn.Sequential(nn.Linear(hidden_size_y, hidden_size_y), act_fn_out,
+                                       nn.Linear(hidden_size_y, self.out_dim_y))
+
+    def initialize_parameters(self, seed=None):
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        # Initialize transformer layers:
+        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, noisy_data, true_X, true_E, lw_X, lw_E):
+        pred = self.forward(noisy_data)
+        
+        # Reshape predictions and targets
+        true_X = torch.reshape(true_X, (-1, true_X.size(-1)))  # (bs * n, dx)
+        true_E = torch.reshape(true_E, (-1, true_E.size(-1)))  # (bs * n * n, de)
+        masked_pred_X = torch.reshape(pred.X, (-1, pred.X.size(-1)))  # (bs * n, dx)
+        masked_pred_E = torch.reshape(pred.E, (-1, pred.E.size(-1)))   # (bs * n * n, de)
+
+        # Remove masked rows
+        mask_X = (true_X != 0.).any(dim=-1)
+        mask_E = (true_E != 0.).any(dim=-1)
+
+        flat_true_X = true_X[mask_X, :]
+        flat_pred_X = masked_pred_X[mask_X, :]
+        flat_true_E = true_E[mask_E, :]
+        flat_pred_E = masked_pred_E[mask_E, :]
+        
+        # Calculate node and edge losses using cross entropy
+        loss_X = F.cross_entropy(flat_pred_X, torch.argmax(flat_true_X, dim=-1)) if true_X.numel() > 0 else 0.0
+        loss_E = F.cross_entropy(flat_pred_E, torch.argmax(flat_true_E, dim=-1)) if true_E.numel() > 0 else 0.0
+        loss = lw_X * loss_X + lw_E * loss_E
+
+        return loss, loss_X, loss_E
+
+    def forward(self, noisy_data):
+        X, E, y, node_mask = noisy_data['X_t'].float(), noisy_data['E_t'].float(), noisy_data['y_t'].float(), noisy_data['node_mask']
+        bs, n = X.shape[0], X.shape[1]
+
+        diag_mask = torch.eye(n)
+        diag_mask = ~diag_mask.type_as(E).bool()
+        diag_mask = diag_mask.unsqueeze(0).unsqueeze(-1).expand(bs, -1, -1, -1)
+
+        X_to_out = X[..., :self.out_dim_X]
+        E_to_out = E[..., :self.out_dim_E]
+        y_to_out = y[..., :self.out_dim_y]
+
+        new_E = self.mlp_in_E(E)
+        new_E = (new_E + new_E.transpose(1, 2)) / 2
+        after_in = PlaceHolder(X=self.mlp_in_X(X), E=new_E, y=self.mlp_in_y(y)).mask(node_mask)
+        X, E, y = after_in.X, after_in.E, after_in.y
+
+        for layer in self.tf_layers:
+            X, E, y = layer(X, E, y, node_mask)
+
+        X = self.mlp_out_X(X)
+        E = self.mlp_out_E(E)
+        y = self.mlp_out_y(y)
+
+        X = (X + X_to_out)
+        E = (E + E_to_out) * diag_mask
+        y = y + y_to_out
+
+        E = 1/2 * (E + torch.transpose(E, 1, 2))
+
+        return PlaceHolder(X=X, E=E, y=y).mask(node_mask)
+
+class XEyTransformerLayer(nn.Module):
+    """ Transformer that updates node, edge and global features
+        d_x: node features
+        d_e: edge features
+        dz : global features
+        n_head: the number of heads in the multi_head_attention
+        dim_feedforward: the dimension of the feedforward network model after self-attention
+        dropout: dropout probablility. 0 to disable
+        layer_norm_eps: eps value in layer normalizations.
+    """
+    def __init__(self, dx: int, de: int, dy: int, n_head: int, dim_ffX: int = 2048,
+                 dim_ffE: int = 128, dim_ffy: int = 2048, dropout: float = 0.1,
+                 layer_norm_eps: float = 1e-5, device=None, dtype=None) -> None:
+        kw = {'device': device, 'dtype': dtype}
+        super().__init__()
+
+        self.self_attn = NodeEdgeBlock(dx, de, dy, n_head, **kw)
+
+        self.linX1 = Linear(dx, dim_ffX, **kw)
+        self.linX2 = Linear(dim_ffX, dx, **kw)
+        self.normX1 = LayerNorm(dx, eps=layer_norm_eps, **kw)
+        self.normX2 = LayerNorm(dx, eps=layer_norm_eps, **kw)
+        self.dropoutX1 = Dropout(dropout)
+        self.dropoutX2 = Dropout(dropout)
+        self.dropoutX3 = Dropout(dropout)
+
+        self.linE1 = Linear(de, dim_ffE, **kw)
+        self.linE2 = Linear(dim_ffE, de, **kw)
+        self.normE1 = LayerNorm(de, eps=layer_norm_eps, **kw)
+        self.normE2 = LayerNorm(de, eps=layer_norm_eps, **kw)
+        self.dropoutE1 = Dropout(dropout)
+        self.dropoutE2 = Dropout(dropout)
+        self.dropoutE3 = Dropout(dropout)
+
+        self.lin_y1 = Linear(dy, dim_ffy, **kw)
+        self.lin_y2 = Linear(dim_ffy, dy, **kw)
+        self.norm_y1 = LayerNorm(dy, eps=layer_norm_eps, **kw)
+        self.norm_y2 = LayerNorm(dy, eps=layer_norm_eps, **kw)
+        self.dropout_y1 = Dropout(dropout)
+        self.dropout_y2 = Dropout(dropout)
+        self.dropout_y3 = Dropout(dropout)
+
+        self.activation = F.relu
+
+    def forward(self, X: Tensor, E: Tensor, y, node_mask: Tensor):
+        """ Pass the input through the encoder layer.
+            X: (bs, n, d)
+            E: (bs, n, n, d)
+            y: (bs, dy)
+            node_mask: (bs, n) Mask for the src keys per batch (optional)
+            Output: newX, newE, new_y with the same shape.
+        """
+
+        newX, newE, new_y = self.self_attn(X, E, y, node_mask=node_mask)
+
+        newX_d = self.dropoutX1(newX)
+        X = self.normX1(X + newX_d)
+
+        newE_d = self.dropoutE1(newE)
+        E = self.normE1(E + newE_d)
+
+        new_y_d = self.dropout_y1(new_y)
+        y = self.norm_y1(y + new_y_d)
+
+        ff_outputX = self.linX2(self.dropoutX2(self.activation(self.linX1(X))))
+        ff_outputX = self.dropoutX3(ff_outputX)
+        X = self.normX2(X + ff_outputX)
+
+        ff_outputE = self.linE2(self.dropoutE2(self.activation(self.linE1(E))))
+        ff_outputE = self.dropoutE3(ff_outputE)
+        E = self.normE2(E + ff_outputE)
+
+        ff_output_y = self.lin_y2(self.dropout_y2(self.activation(self.lin_y1(y))))
+        ff_output_y = self.dropout_y3(ff_output_y)
+        y = self.norm_y2(y + ff_output_y)
+
+        return X, E, y
+
+
+class NodeEdgeBlock(nn.Module):
+    """ Self attention layer that also updates the representations on the edges. """
+    def __init__(self, dx, de, dy, n_head, **kwargs):
+        super().__init__()
+        assert dx % n_head == 0, f"dx: {dx} -- nhead: {n_head}"
+        self.dx = dx
+        self.de = de
+        self.dy = dy
+        self.df = int(dx / n_head)
+        self.n_head = n_head
+
+        # Attention
+        self.q = Linear(dx, dx)
+        self.k = Linear(dx, dx)
+        self.v = Linear(dx, dx)
+
+        # FiLM E to X
+        self.e_add = Linear(de, dx)
+        self.e_mul = Linear(de, dx)
+
+        # FiLM y to E
+        self.y_e_mul = Linear(dy, dx)           # Warning: here it's dx and not de
+        self.y_e_add = Linear(dy, dx)
+
+        # FiLM y to X
+        self.y_x_mul = Linear(dy, dx)
+        self.y_x_add = Linear(dy, dx)
+
+        # Process y
+        self.y_y = Linear(dy, dy)
+        self.x_y = Xtoy(dx, dy)
+        self.e_y = Etoy(de, dy)
+
+        # Output layers
+        self.x_out = Linear(dx, dx)
+        self.e_out = Linear(dx, de)
+        self.y_out = nn.Sequential(nn.Linear(dy, dy), nn.ReLU(), nn.Linear(dy, dy))
+
+    def forward(self, X, E, y, node_mask):
+        """
+        :param X: bs, n, d        node features
+        :param E: bs, n, n, d     edge features
+        :param y: bs, dz           global features
+        :param node_mask: bs, n
+        :return: newX, newE, new_y with the same shape.
+        """
+        bs, n, _ = X.shape
+        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
+
+        # 1. Map X to keys and queries
+        Q = self.q(X) * x_mask           # (bs, n, dx)
+        K = self.k(X) * x_mask           # (bs, n, dx)
+        assert_correctly_masked(Q, x_mask)
+        # 2. Reshape to (bs, n, n_head, df) with dx = n_head * df
+
+        Q = Q.reshape((Q.size(0), Q.size(1), self.n_head, self.df))
+        K = K.reshape((K.size(0), K.size(1), self.n_head, self.df))
+
+        Q = Q.unsqueeze(2)                              # (bs, 1, n, n_head, df)
+        K = K.unsqueeze(1)                              # (bs, n, 1, n head, df)
+
+        # Compute unnormalized attentions. Y is (bs, n, n, n_head, df)
+        Y = Q * K
+        Y = Y / math.sqrt(Y.size(-1))
+        assert_correctly_masked(Y, (e_mask1 * e_mask2).unsqueeze(-1))
+
+        E1 = self.e_mul(E) * e_mask1 * e_mask2                        # bs, n, n, dx
+        E1 = E1.reshape((E.size(0), E.size(1), E.size(2), self.n_head, self.df))
+
+        E2 = self.e_add(E) * e_mask1 * e_mask2                        # bs, n, n, dx
+        E2 = E2.reshape((E.size(0), E.size(1), E.size(2), self.n_head, self.df))
+
+        # Incorporate edge features to the self attention scores.
+        Y = Y * (E1 + 1) + E2                  # (bs, n, n, n_head, df)
+
+        # Incorporate y to E
+        newE = Y.flatten(start_dim=3)                      # bs, n, n, dx
+        ye1 = self.y_e_add(y).unsqueeze(1).unsqueeze(1)  # bs, 1, 1, de
+        ye2 = self.y_e_mul(y).unsqueeze(1).unsqueeze(1)
+        newE = ye1 + (ye2 + 1) * newE
+
+        # Output E
+        newE = self.e_out(newE) * e_mask1 * e_mask2      # bs, n, n, de
+        assert_correctly_masked(newE, e_mask1 * e_mask2)
+
+        # Compute attentions. attn is still (bs, n, n, n_head, df)
+        softmax_mask = e_mask2.expand(-1, n, -1, self.n_head)    # bs, 1, n, 1
+        attn = masked_softmax(Y, softmax_mask, dim=2)  # bs, n, n, n_head
+
+        V = self.v(X) * x_mask                        # bs, n, dx
+        V = V.reshape((V.size(0), V.size(1), self.n_head, self.df))
+        V = V.unsqueeze(1)                                     # (bs, 1, n, n_head, df)
+
+        # Compute weighted values
+        weighted_V = attn * V
+        weighted_V = weighted_V.sum(dim=2)
+
+        # Send output to input dim
+        weighted_V = weighted_V.flatten(start_dim=2)            # bs, n, dx
+
+        # Incorporate y to X
+        yx1 = self.y_x_add(y).unsqueeze(1)
+        yx2 = self.y_x_mul(y).unsqueeze(1)
+        newX = yx1 + (yx2 + 1) * weighted_V
+
+        # Output X
+        newX = self.x_out(newX) * x_mask
+        assert_correctly_masked(newX, x_mask)
+
+        # Process y based on X axnd E
+        y = self.y_y(y)
+        e_y = self.e_y(E)
+        x_y = self.x_y(X)
+        new_y = y + x_y + e_y
+        new_y = self.y_out(new_y)               # bs, dy
+
+        return newX, newE, new_y
+
+class Xtoy(nn.Module):
+    def __init__(self, dx, dy):
+        """ Map node features to global features """
+        super().__init__()
+        self.lin = nn.Linear(4 * dx, dy)
+
+    def forward(self, X):
+        """ X: bs, n, dx. """
+        m = X.mean(dim=1)
+        mi = X.min(dim=1)[0]
+        ma = X.max(dim=1)[0]
+        std = X.std(dim=1)
+        z = torch.hstack((m, mi, ma, std))
+        out = self.lin(z)
+        return out
+
+
+class Etoy(nn.Module):
+    def __init__(self, d, dy):
+        """ Map edge features to global features. """
+        super().__init__()
+        self.lin = nn.Linear(4 * d, dy)
+
+    def forward(self, E):
+        """ E: bs, n, n, de
+            Features relative to the diagonal of E could potentially be added.
+        """
+        m = E.mean(dim=(1, 2))
+        mi = E.min(dim=2)[0].min(dim=1)[0]
+        ma = E.max(dim=2)[0].max(dim=1)[0]
+        std = torch.std(E, dim=(1, 2))
+        z = torch.hstack((m, mi, ma, std))
+        out = self.lin(z)
+        return out
+
+def masked_softmax(x, mask, **kwargs):
+    if mask.sum() == 0:
+        return x
+    x_masked = x.clone()
+    x_masked[mask == 0] = -float("inf")
+    return torch.softmax(x_masked, **kwargs)

torch_molecule/generator/digress/utils.py

@@ -0,0 +1,196 @@
+import torch
+import json
+from rdkit import Chem
+import numpy as np
+from torch_geometric.utils import to_dense_adj, to_dense_batch, remove_self_loops
+
+def to_dense(x, edge_index, edge_attr, batch, max_num_nodes=None):
+    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)
+    if max_num_nodes is None:
+        max_num_nodes = X.size(1)
+    E = to_dense_adj(edge_index=edge_index, batch=batch, edge_attr=edge_attr, max_num_nodes=max_num_nodes)
+    E = encode_no_edge(E)
+    return PlaceHolder(X=X, E=E, y=None), node_mask
+
+def encode_no_edge(E):
+    assert len(E.shape) == 4
+    if E.shape[-1] == 0:
+        return E
+    no_edge = torch.sum(E, dim=3) == 0
+    first_elt = E[:, :, :, 0]
+    first_elt[no_edge] = 1
+    E[:, :, :, 0] = first_elt
+    diag = torch.eye(E.shape[1], dtype=torch.bool).unsqueeze(0).expand(E.shape[0], -1, -1)
+    E[diag] = 0
+    return E
+
+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}
+    tansition_E = np.zeros((118, 118, 5))
+
+    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/gdss/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_gdss import GDSSMolecularGenerator
+
+__all__ = ['GDSSMolecularGenerator']