Orbital calculations

Orbits.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}

module Orbits where

import Units
import Vec

data OrbitalElements
    = OrbitalElements
        { specificAngularMomentum :: !AngularMomentum
        , inclination             :: Inclination
        , rightAscensionOfAN      :: RightAscension
        , eccentricity            :: !Eccentricity
        , argumentOfPeriapsis     :: ArgumentOfPeriapsis
        , trueAnomaly             :: TrueAnomaly
        }
    deriving Show

stVecToOrb :: GravitationalParam -> Position -> Velocity -> Maybe OrbitalElements
stVecToOrb (GravitationalParam mu) (Position r') (Velocity v') = orb where

    orb
        | v == 0 || r == 0 = Nothing
        | otherwise        = Just $ OrbitalElements
            { specificAngularMomentum = AngularMomentum h'
            , inclination             = Inclination i
            , rightAscensionOfAN      = RightAscension ra
            , eccentricity            = Eccentricity e
            , argumentOfPeriapsis     = ArgumentOfPeriapsis w
            , trueAnomaly             = TrueAnomaly ta
            }

    r = vabs r'
    v = vabs v'

    vr = (r' `dot` v') / r

    h' = r' `cross` v'
    h = vabs h'

    i = let (Vec3 _ _ hz) = h'
        in Angle $ acos (hz / h)

    n' = (Vec3 0 0 1) `cross` h'
    n = vabs n'

    ra =
        let (Vec3 nx ny _) = n'
            a              = acos (nx / n)
        in  Angle $ if ny >= 0
                then a
                else 2*pi - a

    e' =
        (1/mu) `mul` ((v' `cross` h') `sub` (mu `mul` unit r'))

    e =
        let k = mu - r * v * v
        in  (1 / mu) * sqrt
            ( (2 * mu - r * v * v) * r * vr * vr + k * k)

    w =
        let (Vec3 _ _ ez)     = e'
            a                 = acos $ (n' `dot` e') / (n * e)
        in  Angle $ if ez >= 0
                then a
                else 2*pi - a

    ta =
        let a = acos $ (1 / e) * (h * h / (mu * r) - 1)
        in Angle $ if vr >= 0
                then a
                else 2*pi - a

orbToStVec :: GravitationalParam -> OrbitalElements -> (Position, Velocity)
orbToStVec (GravitationalParam mu) (OrbitalElements {..}) = (Position r', Velocity v') where
    AngularMomentum     h' = specificAngularMomentum
    Eccentricity        e  = eccentricity
    RightAscension      ra = rightAscensionOfAN
    Inclination         i  = inclination
    ArgumentOfPeriapsis w  = argumentOfPeriapsis
    TrueAnomaly         ta = trueAnomaly

    h = vabs h'

    rpm = (h*h/mu) * (1/(1 + e * cos' ta))
    rp = Vec3  (rpm * cos' ta) (rpm * sin' ta) 0

    vpm = mu/h
    vp = Vec3 (vpm * (-sin' ta)) (vpm * (e + cos' ta)) 0

    cos'ra = cos' ra
    sin'ra = sin' ra

    cos'i = cos' i
    sin'i = sin' i

    cos'w = cos' w
    sin'w = sin' w

    qpx = Mat33
        (cos'ra*cos'w - sin'ra*sin'w*cos'i) ((-cos'ra)*sin'w - sin'ra*cos'i*cos'w)  (sin'ra*sin'i)
        (sin'ra*cos'w + cos'ra*cos'i*sin'w) ((-sin'ra)*sin'w + cos'ra*cos'i*cos'w)  ((-cos'ra)*sin'i)
        (sin'i*sin'w)                       (sin'i*cos'w)                           (cos'i)

    r' = qpx `mmul` rp
    v' = qpx `mmul` vp


trueToEccentric :: Eccentricity -> TrueAnomaly -> EccentricAnomaly
trueToEccentric (Eccentricity e) (TrueAnomaly ta)
    = EccentricAnomaly $ Angle $ atan2 (sqrt (1 - e*e) * sin' ta) (e + cos' ta)

eccentricToMean :: Eccentricity -> EccentricAnomaly -> MeanAnomaly
eccentricToMean (Eccentricity e) (EccentricAnomaly (Angle ea))
    = MeanAnomaly $ Angle $ ea - e * sin ea

meanToEccentric :: Scalar -> Eccentricity -> MeanAnomaly -> EccentricAnomaly
meanToEccentric epsilon (Eccentricity e) (MeanAnomaly (Angle ma))
    = EccentricAnomaly $ Angle $ ea where

        ea
            | e > 0.8   = newton pi
            | otherwise = newton ma

        newton !x
            | abs (x - x') < epsilon = x'
            | otherwise              = newton x'
            where
                x' = x - (x - e * sin x - ma) / (1 - e * cos x)

eccentricToTrue :: Eccentricity -> EccentricAnomaly -> TrueAnomaly
eccentricToTrue (Eccentricity e) (EccentricAnomaly (Angle ea))
    = TrueAnomaly $ Angle $ 2 * atan2 y x
    where
        x  = sqrt (1 - e) * cos (ea/2)
        y  = sqrt (1 + e) * sin (ea/2)