Commit aba8edc2 authored by Erik Strand's avatar Erik Strand

Add jacobian computations

No idea if they're right, since I haven't even printed them yet.
parent b55f3114
......@@ -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
......@@ -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);
......
......@@ -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_;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment