class_docstr: 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>)
class_docstr_tok: ["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",">",")"]
docstr: 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.
docstr_tok: ["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","."]
code: def forward(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.
      """
      with graph.local_scope():
          radial_pooled_values = self.radial_pooling(distances).to(
              feat
          )  # (K, E, 1)
          if self.features_to_use is not None:
              feat = (feat == self.features_to_use).to(feat)  # (V, T)
          graph.ndata["hv"] = feat
          graph.edata["he"] = radial_pooled_values.transpose(
              1, 0
          ).squeeze(
              -1
          )  # (E, K)
          graph.update_all(msg_func, reduce_func)
  
          return graph.ndata["hv_new"].view(graph.num_nodes(), -1)
  
file_path: dgl/python/dgl/nn/pytorch/conv/atomicconv.py
imports: import numpy
  import torch
  import torch.nn
label_desc: Use image node_id 2 for calling the AtomicConv obj's underlying member method code with example usage: obj.forward(graph, feat, distances) and returns: graph