Computing Kepler's elements with Haskell

kepler.hs

module Main where

import Debug.Trace
import Control.Monad (forM_)

epsilon = 1.0e-10 :: Double

add (x1, y1, z1) (x2, y2, z2) = (x1+x2, y1+y2, z1+z2)
sub (x1, y1, z1) (x2, y2, z2) = (x1-x2, y1-y2, z1-z2)
mul (x1, y1, z1) (x2, y2, z2) = (x1*x2, y1*y2, z1*z2)
div (x1, y1, z1) (x2, y2, z2) = (x1/x2, y1/y2, z1/z2)
scale c (x1, y1, z1) = (c*x1, c*y1, c*z1)

dot (x1, y1, z1) (x2, y2, z2) = x1*x2 + y1*y2 + z1*z2
cross (x1, y1, z1) (x2, y2, z2) = (y1*z2 - z1*y2, -(x1*z2 - z1*x2), x1*y2 - y1*x2)

mag v = sqrt (dot v v)
unit v = scale (1.0/mag v) v

vector_product (x, y, z) v = (dot x v, dot y v, dot z v)

sign x = if x < 0.0 then -1.0 else 1.0

period (_, e, n, _, _, _)
    | e < 1.0 = 2.0 * pi / n
    | otherwise = inf
    where
    inf = 1.0/0.0

meanAnomaly (_, e, n, m0, t0, _) t =
    m0 + n * (t - t0)

meanToEccentricAnomaly e mean
    | e < epsilon = circularAnomaly
    | abs (e - 1.0) < epsilon = parabolicAnomaly
    | e < 1.0 = eccentricAnomaly'
    | e > 1.0 = eccentricAnomaly'
    where
    circularAnomaly =
        mean
    parabolicAnomaly =
        -- D^3 / 3.0 + D - M = 0
        -- the only (always) real-valued root of D is: [wolfram alpha]
        (num / denom)**(1.0/3.0) - (denom / num)**(1.0/3.0)
        where
        num = 2.0
        denom = sqrt (9.0 * mean^2 + 4.0) - 3.0 * mean
    eccentricAnomaly' =
        iterate f fmean mean steps
        where
        steps
            | e < 0.3 = 5
            | e < 0.9 = 6
            | e < 1.0 = 8
            | e > 1.0 = 30
        f x
            | e < 0.3 =
                mean + e * sin x
            | e < 0.9 =
                x + (mean + e * sin x - x) / (1.0 - e * cos x )
            | e < 1.0 =
                let s = e * sin x
                    c = e * cos x
                    f0 = x - s - mean
                    f1 = 1.0 - c
                    f2 = s
                in
                    x + (-5.0) * f0 / (f1 + sign(f1) * sqrt(abs(16.0 * f1 * f1 - 20.0 * f0 * f2)))
            | e > 1.0 =
                let s = e * sinh x
                    c = e * cosh x
                    f0 = s - x - mean
                    f1 = c - 1.0
                    f2 = s
                in
                    x + (-5.0) * f0 / (f1 + sign(f1) * sqrt(abs(16.0 * f1 * f1 - 20.0 * f0 * f2)))
        fmean x
            | e < 1.0 = x + e * sin x
            | e > 1.0 = x + e * sinh x
        sign x
            | x < 0.0 = -1.0
            | otherwise = 1.0
        iterate _ _ x 0 = x
        iterate f fmean x n =
            if abs (fmean x - mean) < epsilon then x else iterate f fmean (f x) (n-1)

eccentricToMeanAnomaly e eccentric
    | e < epsilon =
        eccentric
    | abs (e - 1.0) < epsilon =
        eccentric^3 / 3.0 + eccentric
    | e < 1.0 =
        eccentric - e * sin eccentric
    | e > 1.0 =
        e * sinh eccentric - eccentric

eccentricToTrueAnomaly e eccentric
    | e < epsilon =
        eccentric
    | abs (e - 1.0) < epsilon =
        2.0 * atan eccentric
    | e < 1.0 =
        atan2 (sqrt (1.0 - e^2) * sin eccentric) (cos eccentric - e)
    | e > 1.0 =
        2.0 * atan (sqrt ((e + 1.0) / (e - 1.0)) * tanh (eccentric/2.0) )

trueToEccentricAnomaly e true
    | e < epsilon =
        true
    | abs (e - 1.0) < epsilon =
        atan (true / 2.0)
    | e < 1.0 =
        atan2 (sqrt (1.0 - e^2) * sin true) (cos true + e)
    | e > 1.0 =
        sign true * acosh ((e + cos true) / (1.0 + e * cos true))

meanToTrueAnomaly e = eccentricToTrueAnomaly e . meanToEccentricAnomaly e
trueToMeanAnomaly e = eccentricToMeanAnomaly e . trueToEccentricAnomaly e

anomalies elems@(_, e, _, _, _, _) t =
    (mean, eccentric, true)
    where
    mean = meanAnomaly elems t
    eccentric = meanToEccentricAnomaly e mean
    true = eccentricToTrueAnomaly e eccentric

orbitMatrix (i, an, arg) =
    (
        (
            (cos arg * cos an) - (sin arg * sin an * cos i),
            -(cos arg * sin an * cos i) - (sin arg * cos an),
            (sin an * sin i)
        ),
        (
            (sin arg * cos an * cos i) + (cos arg * sin an),
            (cos arg * cos an * cos i) - (sin arg * sin an),
            -(cos an * sin i)
        ),
        (
            sin arg * sin i,
            cos arg * sin i,
            cos i
        )
    )

planarState (p, e, n, _, _, _) (mean, eccentric, true)
    | e < epsilon = circular
    | abs (e - 1.0) < epsilon = parabolic
    | e < 1.0 = elliptic
    | e > 1.0 = hyperbolic
    where
    circular =
        ((x, y, 0.0), (vx, vy, 0.0))
        where
        x = p * cos eccentric
        y = p * sin eccentric
        vx = -v * sin eccentric
        vy = v * cos eccentric
        v = p * n
    parabolic =
        ((x, y, 0.0), (vx, vy, 0.0))
        where
        r = p / (1 + cos true)
        x = r * cos true
        y = r * sin true
        dx = (-p * sin true) / (1.0 + cos true)^2
        dy = p / (1.0 + cos true)
        d = sqrt (dx^2 + dy^2)
        vx = v * dx / d
        vy = v * dy / d
        v = n * sqrt (p^2 / 2.0 * (1.0 + cos true))
    a = p / (1.0 - e^2)
    elliptic =
        ((x, y, 0.0), (vx, vy, 0.0))
        where
        b = a * sqrt (1.0 - e^2)
        x = a * (cos eccentric - e)
        y = b * sin eccentric
        edot = n / (1.0 - e * cos eccentric)
        vx = -a * sin eccentric * edot
        vy =  b * cos eccentric * edot
    hyperbolic =
        ((x, y, 0.0), (vx, vy, 0.0))
        where
        --x = -a * (e - cosh eccentric)
        --y = -a * sqrt(e^2 - 1.0) * sinh eccentric
        r = p / (1.0 + e * cos true)
        x = r * cos true
        y = r * sin true
        mu = (n^2 * (-p / (1.0 - e^2))^3)
        --v = sqrt ((mu / p) * (1.0 + e^2 + 2.0 * cos true))
        v = sqrt (mu * ((2.0 / r) - (1.0 / a)))
        dx = -p * sin true / (1.0 + e * cos true)^2
        dy = p * (e + cos true) / (1.0 + e * cos true)^2
        d = sqrt (dx^2 + dy^2)
        vx = v * dx / d
        vy = v * dy / d


stateToElements gM gm r v t0 =
    trajectory
    where
    trajectory
        | mag v < epsilon || mag r < epsilon || mag k < epsilon =
            linearTrajectory
        | otherwise =
            conicTrajectory
    linearTrajectory =
        undefined

    -- standard gravitational parameter
    mu = gM + gm
    -- specific relative angular momentum, direction = orbital plane normal
    k = cross r v
    -- eccentricity vector, length = eccentricity, direction = to periapsis
    evec = sub (scale (-1.0/mu) (cross k v)) (unit r)
    e = mag evec
    -- line of nodes, length = zero for equatorial, direction = to ascending node
    -- nodes = cross (0.0, 0.0, 1.0) k
    nodes = let (kx, ky, _) = k in (-ky, kx, 0.0)

    conicTrajectory
        | e < epsilon = circularTrajectory
        | abs (e - 1.0) < epsilon = parabolicTrajectory
        | e < 1.0 = ellipticTrajectory
        | e > 1.0 = hyperbolicTrajectory
        | otherwise = undefined
    circularTrajectory =
        trace "circular" ((radius, e, n, true, t0, (i, an, arg)), (true, true, true))
        where
        radius = mu / (dot v v)
        n = sqrt (mu / radius^3)
        -- true anomaly measured from line of nodes -pi..pi
        true
            | mag nodes < epsilon = -- equatorial orbit
                -- measure true anomaly from x-axis
                atan2 y x
            | otherwise =
                -sign (nodes `dot` v) * acos (unit r `dot` unit nodes)
            where
            (x, y, _) = r
        arg = 0.0
    parabolicTrajectory =
        trace "parabolic" ((p, e, n, mean, t0, (i, an, arg)), (true, mean, parabolic))
        where
        -- periapsis distance
        q = mu * (1.0 + cosT) / dot v v
        -- semi-latus rectum
        p = 2.0 * q
        n = sqrt (mu / (2.0 * q^3))
        parabolic = trueToEccentricAnomaly e true
        mean = eccentricToMeanAnomaly e parabolic
    ellipticTrajectory =
        trace "elliptic" ((p, e, n, mean, t0, (i, an, arg)), (true, mean, eccentric))
        where
        n = sqrt (mu / a^3)
        eccentric = sign true * acos cosE
        mean = eccentricToMeanAnomaly e eccentric
    hyperbolicTrajectory =
        trace "hyperbolic" ((p, e, n, mean, t0, (i, an, arg)), (true, mean, hyperbolic))
        where
        n = sqrt (mu / (-a)^3)
        hyperbolic = sign true * acosh cosE
        mean = eccentricToMeanAnomaly e hyperbolic
    -- semi-major axis
    -- a = 1.0 / (2.0 / mag r - dot v v / mu)
    a = p / (1.0 - e^2)
    -- semi-latus rectum of ellipse or hyperbola
    -- p = a * (1.0 - e^2)
    p = dot k k / mu

    -- cosine of true anomaly (not well defined for circular orbits!)
    cosT = dot (unit evec) (unit r)
    -- true anomaly -pi..pi
    true = -sign (evec `dot` v) * acos cosT

    -- (hyperbolic) cosine of eccentric anomaly
    cosE = (e + cosT) / (1.0 + e * cosT)

    -- inclination 0..pi
    i = let (_, _, kz) = k in acos (kz / mag k)

    -- longitude of ascending node -pi..pi
    an
        | mag nodes < epsilon = 0.0     -- equatorial orbit
        | otherwise = atan2 ny nx
        where
        (nx, ny, _) = nodes

    -- argument of periapsis -pi..pi
    arg
        | mag nodes < epsilon =         -- equatorial orbit
            -- measure argument of periapsis from x axis
            atan2 ey ex
        | otherwise =
            sign ez * acos (unit nodes `dot` unit evec)
        where
        (ex, ey, ez) = evec


test = stateToElements 40.0e13 0.0 (7500.0e3, 0.0, 0.0) (0.0, 7877.0, -2000.0) 0

main = do
    forM_ states write
    where
    mu = 40.0e13
    testElems = [
        --stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, 2.0 * vEsc, 0.0) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, 1.1 * vEsc, 0.0) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, vEsc, 0.0) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, 1.1 * vCirc, 0.0) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, vCirc, 0.0) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, 0.0, vCirc) 0.0,
        stateToElements mu 0.0 (0.0, q, 0.0) (0.0, 0.0, vCirc) 0.0,
        stateToElements mu 0.0 (q, 0.0, 0.0) (0.0, 0.75 * vCirc, 0.0) 0.0
        ]
    q = 7500.0e3
    vEsc = sqrt (2.0 * mu / q)
    vCirc = sqrt (mu / q)
    num = 64
    write ((x, y, z), (vx, vy, vz)) = do
        putStrLn $ show x ++ "\t" ++ show y ++ "\t" ++ show z ++ "\t" ++ show vx ++ "\t" ++ show vy ++ "\t" ++ show vz
    states = concat [map (state elems) times | (elems, _) <- testElems]
    times = [(m / num) | m <- [0..num]]
    state elems@(p, e, n, _, _, rot) t' =
        (vector_product matrix r, vector_product matrix v)
        where
        matrix = orbitMatrix rot
        (r, v) = planarState elems anoms
        t = (if e < 1.0 then t'*period else (t'*period - period / 2))
        anoms = anomalies elems t
        period
            | e < 1.0 = 2.0 * pi / n
            | otherwise = 60.0*3 * 24

-- Sources:
-- Karttunen: Johdatus taivaanmekaniikkaan, URSA
-- http://www.bogan.ca/orbits/kepler/orbteqtn.html
-- http://mmae.iit.edu/~mpeet/Classes/MMAE441/Spacecraft/441Lecture17.pdf
-- NOTE: the two sources have different definitions for parabolic anomaly, mean anomaly
-- NOTE: and mean motion of parabolic trajectories