Commit 00355e04 authored by Erik Strand's avatar Erik Strand

Refactor for clarity

parent 278f11a9
......@@ -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})
#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
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;
}
#endif
#include <iostream>
#include "scalar.h"
#include "vector.h"
#include "glut_grapher.h"
#include "n_body.h"
//--------------------------------------------------------------------------------------------------
GlutGrapher grapher;
......@@ -9,154 +7,6 @@ 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;
}
//--------------------------------------------------------------------------------------------------
// 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
//--------------------------------------------------------------------------------------------------
// 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 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);
// 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;
}
}
return d_state_;
}
//--------------------------------------------------------------------------------------------------
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_;
};
//..................................................................................................
NBody::NBody() {
state_.resize(2 * law_.n() * law_.d);
delta_t_ = 1e-4;
}
//..................................................................................................
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';
}
//..................................................................................................
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);
}
//--------------------------------------------------------------------------------------------------
int main() {
uint32_t n_bodies = 3;
......
#include "n_body.h"
//..................................................................................................
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);
// 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;
}
}
return d_state_;
}
//..................................................................................................
NBody::NBody() {
state_.resize(2 * law_.n() * law_.d);
delta_t_ = 1e-4;
}
//..................................................................................................
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';
}
//..................................................................................................
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);
}
#ifndef N_BODY_H
#define N_BODY_H
#include "glut_grapher.h"
#include "integrators.h"
#include "scalar.h"
#include "vector.h"
//--------------------------------------------------------------------------------------------------
// 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
//--------------------------------------------------------------------------------------------------
// Functor that computes derivatives of phase space coordinates w.r.t. time
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 derivatives of phase space coordinates w.r.t. time at 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_;
};
//--------------------------------------------------------------------------------------------------
// N-Body simulation based on integration of NBodyLaw
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_;
};
#endif
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