Examples.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

-- \| Hamilton example suite
--
-- See: https://github.com/mstksg/hamilton#example-app-runner
--
-- Or just run with:
--
-- > $ hamilton-examples --help
-- > $ hamilton-examples [EXAMPLE] --help

import Control.Concurrent
import Control.Monad
import Data.Bifunctor
import Data.Finite
import Data.Foldable
import Data.IORef
import Data.List
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as M
import qualified Data.Vector as VV
import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable.Sized as VS
import GHC.TypeLits
import Graphics.Vty hiding ((<|>))
import Graphics.Vty.CrossPlatform (mkVty)
import Numeric.Hamilton
import Numeric.LinearAlgebra.Static hiding (dim, (<>))
import Numeric.LinearAlgebra.Static.Vector
import Options.Applicative
import qualified Prettyprinter as PP
import System.Exit
import Text.Printf
import Text.Read

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 = [grVec 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)) = grVec <$> [xs, 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 = [grVec 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)) = grVec <$> [xs, 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 + grVec w]
    c0 :: Config 3
    c0 = Cfg (vec3 0 x0 0) (vec3 1 0 (-0.5))

bezier ::
  forall n.
  KnownNat (1 + n) =>
  V.Vector (1 + n) (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 = [grVec xs]
    c0 :: Config 1
    c0 = Cfg (0.5 :: R 1) (0.25 :: R 1)

newtype 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 defaultConfig

  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")
                      . VS.toList
                      . rVec
                      . phsPositions
                      $ p
                  , printf "d<%s>/dt: <%s>" qVec
                      . intercalate ", "
                      . map (printf "%.4f")
                      . VS.toList
                      . rVec
                      . 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 x
    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
  { rrZero :: Double
  -- ^ Where on the screen (0 to 1) to place the other axis
  , rrRatio :: Double
  -- ^ Ratio of height of a terminal character to width
  }
  deriving (Show)

data PlotRange
  = PXY (Double, Double) (Double, Double)
  | PX (Double, Double) RangeRatio
  | PY RangeRatio (Double, Double)

plot ::
  -- | display bounds
  (Int, Int) ->
  PlotRange ->
  -- | points to plot
  [(V2 Double, (Char, Attr))] ->
  [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 p i = case scale p of
      V2 pX pY ->
        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
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V1 #-}
#endif

type V2 = V.Vector 2
pattern V2 :: a -> a -> V2 a
pattern V2 x y <- (V.toList -> [x, y])
  where
    V2 x y = V.fromTuple (x, y)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V2 #-}
#endif

pattern V3 :: a -> a -> a -> V.Vector 3 a
pattern V3 x y z <- (V.toList -> [x, y, z])
  where
    V3 x y z = V.fromTuple (x, y, z)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V3 #-}
#endif

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 = V.fromTuple (x, y, z, a)
#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE V4 #-}
#endif

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 (1 + n), Applicative f, Num a) =>
  V.Vector (1 + n) (f a) ->
  a ->
  f a
bezierCurve ps t =
  foldl' (liftA2 (+)) (pure 0)
    . V.imap
      ( \i ->
          let i' = fromIntegral i
           in fmap (* (fromIntegral (n' `choose` i') * (1 - t) ^ (n' - i') * t ^ i))
      )
    $ ps
  where
    n' :: Int
    n' = fromIntegral (maxBound :: Finite (1 + n))
    choose :: Int -> Int -> Int
    n `choose` k = factorial n `div` (factorial (n - k) * factorial k)
    factorial :: Int -> Int
    factorial m = product [1 .. m]

deriving instance Ord Color