diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Justin Le (c) 2016 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Justin Le nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,227 @@ +Hamilton +======== + +Simulate physics on arbitrary coordinate systems using [automatic +differentiation][ad] and [Hamiltonian mechanics][]. + +[ad]: http://hackage.haskell.org/package/ad +[Hamiltonian mechanics]: https://en.wikipedia.org/wiki/Hamiltonian_mechanics + +For example, a simulating a [double pendulum system][dps] by simulating the +progression of the angles of each bob: + +[dps]: https://en.wikipedia.org/wiki/Double_pendulum + +[![My name is William Rowan Hamilton](http://i.imgur.com/Vaaa2EC.gif)][gifv] + +[gifv]: http://i.imgur.com/Vaaa2EC.gifv + +You only need: + +1. Your generalized coordinates (in this case, `θ1` and `θ2`), and equations + to convert them to cartesian coordinates of your objects: + + ~~~haskell + x1 = sin θ1 + y1 = -cos θ1 + x2 = sin θ1 + sin θ2 / 2 -- second pendulum is half-length + y2 = -cos θ1 - cos θ2 / 2 + ~~~ + +2. The masses/inertias of each of those cartesian coordinates (`m1` for `x1` + and `y1`, `m2` for `x2` and `y2`) + +3. A potential energy function for your objects: + + ~~~haskell + U = (m1 y1 + m2 y2) * g + ~~~ + +And that's it! Hamiltonian mechanics steps your generalized coordinates (`θ1` +and `θ2`) through time, without needing to do any simulation involving +`x1`/`y1`/`x2`/`y2`! And you don't need to worry about tension or any other +stuff like that. All you need is a description of your coordinate system +itself, and the potential energy! + +~~~haskell +doublePendulum :: System 4 2 +doublePendulum = + mkSystem' (vec4 m1 m1 m2 m2) -- masses + (\(V2 θ1 θ2) -> V4 (sin θ1) (-cos θ1) + (sin θ1 + sin θ2/2) (-cos θ1 - cos θ2/2) + ) -- coordinates + (\(V4 _ y1 _ y2) -> (m1 * y1 + m2 * y2) * g) + -- potential +~~~ + +Thanks to [~~Alexander~~ William Rowan Hamilton][WRH], we can express our +system parameterized by arbitrary coordinates and get back equations of motions +as first-order differential equations. This library solves those first-order +differential equations for you using automatic differentiation and some matrix +manipulation. + +[WRH]: https://www.youtube.com/watch?v=SZXHoWwBcDc + +See [documentation][] and [example runner][]. + +[documentation]: https://mstksg.github.io/hamilton/ +[example runner]: https://github.com/mstksg/hamilton/blob/master/app/Examples.hs + +### Full Exmaple + +Let's turn our double pendulum (with the second pendulum half as long) into an +actual running program. Let's say that `g = 5`, `m1 = 1`, and `m2 = 2`. + +First, the system: + +~~~haskell +import Numeric.LinearAlgebra.Static +import qualified Data.Vector.Sized as V + + +doublePendulum :: System 4 2 +doublePendulum = mkSystem' masses coordinates potential + where + masses :: R 4 + masses = vec4 1 1 2 2 + coordinates + :: Floating a + => V.Vector 2 a + -> V.Vector 4 a + coordinates (V2 θ1 θ2) = V4 (sin θ1) (-cos θ1) + (sin θ1 + sin θ2/2) (-cos θ1 - cos θ2/2) + potential + :: Num a + => V.Vector 4 a + -> a + potential (V4 _ y1 _ y2) = (y1 + 2 * y2) * 5 +~~~ + +Neat! Easy, right? + +Okay, now let's run it. Let's pick a starting configuration (state of the +system) of `θ1` and `θ2`: + +~~~haskell +config0 :: Config 2 +config0 = Cfg (vec2 1 0 ) -- initial positions + (vec2 0 0.5) -- initial velocities +~~~ + +Configurations are nice, but Hamiltonian dynamics is all about motion through +phase space, so let's convert this configuration-space representation of the +state into a phase-space representation of the state: + +~~~haskell +phase0 :: Phase 2 +phase0 = toPhase doublePendulum config0 +~~~ + +And now we can ask for the state of our system at any amount of points in time! + +~~~haskell +ghci> evolveHam doublePendulum phase0 [0,0.1 .. 1] +-- result: state of the system at times 0, 0.1, 0.2, 0.3 ... etc. +~~~ + +Or, if you want to run the system step-by-step: + + +~~~haskell +evolution :: [Phase 2] +evolution = iterate (stepHam 0.1 doublePendulum) phase0 +~~~ + +And you can get the position of the coordinates as: + +~~~haskell +positions :: [R 2] +positions = phsPos <$> evolution +~~~ + +And the position in the underlying cartesian space as: + +~~~hakell +positions' :: [R 4] +positions' = underlyingPos doublePendulum <$> positions +~~~ + +Example App runner +------------------ + +Installation: + +~~~bash +$ git clone https://github.com/mstksg/hamilton +$ cd hamilton +$ stack install +~~~ + +Usage: + +~~~bash +$ hamilton-examples [EXAMPLE] (options) +$ hamilton-examples --help +$ hamilton-examples [EXAMPLE] --help +~~~ + +The example runner is a command line application that plots the progression of +several example system through time. + + +| Example | Description | Coordinates | Options | +|--------------|------------------------------------------------------------|---------------------------------------------------------------------|---------------------------------------------------------------| +| `doublepend` | Double pendulum, described above | `θ1`, `θ2` (angles of bobs) | Masses of each bob | +| `pend` | Single pendulum | `θ` (angle of bob) | Initial angle and velocity of bob | +| `room` | Object bounding around walled room | `x`, `y` | Initial launch angle of object | +| `twobody` | Two gravitationally attracted bodies, described below | `r`, `θ` (distance between bodies, angle of rotation) | Masses of bodies and initial angular veocity | +| `spring` | Spring hanging from a block on a rail, holding up a weight | `r`, `x`, `θ` (position of block, spring compression, spring angle) | Masses of block, weight, spring constant, initial compression | +| `bezier` | Bead sliding at constant velocity along bezier curve | `t` (Bezier time parameter) | Control points for arbitrary bezier curve | + +Call with `--help` (or `[EXAMPLE] --help`) for more information. + +More examples +------------- + +### Two-body system under gravity + +[![The two-body solution](http://i.imgur.com/TDEHTcb.gif)][gifv2] + +[gifv2]: http://i.imgur.com/TDEHTcb.gifv + +1. The generalized coordinates are just: + + * `r`, the distance between the two bodies + * `θ`, the current angle of rotation + + ~~~haskell + x1 = m2/(m1+m2) * r * sin θ -- assuming (0,0) is the center of mass + y1 = m2/(m1+m2) * r * cos θ + x2 = -m1/(m1+m2) * r * sin θ + y2 = -m1/(m1+m2) * r * cos θ + ~~~ + +2. The masses/inertias are again `m1` for `x1` and `y1`, and `m2` for `x2` and + `y2` + +3. The potential energy function is the classic gravitational potential: + + ~~~haskell + U = - m1 * m2 / r + ~~~ + +And...that's all you need! + +Here is the actual code for the two-body system: + +~~~haskell +twoBody :: System 4 2 +twoBody = + mkSystem (vec4 m1 m1 m2 m2) -- masses + (\(V2 r θ) -> let r1 = r * m2 / (m1 + m2) + r2 = - r * m1 / (m1 + m2) + in V4 (r1 * cos θ) (r1 * sin θ) + (r2 * cos θ) (r2 * sin θ) + ) -- coordinates + (\(V2 r _) -> - m1 * m2 / r) -- potential +~~~ diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/app/Examples.hs b/app/Examples.hs new file mode 100644 --- /dev/null +++ b/app/Examples.hs @@ -0,0 +1,550 @@ +{-# 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 + diff --git a/hamilton.cabal b/hamilton.cabal new file mode 100644 --- /dev/null +++ b/hamilton.cabal @@ -0,0 +1,56 @@ +name: hamilton +version: 0.1.0.0 +synopsis: Physics on generalized coordinate systems using Hamiltonian Mechanics and AD +description: See README.md (or read online at ) +homepage: https://github.com/mstksg/hamilton +license: BSD3 +license-file: LICENSE +author: Justin Le +maintainer: justin@jle.im +copyright: (c) Justin Le 2016 +category: Physics +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Numeric.Hamilton + build-depends: base >= 4.7 && < 5 + , ad + , comonad + , free + , hmatrix >= 0.18 + , hmatrix-gsl >= 0.18 + , typelits-witnesses + , vector-sized >= 0.4.1 + ghc-options: -Wall + default-language: Haskell2010 + +executable hamilton-examples + hs-source-dirs: app + main-is: Examples.hs + ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall + build-depends: base + , ansi-wl-pprint + , containers + , hamilton + , hmatrix + , optparse-applicative >= 0.13 + , vector + , vector-sized + , vty + default-language: Haskell2010 + +-- test-suite hamilton-test +-- type: exitcode-stdio-1.0 +-- hs-source-dirs: test +-- main-is: Spec.hs +-- build-depends: base +-- , hamilton +-- ghc-options: -threaded -rtsopts -with-rtsopts=-N +-- default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/mstksg/hamilton diff --git a/src/Numeric/Hamilton.hs b/src/Numeric/Hamilton.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Hamilton.hs @@ -0,0 +1,506 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Numeric.Hamilton +-- Description : Hamiltonian dynamics for physical systems on generalized +-- coordinates using automatic differentiation +-- Copyright : (c) Justin Le 2016 +-- License : BSD-3 +-- Maintainer : justin@jle.im +-- Stability : unstable +-- Portability : portable +-- +-- Simulate physical systems on generalized/arbitrary coordinates using +-- Hamiltonian mechanics and automatic differentiation! +-- +-- See the for more +-- information on usage! +-- + +module Numeric.Hamilton + ( -- * Systems and states + -- ** Systems + System + , mkSystem + , mkSystem' + , underlyingPos + -- ** States + , Config(..) + , Phase(..) + , toPhase + , fromPhase + -- * State functions + , momenta + , velocities + , keC + , keP + , pe + , lagrangian + , hamiltonian + , hamEqs + -- * Simulating hamiltonian dynamics + -- ** Over phase space + , stepHam + , evolveHam + , evolveHam' + -- ** Over configuration space + -- | Convenience wrappers over the normal phase-space + -- steppers/simulators that allow you to provide input and expect + -- output in configuration space instead of in phase space. Note that + -- the simulation itself still runs in phase space, so these all + -- require conversions to and from phase space under the hood. + , stepHamC + , evolveHamC + , evolveHamC' + ) where + +import Control.Monad +import Data.Bifunctor +import Data.Foldable +import Data.Kind +import Data.Maybe +import Data.Proxy +import Data.Type.Equality hiding (sym) +import GHC.Generics (Generic) +import GHC.TypeLits +import GHC.TypeLits.Compare +import Numeric.AD +import Numeric.GSL.ODE +import Numeric.LinearAlgebra.Static +import qualified Control.Comonad as C +import qualified Control.Comonad.Cofree as C +import qualified Data.Vector.Generic.Sized as VG +import qualified Data.Vector.Sized as V +import qualified Numeric.LinearAlgebra as LA + +-- | Represents the full state of a system of @n@ generalized coordinates +-- in configuration space (informally, "positions and velocities") +-- +-- A configuration space representaiton is more directly "physically +-- meaningful" and intuitive/understandable to humans than a phase space +-- representation. However, it's much less mathematically ideal to work +-- with because of the lack of some neat underlying symmetries. +-- +-- You can convert a @'Config' n@ into a @'Phase' n@ (convert from +-- configuration space to phase space) for a given system with 'toPhase'. +-- This allows you to state your system in configuration space and then +-- convert it to phase space before handing it off to the hamiltonian +-- machinery. +data Config :: Nat -> Type where + Cfg :: { -- | The current values ("positions") of each of the @n@ + -- generalized coordinates + cfgPositions :: !(R n) + -- | The current rate of changes ("velocities") of each of the + -- @n@ generalized coordinates + , cfgVelocities :: !(R n) + } + -> Config n + deriving (Generic) + +deriving instance KnownNat n => Show (Config n) + +-- | Represents the full state of a system of @n@ generalized coordinates +-- in phase space (informally, "positions and momentums"). +-- +-- Phase space representations are much nicer to work with mathematically +-- because of some neat underlying symmetries. For one, positions and +-- momentums are "interchangeable" in a system; if you swap every +-- coordinate's positions with their momentums, and also swap them in the +-- equations of motions, you get the same system back. This isn't the case +-- with configuration space representations. +-- +-- A hamiltonian simulation basically describes the trajectory of each +-- coordinate through phase space, so this is the /state/ of the +-- simulation. However, configuration space representations are much more +-- understandable to humans, so it might be useful to give an initial state +-- in configuration space using 'Config', and then convert it to a 'Phase' +-- with 'toPhase'. +data Phase :: Nat -> Type where + Phs :: { -- | The current values ("positions") of each of the @n@ + -- generalized coordinates. + phsPositions :: !(R n) + -- | The current conjugate momenta ("momentums") to each of + -- the @n@ generalized coordinates + , phsMomenta :: !(R n) + } + -> Phase n + deriving (Generic) + +deriving instance KnownNat n => Show (Phase n) + +-- | Represents a physical system in which physics happens. A @'System' +-- m n@ is a system whose state described using @n@ generalized coordinates +-- (an "@n@-dimensional" system), where the underlying cartesian coordinate +-- space is @m@-dimensional. +-- +-- For the most part, you are supposed to be able to ignore @m@. @m@ is +-- only provided because it's useful when plotting/drawing the system with +-- a given state back in rectangular coordinates. (The only function that +-- use the @m@ at the moment is 'underlyingPos') +-- +-- A @'System' m n@'s state is described using a @'Config' n@ (which +-- describes the system in configuration space) or a @'Phase' n@ (which +-- describes the system in phase space). +data System :: Nat -> Nat -> Type where + Sys :: { _sysInertia :: R m + , _sysCoords :: R n -> R m + , _sysJacobian :: R n -> L m n + , _sysJacobian2 :: R n -> V.Vector m (Sym n) + , _sysPotential :: R n -> Double + , _sysPotentialGrad :: R n -> R n + } + -> System m n + +-- coordShift +-- :: (KnownNat m, KnownNat n, KnownNat o) +-- => (R o -> R n) +-- -> (R o -> L n o) +-- -> (R o -> V.Vector n (Sym o)) +-- -> System m n +-- -> System m o +-- coordShift c j j2 = \case +-- Sys i c0 j0 j20 p g -> Sys i (c0 . c) +-- ((<>) <$> j0 . c <*> j) +-- ((\d -> fmap _) <$> j2 <*> j20 . c) +-- p g + +-- | Converts the position of generalized coordinates of a system to the +-- coordinates of the system's underlying cartesian coordinate system. +-- Useful for plotting/drawing the system in cartesian space. +underlyingPos + :: System m n + -> R n + -> R m +underlyingPos = _sysCoords + +-- | The potential energy of a system, given the position in the +-- generalized coordinates of the system. +pe :: System m n + -> R n + -> Double +pe = _sysPotential + +vec2r + :: KnownNat n => V.Vector n Double -> R n +vec2r = fromJust . create . VG.fromSized . VG.convert + +r2vec + :: KnownNat n => R n -> V.Vector n Double +r2vec = VG.convert . fromJust . VG.toSized . extract + +vec2l + :: (KnownNat m, KnownNat n) + => V.Vector m (V.Vector n Double) + -> L m n +vec2l = fromJust . (\rs -> withRows rs exactDims) . toList . fmap vec2r + +-- l2vec +-- :: (KnownNat m, KnownNat n) +-- => L m n +-- -> V.Vector m (V.Vector n Double) +-- l2vec = fromJust . V.fromList . map r2vec . toRows + +-- | Create a system with @n@ generalized coordinates by describing its +-- coordinate space (by a function from the generalized coordinates to the +-- underlying cartesian coordinates), the inertia of each of those +-- underlying coordinates, and the pontential energy function. +-- +-- The potential energy function is expressed in terms of the genearlized +-- coordinate space's positions. +mkSystem + :: forall m n. (KnownNat m, KnownNat n) + => R m -- ^ The "inertia" of each of the @m@ coordinates + -- in the underlying cartesian space of the system. This + -- should be mass for linear coordinates and rotational + -- inertia for angular coordinates. + -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a) + -- ^ Conversion function to convert points in the + -- generalized coordinate space to the underlying cartesian + -- space of the system. + -> (forall a. RealFloat a => V.Vector n a -> a) + -- ^ The potential energy of the system as a function of + -- the generalized coordinate space's positions. + -> System m n +mkSystem m f u = + Sys m + (vec2r . f . r2vec) + (tr . vec2l . jacobianT f . r2vec) + (fmap (sym . vec2l . j2 . C.hoistCofree VG.convert) + . VG.convert + . jacobians f + . r2vec + ) + (u . r2vec) + (vec2r . grad u . r2vec) + where + j2 :: C.Cofree (V.Vector n) Double + -> V.Vector n (V.Vector n Double) + j2 = fmap (fmap C.extract . C.unwrap) . C.unwrap + +-- | Convenience wrapper over 'mkSystem' that allows you to specify the +-- potential energy function in terms of the underlying cartesian +-- coordinate space. +mkSystem' + :: forall m n. (KnownNat m, KnownNat n) + => R m -- ^ The "inertia" of each of the @m@ coordinates + -- in the underlying cartesian space of the system. This + -- should be mass for linear coordinates and rotational + -- inertia for angular coordinates. + -> (forall a. RealFloat a => V.Vector n a -> V.Vector m a) + -- ^ Conversion function to convert points in the + -- generalized coordinate space to the underlying cartesian + -- space of the system. + -> (forall a. RealFloat a => V.Vector m a -> a) + -- ^ The potential energy of the system as a function of + -- the underlying cartesian coordinate space's positions. + -> System m n +mkSystem' m f u = mkSystem m f (u . f) + + +-- | Compute the generalized momenta conjugate to each generalized +-- coordinate of a system by giving the configuration-space state of the +-- system. +-- +-- Note that getting the momenta from a @'Phase' n@ involves just using +-- 'phsMomenta'. +momenta + :: (KnownNat m, KnownNat n) + => System m n + -> Config n + -> R n +momenta Sys{..} Cfg{..} = tr j #> diag _sysInertia #> j #> cfgVelocities + where + j = _sysJacobian cfgPositions + +-- | Convert a configuration-space representaiton of the state of the +-- system to a phase-space representation. +-- +-- Useful because the hamiltonian simulations use 'Phase' as its working +-- state, but 'Config' is a much more human-understandable and intuitive +-- representation. This allows you to state your starting state in +-- configuration space and convert to phase space for your simulation to +-- use. +toPhase + :: (KnownNat m, KnownNat n) + => System m n + -> Config n + -> Phase n +toPhase s = Phs <$> cfgPositions <*> momenta s + +-- | The kinetic energy of a system, given the system's state in +-- configuration space. +keC :: (KnownNat m, KnownNat n) + => System m n + -> Config n + -> Double +keC s = do + vs <- cfgVelocities + ps <- momenta s + return $ (vs <.> ps) / 2 + +-- | The Lagrangian of a system (the difference between the kinetic energy +-- and the potential energy), given the system's state in configuration +-- space. +lagrangian + :: (KnownNat m, KnownNat n) + => System m n + -> Config n + -> Double +lagrangian s = do + t <- keC s + u <- pe s . cfgPositions + return (t - u) + +-- | Compute the rate of change of each generalized coordinate by giving +-- the state of the system in phase space. +-- +-- Note that getting the velocities from a @'Config' n@ involves just using +-- 'cfgVelocities'. +velocities + :: (KnownNat m, KnownNat n) + => System m n + -> Phase n + -> R n +velocities Sys{..} Phs{..} = inv jmj #> phsMomenta + where + j = _sysJacobian phsPositions + jmj = tr j <> diag _sysInertia <> j + +-- | Invert 'toPhase' and convert a description of a system's state in +-- phase space to a description of the system's state in configuration +-- space. +-- +-- Possibly useful for showing the phase space representation of a system's +-- state in a more human-readable/human-understandable way. +fromPhase + :: (KnownNat m, KnownNat n) + => System m n + -> Phase n + -> Config n +fromPhase s = Cfg <$> phsPositions <*> velocities s + +-- | The kinetic energy of a system, given the system's state in +-- phase space. +keP :: (KnownNat m, KnownNat n) + => System m n + -> Phase n + -> Double +keP s = do + ps <- phsMomenta + vs <- velocities s + return $ (vs <.> ps) / 2 + +-- | The Hamiltonian of a system (the sum of kinetic energy and the +-- potential energy), given the system's state in phase space. +hamiltonian + :: (KnownNat m, KnownNat n) + => System m n + -> Phase n + -> Double +hamiltonian s = do + t <- keP s + u <- pe s . phsPositions + return (t + u) + +-- | The "hamiltonian equations" for a given system at a given state in +-- phase space. Returns the rate of change of the positions and +-- conjugate momenta, which can be used to progress the simulation through +-- time. +hamEqs + :: (KnownNat m, KnownNat n) + => System m n + -> Phase n + -> (R n, R n) +hamEqs Sys{..} Phs{..} = (dHdp, -dHdq) + where + mm = diag _sysInertia + j = _sysJacobian phsPositions + trj = tr j + j' = unSym <$> _sysJacobian2 phsPositions + jmj = trj <> mm <> j + ijmj = inv jmj + dTdq = vec2r + . flip fmap (tr2 j') $ \djdq -> + -phsMomenta <.> ijmj #> trj #> mm #> djdq #> ijmj #> phsMomenta + dHdp = ijmj #> phsMomenta + dHdq = dTdq + _sysPotentialGrad phsPositions + +tr2 + :: (KnownNat m, KnownNat n, KnownNat o) + => V.Vector m (L n o) + -> V.Vector n (L m o) +tr2 = fmap (fromJust . (\rs -> withRows rs exactDims) . toList) + . sequenceA + . fmap (fromJust . V.fromList . toRows) + +-- | Step a system through phase space over over a single timestep. +stepHam + :: forall m n. (KnownNat m, KnownNat n) + => Double -- ^ timestep to step through + -> System m n -- ^ system to simulate + -> Phase n -- ^ initial state, in phase space + -> Phase n +stepHam r s p = evolveHam @m @n @2 s p (fromJust $ V.fromList [0, r]) + `V.unsafeIndex` 1 + +-- | Evolve a system using a hamiltonian stepper, with the given initial +-- phase space state. +-- +-- Desired solution times provided as a list instead of a sized 'V.Vector'. +-- The output list should be the same length as the input list. +evolveHam' + :: forall m n. (KnownNat m, KnownNat n) + => System m n -- ^ system to simulate + -> Phase n -- ^ initial state, in phase space + -> [Double] -- ^ desired solution times + -> [Phase n] +evolveHam' _ _ [] = [] +evolveHam' s p0 ts = V.withSizedList (toList ts') $ \(v :: V.Vector s Double) -> + case (Proxy %<=? Proxy) :: (2 :<=? s) of + LE Refl -> (if l1 then tail else id) + . toList + $ evolveHam s p0 v + NLE Refl -> error "evolveHam': Internal error" + where + (l1, ts') = case ts of + [x] -> (True , [0,x]) + _ -> (False, ts ) + +-- | Evolve a system using a hamiltonian stepper, with the given initial +-- phase space state. +evolveHam + :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s) + => System m n -- ^ system to simulate + -> Phase n -- ^ initial state, in phase space + -> V.Vector s Double -- ^ desired solution times + -> V.Vector s (Phase n) +evolveHam s p0 ts = fmap toPs . fromJust . V.fromList . LA.toRows + $ odeSolveV RKf45 hi eps eps (const f) (fromPs p0) ts' + where + hi = (V.unsafeIndex ts 1 - V.unsafeIndex ts 0) / 100 + eps = 1.49012e-08 + f :: LA.Vector Double -> LA.Vector Double + f = uncurry (\p m -> LA.vjoin [p,m]) + . join bimap extract . hamEqs s . toPs + ts' = VG.fromSized . VG.convert $ ts + n = fromInteger $ natVal (Proxy @n) + fromPs :: Phase n -> LA.Vector Double + fromPs p = LA.vjoin . map extract $ [phsPositions p, phsMomenta p] + toPs :: LA.Vector Double -> Phase n + toPs v = Phs pP pM + where + Just [pP, pM] = traverse create . LA.takesV [n, n] $ v + +-- | A convenience wrapper for 'evolveHam'' that works on configuration +-- space states instead of phase space states. +-- +-- Note that the simulation itself still runs in phase space; this function +-- just abstracts over converting to and from phase space for the inputs +-- and outputs. +evolveHamC' + :: forall m n. (KnownNat m, KnownNat n) + => System m n -- ^ system to simulate + -> Config n -- ^ initial state, in configuration space + -> [Double] -- ^ desired solution times + -> [Config n] +evolveHamC' s c0 = fmap (fromPhase s) . evolveHam' s (toPhase s c0) + +-- | A convenience wrapper for 'evolveHam' that works on configuration +-- space states instead of phase space states. +-- +-- Note that the simulation itself still runs in phase space; this function +-- just abstracts over converting to and from phase space for the inputs +-- and outputs. +evolveHamC + :: forall m n s. (KnownNat m, KnownNat n, KnownNat s, 2 <= s) + => System m n -- ^ system to simulate + -> Config n -- ^ initial state, in configuration space + -> V.Vector s Double -- ^ desired solution times + -> V.Vector s (Config n) +evolveHamC s c0 = fmap (fromPhase s) . evolveHam s (toPhase s c0) + +-- | Step a system through configuration space over over a single timestep. +-- +-- Note that the simulation itself still runs in phase space; this function +-- just abstracts over converting to and from phase space for the input +-- and output. +stepHamC + :: forall m n. (KnownNat m, KnownNat n) + => Double -- ^ timestep to step through + -> System m n -- ^ system to simulate + -> Config n -- ^ initial state, in phase space + -> Config n +stepHamC r s = fromPhase s . stepHam r s . toPhase s +