Add jacobian computations

No idea if they're right, since I haven't even printed them yet.

_code/notes/diffsim/integrators.h

@@ -11,10 +11,33 @@
 
 //--------------------------------------------------------------------------------------------------
 // A single step of Euler integration
+// This is commented out since I need to make its interface compatible with diff_euler_step.
+/*
 template <typename F, typename Vector, typename Scalar>
 void euler_step(F const& f, Scalar const delta_t, Scalar& t, Vector& x) {
     x += delta_t * f(t, x);
     t += delta_t;
 }
+*/
+
+//--------------------------------------------------------------------------------------------------
+// A single step of differentiable Euler integration
+template <typename F, typename Vector, typename Scalar>
+void diff_euler_step(
+    F& f,
+    Scalar const delta_t,
+    Scalar& t,
+    Vector& state,
+    MatrixX<Scalar>& state_jacobian
+) {
+    uint32_t const rows = state_jacobian.rows();
+    uint32_t const cols = state_jacobian.cols();
+
+    f(t, state);
+    t += delta_t;
+    state += delta_t * f.tangent();
+    state_jacobian =
+        (MatrixX<Scalar>::Identity(rows, cols) + delta_t * f.tangent_jacobian()) * state_jacobian;
+}
 
 #endif

_code/notes/diffsim/n_body.cpp

@@ -13,23 +13,72 @@ void NBodyLaw::set_masses(VectorX<Scalar> const& m) {
 }
 
 //..................................................................................................
-void NBodyLaw::operator()(VectorX<Scalar> const& state, VectorX<Scalar>& gradient) {
+void NBodyLaw::operator()(VectorX<Scalar> const& state, VectorX<Scalar>& tangent) {
     // 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);
+    tangent.segment(0, n_ * d) = state.segment(n_ * d, n_ * d);
 
     // 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();
+    tangent.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);
-            gradient.segment<d>((n_ + i) * d) += m_[j] * force;
-            gradient.segment<d>((n_ + j) * d) -= m_[i] * force;
+            tangent.segment<d>((n_ + i) * d) += m_[j] * force;
+            tangent.segment<d>((n_ + j) * d) -= m_[i] * force;
+        }
+    }
+}
+
+//..................................................................................................
+void NBodyLaw::operator()(
+    VectorX<Scalar> const& state,
+    VectorX<Scalar>& tangent,
+    MatrixX<Scalar>& tangent_jacobian
+) {
+    // Compute derivatives of displacements w.r.t. time, and initialize deriviatives of velocities
+    // w.r.t. time.
+    tangent.segment(0, n_ * d) = state.segment(n_ * d, n_ * d);
+    tangent.segment(n_ * d, n_ * d).setZero();
+
+    // Initialize the jacobian. The upper left and lower right blocks will stay zero; the others are
+    // set below.
+    tangent_jacobian.setZero();
+
+    // Set the upper right block of the jacobian. These are the derivatives of the tangent terms
+    // calculated in the first line of this method.
+    tangent_jacobian.block(0, n_ * d, n_ * d, n_ * d) = MatrixX<Scalar>::Identity(n_ * d, n_ * d);
+
+    // This loop fills in the rest of the tangent and the lower left block of the jacobian.
+    for (uint32_t i = 0; i < n_; ++i) {
+        for (uint32_t j = i + 1; j < n_; ++j) {
+            // Do some shared temp calculations.
+            Eigen::Matrix<Scalar, d, 1> const displacement =
+                state.segment<d>(j * d) - state.segment<d>(i * d);
+            Scalar const distance_squared = displacement.squaredNorm();
+            Scalar const distance = std::sqrt(distance_squared);
+            Scalar const distance_5 = distance_squared * distance_squared * distance;
+
+            // Compute the accelerations as before.
+            Eigen::Matrix<Scalar, d, 1> const force =
+                G * displacement / (distance_squared * distance);
+            tangent.segment<d>((n_ + i) * d) += m_[j] * force;
+            tangent.segment<d>((n_ + j) * d) -= m_[i] * force;
+
+            // Compute shared expression in the dim * dim matrix of partial derivatives of the
+            // accelerations w.r.t. to displacements (of mass i).
+            Eigen::Matrix<Scalar, d, d> dforce;
+            dforce.fill(3 * G * displacement.prod() / distance_5);
+            dforce.diagonal() = (3 * displacement.array() - distance_squared).matrix() / distance_5;
+
+            // Just multiply by masses and set sign to complete partial derivatives of the
+            // accelerations w.r.t. displacement i, and w.r.t. displacement j.
+            tangent_jacobian.block((n_ + i) * d, j * d, d, d) += m_[j] * dforce;
+            tangent_jacobian.block((n_ + j) * d, i * d, d, d) -= m_[i] * dforce;
         }
     }
 }
@@ -45,11 +94,17 @@ void NBody::initialize(
     law_.set_masses(masses);
     delta_t_ = 1e-4;
 
+    // initialize state
     t_ = 0;
     state_ = state;
+    state_jacobian_.resize(state_.size(), state_.size());
+    state_jacobian_.setIdentity();
 
-    d_state_.resize(state_.size());
-    d_state_.setZero();
+    // initialize working memory
+    tangent_.resize(state_.size());
+    tangent_.setZero();
+    tangent_jacobian_.resize(state_.size(), state_.size());
+    tangent_jacobian_.setZero();
 
     points_.resize(law_.n() * law_.d);
     points_ = state_.segment(0, law_.n() * law_.d);
@@ -60,11 +115,7 @@ void NBody::initialize(
 
 //..................................................................................................
 void NBody::step() {
-    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_);
+    diff_euler_step(*this, delta_t_, t_, state_, state_jacobian_);
 
     // hack
     points_ = state_.segment(0, law_.n() * law_.d);

_code/notes/diffsim/n_body.h

@@ -25,8 +25,16 @@ public:
     void set_masses(VectorX<Scalar> const& m);
 
     uint32_t n() const { return n_; }
-    // computes the gradient of phase space coordinates w.r.t. time at a give state
-    void operator()(VectorX<Scalar> const& state, VectorX<Scalar>& gradient);
+    // Computes the derivatives of the phase space coordinates w.r.t. time at a given state. This is
+    // a vector in the tangent bundle of phase space at the specified state, hence the name.
+    void operator()(VectorX<Scalar> const& state, VectorX<Scalar>& tangent);
+    // Computes the derivatives as above, as well as the jacobian of this tangent w.r.t. the phase
+    // space coordinates at a given state.
+    void operator()(
+        VectorX<Scalar> const& state,
+        VectorX<Scalar>& tangent,
+        MatrixX<Scalar>& jacobian_tangent
+    );
 
     static constexpr uint32_t d = 2;
     static constexpr Scalar G = 1e-1;
@@ -49,6 +57,13 @@ public:
     VectorX<Scalar> const& points() const final { return points_; }
     VectorX<Scalar> const& state() const { return state_; }
 
+    // These methods make NBody compatible with the differentiable integrators.
+    void operator()(Scalar, VectorX<Scalar> const& state) {
+        law_(state, tangent_, tangent_jacobian_);
+    }
+    VectorX<Scalar> const& tangent() const { return tangent_; }
+    MatrixX<Scalar> const& tangent_jacobian() const { return tangent_jacobian_; }
+
 private:
     // configuration
     NBodyLaw law_;
@@ -56,11 +71,16 @@ private:
 
     // state
     Scalar t_;
+    // Holds the current state (displacements and velocities).
     VectorX<Scalar> state_;
+    // Holds the jacobian of current state variables w.r.t. initial state variables.
+    MatrixX<Scalar> state_jacobian_;
 
-    // working memory
-    // Used for evaluation of gradient of state w.r.t time (i.e. velocities and accelerations).
-    VectorX<Scalar> d_state_;
+    // working memory (temporary)
+    // Holds the tangent of state w.r.t time (i.e. velocities and accelerations).
+    VectorX<Scalar> tangent_;
+    // Holds the jacobian of tangent_ w.r.t. state variables (displacements and velocities).
+    MatrixX<Scalar> tangent_jacobian_;
 
     // copy of first half of state_, just to make visualization easy
     VectorX<Scalar> points_;