MD Integrators

Integrator.hs

-- 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
import           Data.Array.Accelerate.Control.Lens

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
type Duration = 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 (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 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 (unit (the d + dt), r, r', v')

data Verlet = Verlet

instance Integrator Verlet where
  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')