Full Code of mstksg/hamilton for AI

Repository: mstksg/hamilton
Branch: master
Commit: 618c8d193c85
Files: 15
Total size: 59.1 KB

Directory structure:
gitextract_9tsek9fv/

├── .envrc
├── .git-blame-ignore-revs
├── .github/
│   └── workflows/
│       └── flake-ci.yml
├── .gitignore
├── .travis.yml
├── CHANGELOG.md
├── LICENSE
├── README.md
├── Setup.hs
├── app/
│   └── Examples.hs
├── flake.nix
├── fourmolu.yaml
├── hamilton.cabal
├── src/
│   └── Numeric/
│       └── Hamilton.hs
└── test/
    └── Spec.hs

================================================
FILE CONTENTS
================================================

================================================
FILE: .envrc
================================================
watch_file *.cabal
use flake


================================================
FILE: .git-blame-ignore-revs
================================================
# fourmolu
2f6ca862e31ddf33c86c8fb2660d7e0194b6fc2e


================================================
FILE: .github/workflows/flake-ci.yml
================================================
name: "Flake CI"
on:
  pull_request:
  push:
jobs:
  checks:
    runs-on: ubuntu-latest
    steps:
    - name: Free Disk Space
      uses: insightsengineering/free-disk-space@v1.1.0
    - uses: actions/checkout@v3
    - uses: webfactory/ssh-agent@v0.9.0
      with:
        ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
    - uses: cachix/install-nix-action@v22
      with:
        nix_path: nixpkgs=channel:nixos-unstable
        github_access_token: ${{ secrets.GITHUB_TOKEN }}
        extra_nix_config: |
          trusted-public-keys = hydra.iohk.io:f/Ea+s+dFdN+3Y/G+FDgSq+a5NEWhJGzdjvKNGv0/EQ= cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY= loony-tools:pr9m4BkM/5/eSTZlkQyRt57Jz7OMBxNSUiMC4FkcNfk=
          allow-import-from-derivation = true
          auto-optimise-store = true
          substituters = https://hydra.iohk.io https://cache.nixos.org/ https://cache.iog.io https://cache.zw3rk.com https://mstksg.cachix.org
    - uses: cachix/cachix-action@v13
      with:
        name: mstksg
        authToken: '${{ secrets.CACHIX_AUTH_TOKEN }}'
    - run: nix flake check --show-trace

  cache:
    runs-on: ubuntu-latest
    steps:
    - name: Free Disk Space
      uses: insightsengineering/free-disk-space@v1.1.0
    - uses: actions/checkout@v4.1.1
    - uses: webfactory/ssh-agent@v0.9.0
      with:
        ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
    - uses: cachix/install-nix-action@v22
      with:
        nix_path: nixpkgs=channel:nixos-unstable
        github_access_token: ${{ secrets.GITHUB_TOKEN }}
        extra_nix_config: |
          trusted-public-keys = hydra.iohk.io:f/Ea+s+dFdN+3Y/G+FDgSq+a5NEWhJGzdjvKNGv0/EQ= cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY= loony-tools:pr9m4BkM/5/eSTZlkQyRt57Jz7OMBxNSUiMC4FkcNfk=
          allow-import-from-derivation = true
          auto-optimise-store = true
          substituters = https://hydra.iohk.io https://cache.nixos.org/ https://cache.iog.io https://cache.zw3rk.com https://mstksg.cachix.org
    - uses: cachix/cachix-action@v13
      with:
        name: mstksg
        authToken: '${{ secrets.CACHIX_AUTH_TOKEN }}'
    - run: nix build --show-trace
    - run: nix develop --show-trace

  every-compiler:
    runs-on: ubuntu-latest
    steps:
    - name: Free Disk Space
      uses: insightsengineering/free-disk-space@v1.1.0
    - uses: actions/checkout@v3
    - uses: webfactory/ssh-agent@v0.9.0
      with:
        ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
    - uses: cachix/install-nix-action@v22
      with:
        nix_path: nixpkgs=channel:nixos-unstable
        github_access_token: ${{ secrets.GITHUB_TOKEN }}
        extra_nix_config: |
          trusted-public-keys = hydra.iohk.io:f/Ea+s+dFdN+3Y/G+FDgSq+a5NEWhJGzdjvKNGv0/EQ= cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY= loony-tools:pr9m4BkM/5/eSTZlkQyRt57Jz7OMBxNSUiMC4FkcNfk=
          allow-import-from-derivation = true
          auto-optimise-store = true
          substituters = https://hydra.iohk.io https://cache.nixos.org/ https://cache.iog.io https://cache.zw3rk.com https://mstksg.cachix.org
    - uses: cachix/cachix-action@v13
      with:
        name: mstksg
        authToken: '${{ secrets.CACHIX_AUTH_TOKEN }}'
    - run: nix build .#everyCompiler



================================================
FILE: .gitignore
================================================
/.stack-work
/src/highlight.js
/src/style.css
/dist-newstyle
/.direnv


================================================
FILE: .travis.yml
================================================
script:
- |
  set -ex
  case "$BUILD" in
    stack)
      stack --no-terminal $ARGS test --bench --no-run-benchmarks --haddock --no-haddock-deps
      ;;
    cabal)
      cabal install --enable-tests --enable-benchmarks --force-reinstalls --ghc-options=-O0 --reorder-goals --max-backjumps=-1 $CABALARGS $PACKAGES

      ORIGDIR=$(pwd)
      for dir in $PACKAGES
      do
        cd $dir
        cabal check || [ "$CABALVER" == "1.16" ]
        cabal sdist
        PKGVER=$(cabal info . | awk '{print $2;exit}')
        SRC_TGZ=$PKGVER.tar.gz
        cd dist
        tar zxfv "$SRC_TGZ"
        cd "$PKGVER"
        cabal configure --enable-tests --ghc-options -O0
        cabal build
        if [ "$CABALVER" = "1.16" ] || [ "$CABALVER" = "1.18" ]; then
          cabal test
        else
          cabal test --show-details=streaming --log=/dev/stdout
        fi
        cd $ORIGDIR
      done
      ;;
  esac
  set +ex
install:
- echo "$(ghc --version) [$(ghc --print-project-git-commit-id 2> /dev/null || echo
  '?')]"
- if [ -f configure.ac ]; then autoreconf -i; fi
- |
  set -ex
  case "$BUILD" in
    stack)
      # Add in extra-deps for older snapshots, as necessary
      #
      # This is disabled by default, as relying on the solver like this can
      # make builds unreliable. Instead, if you have this situation, it's
      # recommended that you maintain multiple stack-lts-X.yaml files.

      #stack --no-terminal --install-ghc $ARGS test --bench --dry-run || ( \
      #  stack --no-terminal $ARGS build cabal-install && \
      #  stack --no-terminal $ARGS solver --update-config)

      # Build the dependencies
      stack --no-terminal --install-ghc $ARGS test --bench --only-dependencies
      ;;
    cabal)
      cabal --version
      travis_retry cabal update

      # Get the list of packages from the stack.yaml file. Note that
      # this will also implicitly run hpack as necessary to generate
      # the .cabal files needed by cabal-install.
      PACKAGES=$(stack --install-ghc query locals | grep '^ *path' | sed 's@^ *path:@@')

      cabal install --only-dependencies --enable-tests --enable-benchmarks --force-reinstalls --ghc-options=-O0 --reorder-goals --max-backjumps=-1 $CABALARGS $PACKAGES
      ;;
  esac
  set +ex
jobs:
  include:
  - env: BUILD=cabal GHCVER=8.0.2 CABALVER=1.24 HAPPYVER=1.19.5 ALEXVER=3.1.7
    addons:
      apt:
        sources:
        - hvr-ghc
        packages:
        - cabal-install-1.24
        - ghc-8.0.2
        - happy-1.19.5
        - alex-3.1.7
    compiler: ': #GHC 8.0.2'
  - env: BUILD=cabal GHCVER=8.2.2 CABALVER=2.0 HAPPYVER=1.19.5 ALEXVER=3.1.7
    addons:
      apt:
        sources:
        - hvr-ghc
        packages:
        - cabal-install-2.0
        - ghc-8.2.2
        - happy-1.19.5
        - alex-3.1.7
    compiler: ': #GHC 8.2.2'
  - env: BUILD=cabal GHCVER=8.4.4 CABALVER=2.2 HAPPYVER=1.19.5 ALEXVER=3.1.7
    addons:
      apt:
        sources:
        - hvr-ghc
        packages:
        - cabal-install-2.2
        - ghc-8.4.4
        - happy-1.19.5
        - alex-3.1.7
    compiler: ': #GHC 8.4.4'
  - env: BUILD=cabal GHCVER=8.6.5 CABALVER=2.4 HAPPYVER=1.19.5 ALEXVER=3.1.7
    addons:
      apt:
        sources:
        - hvr-ghc
        packages:
        - cabal-install-2.4
        - ghc-8.6.5
        - happy-1.19.5
        - alex-3.1.7
    compiler: ': #GHC 8.6.5'
  - env: BUILD=cabal GHCVER=head  CABALVER=head HAPPYVER=1.19.5 ALEXVER=3.1.7
    addons:
      apt:
        sources:
        - hvr-ghc
        packages:
        - cabal-install-head
        - ghc-head
        - happy-1.19.5
        - alex-3.1.7
    compiler: ': #GHC HEAD'
  - env: BUILD=stack ARGS=""
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack default'
  - env: BUILD=stack ARGS="--resolver lts-9"
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack 8.0.2'
  - env: BUILD=stack ARGS="--resolver lts-11"
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack 8.2.2'
  - env: BUILD=stack ARGS="--resolver lts-12"
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack 8.4.4'
  - env: BUILD=stack ARGS="--resolver lts-14"
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack 8.6.5'
  - env: BUILD=stack ARGS="--resolver nightly"
    addons:
      apt:
        packages:
        - libgmp-dev
    compiler: ': #stack nightly'
  - env: BUILD=stack ARGS=""
    os: osx
    compiler: ': #stack default osx'
  - env: BUILD=stack ARGS="--resolver lts-9"
    os: osx
    compiler: ': #stack 8.0.2 osx'
  - env: BUILD=stack ARGS="--resolver lts-11"
    os: osx
    compiler: ': #stack 8.2.2 osx'
  - env: BUILD=stack ARGS="--resolver lts-12"
    os: osx
    compiler: ': #stack 8.4.4 osx'
  - env: BUILD=stack ARGS="--resolver lts-14"
    os: osx
    compiler: ': #stack 8.6.5 osx'
  - env: BUILD=stack ARGS="--resolver nightly"
    os: osx
    compiler: ': #stack nightly osx'
  allow_failures:
  - env: BUILD=cabal GHCVER=head  CABALVER=head HAPPYVER=1.19.5 ALEXVER=3.1.7
  - env: BUILD=stack ARGS="--resolver nightly"
os: linux
cache:
  directories:
  - $HOME/.ghc
  - $HOME/.cabal
  - $HOME/.stack
  - $TRAVIS_BUILD_DIR/.stack-work
before_install:
- unset CC
- CABALARGS=""
- if [ "x$GHCVER" = "xhead" ]; then CABALARGS=--allow-newer; fi
- export PATH=/opt/ghc/$GHCVER/bin:/opt/cabal/$CABALVER/bin:$HOME/.local/bin:/opt/alex/$ALEXVER/bin:/opt/happy/$HAPPYVER/bin:$HOME/.cabal/bin:$PATH
- mkdir -p ~/.local/bin
- |
  if [ `uname` = "Darwin" ]
  then
    travis_retry curl --insecure -L https://get.haskellstack.org/stable/osx-x86_64.tar.gz | tar xz --strip-components=1 --include '*/stack' -C ~/.local/bin
  else
    travis_retry curl -L https://get.haskellstack.org/stable/linux-x86_64.tar.gz | tar xz --wildcards --strip-components=1 -C ~/.local/bin '*/stack'
  fi

  # Use the more reliable S3 mirror of Hackage
  mkdir -p $HOME/.cabal
  echo 'remote-repo: hackage.haskell.org:http://hackage.fpcomplete.com/' > $HOME/.cabal/config
  echo 'remote-repo-cache: $HOME/.cabal/packages' >> $HOME/.cabal/config
- |
  if [ `uname` = "Darwin" ]
  then
    brew update
    brew install gsl
  fi
language: generic



================================================
FILE: CHANGELOG.md
================================================
Changelog
=========

Version 0.1.0.4
---------------

*Unreleased*

<https://github.com/mstksg/hamilton/releases/tag/v0.1.0.4>

*Internal*

*   Rewrote more internal functions using *hmatrix-vector-sized*, which should
    yield performance benefits.

Version 0.1.0.3
---------------

*Mar 20, 2018*

<https://github.com/mstksg/hamilton/releases/tag/v0.1.0.3>

*   Compatibility with *base-4.11.0.0* and GHC 8.4
*   Compatibility with *vector-sized-1.0.0.0*

*Internal*

*   Internal conversion functions refactored using *hmatrix-vector-sized*,
    *hessianF*.

Version 0.1.0.2
---------------

*Jan 21, 2018*

<https://github.com/mstksg/hamilton/releases/tag/v0.1.0.2>

*   Compatibility with *typelits-witneses-0.3.0.0*

Version 0.1.0.1
---------------

*Aug 17, 2017*

<https://github.com/mstksg/hamilton/releases/tag/v0.1.0.1>

*   Removed `Num` instance in the examples file, to account for
    *vector-sized*'s new `Num` instances.
*   COMPLETE pragmas for examples file.

Version 0.1.0.0
---------------

*Nov 27, 2016*

<https://github.com/mstksg/hamilton/releases/tag/v0.1.0.0>

*   Initial release.


================================================
FILE: LICENSE
================================================
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.

================================================
FILE: README.md
================================================
Hamilton
========

[![Build Status](https://travis-ci.org/mstksg/hamilton.svg?branch=master)](https://travis-ci.org/mstksg/hamilton)

Simulate physics on arbitrary coordinate systems using [automatic
differentiation][ad] and [Hamiltonian mechanics][].  State only an arbitrary
parameterization of your system and a potential energy function!

[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 a [blog post][] I wrote on this, and also the [hackage documentation][] and the
[example runner user guide][user guide] (and its [source][example runner]).

[blog post]: https://blog.jle.im/entry/introducing-the-hamilton-library.html
[hackage documentation]: http://hackage.haskell.org/package/hamilton
[example runner]: https://github.com/mstksg/hamilton/blob/master/app/Examples.hs
[user guide]: https://github.com/mstksg/hamilton#example-app-runner

### Full Example

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


-- some helper patterns to pattern match on sized vectors
pattern V2 :: a -> a -> V.Vector 2 a
pattern V2 x y <- (V.toList->[x,y])
  where
    V2 x y = fromJust (V.fromList [x,y])

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])
~~~

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 = phsPositions <$> evolution
~~~

And the position in the underlying cartesian space as:

~~~haskell
positions' :: [R 4]
positions' = underlyingPos doublePendulum <$> positions
~~~

Example App runner
------------------

*([Source][example 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, assuming `m1` is `100` and
`m2` is `1`:

~~~haskell
twoBody :: System 4 2
twoBody = mkSystem masses coordinates potential
  where
    masses :: R 4
    masses = vec4 100 100 1 1
    coordinates
        :: Floating a
        => V.Vector 2 a
        -> V.Vector 4 a
    coordinates (V2 r θ) = V4 (r1 * cos θ) (r1 * sin θ)
                              (r2 * cos θ) (r2 * sin θ)
      where
        r1 =   r *   1 / 101
        r2 = - r * 100 / 101
    potential
        :: Num a
        => V.Vector 4 a
        -> a
    potential (V2 r _) = - 100 / r
~~~

Potential improvements
----------------------

*   **Time-dependent systems**:  Shouldn't be an problem in theory/math; just
    add a time parameter before all of the functions.  This opens a lot of
    doors, like deriving inertial forces for free (like the famous Coriolis
    force and centrifugal force).

    The only thing is that it makes the API pretty inconvenient, because it'd
    require all of the functions to also take a time parameter.  Of course, the
    easy way out/ugly solution would be to just offer two versions of the same
    function (one for time-independent systems and one for time-dependent
    systems.  But this is un-ideal.

*   Velocity-dependent potentials:  Would give us the ability to model systems
    with velocity-dependent Lagrangians like a charged particle in an
    electromagnetic field, and also dissipative systems, like systems with
    friction (dependent on `signum v`) and linear & quadratic wind resistance.

    This issue is much harder, theoretically.  It involves inverting arbitrary
    functions `forall a. RealFloat a => V.Vector n a -> V.Vector m a`.  It
    might be possible with the help of some
    [bidirectionalization techniques][bff-pearl], but I can't get the [bff][]
    package to compile, and I'm not sure how to get [bff-mono][] to work with
    numeric functions.

    If anyone is familiar with bidirectionalization techniques and is willing
    to help out, please send me a message or open an issue! :)

[bff-pearl]: https://pdfs.semanticscholar.org/5f0d/ef02dbd96e102be9104d2ceb728d2a2a5beb.pdf
[bff]: http://hackage.haskell.org/package/bff
[bff-mono]: http://hackage.haskell.org/package/bff-mono



================================================
FILE: Setup.hs
================================================
import Distribution.Simple

main = defaultMain


================================================
FILE: app/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


================================================
FILE: flake.nix
================================================
{
  description = "Basic Haskell Project Flake";
  inputs = {
    haskellProjectFlake.url = "github:mstksg/haskell-project-flake";
    nixpkgs.follows = "haskellProjectFlake/nixpkgs";
  };
  outputs =
    { self
    , nixpkgs
    , flake-utils
    , haskellProjectFlake
    }:
    flake-utils.lib.eachDefaultSystem (system:
    let
      name = "hamilton";
      pkgs = import nixpkgs {
        inherit system;
        overlays = [ haskellProjectFlake.overlays."${system}".default ];
      };
      project-flake = pkgs.haskell-project-flake
        {
          inherit name;
          src = ./.;
          excludeCompilerMajors = [ "ghc911" ];
          defaultCompiler = "ghc982";
        };
    in
    {
      packages = project-flake.packages;
      apps = project-flake.apps;
      checks = project-flake.checks;
      devShells = project-flake.devShells;
      legacyPackages."${name}" = project-flake;
    }
    );
}



================================================
FILE: fourmolu.yaml
================================================
column-limit: 100
comma-style: leading
fixities: []
function-arrows: trailing
haddock-style: single-line
haddock-style-module: null
import-export-style: diff-friendly
in-style: right-align
indent-wheres: true
indentation: 2
let-style: inline
newlines-between-decls: 1
record-break-space: true
reexports: []
respectful: true
single-constraint-parens: never
unicode: detect


================================================
FILE: hamilton.cabal
================================================
cabal-version:      1.12

-- This file has been generated from package.yaml by hpack version 0.36.1.
--
-- see: https://github.com/sol/hpack

name:               hamilton
version:            0.1.0.4
synopsis:
  Physics on generalized coordinate systems using Hamiltonian Mechanics and AD

description:
  See README.md (or read online at <https://github.com/mstksg/hamilton#readme>)

category:           Physics
homepage:           https://github.com/mstksg/hamilton#readme
bug-reports:        https://github.com/mstksg/hamilton/issues
author:             Justin Le
maintainer:         justin@jle.im
copyright:          (c) Justin Le 2024
license:            BSD3
license-file:       LICENSE
build-type:         Simple
tested-with:        GHC >=8.10
extra-source-files:
  CHANGELOG.md
  README.md

source-repository head
  type:     git
  location: https://github.com/mstksg/hamilton

library
  exposed-modules:  Numeric.Hamilton
  other-modules:    Paths_hamilton
  hs-source-dirs:   src
  ghc-options:      -Wall -O2
  build-depends:
      ad
    , base                       >=4.9   && <5
    , ghc-typelits-natnormalise
    , hmatrix                    >=0.18
    , hmatrix-gsl                >=0.18
    , hmatrix-vector-sized       >=0.1.3
    , typelits-witnesses         >=0.2.3
    , vector-sized               >=1.0

  default-language: Haskell2010

executable hamilton-examples
  main-is:          Examples.hs
  other-modules:    Paths_hamilton
  hs-source-dirs:   app
  ghc-options:      -Wall -O2 -threaded -rtsopts -with-rtsopts=-N
  build-depends:
      base                   >=4.9   && <5
    , containers
    , finite-typelits
    , ghc-typelits-knownnat
    , hamilton
    , hmatrix                >=0.18
    , hmatrix-vector-sized   >=0.1.3
    , optparse-applicative   >=0.13
    , prettyprinter
    , vector
    , vector-sized           >=1.0
    , vty
    , vty-crossplatform

  default-language: Haskell2010


================================================
FILE: src/Numeric/Hamilton.hs
================================================
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

-- |
-- Module      : Numeric.Hamilton
-- Description : Hamiltonian dynamics for physical systems on generalized
--               coordinates using automatic differentiation
-- Copyright   : (c) Justin Le 2024
-- 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 <https://github.com/mstksg/hamilton#readme> 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
import qualified Data.Vector.Generic.Sized as VG
import qualified Data.Vector.Sized as V
import GHC.Generics (Generic)
import GHC.TypeLits
import GHC.TypeLits.Compare
import Numeric.AD
import Numeric.GSL.ODE
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Static as H
import Numeric.LinearAlgebra.Static.Vector

-- | 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 ::
    { cfgPositions :: !(R n)
    -- ^ The current values ("positions") of each of the @n@
    -- generalized coordinates
    , cfgVelocities :: !(R n)
    -- ^ The current rate of changes ("velocities") of each of the
    -- @n@ generalized coordinates
    } ->
    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 ::
    { phsPositions :: !(R n)
    -- ^ The current values ("positions") of each of the @n@
    -- generalized coordinates.
    , phsMomenta :: !(R n)
    -- ^ The current conjugate momenta ("momentums") to each of
    -- the @n@ generalized coordinates
    } ->
    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
    , _sysHessian :: R n -> V.Vector n (L m n)
    , _sysPotential :: R n -> Double
    , _sysPotentialGrad :: R n -> R n
    } ->
    System m n

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

vec2l ::
  (KnownNat m, KnownNat n) =>
  V.Vector m (V.Vector n Double) ->
  L m n
vec2l = rowsL . fmap gvecR

-- | 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) =>
  -- | 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.
  R m ->
  -- | 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 -> V.Vector m a) ->
  -- | The potential energy of the system as a function of
  -- the generalized coordinate space's positions.
  (forall a. RealFloat a => V.Vector n a -> a) ->
  System m n
mkSystem m f u =
  Sys
    { _sysInertia = m
    , _sysCoords = gvecR . f . grVec
    , _sysJacobian = tr . vec2l . jacobianT f . grVec
    , _sysHessian = tr2 . fmap vec2l . hessianF f . grVec
    , _sysPotential = u . grVec
    , _sysPotentialGrad = gvecR . grad u . grVec
    }
  where
    tr2 ::
      forall o.
      (KnownNat n, KnownNat o) =>
      V.Vector m (L n o) ->
      V.Vector n (L m o)
    tr2 = fmap rowsL . traverse lRows
    {-# INLINE tr2 #-}

-- | 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) =>
  -- | 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.
  R m ->
  -- | 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 -> V.Vector m a) ->
  -- | The potential energy of the system as a function of
  -- the underlying cartesian coordinate space's positions.
  (forall a. RealFloat a => V.Vector m a -> a) ->
  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 H.<> diag _sysInertia H.<> 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.
--
-- Computed using the maths derived in
-- <https://blog.jle.im/entry/hamiltonian-dynamics-in-haskell.html>.
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
    jmj = trj H.<> mm H.<> j
    ijmj = inv jmj
    dTdq = gvecR
      . flip fmap (_sysHessian phsPositions)
      $ \djdq ->
        -(phsMomenta <.> ijmj #> trj #> mm #> djdq #> ijmj #> phsMomenta)
    dHdp = ijmj #> phsMomenta
    dHdq = dTdq + _sysPotentialGrad phsPositions

-- | Step a system through phase space over over a single timestep.
stepHam ::
  forall m n.
  (KnownNat m, KnownNat n) =>
  -- | timestep to step through
  Double ->
  -- | system to simulate
  System m n ->
  -- | initial state, in phase space
  Phase n ->
  Phase n
stepHam r s p =
  evolveHam @m @n @2 s p (V.fromTuple (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 to simulate
  System m n ->
  -- | initial state, in phase space
  Phase n ->
  -- | desired solution times
  [Double] ->
  [Phase n]
evolveHam' _ _ [] = []
evolveHam' s p0 ts = V.withSizedList (toList ts') $ \(v :: V.Vector s Double) ->
  case Proxy @2 %<=? Proxy @s of
    LE Refl ->
      (if l1 then toList . V.tail @(s - 1) else toList) $
        evolveHam s p0 v
    NLE{} -> 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 to simulate
  System m n ->
  -- | initial state, in phase space
  Phase n ->
  -- | desired solution times
  V.Vector s Double ->
  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 =
      (\(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 = case traverse create . LA.takesV [n, n] $ v of
      Just [pP, pM] -> Phs pP pM
      _ -> error "evolveHam: internal error"

-- | 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 to simulate
  System m n ->
  -- | initial state, in configuration space
  Config n ->
  -- | desired solution times
  [Double] ->
  [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 to simulate
  System m n ->
  -- | initial state, in configuration space
  Config n ->
  -- | desired solution times
  V.Vector s Double ->
  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) =>
  -- | timestep to step through
  Double ->
  -- | system to simulate
  System m n ->
  -- | initial state, in phase space
  Config n ->
  Config n
stepHamC r s = fromPhase s . stepHam r s . toPhase s


================================================
FILE: test/Spec.hs
================================================
main :: IO ()
main = putStrLn "Test suite not yet implemented"