github.com/coyove/nj@v0.0.0-20221110084952-c7f8db1065c3/tests/bench/n-body.lua (about)

     1  -- The Computer Language Benchmarks Game
     2  -- https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
     3  -- contributed by Mike Pall
     4  -- modified by Geoff Leyland
     5  -- modified by Mario Pernici
     6  
     7  sun = {}
     8  jupiter = {}
     9  saturn = {}
    10  uranus = {}
    11  neptune = {}
    12  
    13  local sqrt = math.sqrt
    14  
    15  local PI = 3.141592653589793
    16  local SOLAR_MASS = 4 * PI * PI
    17  local DAYS_PER_YEAR = 365.24
    18  sun.x = 0.0
    19  sun.y = 0.0
    20  sun.z = 0.0
    21  sun.vx = 0.0
    22  sun.vy = 0.0
    23  sun.vz = 0.0
    24  sun.mass = SOLAR_MASS
    25  jupiter.x = 4.84143144246472090e+00
    26  jupiter.y = -1.16032004402742839e+00
    27  jupiter.z = -1.03622044471123109e-01
    28  jupiter.vx = 1.66007664274403694e-03 * DAYS_PER_YEAR
    29  jupiter.vy = 7.69901118419740425e-03 * DAYS_PER_YEAR
    30  jupiter.vz = -6.90460016972063023e-05 * DAYS_PER_YEAR
    31  jupiter.mass = 9.54791938424326609e-04 * SOLAR_MASS
    32  saturn.x = 8.34336671824457987e+00
    33  saturn.y = 4.12479856412430479e+00
    34  saturn.z = -4.03523417114321381e-01
    35  saturn.vx = -2.76742510726862411e-03 * DAYS_PER_YEAR
    36  saturn.vy = 4.99852801234917238e-03 * DAYS_PER_YEAR
    37  saturn.vz = 2.30417297573763929e-05 * DAYS_PER_YEAR
    38  saturn.mass = 2.85885980666130812e-04 * SOLAR_MASS
    39  uranus.x = 1.28943695621391310e+01
    40  uranus.y = -1.51111514016986312e+01
    41  uranus.z = -2.23307578892655734e-01
    42  uranus.vx = 2.96460137564761618e-03 * DAYS_PER_YEAR
    43  uranus.vy = 2.37847173959480950e-03 * DAYS_PER_YEAR
    44  uranus.vz = -2.96589568540237556e-05 * DAYS_PER_YEAR
    45  uranus.mass = 4.36624404335156298e-05 * SOLAR_MASS
    46  neptune.x = 1.53796971148509165e+01
    47  neptune.y = -2.59193146099879641e+01
    48  neptune.z = 1.79258772950371181e-01
    49  neptune.vx = 2.68067772490389322e-03 * DAYS_PER_YEAR
    50  neptune.vy = 1.62824170038242295e-03 * DAYS_PER_YEAR
    51  neptune.vz = -9.51592254519715870e-05 * DAYS_PER_YEAR
    52  neptune.mass = 5.15138902046611451e-05 * SOLAR_MASS
    53  
    54  local bodies = {sun,jupiter,saturn,uranus,neptune}
    55  
    56  local function advance(bodies, nbody, dt)
    57    for i=1,nbody do
    58      local bi = bodies[i]
    59      local bix, biy, biz, bimass = bi.x, bi.y, bi.z, bi.mass
    60      local bivx, bivy, bivz = bi.vx, bi.vy, bi.vz
    61      for j=i+1,nbody do
    62        local bj = bodies[j]
    63        local dx, dy, dz = bix-bj.x, biy-bj.y, biz-bj.z
    64        local dist2 = dx*dx + dy*dy + dz*dz
    65        local mag = sqrt(dist2)
    66        mag = dt / (mag * dist2)
    67        local bm = bj.mass*mag
    68        bivx = bivx - (dx * bm)
    69        bivy = bivy - (dy * bm)
    70        bivz = bivz - (dz * bm)
    71        bm = bimass*mag
    72        bj.vx = bj.vx + (dx * bm)
    73        bj.vy = bj.vy + (dy * bm)
    74        bj.vz = bj.vz + (dz * bm)
    75      end
    76      bi.vx = bivx
    77      bi.vy = bivy
    78      bi.vz = bivz
    79      bi.x = bix + dt * bivx
    80      bi.y = biy + dt * bivy
    81      bi.z = biz + dt * bivz
    82    end
    83  end
    84  
    85  local function energy(bodies, nbody)
    86    local e = 0
    87    for i=1,nbody do
    88      local bi = bodies[i]
    89      local vx, vy, vz, bim = bi.vx, bi.vy, bi.vz, bi.mass
    90      e = e + (0.5 * bim * (vx*vx + vy*vy + vz*vz))
    91      for j=i+1,nbody do
    92        local bj = bodies[j]
    93        local dx, dy, dz = bi.x-bj.x, bi.y-bj.y, bi.z-bj.z
    94        local distance = sqrt(dx*dx + dy*dy + dz*dz)
    95        e = e - ((bim * bj.mass) / distance)
    96      end
    97    end
    98    return e
    99  end
   100  
   101  local function offsetMomentum(b, nbody)
   102    local px, py, pz = 0, 0, 0
   103    for i=1,nbody do
   104      local bi = b[i]
   105      local bim = bi.mass
   106      px = px + (bi.vx * bim)
   107      py = py + (bi.vy * bim)
   108      pz = pz + (bi.vz * bim)
   109    end
   110    b[1].vx = -px / SOLAR_MASS
   111    b[1].vy = -py / SOLAR_MASS
   112    b[1].vz = -pz / SOLAR_MASS
   113  end
   114  
   115  local N = tonumber(arg and arg[1]) or 1000
   116  local nbody = #bodies
   117  
   118  offsetMomentum(bodies, nbody)
   119  io.write( string.format("%0.9f",energy(bodies, nbody)), "\n")
   120  for i=1,N do advance(bodies, nbody, 0.01) end
   121  io.write( string.format("%0.9f",energy(bodies, nbody)), "\n")