Commit b55f3114 authored by Erik Strand's avatar Erik Strand

Pull state back as far as possible

This puts it all in one place, and makes it clear how much there is.
parent a029ac50
......@@ -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
......@@ -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,
......
......@@ -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);
......
......@@ -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_;
......
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