WiP orbital mechanics code

gistfile1.txt

#include <iostream>
#include <cmath>
#include <utility>

#include <glm/glm.hpp>
#include <glm/gtc/type_ptr.hpp> 

const double G = 6.674e-11; // gravitational constant
const double PI = 3.14159265358979323846;

double solve_E(double M, double e)
{
    double E = M;
    double diff = 1;
    double epsilon = 1.e-13;
    while(std::abs(diff/E) > epsilon)
    {
        diff = (E - e*std::sin(E) - M)/(1-e*std::cos(E));
        E -= diff;
    }
    return E;
}

double solve_H(double M, double e)
{
    double H = M;
    double diff = 1;
    double epsilon = 1.e-13;
    while(std::abs(diff/H) > epsilon)
    {
        diff = (e*std::sinh(H) - H - M)/(e*std::cosh(H)-1);
        H -= diff;
    }
    return H;
}

struct Orbit {
    double p, e, a, n, mu, T, t0, theta0;
    glm::dmat3 coordinates;
    
    double t2theta(double t)
    {
        t /= 2*PI*n;
        t -= std::floor(t);
        double M = 2*PI*t;
        double theta;
        if(e<1)
        {
            double E = solve_E(M, e);
            theta = 2*atan2(std::sqrt(1+e)*std::sin(E/2), std::sqrt(1-e)*std::cos(E/2));
        }
        else
        {
            double H = solve_H(M, e);
            theta = 2*atan(std::sqrt((e+1)/(e-1))*std::tanh(H/2));
        }
        return theta;
    }

    double theta2t(double theta)
    {
        double M;
        if(e<1)
        {
            double E = 2*atan2(std::sqrt(1-e)*std::sin(theta/2), std::sqrt(1+e)*std::cos(theta/2));
            M = E - e*std::sin(E);
        }
        else
        {
            double H = 2*atanh(std::sqrt((e-1)/(e+1))*std::tan(theta/2));
            M = e*std::sinh(H)-H;
        }
        double t = M*n;
        return t;
    }
    
    glm::dvec3 r(double t)
    {
        double theta = t2theta(t0+t);
        double r = p/(1 + e*cos(theta));
        return coordinates*glm::dvec3(r*std::cos(theta), r*std::sin(theta), 0);
    }

    glm::dvec3 v(double t)
    {
        double theta = t2theta(t0+t);
        double r = p/(1 + e*cos(theta));
        double v = std::sqrt(mu*(2/r-1/a));
        return coordinates*glm::dvec3(-v*std::sin(theta), v*std::cos(theta), 0);
    }
    
    std::pair<glm::dvec3, glm::dvec3> state(double t)
    {
        double theta = t2theta(t0+t);
        double r = p/(1 + e*cos(theta));
        double v = std::sqrt(mu*(2/r-1/a));
        return std::make_pair(
            coordinates*glm::dvec3(r*std::cos(theta), r*std::sin(theta), 0),
            coordinates*glm::dvec3(-v*std::sin(theta), v*std::cos(theta), 0)
        );
    }
};

Orbit calcElements(glm::dvec3 r_vector, glm::dvec3 v_vector, double M)
{
    // the orbited mass (M) is assumed at the coordinate origin (0, 0)
    double mu = G*M;
    
    double r = glm::length(r_vector); // radius
    glm::dvec3 n = r_vector/r; // normal from origin (x)
    double vr = glm::dot(n, v_vector); // radial velocity
    glm::dvec3 t = v_vector-n*vr; // tangent
    double vt = glm::length(t); // tangential velocity
    t /= vt;
    
    double p = (r*r*vt*vt)/mu; // semi-latus rectum
    
    double v0 = std::sqrt(mu/p);
    
    double ec = vt/v0-1; // e*cos(theta)
    double es = vr/v0;   // e*sin(theta)
    double e = std::sqrt(ec*ec + es*es);
    double theta = std::atan2(es, ec);
    
    double c = ec/e;
    double s = es/e;

    glm::dvec3 x = c*n - s*t;
    glm::dvec3 y = s*n + c*t;
    glm::dvec3 z = glm::cross(n,t);
    
    Orbit result;
    result.p = p;
    result.e = e;
    result.a = p/(1-e*e);
    result.mu = mu;
    
    result.n = std::sqrt((std::abs(result.a*result.a*result.a))/mu);
    result.T = 2*PI*result.n;
    result.theta0 = theta;
    result.t0 = result.theta2t(theta);

    result.coordinates = glm::dmat3(x,y,z);
    
    return result;
}

std::pair<glm::dvec3, glm::dvec3> velocity_verlet_step(std::pair<glm::dvec3, glm::dvec3> state, double M, double dt)
{
    double mu = G*M;
    glm::dvec3 r = state.first;
    glm::dvec3 v = state.second;
    
    double rl = glm::length(r);
    
    glm::dvec3 a0 = -mu*r/(rl*rl*rl);
    r += dt*v + 0.5*dt*dt*a0;
    glm::dvec3 a1 = -mu*r/(rl*rl*rl);
    v += 0.5*dt*(a0+a1);
    
    return std::make_pair(r, v);
}

double energy(std::pair<glm::dvec3, glm::dvec3> state, double M)
{
    double r = glm::length(state.first);
    double v2 = glm::dot(state.second, state.second);
    return 0.5*v2 - G*M/r;
}




static double mad(double a, double b, double c)
{
    return a * b + c;
}
 
static void step1(const double *vel0, const double *a, double *d, double *v, double *delta, double *vel, double dt, size_t n)
{
    for(size_t i = 0; i < n; ++i)
    {
        double k1 = a[i] * 0.5 * dt;
        v[i] = vel0[i] + k1;
        d[i] = mad(k1, 0.5, vel0[i]) * 0.5 * dt; // d = K
 
        delta[i] = mad(k1, 1.0/3.0, vel0[i]); // delta = vel0 + (1/3)*k1
        vel[i] = mad(k1, 1.0/3.0, vel0[i]); // vel = vel0 + (1/3)*k1
    }
}
 
static void step2(const double *vel0, const double *a, double *v, double *delta, double *vel, double dt, size_t n)
{
    for(size_t i = 0; i < n; ++i)
    {
        double k2 = a[i] * 0.5 * dt;
 
        v[i] = vel0[i] + k2; // v = vel0 + k2
 
        delta[i] = mad(k2, 1.0/3.0, delta[i]); // delta = vel0 + (1/3)(k1+k2)
        vel[i] = mad(k2, 2.0/3.0, vel[i]); // vel = (1/3)*k1 + (2/3)*k2
    }
}
 
static void step3(const double *vel0, const double *a, double *d, double *v, double *delta, double *vel, double dt, size_t n)
{
    for(size_t i = 0; i < n; ++i)
    {
        double k3 = a[i] * 0.5 * dt;
 
        d[i] = (vel[i] + k3) * dt; // d = L
        v[i] = mad(k3, 2.0, vel0[i]);
 
        delta[i] = mad(k3, 1.0/3.0, delta[i]); // delta = vel0 + (1/3)*(k1+k2+k3)
        vel[i] = mad(k3, 2.0/3.0, vel[i]); // vel = (1/3)*k1+(2/3)*k2+(2/3)*k3
    }
}
 
static void step4(const double *a, double *delta, double *vel, double dt, size_t n)
{
    for(size_t i = 0; i < n; ++i)
    {
        double k4 = a[i] * 0.5 * dt;
 
        vel[i] = mad(k4, 1.0/3.0, vel[i]); // vel = vel0 + (1/3)*k1+(2/3)*k2+(2/3)*k3+(1/3)*k4
        delta[i] = delta[i] * dt; // delta = dt * (vel0 + (1/3)*(k1+k2+k3))
    }
}
 
static void accelerate(const double *s, const double *vel, double *a, double t, size_t n)
{
    double M = 1.9891e30;
    double mu = G*M;
    double l = 0;
    for(size_t i = 0; i < n; ++i)
        l += s[i]*s[i];
    l = std::sqrt(l);
    for(size_t i = 0; i < n; ++i)
        a[i] = -mu*s[i]/(l*l*l);
    
}
 
static void update(const double *s0, const double *delta, double *s, size_t n)
{
    for(size_t i = 0; i < n; ++i)
        s[i] = s0[i] + delta[i];
}
 
static void sim_step(const double *s0, const double *vel0, double *s, double *vel, double t0, double dt, size_t n)
{
    double d_temp[n], v_temp[n];
    double a[n];
    double delta[n];
 
    accelerate(s0, vel0, a, t0, n);
    step1(vel0, a, d_temp, v_temp, delta, vel, dt, n);
    update(s0, d_temp, s, n);
 
    accelerate(s, v_temp, a, t0 + 0.5 * dt, n);
    step2(vel0, a, v_temp, delta, vel, dt, n);
    accelerate(s, v_temp, a, t0 + 0.5 * dt, n);
    step3(vel0, a, d_temp, v_temp, delta, vel, dt, n);
    update(s0, d_temp, s, n);
 
    accelerate(s, v_temp, a, t0 + dt, n);
    step4(a, delta, vel, dt, n);
    update(s0, delta, s, n);
 
    accelerate(s, vel, a, t0+dt, n);
    //~ printf("%f\t%f\t%f\t%f\n", t0+dt, s[0], vel[0], a[0]);
}






int main()
{
    double M_sun = 1.9891e30;
    double v_earth = 29.29e3;
    double r_earth = 152.10e9;

    auto state = std::make_pair(glm::dvec3(r_earth, 0, 0), glm::dvec3(0, 1.0*v_earth, 0));
    
    double E0 = energy(state, M_sun);
    
    Orbit orbit = calcElements(state.first, state.second, M_sun);
    int steps = 0;
    double dt = 1;
    double t = 0;
    for(;t<1000*orbit.T;t+=dt, ++steps)
    {
        if(steps%100000 == 0)
        {
            auto state2 = orbit.state(t);
            
           
            std::cout << t/orbit.T << ' ' << energy(state, M_sun) << ' ' << energy(state2, M_sun) << ' ' << energy(state, M_sun) - E0 << ' ' << energy(state2, M_sun) - E0 << std::endl;
        }
        
        //~ state = velocity_verlet_step(state, M_sun, dt);
        auto newstate = state;
        sim_step(glm::value_ptr(state.first), glm::value_ptr(state.second), glm::value_ptr(newstate.first), glm::value_ptr(newstate.second), t, dt, 3);
        state = newstate;        
    }
    
    
    std::cout << dt/orbit.T << std::endl;
}