Skip to content

Instantly share code, notes, and snippets.

@zmactep

zmactep/Integrator.hs

Last active October 31, 2017 12:06
Download ZIP

Revisions

  1. zmactep revised this gist Oct 31, 2017. 1 changed file with 38 additions and 11 deletions.
    49 changes: 38 additions & 11 deletions Integrator.hs
    View file
    Original file line number Diff line number Diff line change
    @@ -5,6 +5,7 @@ module Integrator where

    import Data.Array.Accelerate
    import Data.Array.Accelerate.Linear
    import Data.Array.Accelerate.Control.Lens

    type R = Float
    type V3R = V3 R
    @@ -13,6 +14,7 @@ type Velocity = V3 R
    type Force = V3 R
    type Mass = R
    type Timestep = R
    type Duration = R

    class ForceField f where
    forces :: f -> Acc (Vector Mass)
    @@ -24,25 +26,50 @@ class Integrator i where
    integrate :: ForceField f => i -> f
    -> Exp Timestep
    -> Acc (Vector Mass)
    -> Acc (Vector Position)
    -> Acc (Vector Position)
    -> Acc (Vector Velocity)
    -> Acc (Vector Position, Vector Velocity)
    -> Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity)
    -> Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity)

    simulate :: (Integrator i, ForceField f) => i -> f
    -> Timestep
    -> Duration
    -> Acc (Vector Mass)
    -> Acc (Vector Position, Vector Position, Vector Velocity)
    -> Acc (Vector Position, Vector Position, Vector Velocity)
    simulate i f step duration mass state = let result = awhile cond body init
    ro = result ^. _2
    r = result ^. _3
    v = result ^. _4
    in lift (ro, r, v)
    where cond :: Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity) -> Acc (Scalar Bool)
    cond state = let d = state ^. _1
    in unit (the d < constant duration)

    body :: Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity)
    -> Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity)
    body = integrate i f (constant step) mass

    init :: Acc (Scalar Duration, Vector Position, Vector Position, Vector Velocity)
    init = lift (unit (constant 0), state ^. _1, state ^. _2, state ^. _3)


    data Euler = Euler

    instance Integrator Euler where
    integrate _ ff dt m _ r v = let f = forces ff m r v
    integrate _ ff dt m state = let d = state ^. _1
    r = state ^. _3
    v = state ^. _4
    f = forces ff m r v
    a' = zipWith (^/) f m
    r' = zipWith3 (\a b c -> a + b + c) r (map (dt *^) v) (map (dt * dt / 2 *^) a')
    v' = zipWith (+) v (map (dt *^) a')
    in lift (r', v')
    in lift (unit (the d + dt), r, r', v')

    data Verlet = Verlet

    instance Integrator Verlet where
    integrate _ ff dt m ro r v = let f = forces ff m r v
    a' = zipWith (^/) f m
    r' = zipWith3 (\a b c -> a - b + c) (map (2 *^) r) ro (map (dt * dt *^) a')
    v' = map (^/ (2 * dt)) (zipWith (-) r' ro)
    in lift (r', v')
    integrate _ ff dt m state = let (d, ro, r, v) = unlift state
    f = forces ff m r v
    a' = zipWith (^/) f m
    r' = zipWith3 (\a b c -> a - b + c) (map (2 *^) r) ro (map (dt * dt *^) a')
    v' = map (^/ (2 * dt)) (zipWith (-) r' ro)
    in lift (unit (the d + dt), r, r', v')
  2. zmactep created this gist Oct 31, 2017.
    48 changes: 48 additions & 0 deletions Integrator.hs
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,48 @@
    -- Based on http://www2.mpip-mainz.mpg.de/~andrienk/journal_club/integrators.pdf
    {-# LANGUAGE NoImplicitPrelude #-}

    module Integrator where

    import Data.Array.Accelerate
    import Data.Array.Accelerate.Linear

    type R = Float
    type V3R = V3 R
    type Position = V3 R
    type Velocity = V3 R
    type Force = V3 R
    type Mass = R
    type Timestep = R

    class ForceField f where
    forces :: f -> Acc (Vector Mass)
    -> Acc (Vector Position)
    -> Acc (Vector Velocity)
    -> Acc (Vector Force)

    class Integrator i where
    integrate :: ForceField f => i -> f
    -> Exp Timestep
    -> Acc (Vector Mass)
    -> Acc (Vector Position)
    -> Acc (Vector Position)
    -> Acc (Vector Velocity)
    -> Acc (Vector Position, Vector Velocity)

    data Euler = Euler

    instance Integrator Euler where
    integrate _ ff dt m _ r v = let f = forces ff m r v
    a' = zipWith (^/) f m
    r' = zipWith3 (\a b c -> a + b + c) r (map (dt *^) v) (map (dt * dt / 2 *^) a')
    v' = zipWith (+) v (map (dt *^) a')
    in lift (r', v')

    data Verlet = Verlet

    instance Integrator Verlet where
    integrate _ ff dt m ro r v = let f = forces ff m r v
    a' = zipWith (^/) f m
    r' = zipWith3 (\a b c -> a - b + c) (map (2 *^) r) ro (map (dt * dt *^) a')
    v' = map (^/ (2 * dt)) (zipWith (-) r' ro)
    in lift (r', v')