BSD/MIT licensed

Rocket.hs

-- This is a Haskell implementation of aphyr's
-- "Clojure from the ground up - modeling"
--
-- http://aphyr.com/posts/312-clojure-from-the-ground-up-modeling
--
-- Compile this with:
-- 
--   $ ghc -O Rocket.hs

import GHC.Exts

data Cartesian = Cartesian { x :: Double,
                             y :: Double,
                             z :: Double } deriving (Show)

instance Num Cartesian where
    (Cartesian x1 y1 z1) + (Cartesian x2 y2 z2) = Cartesian (x1 + x2) (y1 + y2) (z1 + z2)
    (Cartesian x1 y1 z1) - (Cartesian x2 y2 z2) = Cartesian (x1 - x2) (y1 - y2) (z1 - z2)
    (Cartesian x1 y1 z1) * (Cartesian x2 y2 z2) = Cartesian (x1 * x2) (y1 * y2) (z1 * z2)
    abs (Cartesian x y z) = Cartesian (abs x) (abs y) (abs z)
    signum = id
    fromInteger _ = Cartesian 0 0 0

data Spherical = Spherical { r :: Double,
                             theta :: Double,
                             phi :: Double } deriving (Show)

data Craft = Craft {
    time :: Int,
    dryMass :: Double,
    fuelMass :: Double,
    isp :: Double,
    maxFuelRate :: Double,
    velocity :: Velocity,
    position :: Position,
    nextStage :: Maybe Craft } deriving (Show)

type Position = Cartesian 
type Velocity = Cartesian 
type Trajectory = [Craft]

initialSpaceCenterPosition :: Position
initialSpaceCenterPosition = Cartesian earthEquatorialRadius 0 0

initialSpaceCenterVelocity :: Velocity
initialSpaceCenterVelocity = Cartesian 0 earthEquatorialSpeed 0

magnitude :: Cartesian -> Double
magnitude (Cartesian x y z) = sqrt (x ^ 2 + y ^ 2 + z ^ 2)

scale :: Double -> Cartesian -> Cartesian
scale f (Cartesian x y z) = Cartesian (x * f) (y * f) (z * f)

unitVector :: Cartesian -> Cartesian
unitVector c = scale (1 / magnitude c) c

dotProduct :: Cartesian -> Cartesian -> Double
dotProduct (Cartesian x1 y1 z1) (Cartesian x2 y2 z2) =
    x1 * x2 + y1 * y2 + z1 * z2

projection :: Cartesian -> Cartesian -> Cartesian
projection a b = scale (dotProduct a c) c
    where c = unitVector b

rejection :: Cartesian -> Cartesian -> Cartesian
rejection a b = a - projection a b

cartesianToSpherical :: Cartesian -> Spherical
cartesianToSpherical c@(Cartesian x y z) =
    Spherical r (atan (y / x)) (acos (z / r))
        where r = magnitude c

sphericalToCartesian :: Spherical -> Cartesian
sphericalToCartesian (Spherical r theta phi) =
    Cartesian
       (r * cos theta * sin phi)
       (r * sin theta * sin phi)
       (r * cos phi)

-- Radius of the earth, in meters
earthEquatorialRadius = 6378137

-- Length of an earth day, in seconds.
earthDay = 86400

earthEquatorialSpeed = (2 * pi * earthEquatorialRadius) / earthDay

-- Acceleration of gravity in meters/s^2
g = -9.8

centaur = Craft {
    time=0,
    dryMass=2361,
    fuelMass=13897,
    isp=4354,
    maxFuelRate=13897 / 470,
    velocity=initialSpaceCenterVelocity,
    position=initialSpaceCenterPosition,
    nextStage=Nothing}

atlasV :: Craft -> Craft
atlasV next =
    Craft {
        time=0,
        dryMass=50050,
        fuelMass=284450,
        isp=3050,
        maxFuelRate=284450 / 252,
        velocity=initialSpaceCenterVelocity,
        position=initialSpaceCenterPosition,
        nextStage=Just next}

ascent = [0, 300]
circularization = [400, 1000]

orientation :: Craft -> Cartesian
orientation (Craft {time = t, velocity = v, position = p})
    | (head ascent <= t) && (t <= last ascent) = p
    | (head circularization <= t) && (t <= last circularization) = rejection v p
    | otherwise = v

fuelRate :: Craft -> Double
fuelRate (Craft {time = t, fuelMass = f, maxFuelRate = m})
    | f <= 0 = 0
    | (head ascent <= t) && (t <= last ascent) = m
    | (head circularization <= t) && (t <= last circularization) = m
    | otherwise = 0

stage :: Craft -> Craft
stage c@(Craft t _ f _ _ v p n)
    | f > 0 = c
    | otherwise = case n of
        Nothing -> c
        (Just (Craft _ dm fm isp mfr _ _ n')) ->
            Craft t dm fm isp mfr v p n'

thrust :: Craft -> Double
thrust c = fuelRate c * isp c

mass :: Craft -> Double
mass (Craft {dryMass = d, fuelMass = f}) = d + f

gravityForce :: Craft -> Cartesian
gravityForce c@(Craft {position = p}) =
    sphericalToCartesian (Spherical r theta phi)
        where
            (Spherical _ theta phi) = cartesianToSpherical p
            r = mass c * g

altitude :: Craft -> Double
altitude (Craft {position = p}) = r - earthEquatorialRadius
    where (Spherical r _ _) = cartesianToSpherical p

engineForce :: Craft -> Cartesian
engineForce c = scale (thrust c) (unitVector (orientation c)) 

totalForce :: Craft -> Cartesian
totalForce c = engineForce c + gravityForce c

acceleration :: Craft -> Cartesian
acceleration c = scale (1 / mass c) (totalForce c)

step :: Craft -> Int -> Craft
step c@(Craft t dm fm isp mfr v p n) dt =
    Craft t' dm fm' isp mfr v' p' n
        where
            t' = t + dt
            dt' = fromIntegral dt
            fm' = fm - (dt' * fuelRate c)
            p' = p + scale dt' v
            v' = v + scale dt' (acceleration c)

stagedStep :: Craft -> Int -> Craft
stagedStep c = step (stage c)

stepBackwards :: Int -> Craft -> Craft
stepBackwards dt c = stagedStep c dt

trajectory :: Int -> Craft -> Trajectory
trajectory dt = iterate (stepBackwards 1)

isAboveGround :: Craft -> Bool
isAboveGround c = 0 <= altitude c

flight :: Trajectory -> Trajectory
flight = takeWhile isAboveGround

crashed :: Trajectory -> Bool
crashed t = not (all isAboveGround (takeWhile f t))
    where
        f c = time c <= (10 * 3600)

crashTime :: Trajectory -> Int
crashTime t = case reverse (flight t) of
                (x:_) -> time x
                _ -> 0

apoapsis :: Trajectory -> Double
apoapsis t = maximum (map altitude (flight t))

apoapsisTime :: Trajectory -> Int
apoapsisTime [] = 0
apoapsisTime [t] = time t
apoapsisTime t = time (last (sortWith altitude (flight t)))

main = do
    let trj = trajectory 1 (atlasV centaur)
    let hasCrashed = crashed trj

    if hasCrashed then do
        putStrLn $ "Crashed at " ++ show (crashTime trj) ++ " seconds"
        putStrLn $ "Maximum altitude " ++ show (apoapsis trj)
            ++ " meters at " ++ show (apoapsisTime trj) ++ " seconds"
    else
      putStrLn "Reached orbit!"