code: def init(total_state, n_body, fea_num, orbit):
      data = np.zeros((total_state, n_body, fea_num), dtype=float)
      if orbit:
          data[0][0][0] = 100
          data[0][0][1:5] = 0.0
          # The position are initialized randomly.
          for i in range(1, n_body):
              data[0][i][0] = np.random.rand() * 8.98 + 0.02
              distance = np.random.rand() * 90.0 + 10.0
              theta = np.random.rand() * 360
              theta_rad = pi / 2 - radians(theta)
              data[0][i][1] = distance * cos(theta_rad)
              data[0][i][2] = distance * sin(theta_rad)
              data[0][i][3] = (
                  -1
                  * data[0][i][2]
                  / norm(data[0][i][1:3])
                  * (G * data[0][0][0] / norm(data[0][i][1:3]) ** 2)
                  * distance
                  / 1000
              )
              data[0][i][4] = (
                  data[0][i][1]
                  / norm(data[0][i][1:3])
                  * (G * data[0][0][0] / norm(data[0][i][1:3]) ** 2)
                  * distance
                  / 1000
              )
      else:
          for i in range(n_body):
              data[0][i][0] = np.random.rand() * 8.98 + 0.02
              distance = np.random.rand() * 90.0 + 10.0
              theta = np.random.rand() * 360
              theta_rad = pi / 2 - radians(theta)
              data[0][i][1] = distance * cos(theta_rad)
              data[0][i][2] = distance * sin(theta_rad)
              data[0][i][3] = np.random.rand() * 6.0 - 3.0
              data[0][i][4] = np.random.rand() * 6.0 - 3.0
      return data
  
file_path: dgl/examples/pytorch/graphsim/n_body_sim.py
imports: from __future__ import absolute_import, division, print_function
  import argparse
  import os
  from math import cos, pi, radians, sin
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: init(total_state, n_body, fea_num, orbit) and returns: data