Experiment with three particles

_code/notes/diffsim/main.cpp

@@ -16,6 +16,27 @@ Vector euler_step(F const& f, Scalar t, Vector const& x, Scalar delta_t) {
     return x + delta_t * f(t, x);
 }
 
+//--------------------------------------------------------------------------------------------------
+// A single step of fourth order Runge-Kutta
+template <typename F, typename Vector, typename Scalar>
+Vector rk4_step(F const& f, Scalar t, Vector const& x, Scalar delta_t) {
+    Scalar const half_delta_t = static_cast<Scalar>(0.5) * delta_t;
+    Vector result = x;
+    // k1
+    Vector tmp = delta_t * f(t, x);
+    result += 0.16666666666666 * tmp;
+    // k2
+    tmp = delta_t * f(half_delta_t, x + static_cast<Scalar>(0.5) * tmp);
+    result += 0.33333333333333 * tmp;
+    // k3
+    tmp = delta_t * f(half_delta_t, x + static_cast<Scalar>(0.5) * tmp);
+    result += 0.33333333333333 * tmp;
+    // k4
+    tmp = delta_t * f(t + delta_t, x + tmp);
+    result += 0.16666666666666 * tmp;
+    return result;
+}
+
 //--------------------------------------------------------------------------------------------------
 // state vectors (i.e. phase space points) are packed as follows:
 // 0, ..., d - 1 give (up to) x, y, z for mass zero
@@ -30,45 +51,48 @@ Vector euler_step(F const& f, Scalar t, Vector const& x, Scalar delta_t) {
 // Functor that computes phase space derivatives
 class NBodyLaw {
 public:
-    NBodyLaw();
-    void set_masses(VectorX<Scalar> const& m) { m_ = m; }
+    NBodyLaw() {}
+    void resize(uint32_t n);
+    void set_masses(VectorX<Scalar> const& m) { resize(m.size()); m_ = m; }
 
+    uint32_t n() const { return n_; }
     // returns derivative of state for a give state
     VectorX<Scalar> const& operator()(VectorX<Scalar> const& state);
 
-    static constexpr uint32_t n = 2;
     static constexpr uint32_t d = 2;
-    static constexpr Scalar G = 1e-3;
+    static constexpr Scalar G = 1e-1;
 
 private:
     void compute_force();
 
+    uint32_t n_;
     VectorX<Scalar> m_;
     // holds velocities and derivatives
     VectorX<Scalar> d_state_;
 };
 
 //..................................................................................................
-NBodyLaw::NBodyLaw() {
-    m_.resize(n);
-    d_state_.resize(2 * n * d);
+void NBodyLaw::resize(uint32_t n) {
+    n_ = n;
+    m_.resize(n_);
+    d_state_.resize(2 * n_ * d);
 }
 
 //..................................................................................................
 VectorX<Scalar> const& NBodyLaw::operator()(VectorX<Scalar> const& state) {
     // Copy in velocities
-    d_state_.segment<n * d>(0) = state.segment<n * d>(n * d);
+    d_state_.segment(0, n_ * d) = state.segment(n_ * d, n_ * d);
 
     // Compute forces
-    d_state_.segment<n * d>(n * d).setZero();
-    for (uint32_t i = 0; i < n; ++i) {
-        for (uint32_t j = i + 1; j < n; ++j) {
+    d_state_.segment(n_ * d, n_ * d).setZero();
+    for (uint32_t i = 0; i < n_; ++i) {
+        for (uint32_t j = i + 1; j < n_; ++j) {
             Eigen::Matrix<Scalar, d, 1> force = state.segment<d>(j * d) - state.segment<d>(i * d);
             Scalar const distance_squared = force.squaredNorm();
             Scalar const distance = std::sqrt(distance_squared);
             force *= G / (distance_squared * distance);
-            d_state_.segment<d>((n + i) * d) += m_[j] * force;
-            d_state_.segment<d>((n + j) * d) -= m_[i] * force;
+            d_state_.segment<d>((n_ + i) * d) += m_[j] * force;
+            d_state_.segment<d>((n_ + j) * d) -= m_[i] * force;
         }
     }
 
@@ -97,8 +121,8 @@ private:
 
 //..................................................................................................
 NBody::NBody() {
-    state_.resize(2 * law_.n * law_.d);
-    delta_t_ = 1e-3;
+    state_.resize(2 * law_.n() * law_.d);
+    delta_t_ = 1e-4;
 }
 
 //..................................................................................................
@@ -109,11 +133,11 @@ void NBody::initialize(
     t_ = 0;
     law_.set_masses(m);
     state_ = state;
-    //state_.resize(2 * law_.n * law_.d);
-    //state_.segment(0, law_.n * law_.d) = x0;
-    //state_.segment(law_.n * law_.d, law_.n * law_.d) = v0;
-    points_.resize(law_.n * law_.d);
-    points_ = state_.segment(0, law_.n * law_.d);
+    //state_.resize(2 * law_.n() * law_.d);
+    //state_.segment(0, law_.n() * law_.d) = x0;
+    //state_.segment(law_.n() * law_.d, law_.n() * law_.d) = v0;
+    points_.resize(law_.n() * law_.d);
+    points_ = state_.segment(0, law_.n() * law_.d);
 
     std::cout << "state.size() = " << state_.size() << '\n';
     std::cout << "points.size() = " << points_.size() << '\n';
@@ -121,7 +145,7 @@ void NBody::initialize(
 
 //..................................................................................................
 void NBody::step() {
-    state_ = euler_step(
+    state_ = rk4_step(
         [this](Scalar, VectorX<Scalar> const& state) { return law_(state); },
         t_,
         state_,
@@ -130,21 +154,27 @@ void NBody::step() {
     t_ += delta_t_;
 
     // hack
-    points_ = state_.segment(0, law_.n * law_.d);
+    points_ = state_.segment(0, law_.n() * law_.d);
 }
 
 //--------------------------------------------------------------------------------------------------
 int main() {
-    VectorX<Scalar> masses(2);
-    masses[0] = 1e3; // sun
+    uint32_t n_bodies = 3;
+    VectorX<Scalar> masses(n_bodies);
+    masses[0] = 1.5; // sun
     masses[1] = 1; // earth
-    VectorX<Scalar> initial_conditions(8);
-    Scalar const earth_v = -1;
-    Scalar const sun_v = -earth_v * masses[1] / masses[0];
+    masses[2] = 0.2; // moon
+    VectorX<Scalar> initial_conditions(2 * n_bodies * NBodyLaw::d);
+    Scalar const earth_v_y = -0.3;
+    Scalar const moon_v_x = 0.3;
+    Scalar const sun_v_y = -earth_v_y * masses[1] / masses[0];
+    Scalar const sun_v_x = -moon_v_x * masses[2] / masses[0];
     initial_conditions.segment<2>(0) = Vector2<Scalar>(0, 0);       // position of sun
     initial_conditions.segment<2>(2) = Vector2<Scalar>(1, 0);       // position of earth
-    initial_conditions.segment<2>(4) = Vector2<Scalar>(0, sun_v);   // velocity of sun
-    initial_conditions.segment<2>(6) = Vector2<Scalar>(0, earth_v); // velocity of earth
+    initial_conditions.segment<2>(4) = Vector2<Scalar>(0, 0.8);       // position of moon
+    initial_conditions.segment<2>(6) = Vector2<Scalar>(sun_v_x, sun_v_y);   // velocity of sun
+    initial_conditions.segment<2>(8) = Vector2<Scalar>(0, earth_v_y); // velocity of earth
+    initial_conditions.segment<2>(10) = Vector2<Scalar>(moon_v_x, 0); // velocity of moon
 
     NBody nbody;
     nbody.initialize(masses, initial_conditions);