"""AttentiveFP""" # pylint: disable= no-member, arguments-differ, invalid-name import dgl.function as fn import torch import torch.nn as nn import torch.nn.functional as F from dgl.nn.pytorch import edge_softmax __all__ = ['AttentiveFPGNN'] # pylint: disable=W0221, C0103, E1101 class AttentiveGRU1(nn.Module): """Update node features with attention and GRU. This will be used for incorporating the information of edge features into node features for message passing. Parameters ---------- node_feat_size : int Size for the input node features. edge_feat_size : int Size for the input edge (bond) features. edge_hidden_size : int Size for the intermediate edge (bond) representations. dropout : float The probability for performing dropout. """ def __init__(self, node_feat_size, edge_feat_size, edge_hidden_size, dropout): super(AttentiveGRU1, self).__init__() self.edge_transform = nn.Sequential( nn.Dropout(dropout), nn.Linear(edge_feat_size, edge_hidden_size) ) self.gru = nn.GRUCell(edge_hidden_size, node_feat_size) def forward(self, g, edge_logits, edge_feats, node_feats): """Update node representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs edge_logits : float32 tensor of shape (E, 1) The edge logits based on which softmax will be performed for weighting edges within 1-hop neighborhoods. E represents the number of edges. edge_feats : float32 tensor of shape (E, edge_feat_size) Previous edge features. node_feats : float32 tensor of shape (V, node_feat_size) Previous node features. V represents the number of nodes. Returns ------- float32 tensor of shape (V, node_feat_size) Updated node features. """ g = g.local_var() g.edata['e'] = edge_softmax(g, edge_logits) * self.edge_transform(edge_feats) g.update_all(fn.copy_edge('e', 'm'), fn.sum('m', 'c')) context = F.elu(g.ndata['c']) return F.relu(self.gru(context, node_feats)) class AttentiveGRU2(nn.Module): """Update node features with attention and GRU. This will be used in GNN layers for updating node representations. Parameters ---------- node_feat_size : int Size for the input node features. edge_hidden_size : int Size for the intermediate edge (bond) representations. dropout : float The probability for performing dropout. """ def __init__(self, node_feat_size, edge_hidden_size, dropout): super(AttentiveGRU2, self).__init__() self.project_node = nn.Sequential( nn.Dropout(dropout), nn.Linear(node_feat_size, edge_hidden_size) ) self.gru = nn.GRUCell(edge_hidden_size, node_feat_size) def forward(self, g, edge_logits, node_feats): """Update node representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs edge_logits : float32 tensor of shape (E, 1) The edge logits based on which softmax will be performed for weighting edges within 1-hop neighborhoods. E represents the number of edges. node_feats : float32 tensor of shape (V, node_feat_size) Previous node features. V represents the number of nodes. Returns ------- float32 tensor of shape (V, node_feat_size) Updated node features. """ g = g.local_var() g.edata['a'] = edge_softmax(g, edge_logits) g.ndata['hv'] = self.project_node(node_feats) g.update_all(fn.src_mul_edge('hv', 'a', 'm'), fn.sum('m', 'c')) context = F.elu(g.ndata['c']) return F.relu(self.gru(context, node_feats)) class GetContext(nn.Module): """Generate context for each node by message passing at the beginning. This layer incorporates the information of edge features into node representations so that message passing needs to be only performed over node representations. Parameters ---------- node_feat_size : int Size for the input node features. edge_feat_size : int Size for the input edge (bond) features. graph_feat_size : int Size of the learned graph representation (molecular fingerprint). dropout : float The probability for performing dropout. """ def __init__(self, node_feat_size, edge_feat_size, graph_feat_size, dropout): super(GetContext, self).__init__() self.project_node = nn.Sequential( nn.Linear(node_feat_size, graph_feat_size), nn.LeakyReLU() ) self.project_edge1 = nn.Sequential( nn.Linear(node_feat_size + edge_feat_size, graph_feat_size), nn.LeakyReLU() ) self.project_edge2 = nn.Sequential( nn.Dropout(dropout), nn.Linear(2 * graph_feat_size, 1), nn.LeakyReLU() ) self.attentive_gru = AttentiveGRU1(graph_feat_size, graph_feat_size, graph_feat_size, dropout) def apply_edges1(self, edges): """Edge feature update. Parameters ---------- edges : EdgeBatch Container for a batch of edges Returns ------- dict Mapping ``'he1'`` to updated edge features. """ return {'he1': torch.cat([edges.src['hv'], edges.data['he']], dim=1)} def apply_edges2(self, edges): """Edge feature update. Parameters ---------- edges : EdgeBatch Container for a batch of edges Returns ------- dict Mapping ``'he2'`` to updated edge features. """ return {'he2': torch.cat([edges.dst['hv_new'], edges.data['he1']], dim=1)} def forward(self, g, node_feats, edge_feats): """Incorporate edge features and update node representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs. node_feats : float32 tensor of shape (V, node_feat_size) Input node features. V for the number of nodes. edge_feats : float32 tensor of shape (E, edge_feat_size) Input edge features. E for the number of edges. Returns ------- float32 tensor of shape (V, graph_feat_size) Updated node features. """ g = g.local_var() g.ndata['hv'] = node_feats g.ndata['hv_new'] = self.project_node(node_feats) g.edata['he'] = edge_feats g.apply_edges(self.apply_edges1) g.edata['he1'] = self.project_edge1(g.edata['he1']) g.apply_edges(self.apply_edges2) logits = self.project_edge2(g.edata['he2']) return self.attentive_gru(g, logits, g.edata['he1'], g.ndata['hv_new']) class GNNLayer(nn.Module): """GNNLayer for updating node features. This layer performs message passing over node representations and update them. Parameters ---------- node_feat_size : int Size for the input node features. graph_feat_size : int Size for the graph representations to be computed. dropout : float The probability for performing dropout. """ def __init__(self, node_feat_size, graph_feat_size, dropout): super(GNNLayer, self).__init__() self.project_edge = nn.Sequential( nn.Dropout(dropout), nn.Linear(2 * node_feat_size, 1), nn.LeakyReLU() ) self.attentive_gru = AttentiveGRU2(node_feat_size, graph_feat_size, dropout) def apply_edges(self, edges): """Edge feature generation. Generate edge features by concatenating the features of the destination and source nodes. Parameters ---------- edges : EdgeBatch Container for a batch of edges. Returns ------- dict Mapping ``'he'`` to the generated edge features. """ return {'he': torch.cat([edges.dst['hv'], edges.src['hv']], dim=1)} def forward(self, g, node_feats): """Perform message passing and update node representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs. node_feats : float32 tensor of shape (V, node_feat_size) Input node features. V for the number of nodes. Returns ------- float32 tensor of shape (V, graph_feat_size) Updated node features. """ g = g.local_var() g.ndata['hv'] = node_feats g.apply_edges(self.apply_edges) logits = self.project_edge(g.edata['he']) return self.attentive_gru(g, logits, node_feats) class AttentiveFPGNN(nn.Module): """`Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism `__ This class performs message passing in AttentiveFP and returns the updated node representations. Parameters ---------- node_feat_size : int Size for the input node features. edge_feat_size : int Size for the input edge features. num_layers : int Number of GNN layers. Default to 2. graph_feat_size : int Size for the graph representations to be computed. Default to 200. dropout : float The probability for performing dropout. Default to 0. """ def __init__(self, node_feat_size, edge_feat_size, num_layers=2, graph_feat_size=200, dropout=0.): super(AttentiveFPGNN, self).__init__() self.init_context = GetContext(node_feat_size, edge_feat_size, graph_feat_size, dropout) self.gnn_layers = nn.ModuleList() for _ in range(num_layers - 1): self.gnn_layers.append(GNNLayer(graph_feat_size, graph_feat_size, dropout)) def forward(self, g, node_feats, edge_feats): """Performs message passing and updates node representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs. node_feats : float32 tensor of shape (V, node_feat_size) Input node features. V for the number of nodes. edge_feats : float32 tensor of shape (E, edge_feat_size) Input edge features. E for the number of edges. Returns ------- node_feats : float32 tensor of shape (V, graph_feat_size) Updated node representations. """ node_feats = self.init_context(g, node_feats, edge_feats) for gnn in self.gnn_layers: node_feats = gnn(g, node_feats) return node_feats