Commit 9640c329 authored by Erik Strand's avatar Erik Strand

Find a stable orbit in a differentiable simulation

parents 278f11a9 c0d97152
[submodule "_code/extern/optimization"]
path = _code/extern/optimization
url = ssh://git@gitlab.cba.mit.edu:846/erik/optimization.git
# Clara
add_library(clara INTERFACE)
target_include_directories(clara INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}/clara)
# optimization
# I should configure the project to place nice as a subproject, but for now I just want headers.
add_library(optimization INTERFACE)
target_include_directories(optimization INTERFACE optimization/optimization)
# should also include Eigen
target_link_libraries(optimization INTERFACE shared_settings)
Subproject commit c87a5e0e10d42eab8a7a74191af8127301a0dbac
......@@ -3,6 +3,7 @@ find_package(GLUT REQUIRED)
add_executable(nbody
main.cpp
n_body.cpp
glut_grapher.cpp
)
target_link_libraries(nbody shared_code ${OPENGL_LIBRARIES} ${GLUT_LIBRARY})
target_link_libraries(nbody shared_code optimization ${OPENGL_LIBRARIES} ${GLUT_LIBRARY})
from sympy import *
from sympy.assumptions.assume import global_assumptions
import sympy as sp
#from sympy.assumptions.assume import global_assumptions
def print_expr(name, expr):
print(name + ": ")
print(latex(expr))
print(sp.latex(expr))
print("")
t = Symbol('t', real=True)
dt = Symbol('\Delta t', real=True)
x = Symbol('x', real=True)
xn = Function('x_n', real=True)
f = Function('f', real=True)
t = sp.Symbol('t', real=True)
dt = sp.Symbol('\Delta t', real=True)
x = sp.Symbol('x', real=True)
xn = sp.Function('x_n', real=True)
f = sp.Function('f', real=True)
x_1, y_1, x_2, y_2 = sp.symbols("x_1 y_1 x_2 y_2", real=True)
G, m_2 = sp.symbols("G m_2", real=True)
def euler():
expr = xn(x) + dt * f(t, xn(x))
return simplify(Derivative(expr, x))
return sp.simplify(sp.diff(expr, x))
def rk4():
k1 = dt * f(t, xn(x))
......@@ -22,7 +25,24 @@ def rk4():
k3 = dt * f(t + dt/2, xn(x) + k2/2)
k4 = dt * f(t + dt, xn(x) + k3)
expr = xn(x) + k1/6 + k2/3 + k3/3 + k4/6
return simplify(Derivative(expr, x))
return sp.simplify(sp.diff(expr, x))
def nbody():
distance = sp.sqrt((x_2 - x_1)**2 + (y_2 - y_1)**2)
f12_x = G * (x_2 - x_1) / distance**3
f12_y = G * (y_2 - y_1) / distance**3
df12_x_x1 = sp.simplify(sp.diff(f12_x, x_1))
df12_x_y1 = sp.simplify(sp.diff(f12_x, y_1))
df12_y_x1 = sp.simplify(sp.diff(f12_y, x_1))
df12_y_y1 = sp.simplify(sp.diff(f12_y, y_1))
print_expr("f12_x", f12_x)
#print_expr("f12_x", f12_y)
print_expr("df12_x_x1", df12_x_x1)
print_expr("df12_x_y1", df12_x_y1)
print_expr("df12_y_x1", df12_y_x1)
print_expr("df12_y_y1", df12_y_y1)
print_expr("euler", euler())
print_expr("rk4", rk4())
#print_expr("euler", euler())
#print_expr("rk4", rk4())
nbody()
#ifndef INTEGRATORS_H
#define INTEGRATORS_H
#include <cmath>
#include <utility>
//--------------------------------------------------------------------------------------------------
// All these integrators are for numerically solving problems of the form x'(t) = f(t, x). Note
// that x can be vector valued, so higher order problems can be solved by converting to a system of
// first order problems.
//--------------------------------------------------------------------------------------------------
// 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
#include <iostream>
#include "scalar.h"
#include "vector.h"
#include "glut_grapher.h"
#include "n_body.h"
#include "optimizers/nelder_mead/nelder_mead.h"
#include "optimizers/conjugate_gradient_descent/conjugate_gradient_descent.h"
#include "objectives/rosenbrock.h"
//--------------------------------------------------------------------------------------------------
GlutGrapher grapher;
......@@ -9,188 +11,149 @@ void display() { grapher.display(); }
void idle() { grapher.idle(); }
void mouse(int, int, int, int) { exit(0); }
//--------------------------------------------------------------------------------------------------
// 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;
}
struct LaunchObjective {
LaunchObjective() {
masses_.resize(n_bodies_);
masses_[0] = 1e1; // planet
masses_[1] = 1e-1; // satellite
initial_conditions_.resize(2 * dim_ * n_bodies_);
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); // velocity of sun
initial_conditions_.segment<dim_>(3 * dim_) = Vector2<Scalar>(0, 0); // velocity of earth
}
//--------------------------------------------------------------------------------------------------
// state vectors (i.e. phase space points) are packed as follows:
// 0, ..., d - 1 give (up to) x, y, z for mass zero
// d, ..., 2 * d - 1 give (up to) x, y, z for mass one
// ...
// (n - 1) * d, ..., n * d - 1 give (up to) x, y, z for mass n - 1
// n * d, ..., (n + 1) * d give (up to) v_x, v_y, v_z for mass zero
// ...
// (2 * n - 1) * d, ..., 2 * n * d - 1 give (up to) v_x, v_y, v_z for mass n - 1
NBody& nbody() { return nbody_; }
//--------------------------------------------------------------------------------------------------
// Functor that computes phase space derivatives
class NBodyLaw {
public:
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 d = 2;
static constexpr Scalar G = 1e-1;
private:
void compute_force();
uint32_t n_;
VectorX<Scalar> m_;
// holds velocities and derivatives
VectorX<Scalar> d_state_;
};
void prep_initial_conditions(VectorX<Scalar> const& velocity) {
initial_conditions_.segment<dim_>(3 * dim_) = velocity;
nbody_.initialize(dim_, masses_, initial_conditions_, delta_t_);
}
//..................................................................................................
void NBodyLaw::resize(uint32_t n) {
n_ = n;
m_.resize(n_);
d_state_.resize(2 * n_ * d);
}
void operator()(VectorX<Scalar> const& velocity, Scalar& loss) {
// prepare initial conditions
initial_conditions_.segment<dim_>(3 * dim_) = velocity;
//..................................................................................................
VectorX<Scalar> const& NBodyLaw::operator()(VectorX<Scalar> const& state) {
// Copy in velocities
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) {
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;
// run the simulation
nbody_.initialize(dim_, masses_, initial_conditions_, delta_t_);
for (uint32_t i = 0; i < n_steps_; ++i) {
nbody_.step();
}
// compute loss
Vector2<Scalar> const goal = Vector2<Scalar>(0, 1);
Vector2<Scalar> const difference = nbody_.state().segment<dim_>(dim_) - goal;
loss = difference.squaredNorm();
// print stuff
std::cout << "loss = " << loss << "\n";
}
return d_state_;
}
void operator()(VectorX<Scalar> const& velocity, Scalar& loss, VectorX<Scalar>& gradient) {
// prepare initial conditions
initial_conditions_.segment<dim_>(3 * dim_) = velocity;
//--------------------------------------------------------------------------------------------------
class NBody : public SimulationInterface {
public:
NBody();
void initialize(VectorX<Scalar> const& m, 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_;
Scalar delta_t_;
VectorX<Scalar> state_;
NBodyLaw law_;
// copy of first half of state_, just to make visualization easy
VectorX<Scalar> points_;
};
// run the simulation
nbody_.initialize(dim_, masses_, initial_conditions_, delta_t_);
for (uint32_t i = 0; i < n_steps_; ++i) {
nbody_.step();
}
//..................................................................................................
NBody::NBody() {
state_.resize(2 * law_.n() * law_.d);
delta_t_ = 1e-4;
}
// compute loss
Vector2<Scalar> const goal = Vector2<Scalar>(0, 1);
Vector2<Scalar> const difference = nbody_.state().segment<dim_>(dim_) - goal;
loss = difference.squaredNorm();
// compute gradient of loss
// loss = difference^T * difference
// d_loss = 2 * difference^T * d_difference
gradient = 2 * difference.transpose() * nbody_.state_jacobian().block<dim_, dim_>(dim_, 3 * dim_);
// print stuff
std::cout << "loss = " << loss << "\n";
std::cout << "loss gradient\n";
std::cout << gradient << '\n';
//std::cout << "state jacobian\n";
//std::cout << nbody_.state_jacobian() << '\n';
}
//..................................................................................................
void NBody::initialize(
VectorX<Scalar> const& m,
VectorX<Scalar> const& state
) {
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);
std::cout << "state.size() = " << state_.size() << '\n';
std::cout << "points.size() = " << points_.size() << '\n';
}
constexpr static uint32_t dim_ = 2;
constexpr static uint32_t n_bodies_ = 2;
constexpr static Scalar sim_time_ = 1.4;
constexpr static Scalar delta_t_ = 1e-3;
constexpr static uint32_t n_steps_ = static_cast<uint32_t>(sim_time_ / delta_t_);
//..................................................................................................
void NBody::step() {
state_ = rk4_step(
[this](Scalar, VectorX<Scalar> const& state) { return law_(state); },
t_,
state_,
delta_t_
);
t_ += delta_t_;
// hack
points_ = state_.segment(0, law_.n() * law_.d);
}
VectorX<Scalar> masses_;
VectorX<Scalar> initial_conditions_;
NBody nbody_;
};
//--------------------------------------------------------------------------------------------------
int main() {
uint32_t 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);
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
NBody nbody;
nbody.initialize(masses, initial_conditions);
grapher.initialize(
&nbody,
BoundingBox<Vector2<Scalar>>(
{-2, -2},
{2, 2}
)
);
grapher.init_glut();
glutDisplayFunc(display);
glutMouseFunc(mouse);
glutIdleFunc(idle);
grapher.run();
LaunchObjective objective;
VectorX<Scalar> satellite_velocity;
satellite_velocity.resize(2);
satellite_velocity[0] = 0;
satellite_velocity[1] = 0.4;
//Scalar loss;
//VectorX<Scalar> gradient;
//objective(satellite_velocity, loss, gradient);
//std::cout << '\n';
bool do_nm = false;
bool do_cgd = true;
bool do_vis = true;
if (do_nm) {
optimization::NelderMead<-1> nelder_mead(
10000, // max evaluations
1e-6 // relative y tolerance
);
nelder_mead.optimize(objective, satellite_velocity, 0.5);
satellite_velocity = nelder_mead.point();
std::cout << "n evaluations: " << nelder_mead.n_evaluations() << '\n';
std::cout << "final point: " << nelder_mead.point() << '\n';
std::cout << "final loss: " << nelder_mead.value() << '\n';
std::cout << '\n';
}
if (do_cgd) {
optimization::ConjugateGradientDescent<-1> cgd(
1e-3, // gradient threshold
-1, // abs threshold
100000, // max evaluations
10000 // max iterations
);
cgd.optimize(objective, satellite_velocity);
satellite_velocity = cgd.point();
std::cout << "n evaluations: " << cgd.n_evaluations() << '\n';
std::cout << "n iterations: " << cgd.n_iterations() << '\n';
std::cout << "final point: " << cgd.point() << '\n';
std::cout << '\n';
}
if (do_vis) {
objective.prep_initial_conditions(satellite_velocity);
grapher.initialize(
&objective.nbody(),
BoundingBox<Vector2<Scalar>>(
{-2, -2},
{2, 2}
)
);
grapher.init_glut();
glutDisplayFunc(display);
glutMouseFunc(mouse);
glutIdleFunc(idle);
grapher.run();
}
return 0;
}
#include "n_body.h"
//..................................................................................................
void NBodyLaw::resize(uint32_t n) {
n_ = n;
m_.resize(n_);
}
//..................................................................................................
void NBodyLaw::set_masses(VectorX<Scalar> const& m) {
n_ = m.size();
m_ = m;
}
//..................................................................................................
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.
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.
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);
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;
}
}
}
//..................................................................................................
void NBody::initialize(
uint32_t dim,
VectorX<Scalar> const& masses,
VectorX<Scalar> const& state,
Scalar delta_t
) {
assert(2 * dim * masses.size() == state.size());
law_.set_masses(masses);
delta_t_ = delta_t;
// initialize state
t_ = 0;
state_ = state;
state_jacobian_.resize(state_.size(), state_.size());
state_jacobian_.setIdentity();
// 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);
//std::cout << "state.size() = " << state_.size() << '\n';
//std::cout << "points.size() = " << points_.size() << '\n';
}
//..................................................................................................
void NBody::step() {
diff_euler_step(*this, delta_t_, t_, state_, state_jacobian_);
// hack
points_ = state_.segment(0, law_.n() * law_.d);
}
#ifndef N_BODY_H
#define N_BODY_H
#include "glut_grapher.h"
#include "integrators.h"
#include "scalar.h"
#include "vector.h"
//--------------------------------------------------------------------------------------------------