"""Torch Module for Atomic Convolution Layer"""# pylint: disable= no-member, arguments-differ, invalid-nameimportnumpyasnpimporttorchasthimporttorch.nnasnnclassRadialPooling(nn.Module):r"""Radial pooling from `Atomic Convolutional Networks for Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__ We denote the distance between atom :math:`i` and :math:`j` by :math:`r_{ij}`. A radial pooling layer transforms distances with radial filters. For radial filter indexed by :math:`k`, it projects edge distances with .. math:: h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2) If :math:`r_{ij} < c_k`, .. math:: f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1), else, .. math:: f_{ij}^{k} = 0. Finally, .. math:: e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k} Parameters ---------- interaction_cutoffs : float32 tensor of shape (K) :math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs and two atoms are considered as connected if the distance between them is smaller than the cutoffs. K for the number of radial filters. rbf_kernel_means : float32 tensor of shape (K) :math:`r_k` in the equations above. K for the number of radial filters. rbf_kernel_scaling : float32 tensor of shape (K) :math:`\gamma_k` in the equations above. K for the number of radial filters. """def__init__(self,interaction_cutoffs,rbf_kernel_means,rbf_kernel_scaling):super(RadialPooling,self).__init__()self.interaction_cutoffs=nn.Parameter(interaction_cutoffs.reshape(-1,1,1),requires_grad=True)self.rbf_kernel_means=nn.Parameter(rbf_kernel_means.reshape(-1,1,1),requires_grad=True)self.rbf_kernel_scaling=nn.Parameter(rbf_kernel_scaling.reshape(-1,1,1),requires_grad=True)defforward(self,distances):""" Description ----------- Apply the layer to transform edge distances. Parameters ---------- distances : Float32 tensor of shape (E, 1) Distance between end nodes of edges. E for the number of edges. Returns ------- Float32 tensor of shape (K, E, 1) Transformed edge distances. K for the number of radial filters. """scaled_euclidean_distance=(-self.rbf_kernel_scaling*(distances-self.rbf_kernel_means)**2)# (K, E, 1)rbf_kernel_results=th.exp(scaled_euclidean_distance)# (K, E, 1)cos_values=0.5*(th.cos(np.pi*distances/self.interaction_cutoffs)+1)# (K, E, 1)cutoff_values=th.where(distances<=self.interaction_cutoffs,cos_values,th.zeros_like(cos_values),)# (K, E, 1)# Note that there appears to be an inconsistency between the paper and# DeepChem's implementation. In the paper, the scaled_euclidean_distance first# gets multiplied by cutoff_values, followed by exponentiation. Here we follow# the practice of DeepChem.returnrbf_kernel_results*cutoff_valuesdefmsg_func(edges):""" Description ----------- Send messages along edges. Parameters ---------- edges : EdgeBatch A batch of edges. Returns ------- dict mapping 'm' to Float32 tensor of shape (E, K * T) Messages computed. E for the number of edges, K for the number of radial filters and T for the number of features to use (types of atomic number in the paper). """return{"m":th.einsum("ij,ik->ijk",edges.src["hv"],edges.data["he"]).view(len(edges),-1)}defreduce_func(nodes):""" Description ----------- Collect messages and update node representations. Parameters ---------- nodes : NodeBatch A batch of nodes. Returns ------- dict mapping 'hv_new' to Float32 tensor of shape (V, K * T) Updated node representations. V for the number of nodes, K for the number of radial filters and T for the number of features to use (types of atomic number in the paper). """return{"hv_new":nodes.mailbox["m"].sum(1)}
[docs]classAtomicConv(nn.Module):r"""Atomic Convolution Layer from `Atomic Convolutional Networks for Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__ Denoting the type of atom :math:`i` by :math:`z_i` and the distance between atom :math:`i` and :math:`j` by :math:`r_{ij}`. **Distance Transformation** An atomic convolution layer first transforms distances with radial filters and then perform a pooling operation. For radial filter indexed by :math:`k`, it projects edge distances with .. math:: h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2) If :math:`r_{ij} < c_k`, .. math:: f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1), else, .. math:: f_{ij}^{k} = 0. Finally, .. math:: e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k} **Aggregation** For each type :math:`t`, each atom collects distance information from all neighbor atoms of type :math:`t`: .. math:: p_{i, t}^{k} = \sum_{j\in N(i)} e_{ij}^{k} * 1(z_j == t) Then concatenate the results for all RBF kernels and atom types. Parameters ---------- interaction_cutoffs : float32 tensor of shape (K) :math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs and two atoms are considered as connected if the distance between them is smaller than the cutoffs. K for the number of radial filters. rbf_kernel_means : float32 tensor of shape (K) :math:`r_k` in the equations above. K for the number of radial filters. rbf_kernel_scaling : float32 tensor of shape (K) :math:`\gamma_k` in the equations above. K for the number of radial filters. features_to_use : None or float tensor of shape (T) In the original paper, these are atomic numbers to consider, representing the types of atoms. T for the number of types of atomic numbers. Default to None. Note ---- * This convolution operation is designed for molecular graphs in Chemistry, but it might be possible to extend it to more general graphs. * There seems to be an inconsistency about the definition of :math:`e_{ij}^{k}` in the paper and the author's implementation. We follow the author's implementation. In the paper, :math:`e_{ij}^{k}` was defined as :math:`\exp(-\gamma_{k}|r_{ij}-r_{k}|^2 * f_{ij}^{k})`. * :math:`\gamma_{k}`, :math:`r_k` and :math:`c_k` are all learnable. Example ------- >>> import dgl >>> import numpy as np >>> import torch as th >>> from dgl.nn import AtomicConv >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])) >>> feat = th.ones(6, 1) >>> edist = th.ones(6, 1) >>> interaction_cutoffs = th.ones(3).float() * 2 >>> rbf_kernel_means = th.ones(3).float() >>> rbf_kernel_scaling = th.ones(3).float() >>> conv = AtomicConv(interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling) >>> res = conv(g, feat, edist) >>> res tensor([[0.5000, 0.5000, 0.5000], [0.5000, 0.5000, 0.5000], [0.5000, 0.5000, 0.5000], [1.0000, 1.0000, 1.0000], [0.5000, 0.5000, 0.5000], [0.0000, 0.0000, 0.0000]], grad_fn=<ViewBackward>) """def__init__(self,interaction_cutoffs,rbf_kernel_means,rbf_kernel_scaling,features_to_use=None,):super(AtomicConv,self).__init__()self.radial_pooling=RadialPooling(interaction_cutoffs=interaction_cutoffs,rbf_kernel_means=rbf_kernel_means,rbf_kernel_scaling=rbf_kernel_scaling,)iffeatures_to_useisNone:self.num_channels=1self.features_to_use=Noneelse:self.num_channels=len(features_to_use)self.features_to_use=nn.Parameter(features_to_use,requires_grad=False)
[docs]defforward(self,graph,feat,distances):""" Description ----------- Apply the atomic convolution layer. Parameters ---------- graph : DGLGraph Topology based on which message passing is performed. feat : Float32 tensor of shape :math:`(V, 1)` Initial node features, which are atomic numbers in the paper. :math:`V` for the number of nodes. distances : Float32 tensor of shape :math:`(E, 1)` Distance between end nodes of edges. E for the number of edges. Returns ------- Float32 tensor of shape :math:`(V, K * T)` Updated node representations. :math:`V` for the number of nodes, :math:`K` for the number of radial filters, and :math:`T` for the number of types of atomic numbers. """withgraph.local_scope():radial_pooled_values=self.radial_pooling(distances).to(feat)# (K, E, 1)ifself.features_to_useisnotNone:feat=(feat==self.features_to_use).to(feat)# (V, T)graph.ndata["hv"]=featgraph.edata["he"]=radial_pooled_values.transpose(1,0).squeeze(-1)# (E, K)graph.update_all(msg_func,reduce_func)returngraph.ndata["hv_new"].view(graph.num_nodes(),-1)# (V, K * T)