Implementation of the N-Body Problem.

nbody_problem/cpp/main.cpp

@@ -0,0 +1,170 @@
+#include <iostream>
+#include <math.h>
+#include <numeric>
+#include <vector>
+
+typedef double real;
+using namespace std;
+
+class Star {
+public:
+  real m;
+  vector<real> r;
+  vector<real> v;
+  vector<real> a, a0;
+
+  Star() {
+    r.assign(3, 0);
+    v.assign(3, 0);
+    a.assign(3, 0);
+    a0.assign(3, 0);
+  }
+
+  Star(real mass, vector<real> pos, vector<real> vel) {
+    m = mass;
+    r = pos;
+    v = vel;
+  }
+
+  friend ostream &operator<<(ostream &so, const Star &si) {
+    so << si.m << " " << si.r[0] << " " << si.r[1] << " " << si.r[2] << " "
+       << si.v[0] << " " << si.v[1] << " " << si.v[2] << endl;
+    return so;
+  }
+};
+
+class Cluster : public Star {
+protected:
+public:
+  vector<Star> s;
+  Cluster() : Star() {}
+
+  void acceleration() {
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si)
+      si->a.assign(3, 0);
+
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si) {
+      vector<real> rij(3);
+      real init = 0.0;
+
+      for (vector<Star>::iterator sj = s.begin(); sj != s.end(); ++sj) {
+        if (si != sj) {
+
+          for (int i = 0; i != 3; ++i)
+            rij[i] = si->r[i] - sj->r[i];
+
+          real RdotR = inner_product(rij.begin(), rij.end(), rij.begin(), init);
+          real apre = 1. / sqrt(RdotR * RdotR * RdotR);
+
+          for (int i = 0; i != 3; ++i) {
+            si->a[i] = sj->m * -1 * apre * rij[i];
+          }
+        }
+      }
+    }
+  }
+
+  void updatePositions(real dt) {
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si) {
+      si->a0 = si->a;
+      for (int i = 0; i != 3; ++i)
+        si->r[i] += dt * si->v[i] + 0.5 * dt * dt * si->a0[i];
+    }
+  }
+
+  void updateVelocities(real dt) {
+
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si) {
+      for (int i = 0; i != 3; ++i)
+        si->v[i] += 0.5 * dt * (si->a0[i] + si->a[i]);
+      si->a0 = si->a;
+    }
+  }
+
+  vector<real> energies() {
+    real init = 0;
+    vector<real> E(3), rij(3);
+    E.assign(3, 0);
+
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si)
+      E[1] += 0.5 * si->m *
+              inner_product(si->v.begin(), si->v.end(), si->v.begin(), init);
+
+    // Potential energy
+    for (vector<Star>::iterator si = s.begin(); si != s.end(); ++si) {
+      for (vector<Star>::iterator sj = si + 1; sj != s.end(); ++sj) {
+        for (int i = 0; i != 3; ++i)
+          rij[i] = si->r[i] - sj->r[i];
+        E[2] -= si->m * sj->m /
+                sqrt(inner_product(rij.begin(), rij.end(), rij.begin(), init));
+      }
+    }
+    E[0] = E[1] + E[2];
+    return E;
+  }
+
+  // Print function
+  friend ostream &operator<<(ostream &so, Cluster &cl) {
+    for (vector<Star>::iterator si = cl.s.begin(); si != cl.s.end(); ++si)
+      so << *si;
+    return so;
+  }
+};
+
+int main(int argc, char* argv[]) {
+
+  Cluster cl;
+  real m;
+  int dummy;
+  vector<real> r(3), v(3);
+
+  // Read input data from the command line (makeplummer | dumbp)
+  do {
+    cin >> dummy;
+    cin >> m;
+    for (int i = 0; i != 3; ++i)
+      cin >> r[i];
+    for (int i = 0; i != 3; ++i)
+      cin >> v[i];
+    cl.s.push_back(Star(m, r, v));
+  } while (!cin.eof());
+
+  cl.s.pop_back();
+
+  // Compute initial energu of the system
+  vector<real> E(3), E0(3);
+  E0 = cl.energies();
+
+  // Start time, end time and simulation step
+  real t = 0.0;
+  real tend;
+
+  if (argc > 1)
+    tend = strtod(argv[1], NULL);
+  else
+    tend = 10.0;
+
+  real dt = 1e-3;
+  int k = 0;
+
+  for (int i = 0; i < 50; i++) {
+    cl.acceleration();
+
+    while (t < tend) {
+      cl.updatePositions(dt);
+
+      cl.acceleration();
+
+      cl.updateVelocities(dt);
+
+      t += dt;
+      k += 1;
+      if (k % 100 == 0) {
+        E = cl.energies();
+      }
+    }
+    t = 0;
+  }
+
+  return 0;
+}

nbody_problem/python/compiled_methods.py

@@ -0,0 +1,108 @@
+import sys
+import math
+import timeit
+
+import numpy as np
+import pandas as pd
+
+from transonic import jit
+
+def load_input_data(path):
+
+    df = pd.read_csv(
+        path, names = ["mass", "x", "y", "z", "vx", "vy", "vz"], delimiter=r"\s+"
+    )
+
+    masses = df["mass"].values.copy()
+    positions = df.loc[:, ["x", "y", "z"]].values.copy()
+    velocities = df.loc[:, ["vx", "vy", "vz"]].values.copy()
+
+    return masses, positions, velocities
+
+def compute_accelerations(accelerations, masses, positions):
+    nb_particles = masses.size
+
+    for index_p0 in range(nb_particles - 1):
+        position0 = positions[index_p0]
+        mass0 = masses[index_p0]
+
+        # TODO: Use compiled methods like Numba & Pythran in vectorized approach.
+        # Issue: https://github.com/khushi-411/numpy-benchmarks/issues/4
+        for index_p1 in range(index_p0 + 1, nb_particles):
+            mass1 = masses[index_p1]
+
+            vectors = position0 - positions[index_p1]
+
+            distances = (vectors**2).sum()
+            coefs = 1./distances**1.5
+
+            accelerations[index_p0] += mass1 * -1 * vectors * coefs
+            accelerations[index_p1] += mass0 * vectors * coefs
+
+    return accelerations
+
+@jit
+def loop(
+        time_step, nb_steps, masses, positions,
+        velocities):
+
+    accelerations = np.zeros_like(positions)
+    accelerations1 = np.zeros_like(positions)
+
+    accelerations = compute_accelerations(accelerations, masses, positions)
+
+    time = 0.0
+    energy0, _, _ = compute_energies(masses, positions, velocities)
+    energy_previous = energy0
+
+    for step in range(nb_steps):
+        positions += time_step*velocities + 0.5*accelerations*time_step**2
+
+        accelerations, accelerations1 = accelerations1, accelerations
+        accelerations.fill(0)
+        accelerations = compute_accelerations(accelerations, masses, positions)
+
+        velocities += 0.5*time_step*(accelerations+accelerations1)
+
+        time += time_step
+
+        if not step % 100:
+            energy, _, _ = compute_energies(masses, positions, velocities)
+            energy_previous = energy
+
+    return energy, energy0
+
+def compute_energies(masses, positions, velocities):
+
+    ke = 0.5 * masses @ np.sum(velocities**2, axis=1)
+
+    nb_particles = masses.size
+    pe = 0.0
+    for index_p0 in range(nb_particles - 1):
+
+        mass0 = masses[index_p0]
+        for index_p1 in range(index_p0 + 1, nb_particles):
+
+            mass1 = masses[index_p1]
+            vector = positions[index_p0] - positions[index_p1]
+
+            distance = math.sqrt((vector**2).sum())
+
+            pe -= (mass0*mass1) / distance**2
+
+    return ke+pe, ke, pe
+
+if __name__ == "__main__":
+
+    try:
+        time_end = float(sys.argv[2])
+    except IndexError:
+        time_end = 10.0
+
+    time_step = 0.001
+    nb_steps = int(time_end/time_step) + 1
+
+    path_input = sys.argv[1]
+    masses, positions, velocities = load_input_data(path_input)
+
+    print('time taken:', timeit.timeit('loop(time_step, nb_steps, masses, positions, velocities)', globals=globals(), number=50))

nbody_problem/python/numpy_python.py

@@ -0,0 +1,105 @@
+import sys
+import math
+import timeit
+
+import numpy as np
+import pandas as pd
+
+def load_input_data(path):
+
+    df = pd.read_csv(
+        path, names = ["mass", "x", "y", "z", "vx", "vy", "vz"], delimiter=r"\s+"
+    )
+
+    masses = df["mass"].values.copy()
+    positions = df.loc[:, ["x", "y", "z"]].values.copy()
+    velocities = df.loc[:, ["vx", "vy", "vz"]].values.copy()
+
+    return masses, positions, velocities
+
+def compute_accelerations(accelerations, masses, positions):
+    nb_particles = masses.size
+
+    for index_p0 in range(nb_particles - 1):
+        position0 = positions[index_p0]
+        mass0 = masses[index_p0]
+
+        vectors = position0 - positions[index_p0 + 1: nb_particles]
+
+        distances = (vectors**2).sum(axis=1)
+        coefs = 1./distances**1.5
+
+        accelerations[index_p0] += np.sum((masses[index_p0 + 1: nb_particles] * -1 * vectors.T * coefs).T, axis=0)
+        accelerations[index_p0 + 1: nb_particles] += (mass0 * vectors.T * coefs).T
+
+    return accelerations
+
+def numpy_loop(
+        time_step: float,
+        nb_steps: int,
+        masses: "float[]",
+        positions: "float[:,:]",
+        velocities: "float[:,:]",
+    ):
+
+    accelerations = np.zeros_like(positions)
+    accelerations1 = np.zeros_like(positions)
+
+    accelerations = compute_accelerations(accelerations, masses, positions)
+
+    time = 0.0
+
+    energy0, _, _ = compute_energies(masses, positions, velocities)
+    energy_previous = energy0
+
+    for step in range(nb_steps):
+        positions += time_step*velocities + 0.5*accelerations*time_step**2
+
+        accelerations, accelerations1 = accelerations1, accelerations
+        accelerations.fill(0)
+        accelerations = compute_accelerations(accelerations, masses, positions)
+
+        velocities += 0.5*time_step*(accelerations+accelerations1)
+
+        time += time_step
+
+        if not step%100:
+            energy, _, _ = compute_energies(masses, positions, velocities)
+            energy_previous = energy
+
+    return energy, energy0
+
+def compute_energies(masses, positions, velocities):
+
+    ke = 0.5 * masses @ np.sum(velocities**2, axis=1)
+
+    nb_particles = masses.size
+    pe = 0.0
+    for index_p0 in range(nb_particles - 1):
+
+        mass0 = masses[index_p0]
+        for index_p1 in range(index_p0 + 1, nb_particles):
+
+            mass1 = masses[index_p1]
+            vector = positions[index_p0] - positions[index_p1]
+
+            distance = math.sqrt((vector**2).sum())
+
+            pe -= (mass0*mass1) / distance**2
+
+    return ke+pe, ke, pe
+
+if __name__ == "__main__":
+
+    try:
+        time_end = float(sys.argv[2])
+    except IndexError:
+        time_end = 10.0
+
+    time_step = 0.001
+    nb_steps = int(time_end/time_step) + 1
+
+    path_input = sys.argv[1]
+    masses, positions, velocities = load_input_data(path_input)
+
+    print('time taken:', timeit.timeit('numpy_loop(time_step, nb_steps, masses, positions, velocities)', globals=globals(), number=50))