class_docstr:  The ComENet from the `"ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs" <https://arxiv.org/abs/2206.08515>`_ paper.
  
  Args:
      cutoff (float, optional): Cutoff distance for interatomic interactions. (default: :obj:`8.0`)
      num_layers (int, optional): Number of building blocks. (default: :obj:`4`)
      hidden_channels (int, optional): Hidden embedding size. (default: :obj:`256`)
      middle_channels (int, optional): Middle embedding size for the two layer linear block. (default: :obj:`256`)
      out_channels (int, optional): Size of each output sample. (default: :obj:`1`)
      num_radial (int, optional): Number of radial basis functions. (default: :obj:`3`)
      num_spherical (int, optional): Number of spherical harmonics. (default: :obj:`2`)
      num_output_layers (int, optional): Number of linear layers for the output blocks. (default: :obj:`3`)
code: def _forward(self, data):
      batch = data.batch
      z = data.z.long()
      pos = data.pos
      num_nodes = z.size(0)
  
      edge_index = radius_graph(pos, r=self.cutoff, batch=batch)
      j, i = edge_index
  
      vecs = pos[j] - pos[i]
      dist = vecs.norm(dim=-1)
  
      # Embedding block.
      x = self.emb(z)
  
      # Calculate distances.
      _, argmin0 = scatter_min(dist, i, dim_size=num_nodes)
      argmin0[argmin0 >= len(i)] = 0
      n0 = j[argmin0]
      add = torch.zeros_like(dist).to(dist.device)
      add[argmin0] = self.cutoff
      dist1 = dist + add
  
      _, argmin1 = scatter_min(dist1, i, dim_size=num_nodes)
      argmin1[argmin1 >= len(i)] = 0
      n1 = j[argmin1]
      # --------------------------------------------------------
  
      _, argmin0_j = scatter_min(dist, j, dim_size=num_nodes)
      argmin0_j[argmin0_j >= len(j)] = 0
      n0_j = i[argmin0_j]
  
      add_j = torch.zeros_like(dist).to(dist.device)
      add_j[argmin0_j] = self.cutoff
      dist1_j = dist + add_j
  
      # i[argmin] = range(0, num_nodes)
      _, argmin1_j = scatter_min(dist1_j, j, dim_size=num_nodes)
      argmin1_j[argmin1_j >= len(j)] = 0
      n1_j = i[argmin1_j]
  
      # ----------------------------------------------------------
  
      # n0, n1 for i
      n0 = n0[i]
      n1 = n1[i]
  
      # n0, n1 for j
      n0_j = n0_j[j]
      n1_j = n1_j[j]
  
      # tau: (iref, i, j, jref)
      # when compute tau, do not use n0, n0_j as ref for i and j,
      # because if n0 = j, or n0_j = i, the computed tau is zero
      # so if n0 = j, we choose iref = n1
      # if n0_j = i, we choose jref = n1_j
      mask_iref = n0 == j
      iref = torch.clone(n0)
      iref[mask_iref] = n1[mask_iref]
      idx_iref = argmin0[i]
      idx_iref[mask_iref] = argmin1[i][mask_iref]
  
      mask_jref = n0_j == i
      jref = torch.clone(n0_j)
      jref[mask_jref] = n1_j[mask_jref]
      idx_jref = argmin0_j[j]
      idx_jref[mask_jref] = argmin1_j[j][mask_jref]
  
      pos_ji, pos_in0, pos_in1, pos_iref, pos_jref_j = (
          vecs,
          vecs[argmin0][i],
          vecs[argmin1][i],
          vecs[idx_iref],
          vecs[idx_jref],
      )
  
      # Calculate angles.
      a = ((-pos_ji) * pos_in0).sum(dim=-1)
      b = torch.cross(-pos_ji, pos_in0).norm(dim=-1)
      theta = torch.atan2(b, a)
      theta[theta < 0] = theta[theta < 0] + math.pi
  
      # Calculate torsions.
      dist_ji = pos_ji.pow(2).sum(dim=-1).sqrt()
      plane1 = torch.cross(-pos_ji, pos_in0)
      plane2 = torch.cross(-pos_ji, pos_in1)
      a = (plane1 * plane2).sum(
          dim=-1
      )  # cos_angle * |plane1| * |plane2|
      b = (torch.cross(plane1, plane2) * pos_ji).sum(dim=-1) / dist_ji
      phi = torch.atan2(b, a)
      phi[phi < 0] = phi[phi < 0] + math.pi
  
      # Calculate right torsions.
      plane1 = torch.cross(pos_ji, pos_jref_j)
      plane2 = torch.cross(pos_ji, pos_iref)
      a = (plane1 * plane2).sum(
          dim=-1
      )  # cos_angle * |plane1| * |plane2|
      b = (torch.cross(plane1, plane2) * pos_ji).sum(dim=-1) / dist_ji
      tau = torch.atan2(b, a)
      tau[tau < 0] = tau[tau < 0] + math.pi
  
      feature1 = self.feature1(dist, theta, phi)
      feature2 = self.feature2(dist, tau)
  
      # Interaction blocks.
      for interaction_block in self.interaction_blocks:
          x = interaction_block(
              x, feature1, feature2, edge_index, batch
          )
  
      for lin in self.lins:
          x = self.act(lin(x))
      x = self.lin_out(x)
  
      energy = scatter(x, batch, dim=0)
      return energy
  
imports: from torch_cluster import radius_graph
  from torch_geometric.nn import GraphConv, GraphNorm
  from torch_geometric.nn import inits
  from .features import angle_emb, torsion_emb
  from torch_scatter import scatter, scatter_min
  from torch.nn import Embedding
  import torch
  from torch import nn
  from torch import Tensor
  import torch.nn.functional
  import math
  from math import sqrt
label_desc: Use image node_id 3 for calling the ComENet obj's underlying member method code with example usage: obj._forward(data) and returns: energy