kepler.cpp

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

#include <glm/glm.hpp>

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

// solvers for Kepler's equation (Newton's method)
double solve_E(double M, double e)
{
    double E = M;
    double diff = 1.0;
    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.0;
    double epsilon = 1.e-13;
    while(std::abs(diff/H) > epsilon)
    {
        diff = (e*std::sinh(H) - H - M)/(e*std::cosh(H)-1.0);
        H -= diff;
    }
    return H;
}

struct Orbit {
    double p, e, a, n, mu, T, t0, theta0;
    glm::dmat3 coordinates;

    double t2theta(double t)
    {
        t /= 2.0*PI*n;
        t -= std::floor(t);
        double M = 2.0*PI*t;
        double theta;
        if(e<1)
        {
            double E = solve_E(M, e);
            theta = 2.0*atan2(std::sqrt(1.0+e)*std::sin(E/2.0), std::sqrt(1.0-e)*std::cos(E/2.0));
        }
        else
        {
            double H = solve_H(M, e);
            theta = 2.0*atan(std::sqrt((e+1.0)/(e-1.0))*std::tanh(H/2.0));
        }
        return theta;
    }

    double theta2t(double theta)
    {
        double M;
        if(e<1)
        {
            double E = 2.0*atan2(std::sqrt(1-e)*std::sin(theta/2.0), std::sqrt(1.0+e)*std::cos(theta/2.0));
            M = E - e*std::sin(E);
        }
        else
        {
            double H = 2.0*atanh(std::sqrt((e-1.0)/(e+1.0))*std::tan(theta/2.0));
            M = e*std::sinh(H)-H;
        }
        double t = M*n;
        return t;
    }

    // position at time t
    glm::dvec3 r(double t)
    {
        double theta = t2theta(t0+t);
        double r = p/(1.0 + e*cos(theta));
        return coordinates*glm::dvec3(r*std::cos(theta), r*std::sin(theta), 0.0);
    }

    // velocity at time t
    glm::dvec3 v(double t)
    {
        double theta = t2theta(t0+t);
        double r = p/(1 + e*cos(theta));
        double v = std::sqrt(mu*(2.0/r-1.0/a));
        return coordinates*glm::dvec3(-v*std::sin(theta), v*std::cos(theta), 0.0);
    }

    // position and velocity at time t (solves Kepler's equation only once)
    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.0/r-1.0/a));
        return std::make_pair(
            coordinates*glm::dvec3(r*std::cos(theta), r*std::sin(theta), 0.0),
            coordinates*glm::dvec3(-v*std::sin(theta), v*std::cos(theta), 0.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.0; // 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;

    // find basis vector of our coordinate frame
    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.0-e*e);
    result.mu = mu;

    result.n = std::sqrt((std::abs(result.a*result.a*result.a))/mu);
    result.T = 2.0*PI*result.n;
    result.theta0 = theta;
    result.t0 = result.theta2t(theta);

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

    return result;
}

// calculates energy (which should be preserved)
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;
}

// Runge-Kutta-Nyström integrator
template<class F>
std::pair<glm::dvec3, glm::dvec3> rknv(std::pair<glm::dvec3, glm::dvec3> state, const F &f, double h)
{
    glm::dvec3 r = state.first;
    glm::dvec3 v = state.second;
    glm::dvec3 k[4];

    k[0] = f(r, v);
    k[1] = f(r + 1./2.*h*v + h*h*1./8.*k[0], v + 1./2.*h*k[0]);
    k[2] = f(r + 1./2.*h*v + h*h*1./8.*k[0], v + 1./2.*h*k[1]);
    k[3] = f(r + h*v + h*h*1./2.*k[2], v + h*k[2]);

    r += h*v + h*h*1./6.*(k[0] + k[1] + k[2]);
    v += h*1./6.*(k[0] + 2.*k[1] + 2.*k[2] + k[3]);

    return std::make_pair(r, v);
}

// acceleration ftor for the numerical integrator
struct acceleration {
    acceleration(double M) : M(M) { }
    glm::dvec3 operator()(glm::dvec3 r, glm::dvec3) const
    {
        double mu = G*M;
        double rr = glm::length(r);
        return -mu*r/(rr*rr*rr);
    }
private:
    double M;
};

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

    std::cout << std::setw(20) << "dt [s]" << std::setw(20) << "distance [m]" << std::setw(20) << "rel energy error" << std::endl;

    // sweep over timesteps
    for(double dt =7*24*60*60;dt>1.0;dt/=2.)
    {
        auto state = std::make_pair(glm::dvec3(r_earth, 0, 0), glm::dvec3(0, 1.0*v_earth, 0));
        Orbit orbit = calcElements(state.first, state.second, M_sun);

        double E0 = energy(state, M_sun);

        double t = 0.;
        acceleration a(M_sun);
        for(;t<orbit.T;t+=dt) // simulate for one orbital period (orbit.T)
        {
            state = rknv(state, a, dt);
        }
        auto state2 = orbit.state(t);

        // compare numeric solution with closed form orbit
        std::cout << std::setw(20) << dt << std::setw(20) << glm::distance(state.first, state2.first) << std::setw(20) << (energy(state, M_sun) - E0)/E0 << std::endl;
    }
}