Pull state back as far as possible

This puts it all in one place, and makes it clear how much there is.

_code/notes/diffsim/integrators.h

@@ -12,29 +12,9 @@
 //--------------------------------------------------------------------------------------------------
 // A single step of Euler integration
 template <typename F, typename Vector, typename Scalar>
-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;
+void euler_step(F const& f, Scalar const delta_t, Scalar& t, Vector& x) {
+    x += delta_t * f(t, x);
+    t += delta_t;
 }
 
 #endif

_code/notes/diffsim/main.cpp

@@ -9,25 +9,28 @@ void mouse(int, int, int, int) { exit(0); }
 
 //--------------------------------------------------------------------------------------------------
 int main() {
-    uint32_t n_bodies = 3;
+    constexpr uint32_t dim = 2;
+    uint32_t const n_bodies = 3;
+
     VectorX<Scalar> masses(n_bodies);
     masses[0] = 1.5; // sun
     masses[1] = 1; // earth
     masses[2] = 0.2; // moon
-    VectorX<Scalar> initial_conditions(2 * n_bodies * NBodyLaw::d);
+
+    VectorX<Scalar> initial_conditions(dim * 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, 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
+    initial_conditions.segment<dim>(0 * dim) = Vector2<Scalar>(0, 0);         // position of sun
+    initial_conditions.segment<dim>(1 * dim) = Vector2<Scalar>(1, 0);         // position of earth
+    initial_conditions.segment<dim>(2 * dim) = Vector2<Scalar>(0, 0.8);       // position of moon
+    initial_conditions.segment<dim>(3 * dim) = Vector2<Scalar>(sun_v_x, sun_v_y); // velocity of sun
+    initial_conditions.segment<dim>(4 * dim) = Vector2<Scalar>(0, earth_v_y); // velocity of earth
+    initial_conditions.segment<dim>(5 * dim) = Vector2<Scalar>(moon_v_x, 0);  // velocity of moon
 
     NBody nbody;
-    nbody.initialize(masses, initial_conditions);
+    nbody.initialize(dim, masses, initial_conditions);
 
     grapher.initialize(
         &nbody,

_code/notes/diffsim/n_body.cpp

@@ -4,47 +4,53 @@
 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(0, n_ * d) = state.segment(n_ * d, n_ * d);
+void NBodyLaw::set_masses(VectorX<Scalar> const& m) {
+    n_ = m.size();
+    m_ = m;
+}
+
+//..................................................................................................
+void NBodyLaw::operator()(VectorX<Scalar> const& state, VectorX<Scalar>& gradient) {
+    // The derivatives of locations with respect to time are just the velocities, which we know.
+    gradient.segment(0, n_ * d) = state.segment(n_ * d, n_ * d);
 
-    // Compute forces
-    d_state_.segment(n_ * d, n_ * d).setZero();
+    // The derivatives of velocities with respect to time are accelerations. Here we use Newton's
+    // law of universal gravitation to compute forces, then divide out the relevant mass to get
+    // accelerations (or more literally, we never multiply by it in the first place). This involves
+    // a sum over all pairs of particles, which could be truncated for distant particles if needed.
+    gradient.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;
+            gradient.segment<d>((n_ + i) * d) += m_[j] * force;
+            gradient.segment<d>((n_ + j) * d) -= m_[i] * force;
         }
     }
-
-    return d_state_;
-}
-
-//..................................................................................................
-NBody::NBody() {
-    state_.resize(2 * law_.n() * law_.d);
-    delta_t_ = 1e-4;
 }
 
 //..................................................................................................
 void NBody::initialize(
-    VectorX<Scalar> const& m,
+    uint32_t dim,
+    VectorX<Scalar> const& masses,
     VectorX<Scalar> const& state
 ) {
+    assert(2 * dim * masses.size() == state.size());
+
+    law_.set_masses(masses);
+    delta_t_ = 1e-4;
+
     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;
+
+    d_state_.resize(state_.size());
+    d_state_.setZero();
+
     points_.resize(law_.n() * law_.d);
     points_ = state_.segment(0, law_.n() * law_.d);
 
@@ -54,13 +60,11 @@ void NBody::initialize(
 
 //..................................................................................................
 void NBody::step() {
-    state_ = rk4_step(
-        [this](Scalar, VectorX<Scalar> const& state) { return law_(state); },
-        t_,
-        state_,
-        delta_t_
-    );
-    t_ += delta_t_;
+    auto const f = [this](Scalar, VectorX<Scalar> const& state) -> VectorX<Scalar> const& {
+        law_(state, d_state_);
+        return d_state_;
+    };
+    euler_step(f, delta_t_, t_, state_);
 
     // hack
     points_ = state_.segment(0, law_.n() * law_.d);

_code/notes/diffsim/n_body.h

@@ -22,11 +22,11 @@ class NBodyLaw {
 public:
     NBodyLaw() {}
     void resize(uint32_t n);
-    void set_masses(VectorX<Scalar> const& m) { resize(m.size()); m_ = m; }
+    void set_masses(VectorX<Scalar> const& m);
 
     uint32_t n() const { return n_; }
-    // returns derivatives of phase space coordinates w.r.t. time at a give state
-    VectorX<Scalar> const& operator()(VectorX<Scalar> const& state);
+    // computes the gradient of phase space coordinates w.r.t. time at a give state
+    void operator()(VectorX<Scalar> const& state, VectorX<Scalar>& gradient);
 
     static constexpr uint32_t d = 2;
     static constexpr Scalar G = 1e-1;
@@ -34,28 +34,33 @@ public:
 private:
     void compute_force();
 
-    uint32_t n_;
+    uint32_t n_; // number of particles
     VectorX<Scalar> m_;
-    // holds velocities and derivatives
-    VectorX<Scalar> d_state_;
 };
 
 //--------------------------------------------------------------------------------------------------
 // N-Body simulation based on integration of NBodyLaw
 class NBody : public SimulationInterface {
 public:
-    NBody();
-    void initialize(VectorX<Scalar> const& m, VectorX<Scalar> const& state);
+    NBody() {}
+    void initialize(uint32_t dim, VectorX<Scalar> const& masses, VectorX<Scalar> const& state);
     void step() final;
 
     VectorX<Scalar> const& points() const final { return points_; }
     VectorX<Scalar> const& state() const { return state_; }
 
 private:
-    Scalar t_;
+    // configuration
+    NBodyLaw law_;
     Scalar delta_t_;
+
+    // state
+    Scalar t_;
     VectorX<Scalar> state_;
-    NBodyLaw law_;
+
+    // working memory
+    // Used for evaluation of gradient of state w.r.t time (i.e. velocities and accelerations).
+    VectorX<Scalar> d_state_;
 
     // copy of first half of state_, just to make visualization easy
     VectorX<Scalar> points_;