OpenDynamical.hs

{-# LANGUAGE Arrows #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE ScopedTypeVariables #-}
#define OD_DEBUG 0
module OpenDynamical (
  OD(..), Aff, Add, Finite, Size, KnownNat,
  showOD, integral, integralV, integralN,
  Integrator, odStep, solve, solve', euler, rk4,
#ifdef OD_DEBUG
  module Linear.V0, module Linear.V1,
#endif
  sine, sine',
  ) where


import Linear.Affine
#ifdef OD_DEBUG
import Data.Functor.Classes
#endif
import Linear.Vector
import Linear.V hiding (int)
import Linear.V0
import Linear.V1
import qualified Data.Vector as DV
import GHC.TypeLits
import GHC.TypeLits.Witnesses
import Data.Proxy
import qualified Control.Category as C
import Control.Arrow
import Control.Applicative

#ifdef OD_DEBUG
type Aff v = (Affine v, Show1 v)
type Add v = (Additive v, Show1 v)
#else
type Aff v = Affine v
type Add v = Additive v
#endif
data OD s a b where
  OD :: KnownNat n =>
    !(V n s) -> !(a -> V n s -> V n s) ->
      !(a -> V n s -> b) -> OD s a b

showOD :: Show s => OD s a b -> String
#ifdef OD_DEBUG
showOD (OD v0 _ _) = showsPrec1 0 v0 ""
#else
showOD _ = ""
#endif

integral :: (Aff v, Finite (Diff v), KnownNat (Size (Diff v)), Num s) =>
  v s -> OD s (Diff v s) (v s)
integral v0 = OD zero d o
  where d dv _ = toV dv
        o _ v = v0 .+^ fromV v

integralV :: (Add v, Finite v, KnownNat (Size v)) => v s -> OD s (v s) (v s)
integralV v0 = OD (toV v0) d o
  where d dv _ = toV dv
        o _ v = fromV v

integralN :: s -> OD s s s
integralN v0 = OD (toV (V1 v0)) d o
  where d dv _ = toV (V1 dv)
        o _ v = case fromV v of V1 v' -> v'

cut :: forall n m s. (KnownNat n, KnownNat m) => V (n + m) s -> (V n s, V m s)
cut (V ab) = let (a, b) = DV.splitAt n ab in (V a, V b)
  where n = fromInteger (natVal (Proxy :: Proxy n))

paste :: V n s -> V m s -> V (n + m) s
paste (V a) (V b) = V (a DV.++ b)

slen :: KnownNat n => V n s -> SNat n
slen _ = SNat

nil :: V 0 a
nil = toV V0

odPaste :: (KnownNat n, KnownNat m) =>
  V n s -> V m s ->
    (a -> V n s -> V m s -> (V n s, V m s)) ->
    (a -> V n s -> V m s -> b) ->
    OD s a b
odPaste v0 w0 f q =
  case slen v0 %+ slen w0 of SNat -> OD (paste v0 w0) f' q'
  where f' a vw = let (v, w) = cut vw
                      (dv, dw) = f a v w
                  in paste dv dw
        q' a vw = let (v, w) = cut vw
                  in q a v w

instance C.Category (OD s) where
  id = OD nil (\_ _ -> nil) const
  OD w0 e p . OD v0 d o = odPaste v0 w0 f q
    where f a v w =
            let (dv, b) = (d a v, o a v)
                dw = e b w
            in (dv, dw)
          q a v w =
            let b = o a v
                c = p b w
            in c

instance Arrow (OD s) where
  arr f = OD nil (\_ _ -> nil) (\a _ -> f a)
  OD v0 d o *** OD w0 e p = odPaste v0 w0 f q
    where f (a, a') v w = (d a v, e a' w)
          q (a, a') v w = (o a v, p a' w)

instance Functor (OD s a) where
  fmap f (OD v0 d o) = OD v0 d ((fmap . fmap) f o)

instance Applicative (OD s a) where
  pure = arr . const
  liftA2 f od1 od2 = uncurry f <$> (od1 &&& od2)

instance Num s => ArrowChoice (OD s) where
  OD v0 d o +++ OD w0 e p = odPaste v0 w0 f q
    where f (Left a)   v _ = (d a v, zero)
          f (Right a') _ w = (zero, e a' w)
          q (Left a)   v _ = Left (o a v)
          q (Right a') _ w = Right (p a' w)

instance ArrowLoop (OD s) where
  loop (OD v0 f o) = OD v0 f' o'
    where f' b v =
            let (_, d) = o (b, d) v
            in f (b, d) v
          o' b v =
            let (c, d) = o (b, d) v
            in c


type Integrator s a b = a -> s -> OD s a b -> OD s a b

odStep :: Integrator s a b -> a -> s -> OD s a b -> (b, OD s a b)
odStep int a h od@(OD v0 _ o) = (o a v0, int a h od)

solve :: Integrator s a b -> [(a, s)] -> OD s a b -> [b]
solve int = foldr step (const [])
  where step (a, h) k od = let (b, od') = odStep int a h od
                           in b : k od'

solve' :: Integrator s a b -> [a] -> s -> OD s a b -> [b]
solve' int as h = solve int [(a, h) | a <- as]


euler :: Num s => Integrator s a b
euler a h (OD v0 d o) = OD v1 d o
  where dv = d a v0
        v1 = v0 ^+^ h *^ dv

-- TODO is it ok to keep using the same a... ?
rk4 :: Fractional s => Integrator s a b
rk4 a h (OD v0 d o) = OD v1 d o
  where dv1 = d a v0
        dv2 = d a (v0 ^+^ (h/2) *^ dv1)
        dv3 = d a (v0 ^+^ (h/2) *^ dv2)
        dv4 = d a (v0 ^+^ h *^ dv3)
        v1 = v0 ^+^ (h/6) *^ (dv1 ^+^ 2 *^ dv2 ^+^ 2 *^ dv3 ^+^ dv4)


sine :: Num s => OD s () s
sine = proc () -> do
  rec
    let y'' = -y
    y' <- integralN 1 -< y''
    y <- integralN 0 -< y'
  returnA -< y

sine' :: forall s. Num s => OD s () s
sine' = OD (V [0, 1]) d o
  where d :: () -> V 2 s -> V 2 s
        d _ (V [y, y']) = V [y', -y]
        d _ _ = error "impossible"
        o :: () -> V 2 s -> s
        o _ (V [y, _]) = y
        o _ _ = error "impossible"

{-
main :: IO ()
main = do
  print . sum . take 100000 $ solve' rk4 (repeat ()) (0.01 :: Double) sine'
  print . sum . take 100000 $ solve' rk4 (repeat ()) (0.01 :: Double) sine
  -}