Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch_geometric.nn import SAGEConv, Set2Set | |
| class EncodeBeads(nn.Module): | |
| ''' | |
| The model is partially taken from: Chem. Sci., 2022, 13, 4498-4511 | |
| It is a Regularized Autoencoder with a message passing layer. | |
| ''' | |
| def __init__(self, dim_ae): | |
| super(EncodeBeads,self).__init__() | |
| self.dim_ae = dim_ae | |
| self.node_feature_dim = self.dim_ae['m'] #Should it be m*n? | |
| self.embedding_dim = self.dim_ae['embedding_dim'] | |
| self.encoder_dim = self.dim_ae['encoder_dim'] | |
| self.latent_dim = self.dim_ae['latent_dim'] | |
| self.mp_steps = self.dim_ae['mp_steps'] | |
| #Encoding the beads | |
| self.embedding = nn.Linear(self.node_feature_dim, self.embedding_dim) | |
| self.lin0 = nn.Linear(self.embedding_dim ,self.encoder_dim) | |
| self.conv = SAGEConv(self.encoder_dim, self.encoder_dim, aggr="add", root_weight=True, project=False) | |
| self.gru = nn.GRU(self.encoder_dim, self.encoder_dim) | |
| self.set2set = Set2Set(self.encoder_dim, processing_steps=3) | |
| self.lin1 = nn.Linear(2 * self.encoder_dim, self.encoder_dim) | |
| self.final = nn.Linear(self.encoder_dim, self.latent_dim) | |
| def forward(self, data): | |
| x, edge_index, batch = data.x, data.edge_index, data.batch | |
| #Embedding | |
| emb = self.embedding(x) | |
| emb = F.leaky_relu(emb) | |
| #Linear layer | |
| out = F.leaky_relu(self.lin0(emb)) | |
| h = out.unsqueeze(0) | |
| for _ in range(self.mp_steps): | |
| m = F.relu(self.conv(out, edge_index)) | |
| out, h = self.gru(m.unsqueeze(0),h) | |
| out = out.squeeze(0) | |
| out_s2s = self.set2set(out, batch) | |
| out = F.leaky_relu(self.lin1(out_s2s)) | |
| return self.final(out) | |
| class DecodeBeads(nn.Module): | |
| #Using a MLP on the dimensions of the encoder and latent. | |
| def __init__(self,dim_ae): | |
| super(DecodeBeads,self).__init__() | |
| self.dim_ae = dim_ae | |
| self.n = self.dim_ae['n'] #Number of beads Max 3 for the moment number of nodes | |
| self.m = self.dim_ae['m'] #node feature | |
| self.latent_dim = self.dim_ae['latent_dim'] | |
| self.act_fn = nn.LeakyReLU() | |
| # Keep positive values | |
| # self.sm_nodes = nn.Softmax(dim=1) | |
| # Dimension of the decoder | |
| self.node_hidden = self.dim_ae['node_hidden'] | |
| self.edge_prelayer_dim = self.dim_ae['edge_prelay'] | |
| self.edge_hidden = self.dim_ae['edge_hidden'] | |
| self.encoded_node = self.dim_ae['encoder_node'] | |
| #Decode layers | |
| # self.triu_mask = torch.ones(self.n, self.n).triu_().bool() | |
| row, col = torch.triu_indices(self.n, self.n, offset=1) | |
| self.edge_iterator = torch.stack([row, col], dim=1) | |
| # Node decoding | |
| self.node_reconst = self.construct_linear_layers(self.latent_dim, | |
| self.node_hidden, | |
| (self.m+1)*self.n,) | |
| # Edge decoding | |
| self.reconst_node_encod = self.construct_linear_layers(self.m, | |
| self.encoded_node[:-1], | |
| self.encoded_node[-1]) | |
| self.edge_prelayer = self.construct_linear_layers(self.latent_dim, | |
| self.edge_prelayer_dim[:-1], | |
| self.edge_prelayer_dim[-1]) | |
| self.edge_reconst = self.construct_linear_layers( | |
| self.encoded_node[-1]*self.n + self.edge_prelayer_dim[-1], | |
| self.edge_hidden, | |
| (self.n*(self.n-1))//2, | |
| ) | |
| def construct_linear_layers( | |
| input_dim, | |
| hidden_dims, | |
| output_dim, | |
| nonlinearity=nn.LeakyReLU(negative_slope=0.1)): | |
| """ | |
| Constructs a torch.nn module with sequential linear layers with the specified input and | |
| output dimensions and the specified hidden dimensions. | |
| The nonlinearity is applied to all but the last layer. | |
| :param input_dim: Input dimension of the first layer | |
| :param hidden_dims: List of hidden dimensions | |
| :param output_dim: Output dimension of the last layer | |
| :param nonlinearity: Nonlinearity to apply to all but the last layer | |
| :return: List of linear layers | |
| """ | |
| hidden_dims = [input_dim] + hidden_dims + [output_dim] | |
| layers = [] | |
| for i in range(len(hidden_dims) - 1): | |
| linear = nn.Linear(hidden_dims[i], hidden_dims[i + 1]) | |
| layers.append(linear) | |
| if i < len(hidden_dims) - 2: | |
| ### Add nonlinearity to all but the last layer ### | |
| layers.append(nonlinearity) | |
| return nn.Sequential(*layers) | |
| def build_dense_adj(self, upper_adj_mat,bs): | |
| adj = torch.zeros(bs, self.n, self.n, device=upper_adj_mat.device, dtype=upper_adj_mat.dtype) | |
| i = self.edge_iterator[:, 0] | |
| j = self.edge_iterator[:, 1] | |
| # fill symmetric entries | |
| adj[:, i, j] = upper_adj_mat | |
| adj[:, j, i] = upper_adj_mat | |
| return adj | |
| def forward(self,x): | |
| # Nodes | |
| nodes = self.node_reconst(x).view(-1,self.n,self.m+1) | |
| # Edges | |
| encoded_latent_space = self.act_fn(self.edge_prelayer(x)) | |
| encoded_nodes = self.reconst_node_encod(nodes.view(-1,self.m+1)[:,:-1].contiguous()) | |
| encoded_nodes = encoded_nodes.view(-1,self.n*encoded_nodes.shape[-1]) | |
| edge_rep = torch.cat([encoded_nodes, encoded_latent_space], dim=-1) | |
| upper_adj_mat = self.edge_reconst(edge_rep) | |
| # Get adjacency matrix | |
| adj = self.build_dense_adj(upper_adj_mat,x.shape[0]) | |
| return nodes, adj | |
| class Autoencoder(nn.Module): | |
| """ | |
| Autoencoder for molecule encoding | |
| """ | |
| def __init__(self,dim_ae): | |
| super().__init__() | |
| self.dim_ae = dim_ae | |
| self.encoder = EncodeBeads(self.dim_ae) | |
| self.decoder = DecodeBeads(self.dim_ae) | |
| def forward(self, data, only_latent_space=True): | |
| ### Encode data ### | |
| z = self.encoder(data) | |
| if only_latent_space: | |
| return z | |
| else: | |
| ### Decode latent space ### | |
| nodes, adj_mat = self.decoder(z) | |
| return nodes, adj_mat, z | |