Added nbodies example

resources/examples/nbodies.c3

@@ -0,0 +1,161 @@
+module nbodies;
+
+const PI = 3.141592653589793;
+const SOLAR_MASS = 4 * PI * PI;
+const DAYS_PER_YEAR = 365.24;
+const DT = 1e-2;
+const RECIP_DT = (1.0/DT);
+
+struct Planet
+{
+  double x, y, z;
+  double vx, vy, vz;
+  double mass;
+}
+
+func void advance(Planet[] bodies) @noinline
+{
+    usize nbodies = bodies.len();
+    foreach (i, Planet* &b : bodies)
+    {
+        for (usize j = i + 1; j < nbodies; j++)
+        {
+            Planet* b2 = &bodies[j];
+            double dx = b.x - b2.x;
+            double dy = b.y - b2.y;
+            double dz = b.z - b2.z;
+            double inv_distance = 1.0/sqrt(dx * dx + dy * dy + dz * dz);
+            double mag = inv_distance * inv_distance * inv_distance;
+            b.vx -= dx * b2.mass * mag;
+            b.vy -= dy * b2.mass * mag;
+            b.vz -= dz * b2.mass * mag;
+            b2.vx += dx * b.mass * mag;
+            b2.vy += dy * b.mass * mag;
+            b2.vz += dz * b.mass * mag;
+        }
+    }
+    foreach (&b : bodies)
+    {
+        b.x += b.vx;
+        b.y += b.vy;
+        b.z += b.vz;
+    }
+}
+
+func double energy(Planet[] bodies)
+{
+    double e;
+    usize nbodies = bodies.len();
+
+    foreach (i, Planet* &b : bodies)
+    {
+        e += 0.5 * b.mass * (b.vx * b.vx + b.vy * b.vy + b.vz * b.vz);
+        for (usize j = i + 1; j < nbodies; j++)
+        {
+            Planet* b2 = &bodies[j];
+            double dx = b.x - b2.x;
+            double dy = b.y - b2.y;
+            double dz = b.z - b2.z;
+            double distance = sqrt(dx * dx + dy * dy + dz * dz);
+            e -= (b.mass * b2.mass) / distance;
+        }
+    }
+    return e;
+}
+
+func void offset_momentum(Planet[] bodies)
+{
+    double px;
+    double py;
+    double pz;
+    foreach (&b : bodies)
+    {
+        px += b.vx * b.mass;
+        py += b.vy * b.mass;
+        pz += b.vz * b.mass;
+    }
+    bodies[0].vx = - px / SOLAR_MASS;
+    bodies[0].vy = - py / SOLAR_MASS;
+    bodies[0].vz = - pz / SOLAR_MASS;
+}
+
+Planet[*] planet_bodies = {
+  {                               /* sun */
+    0, 0, 0, 0, 0, 0, SOLAR_MASS
+  },
+  {                               /* jupiter */
+    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
+  },
+  {                               /* saturn */
+    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
+  },
+  {                               /* uranus */
+    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
+  },
+  {                               /* neptune */
+    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
+  }
+};
+
+
+/*
+ * Rescale certain properties of bodies. That allows doing
+ * consequential advance()'s as if dt were equal to 1.0.
+ *
+ * When all advances done, rescale bodies back to obtain correct energy.
+ */
+func void scale_bodies(Planet[] bodies, double scale)
+{
+    foreach (&b : bodies)
+    {
+        b.mass *= scale * scale;
+        b.vx *= scale;
+        b.vy *= scale;
+        b.vz *= scale;
+    }
+}
+
+extern func int atoi(char *s);
+extern func int printf(char *s, ...);
+extern func double sqrt(double);
+
+func int main(int argc, char ** argv)
+{
+    int n = atoi(argv[1]);
+
+    Planet[] bodies = &planet_bodies;
+    offset_momentum(bodies);
+    printf ("%.9f\n", energy(bodies));
+    scale_bodies(bodies, DT);
+    for (int i = 1; i <= n; i++)
+    {
+        advance(bodies);
+    }
+    scale_bodies(bodies, RECIP_DT);
+    printf ("%.9f\n", energy(bodies));
+    return 0;
+}