nbody.hk

//
// nbody.hk
//

import { PI } from numbers;
import { sqrt } from math;

let SOLAR_MASS = 4 * PI * PI;
let DAYS_PER_YEAR = 365.24;

struct Body {
  x,
  y,
  z,
  vx,
  vy,
  vz,
  mass
}

fn sun() => Body { 0, 0, 0, 0, 0, 0, SOLAR_MASS };

fn jupiter() => Body {
  4.84143144246472090e+00,
  -1.16032004402742839e+00,
  -1.03622044471123109e-01,
  1.66007664274403694e-03 * DAYS_PER_YEAR,
  7.69901118419740425e-03 * DAYS_PER_YEAR,
  -6.90460016972063023e-05 * DAYS_PER_YEAR,
  9.54791938424326609e-04 * SOLAR_MASS
};

fn saturn() => Body {
  8.34336671824457987e+00,
  4.12479856412430479e+00,
  -4.03523417114321381e-01,
  -2.76742510726862411e-03 * DAYS_PER_YEAR,
  4.99852801234917238e-03 * DAYS_PER_YEAR,
  2.30417297573763929e-05 * DAYS_PER_YEAR,
  2.85885980666130812e-04 * SOLAR_MASS
};

fn uranus() => Body {
  1.28943695621391310e+01,
  -1.51111514016986312e+01,
  -2.23307578892655734e-01,
  2.96460137564761618e-03 * DAYS_PER_YEAR,
  2.37847173959480950e-03 * DAYS_PER_YEAR,
  -2.96589568540237556e-05 * DAYS_PER_YEAR,
  4.36624404335156298e-05 * SOLAR_MASS
};

fn neptune() => Body {
  1.53796971148509165e+01,
  -2.59193146099879641e+01,
  1.79258772950371181e-01,
  2.68067772490389322e-03 * DAYS_PER_YEAR,
  1.62824170038242295e-03 * DAYS_PER_YEAR,
  -9.51592254519715870e-05 * DAYS_PER_YEAR,
  5.15138902046611451e-05 * SOLAR_MASS
};

fn offset_momentum(bodies) {
  var px = 0;
  var py = 0;
  var pz = 0;
  let size = len(bodies);
  var i = 0;
  while (i < size) {
    let body = bodies[i];
    let mass = body.mass;
    px += body.vx * mass;
    py += body.vy * mass;
    pz += body.vz * mass;
    i++;
  }
  bodies[0].vx = -px / SOLAR_MASS;
  bodies[0].vy = -py / SOLAR_MASS;
  bodies[0].vz = -pz / SOLAR_MASS;
  return bodies;
}

fn advance(bodies, dt) {
  let size = len(bodies);
  var i = 0;
  while (i < size) {
    var bodyi = bodies[i];
    var vxi = bodyi.vx;
    var vyi = bodyi.vy;
    var vzi = bodyi.vz;
    var j = i + 1;
    while (j < size) {
      var bodyj = bodies[j];
      let dx = bodyi.x - bodyj.x;
      let dy = bodyi.y - bodyj.y;
      let dz = bodyi.z - bodyj.z;
      let d2 = dx * dx + dy * dy + dz * dz;
      let mag = dt / (d2 * sqrt(d2));
      let massj = bodyj.mass;
      vxi -= dx * massj * mag;
      vyi -= dy * massj * mag;
      vzi -= dz * massj * mag;
      let massi = bodyi.mass;
      bodyj.vx += dx * massi * mag;
      bodyj.vy += dy * massi * mag;
      bodyj.vz += dz * massi * mag;
      bodies[j] = bodyj;
      j++;
    }
    bodyi.vx = vxi;
    bodyi.vy = vyi;
    bodyi.vz = vzi;
    bodies[i] = bodyi;
    i++;
  }
  i = 0;
  while (i < size) {
    bodies[i].x += dt * bodies[i].vx;
    bodies[i].y += dt * bodies[i].vy;
    bodies[i].z += dt * bodies[i].vz;
    i++;
  }
  return bodies;
}

fn energy(bodies) {
  var e = 0;
  let size = len(bodies);
  var i = 0;
  while (i < size) {
    let bodyi = bodies[i];
    e += 0.5 * bodyi.mass * ( bodyi.vx * bodyi.vx + 
        bodyi.vy * bodyi.vy + bodyi.vz * bodyi.vz );
    var j = i + 1;
    while (j < size) {
      let bodyj = bodies[j];
      let dx = bodyi.x - bodyj.x;
      let dy = bodyi.y - bodyj.y;
      let dz = bodyi.z - bodyj.z;
      let distance = sqrt(dx * dx + dy * dy + dz * dz);
      e -= (bodyi.mass * bodyj.mass) / distance;
      j++;
    }
    i++;
  }
  return e;
}

let n = to_int(args[0]);
var bodies = [ sun(), jupiter(), saturn(), uranus(), neptune() ];
bodies = offset_momentum(bodies);
println(energy(bodies));
var i = 0;
while (i < n) {
  bodies = advance(bodies, 0.01);
  i++;
}
println(energy(bodies));