kepler_improved.py

import numpy as np
from astropy import units as u

from poliastro.bodies import Earth
from poliastro.core.angles import (
    E_to_nu,
    _kepler_equation,
    _kepler_equation_prime,
    nu_to_M,
)
from poliastro.core.elements import rv2coe
from poliastro.examples import iss
from poliastro.twobody.propagation import cowell
from poliastro.twobody.orbit import Orbit


def kepler_improved(k, r0, v0, tof):
    """ Solves the kepler problem by a non iteraiteve method.

    Parameters
    ----------
    k: float
        Gravitational parameter
    r0: vector 1x3
        Initial position vector
    v0: vector 1x3
        Initial velocity vector
    tof: float
        Time of flight

    Returns
    -------
    rf: vector 1x3
        Final position vector
    vf: vector 1x3
        Final velocity vector
    """

    # Convert to international unit system
    k = k.to(u.m ** 3 / u.s ** 2).value
    r0 = r0.to(u.m).value
    v0 = v0.to(u.m / u.s).value
    tof = tof.to(u.s).value

    # Solve first for eccentricity and mean anomaly
    p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0)
    M0 = nu_to_M(nu, ecc)
    a = p / (1 - ecc ** 2)
    n = np.sqrt(k / a ** 3)
    M = M0 + n * tof

    # Range between -pi and pi
    M = M % (2 * np.pi)
    if M > np.pi:
        M = -(2 * np.pi - M)

    # ------ ALGORITHM STARTS -------

    # Equation (20)
    alpha = (3 * np.pi ** 2 + 1.6 * (np.pi - np.abs(M)) / (1 + ecc)) / (np.pi ** 2 - 6)

    # Equation (5)
    d = 3 * (1 - ecc) + alpha * ecc

    # Equation (9)
    q = 2 * alpha * d * (1 - ecc) - M ** 2

    # Equation (10)
    r = 3 * alpha * d * (d - 1 + ecc) * M + M ** 3

    # Equation (14)
    w = (np.abs(r) + np.sqrt(q ** 3 + r ** 2)) ** (2 / 3)

    # Equation (15)
    E1 = (2 * r * w / (w ** 2 + w * q + q ** 2) + M) / d

    # Equation (26)
    f0 = _kepler_equation(E1, M, ecc)
    f1 = _kepler_equation_prime(E1, M, ecc)
    f2 = ecc * np.sin(E1)
    f3 = ecc * np.cos(E1)
    f4 = -f2

    # Equation (22)
    delta3 = -f0 / (f1 - 0.5 * f0 * f2 / f1)
    delta4 = -f0 / (f1 + 0.5 * delta3 * f2 + 1 / 6 * delta3 ** 2 * f3)
    delta5 = -f0 / (
        f1 + 0.5 * delta4 * f2 + 1 / 6 * delta4 ** 2 * f3 + 1 / 24 * delta4 ** 3 * f4
    )

    # Finally
    E5 = E1 + delta5
    nu = E_to_nu(E5, ecc)
    return nu * u.rad

def hard_orbit():
    """ Propagates the hard orbit. """
    r = [8.0e3, 1.0e3, 0.0] * u.km
    v = [-0.5, -0.5, 0.0] * u.km / u.s
    ss = Orbit.from_vectors(Earth, r, v)
    print(ss.classical())
    a, ecc, inc, raan, argp, nu = ss.classical()

    # Cowell propagator (integrates two-body ODE)
    ss_cowell = ss.propagate(1 * u.h, method=cowell)
    nu_expected = ss_cowell.nu
    print("Nu nu_expected:", nu_expected)
    print(ss_cowell.classical())

    # Kepler_improved propagator (Defined in the first script function)
    nu = kepler_improved(ss.attractor.k, ss.r, ss.v, 1 * u.h)
    ss_kepler_improved = Orbit.from_classical(Earth, a, ecc, inc, raan, argp, nu)
    print("Nu obtained:", nu)
    print(ss_kepler_improved.classical())
    
if __name__ == '__main__':
    hard_orbit()