hamilton-0.1.0.0: app/Examples.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Hamilton example suite
--
-- See: https://github.com/mstksg/hamilton#example-app-runner
--
-- Or just run with:
--
-- > $ hamtilton-examples --help
-- > $ hamtilton-examples [EXAMPLE] --help
--
import Control.Concurrent
import Control.Monad
import Data.Bifunctor
import Data.Foldable
import Data.IORef
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Proxy
import GHC.TypeLits
import Graphics.Vty hiding (Config, (<|>))
import Numeric.Hamilton
import Numeric.LinearAlgebra.Static hiding (dim, (<>))
import Options.Applicative
import System.Exit
import Text.Printf
import Text.Read
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import qualified Data.Vector as VV
import qualified Data.Vector.Generic.Sized as VG
import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable as VS
import qualified Text.PrettyPrint.ANSI.Leijen as PP
data SysExample where
SE :: (KnownNat m, KnownNat n)
=> { seName :: String
, seCoords :: V.Vector n String
, seSystem :: System m n
, seDraw :: R m -> [V2 Double]
, seInit :: Phase n
}
-> SysExample
pendulum :: Double -> Double -> SysExample
pendulum θ0 ω0 = SE "Single pendulum" (V1 "θ") s f (toPhase s c0)
where
s :: System 2 1
s = mkSystem' (vec2 1 1 ) -- masses
(\(V1 θ) -> V2 (sin θ) (0.5 - cos θ)) -- coordinates
(\(V2 _ y) -> y ) -- potential
f :: R 2 -> [V2 Double]
f xs = [r2vec xs]
c0 :: Config 1
c0 = Cfg (konst θ0 :: R 1) (konst ω0 :: R 1)
doublePendulum :: Double -> Double -> SysExample
doublePendulum m1 m2 = SE "Double pendulum" (V2 "θ1" "θ2") s f (toPhase s c0)
where
s :: System 4 2
s = mkSystem' (vec4 m1 m1 m2 m2) -- masses
(\(V2 θ1 θ2) -> V4 (sin θ1) (1 - cos θ1)
(sin θ1 + sin θ2/2) (1 - cos θ1 - cos θ2/2)
) -- coordinates
(\(V4 _ y1 _ y2) -> 5 * (realToFrac m1 * y1 + realToFrac m2 * y2))
-- potential
f :: R 4 -> [V2 Double]
f (split->(xs,ys))= [r2vec xs, r2vec ys]
c0 :: Config 2
c0 = Cfg (vec2 (pi/2) 0) (vec2 0 0)
room :: Double -> SysExample
room θ = SE "Room" (V2 "x" "y") s f (toPhase s c0)
where
s :: System 2 2
s = mkSystem (vec2 1 1) -- masses
id -- coordinates
(\(V2 x y) -> sum [ 2 * y -- gravity
, 1 - logistic (-1) 10 0.1 y -- bottom wall
, logistic 1 10 0.1 y -- top wall
, 1 - logistic (-2) 10 0.1 x -- left wall
, logistic 2 10 0.1 x -- right wall
]
) -- potential
f :: R 2 -> [V2 Double]
f xs = [r2vec xs]
c0 :: Config 2
c0 = Cfg (vec2 (-1) 0.25) (vec2 (cos θ) (sin θ))
twoBody :: Double -> Double -> Double -> SysExample
twoBody m1 m2 ω0 = SE "Two-Body" (V2 "r" "θ") s f (toPhase s c0)
where
mT :: Double
mT = m1 + m2
s :: System 4 2
s = mkSystem (vec4 m1 m1 m2 m2) -- masses
-- positions are calculated assuming (0,0) is the center
-- of mass
(\(V2 r θ) -> let r1 = r * realToFrac (-m2 / mT)
r2 = r * realToFrac (m1 / mT)
in V4 (r1 * cos θ) (r1 * sin θ)
(r2 * cos θ) (r2 * sin θ)
) -- coordinates
(\(V2 r _) -> - realToFrac (m1 * m2) / r) -- potential
f :: R 4 -> [V2 Double]
f (split->(xs,ys))= [r2vec xs, r2vec ys]
c0 :: Config 2
c0 = Cfg (vec2 2 0) (vec2 0 ω0)
spring
:: Double -> Double -> Double -> Double -> SysExample
spring mB mW k x0 = SE "Spring hanging from block" (V3 "r" "x" "θ") s f (toPhase s c0)
where
s :: System 3 3
s = mkSystem (vec3 mB mW mW) -- masses
(\(V3 r x θ) -> V3 r (r + (1 + x) * sin θ) ((1 + x) * (-cos θ))) -- coordinates
(\(V3 r x θ) -> realToFrac k * x**2 / 2 -- spring
+ (1 - logistic (-1.5) 25 0.1 r) -- left rail wall
+ ( logistic 1.5 25 0.1 r) -- right rail wall
+ realToFrac mB * ((1 + x) * (-cos θ)) -- gravity
)
f :: R 3 -> [V2 Double]
f (headTail->(b,w)) = [V2 b 1, V2 0 1 + r2vec w]
c0 :: Config 3
c0 = Cfg (vec3 0 x0 0) (vec3 1 0 (-0.5))
bezier
:: forall n. KnownNat n
=> V.Vector (n + 1) (V2 Double)
-> SysExample
bezier ps = SE "Bezier" (V1 "t") s f (toPhase s c0)
where
s :: System 2 1
s = mkSystem (vec2 1 1) -- masses
(\(V1 t) -> bezierCurve (fmap realToFrac <$> ps) t) -- coordinates
(\(V1 t) -> (1 - logistic 0 5 0.05 t) -- left wall
+ logistic 1 5 0.05 t -- right wall
)
f :: R 2 -> [V2 Double]
f xs = [r2vec xs]
c0 :: Config 1
c0 = Cfg (0.5 :: R 1) (0.25 :: R 1)
data ExampleOpts = EO { eoChoice :: SysExampleChoice }
data SysExampleChoice =
SECDoublePend Double Double
| SECPend Double Double
| SECRoom Double
| SECTwoBody Double Double Double
| SECSpring Double Double Double Double
| SECBezier (NE.NonEmpty (V2 Double))
parseEO :: Parser ExampleOpts
parseEO = EO <$> (parseSEC <|> pure (SECDoublePend 1 1))
parseSEC :: Parser SysExampleChoice
parseSEC = subparser . mconcat $
[ command "doublepend" $
info (helper <*> parseDoublePend)
(progDesc "Double pendulum (default)")
, command "pend" $
info (helper <*> parsePend )
(progDesc "Single pendulum")
, command "room" $
info (helper <*> parseRoom )
(progDesc "Ball in room, bouncing off of walls")
, command "twobody" $
info (helper <*> parseTwoBody )
(progDesc "Two-body graviational simulation. Note that bodies will only orbit if H < 0.")
, command "spring" $
info (helper <*> parseSpring )
(progDesc "A spring hanging from a block on a rail, holding up a mass. Block is constrained to bounce between -1.5 and 1.5.")
, command "bezier" $
info (helper <*> parseBezier )
(progDesc "Particle moving along a parameterized bezier curve")
, metavar "EXAMPLE"
]
where
parsePend
= SECPend <$> option auto ( long "angle"
<> short 'a'
<> metavar "ANGLE"
<> help "Intitial rightward angle (in degrees) of bob"
<> value 0
<> showDefault
)
<*> option auto ( long "vel"
<> short 'v'
<> metavar "VELOCITY"
<> help "Initial rightward angular velocity of bob"
<> value 1
<> showDefault
)
parseDoublePend
= SECDoublePend <$> option auto ( long "m1"
<> metavar "MASS"
<> help "Mass of first bob"
<> value 1
<> showDefault
)
<*> option auto ( long "m2"
<> metavar "MASS"
<> help "Mass of second bob"
<> value 1
<> showDefault
)
parseRoom
= SECRoom <$> option auto ( long "angle"
<> short 'a'
<> metavar "ANGLE"
<> help "Initial upward launch angle (in degrees) of object"
<> value 45
<> showDefault
)
parseTwoBody
= SECTwoBody <$> option auto ( long "m1"
<> metavar "MASS"
<> help "Mass of first body"
<> value 5
<> showDefault
)
<*> option auto ( long "m2"
<> metavar "MASS"
<> help "Mass of second body"
<> value 0.5
<> showDefault
)
<*> option auto ( long "vel"
<> short 'v'
<> metavar "VELOCITY"
<> help "Initial angular velocity of system"
<> value 0.5
<> showDefault
)
parseSpring
= SECSpring <$> option auto ( long "block"
<> short 'b'
<> metavar "MASS"
<> help "Mass of block on rail"
<> value 2
<> showDefault
)
<*> option auto ( long "weight"
<> short 'w'
<> metavar "MASS"
<> help "Mass of weight hanging from spring"
<> value 1
<> showDefault
)
<*> option auto ( short 'k'
<> metavar "NUM"
<> help "Spring constant / stiffness of spring"
<> value 10
<> showDefault
)
<*> option auto ( short 'x'
<> metavar "DIST"
<> help "Initial displacement of spring"
<> value 0.1
<> showDefault
)
parseBezier
= SECBezier <$> option f ( long "points"
<> short 'p'
<> metavar "POINTS"
<> help "List of control points (at least one), as tuples"
<> value (V2 (-1) (-1) NE.:| [V2 (-2) 1, V2 0 1, V2 1 (-1), V2 2 1])
<> showDefaultWith (show . map (\(V2 x y) -> (x, y)) . toList)
)
where f = eitherReader $ \s -> do
ps <- maybe (Left "Bad parse") Right
$ readMaybe s
maybe (Left "At least one control point required") Right
$ NE.nonEmpty (uncurry V2 <$> ps)
data SimOpts = SO { soZoom :: Double
, soRate :: Double
, soHist :: Int
}
deriving (Show)
data SimEvt = SEQuit
| SEZoom Double
| SERate Double
| SEHist Int
main :: IO ()
main = do
EO{..} <- execParser $ info (helper <*> parseEO)
( fullDesc
<> header "hamilton-examples - hamilton library example suite"
<> progDescDoc (Just descr)
)
vty <- mkVty =<< standardIOConfig
opts <- newIORef $ SO 0.5 1 25
t <- forkIO . loop vty opts $ case eoChoice of
SECDoublePend m1 m2 -> doublePendulum m1 m2
SECPend d0 ω0 -> pendulum (d0 / 180 * pi) ω0
SECRoom d0 -> room (d0 / 180 * pi)
SECTwoBody m1 m2 ω0 -> twoBody m1 m2 ω0
SECSpring mB mW k x0 -> spring mB mW k x0
SECBezier (p NE.:| ps) -> V.withSized (VV.fromList ps)
(bezier . V.cons p)
forever $ do
e <- nextEvent vty
forM_ (processEvt e) $ \case
SEQuit -> do
killThread t
shutdown vty
exitSuccess
SEZoom s ->
modifyIORef opts $ \o -> o { soZoom = soZoom o * s }
SERate r ->
modifyIORef opts $ \o -> o { soRate = soRate o * r }
SEHist h ->
modifyIORef opts $ \o -> o { soHist = soHist o + h }
where
fps :: Double
fps = 12
screenRatio :: Double
screenRatio = 2.1
ptAttrs :: [(Char, Color)]
ptAttrs = ptChars `zip` ptColors
where
ptColors = cycle [white,yellow,blue,red,green]
ptChars = cycle "o*+~"
loop :: Vty -> IORef SimOpts -> SysExample -> IO ()
loop vty oRef SE{..} = go M.empty seInit
where
qVec = intercalate "," . V.toList $ seCoords
go hists p = do
SO{..} <- readIORef oRef
let p' = stepHam (soRate / fps) seSystem p -- progress the simulation
xb = (- recip soZoom, recip soZoom)
infobox = vertCat . map (string defAttr) $
[ printf "[ %s ]" seName
, printf " <%s> : <%s>" qVec . intercalate ", "
. map (printf "%.4f") . r2list . phsPositions $ p
, printf "d<%s>/dt: <%s>" qVec . intercalate ", "
. map (printf "%.4f") . r2list . velocities seSystem $ p
, printf "KE: %.4f" . keP seSystem $ p
, printf "PE: %.4f" . pe seSystem . phsPositions $ p
, printf "H : %.4f" . hamiltonian seSystem $ p
, " "
, printf "rate: x%.2f <>" $ soRate
, printf "hist: % 5d []" $ soHist
, printf "zoom: x%.2f -+" $ soZoom
]
pts = (`zip` ptAttrs) . seDraw . underlyingPos seSystem . phsPositions
$ p
hists' = foldl' (\h (r, a) -> M.insertWith (addHist soHist) a [r] h) hists pts
dr <- displayBounds $ outputIface vty
update vty . picForLayers . (infobox:) . plot dr (PX xb (RR 0.5 screenRatio)) $
((second . second) (defAttr `withForeColor`) <$> pts)
++ (map (\((_,c),r) -> (r, ('.', defAttr `withForeColor` c)))
. concatMap sequence
. M.toList
$ hists'
)
threadDelay (round (1000000 / fps))
go hists' p'
addHist hl new old = take hl (new ++ old)
descr :: PP.Doc
descr = PP.vcat
[ "Run examples from the hamilton library example suite."
, "Use with [EXAMPLE] --help for more per-example options."
, ""
, "To adjust rate/history/zoom, use keys <>/[]/-+, respectively."
, ""
, "See: https://github.com/mstksg/hamilton#example-app-runner"
]
processEvt
:: Event -> Maybe SimEvt
processEvt = \case
EvKey KEsc [] -> Just SEQuit
EvKey (KChar 'c') [MCtrl] -> Just SEQuit
EvKey (KChar 'q') [] -> Just SEQuit
EvKey (KChar '+') [] -> Just $ SEZoom (sqrt 2)
EvKey (KChar '-') [] -> Just $ SEZoom (sqrt 0.5)
EvKey (KChar '>') [] -> Just $ SERate (sqrt 2)
EvKey (KChar '<') [] -> Just $ SERate (sqrt (1/2))
EvKey (KChar ']') [] -> Just $ SEHist 5
EvKey (KChar '[') [] -> Just $ SEHist (-5)
_ -> Nothing
data RangeRatio = RR { -- | Where on the screen (0 to 1) to place the other axis
rrZero :: Double
-- | Ratio of height of a terminal character to width
, rrRatio :: Double
}
deriving (Show)
data PlotRange = PXY (Double, Double) (Double, Double)
| PX (Double, Double) RangeRatio
| PY RangeRatio (Double, Double)
plot
:: (Int, Int) -- ^ display bounds
-> PlotRange
-> [(V2 Double, (Char, Attr))] -- ^ points to plot
-> [Image]
plot (wd,ht) pr = map (crop wd ht)
. (++ bgs)
. map (\(p, (c, a)) -> place EQ EQ p $ char a c)
where
wd' = fromIntegral wd
ht' = fromIntegral ht
((xmin, xmax), (ymin, ymax)) = mkRange (wd', ht') pr
origin = place EQ EQ (V2 0 0) $ char defAttr '+'
xaxis = place EQ EQ (V2 0 0) $ charFill defAttr '-' wd 1
yaxis = place EQ EQ (V2 0 0) $ charFill defAttr '|' 1 ht
xrange = xmax - xmin
yrange = ymax - ymin
bg = backgroundFill wd ht
scale (V2 pX pY) = V2 x y
where
x = round $ (pX - xmin) * (wd' / xrange)
y = round $ (pY - ymin) * (ht' / yrange)
place aX aY (scale->(V2 pX pY)) i
= translate (fAlign aX (imageWidth i))
(fAlign aY (imageHeight i))
. translate pX pY
$ i
labels = [ place LT EQ (V2 xmin 0) . string defAttr $ printf "%.2f" xmin
, place GT EQ (V2 xmax 0) . string defAttr $ printf "%.2f" xmax
, place EQ LT (V2 0 ymin) . string defAttr $ printf "%.2f" ymin
, place EQ GT (V2 0 ymax) . string defAttr $ printf "%.2f" ymax
]
bgs = labels ++ [origin, xaxis, yaxis, bg]
fAlign = \case
LT -> const 0
EQ -> negate . (`div` 2)
GT -> negate
mkRange
:: (Double, Double)
-> PlotRange
-> ((Double, Double), (Double, Double))
mkRange (wd, ht) = \case
PXY xb yb -> (xb, yb)
PX xb RR{..} ->
let yr = (uncurry (-) xb) * ht / wd * rrRatio
y0 = (rrZero - 1) * yr
in (xb, (y0, y0 + yr))
PY RR{..} yb ->
let xr = (uncurry (-) yb) * wd / ht / rrRatio
x0 = (rrZero - 1) * xr
in ((x0, x0 + xr), yb)
pattern V1 :: a -> V.Vector 1 a
pattern V1 x <- (V.head->x)
where
V1 x = V.singleton x
type V2 = V.Vector 2
pattern V2 :: a -> a -> V2 a
pattern V2 x y <- (V.toList->[x,y])
where
V2 x y = fromJust (V.fromList [x,y])
pattern V3 :: a -> a -> a -> V.Vector 3 a
pattern V3 x y z <- (V.toList->[x,y,z])
where
V3 x y z = fromJust (V.fromList [x,y,z])
pattern V4 :: a -> a -> a -> a -> V.Vector 4 a
pattern V4 x y z a <- (V.toList->[x,y,z,a])
where
V4 x y z a = fromJust (V.fromList [x,y,z,a])
r2list
:: KnownNat n
=> R n
-> [Double]
r2list = VS.toList . extract
r2vec
:: KnownNat n
=> R n
-> V.Vector n Double
r2vec = VG.convert . fromJust . VG.toSized . extract
logistic
:: Floating a => a -> a -> a -> a -> a
logistic pos ht width = \x -> ht / (1 + exp (- beta * (x - pos)))
where
beta = log (0.9 / (1 - 0.9)) / width
bezierCurve
:: forall n f a. (KnownNat n, Applicative f, Num a)
=> V.Vector (n + 1) (f a)
-> a
-> f a
bezierCurve ps t =
foldl' (liftA2 (+)) (pure 0)
. V.imap (\i -> fmap ((* (fromIntegral (n' `choose` i) * (1 - t)^(n' - i) * t^i))))
$ ps
where
n' :: Int
n' = fromInteger $ natVal (Proxy @n)
choose :: Int -> Int -> Int
n `choose` k = factorial n `div` (factorial (n - k) * factorial k)
factorial :: Int -> Int
factorial m = product [1..m]
instance (KnownNat n, Num a) => Num (V.Vector n a) where
(+) = liftA2 (+)
(-) = liftA2 (-)
(*) = liftA2 (*)
negate = fmap negate
abs = fmap abs
signum = fmap signum
fromInteger = pure . fromInteger
instance (KnownNat n, Fractional a) => Fractional (V.Vector n a) where
(/) = liftA2 (/)
recip = fmap recip
fromRational = pure . fromRational
deriving instance Ord Color