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

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