class_docstr:  The ProNet from the "Learning Protein Representations via Complete 3D Graph Networks" paper.
  
  Args:
      level: (str, optional): The level of protein representations. It could be :obj:`aminoacid`, obj:`backbone`, and :obj:`allatom`. (default: :obj:`aminoacid`)
      num_blocks (int, optional): Number of building blocks. (default: :obj:`4`)
      hidden_channels (int, optional): Hidden embedding size. (default: :obj:`128`)
      out_channels (int, optional): Size of each output sample. (default: :obj:`1`)
      mid_emb (int, optional): Embedding size used for geometric features. (default: :obj:`64`)
      num_radial (int, optional): Number of radial basis functions. (default: :obj:`6`)
      num_spherical (int, optional): Number of spherical harmonics. (default: :obj:`2`)
      cutoff (float, optional): Cutoff distance for interatomic interactions. (default: :obj:`10.0`)
      max_num_neighbors (int, optional): Max number of neighbors during graph construction. (default: :obj:`32`)
      int_emb_layers (int, optional): Number of embedding layers in the interaction block. (default: :obj:`3`)
      out_layers (int, optional): Number of layers for features after interaction blocks. (default: :obj:`2`)
      num_pos_emb (int, optional): Number of positional embeddings. (default: :obj:`16`)
      dropout (float, optional): Dropout. (default: :obj:`0`)
      data_augment_eachlayer (bool, optional): Data augmentation tricks. If set to :obj:`True`, will add noise to the node features before each interaction block. (default: :obj:`False`)
      euler_noise (bool, optional): Data augmentation tricks. If set to :obj:`True`, will add noise to Euler angles. (default: :obj:`False`)
      
class_docstr_tok: ["The","ProNet","from","the","``","Learning","Protein","Representations","via","Complete","3D","Graph","Networks","''","paper",".","Args",":","level",":","(","str",",","optional",")",":","The","level","of","protein","representations",".","It","could","be",":","obj",":","`","aminoacid","`",",","obj",":","`","backbone","`",",","and",":","obj",":","`","allatom","`",".","(","default",":",":","obj",":","`","aminoacid","`",")","num_blocks","(","int",",","optional",")",":","Number","of","building","blocks",".","(","default",":",":","obj",":","`","4","`",")","hidden_channels","(","int",",","optional",")",":","Hidden","embedding","size",".","(","default",":",":","obj",":","`","128","`",")","out_channels","(","int",",","optional",")",":","Size","of","each","output","sample",".","(","default",":",":","obj",":","`","1","`",")","mid_emb","(","int",",","optional",")",":","Embedding","size","used","for","geometric","features",".","(","default",":",":","obj",":","`","64","`",")","num_radial","(","int",",","optional",")",":","Number","of","radial","basis","functions",".","(","default",":",":","obj",":","`","6","`",")","num_spherical","(","int",",","optional",")",":","Number","of","spherical","harmonics",".","(","default",":",":","obj",":","`","2","`",")","cutoff","(","float",",","optional",")",":","Cutoff","distance","for","interatomic","interactions",".","(","default",":",":","obj",":","`","10.0","`",")","max_num_neighbors","(","int",",","optional",")",":","Max","number","of","neighbors","during","graph","construction",".","(","default",":",":","obj",":","`","32","`",")","int_emb_layers","(","int",",","optional",")",":","Number","of","embedding","layers","in","the","interaction","block",".","(","default",":",":","obj",":","`","3","`",")","out_layers","(","int",",","optional",")",":","Number","of","layers","for","features","after","interaction","blocks",".","(","default",":",":","obj",":","`","2","`",")","num_pos_emb","(","int",",","optional",")",":","Number","of","positional","embeddings",".","(","default",":",":","obj",":","`","16","`",")","dropout","(","float",",","optional",")",":","Dropout",".","(","default",":",":","obj",":","`","0","`",")","data_augment_eachlayer","(","bool",",","optional",")",":","Data","augmentation","tricks",".","If","set","to",":","obj",":","`","True","`",",","will","add","noise","to","the","node","features","before","each","interaction","block",".","(","default",":",":","obj",":","`","False","`",")","euler_noise","(","bool",",","optional",")",":","Data","augmentation","tricks",".","If","set","to",":","obj",":","`","True","`",",","will","add","noise","to","Euler","angles",".","(","default",":",":","obj",":","`","False","`",")"]
code: def forward(self, batch_data):
      z, pos, batch = (
          torch.squeeze(batch_data.x.long()),
          batch_data.coords_ca,
          batch_data.batch,
      )
      pos_n = batch_data.coords_n
      pos_c = batch_data.coords_c
      bb_embs = batch_data.bb_embs
      side_chain_embs = batch_data.side_chain_embs
  
      device = z.device
  
      if self.level == "aminoacid":
          x = self.embedding(z)
      elif self.level == "backbone":
          x = torch.cat(
              [
                  torch.squeeze(
                      F.one_hot(z, num_classes=num_aa_type).float()
                  ),
                  bb_embs,
              ],
              dim=1,
          )
          x = self.embedding(x)
      elif self.level == "allatom":
          x = torch.cat(
              [
                  torch.squeeze(
                      F.one_hot(z, num_classes=num_aa_type).float()
                  ),
                  bb_embs,
                  side_chain_embs,
              ],
              dim=1,
          )
          x = self.embedding(x)
      else:
          print("No supported model!")
  
      edge_index = radius_graph(
          pos,
          r=self.cutoff,
          batch=batch,
          max_num_neighbors=self.max_num_neighbors,
      )
      pos_emb = self.pos_emb(edge_index, self.num_pos_emb)
      j, i = edge_index
  
      # Calculate distances.
      dist = (pos[i] - pos[j]).norm(dim=1)
  
      num_nodes = len(z)
  
      # Calculate angles theta and phi.
      refi0 = (i - 1) % num_nodes
      refi1 = (i + 1) % num_nodes
  
      a = ((pos[j] - pos[i]) * (pos[refi0] - pos[i])).sum(dim=-1)
      b = torch.cross(pos[j] - pos[i], pos[refi0] - pos[i]).norm(dim=-1)
      theta = torch.atan2(b, a)
  
      plane1 = torch.cross(pos[refi0] - pos[i], pos[refi1] - pos[i])
      plane2 = torch.cross(pos[refi0] - pos[i], pos[j] - pos[i])
      a = (plane1 * plane2).sum(dim=-1)
      b = (torch.cross(plane1, plane2) * (pos[refi0] - pos[i])).sum(
          dim=-1
      ) / ((pos[refi0] - pos[i]).norm(dim=-1))
      phi = torch.atan2(b, a)
  
      feature0 = self.feature0(dist, theta, phi)
  
      if self.level == "backbone" or self.level == "allatom":
          # Calculate Euler angles.
          Or1_x = pos_n[i] - pos[i]
          Or1_z = torch.cross(
              Or1_x, torch.cross(Or1_x, pos_c[i] - pos[i])
          )
          Or1_z_length = Or1_z.norm(dim=1) + 1e-7
  
          Or2_x = pos_n[j] - pos[j]
          Or2_z = torch.cross(
              Or2_x, torch.cross(Or2_x, pos_c[j] - pos[j])
          )
          Or2_z_length = Or2_z.norm(dim=1) + 1e-7
  
          Or1_Or2_N = torch.cross(Or1_z, Or2_z)
  
          angle1 = torch.atan2(
              (torch.cross(Or1_x, Or1_Or2_N) * Or1_z).sum(dim=-1)
              / Or1_z_length,
              (Or1_x * Or1_Or2_N).sum(dim=-1),
          )
          angle2 = torch.atan2(
              torch.cross(Or1_z, Or2_z).norm(dim=-1),
              (Or1_z * Or2_z).sum(dim=-1),
          )
          angle3 = torch.atan2(
              (torch.cross(Or1_Or2_N, Or2_x) * Or2_z).sum(dim=-1)
              / Or2_z_length,
              (Or1_Or2_N * Or2_x).sum(dim=-1),
          )
  
          if self.euler_noise:
              euler_noise = torch.clip(
                  torch.empty(3, len(angle1))
                  .to(device)
                  .normal_(mean=0.0, std=0.025),
                  min=-0.1,
                  max=0.1,
              )
              angle1 += euler_noise[0]
              angle2 += euler_noise[1]
              angle3 += euler_noise[2]
  
          feature1 = torch.cat(
              (
                  self.feature1(dist, angle1),
                  self.feature1(dist, angle2),
                  self.feature1(dist, angle3),
              ),
              1,
          )
  
      elif self.level == "aminoacid":
          refi = (i - 1) % num_nodes
  
          refj0 = (j - 1) % num_nodes
          refj = (j - 1) % num_nodes
          refj1 = (j + 1) % num_nodes
  
          mask = refi0 == j
          refi[mask] = refi1[mask]
          mask = refj0 == i
          refj[mask] = refj1[mask]
  
          plane1 = torch.cross(pos[j] - pos[i], pos[refi] - pos[i])
          plane2 = torch.cross(pos[j] - pos[i], pos[refj] - pos[j])
          a = (plane1 * plane2).sum(dim=-1)
          b = (torch.cross(plane1, plane2) * (pos[j] - pos[i])).sum(
              dim=-1
          ) / dist
          tau = torch.atan2(b, a)
  
          feature1 = self.feature1(dist, tau)
  
      # Interaction blocks.
      for interaction_block in self.interaction_blocks:
          if self.data_augment_eachlayer:
              # add gaussian noise to features
              gaussian_noise = torch.clip(
                  torch.empty(x.shape)
                  .to(device)
                  .normal_(mean=0.0, std=0.025),
                  min=-0.1,
                  max=0.1,
              )
              x += gaussian_noise
          x = interaction_block(
              x, feature0, feature1, pos_emb, edge_index, batch
          )
  
      y = scatter(x, batch, dim=0)
  
      for lin in self.lins_out:
          y = self.relu(lin(y))
          y = self.dropout(y)
      y = self.lin_out(y)
      return y
  
imports: from torch_geometric.nn import inits, MessagePassing
  from torch_geometric.nn import radius_graph
  from .features import d_angle_emb, d_theta_phi_emb
  from torch_scatter import scatter
  from torch_sparse import matmul
  import torch
  from torch import nn
  from torch.nn import Embedding
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 4 for calling the ProNet obj's underlying member method code with example usage: obj.forward(batch_data) and returns: y