{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Safe #-}

{- | 
Module      :  Physics.Learn.RungeKutta
Copyright   :  (c) Scott N. Walck 2012-2019
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Differential equation solving using 4th-order Runge-Kutta
-}

module Physics.Learn.RungeKutta
    ( rungeKutta4
    , integrateSystem
    )
    where

import Physics.Learn.StateSpace
    ( StateSpace(..)
    , Diff
    , Time
    , (.+^)
    )
import Data.VectorSpace
    ( (^+^)
    , (*^)
    , (^/)
    )

-- | Take a single 4th-order Runge-Kutta step
rungeKutta4 :: StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 :: forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 p -> Diff p
f Scalar (Diff p)
dt p
y
    = let k0 :: Diff p
k0 = Scalar (Diff p)
dt forall v. VectorSpace v => Scalar v -> v -> v
*^ p -> Diff p
f p
y
          k1 :: Diff p
k1 = Scalar (Diff p)
dt forall v. VectorSpace v => Scalar v -> v -> v
*^ p -> Diff p
f (p
y forall p. StateSpace p => p -> Diff p -> p
.+^ Diff p
k0 forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ Scalar (Diff p)
2)
          k2 :: Diff p
k2 = Scalar (Diff p)
dt forall v. VectorSpace v => Scalar v -> v -> v
*^ p -> Diff p
f (p
y forall p. StateSpace p => p -> Diff p -> p
.+^ Diff p
k1 forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ Scalar (Diff p)
2)
          k3 :: Diff p
k3 = Scalar (Diff p)
dt forall v. VectorSpace v => Scalar v -> v -> v
*^ p -> Diff p
f (p
y forall p. StateSpace p => p -> Diff p -> p
.+^ Diff p
k2)
      in p
y forall p. StateSpace p => p -> Diff p -> p
.+^ (Diff p
k0 forall v. AdditiveGroup v => v -> v -> v
^+^ Scalar (Diff p)
2 forall v. VectorSpace v => Scalar v -> v -> v
*^ Diff p
k1 forall v. AdditiveGroup v => v -> v -> v
^+^ Scalar (Diff p)
2 forall v. VectorSpace v => Scalar v -> v -> v
*^ Diff p
k2 forall v. AdditiveGroup v => v -> v -> v
^+^ Diff p
k3) forall v s.
(VectorSpace v, s ~ Scalar v, Fractional s) =>
v -> s -> v
^/ Scalar (Diff p)
6

-- | Solve a first-order system of differential equations with 4th-order Runge-Kutta
integrateSystem :: StateSpace p => (p -> Diff p) -> Time p -> p -> [p]
integrateSystem :: forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> [p]
integrateSystem p -> Diff p
systemDerivative Time p
dt
    = forall a. (a -> a) -> a -> [a]
iterate (forall p. StateSpace p => (p -> Diff p) -> Time p -> p -> p
rungeKutta4 p -> Diff p
systemDerivative Time p
dt)



{-
-- | Take a single 4th-order Runge-Kutta step
rungeKutta4 :: (VectorSpace v, Fractional (Scalar v)) => (v -> v) -> Scalar v -> v -> v
rungeKutta4 f h y
    = let k0 = h *^ f y
          k1 = h *^ f (y ^+^ k0 ^/ 2)
          k2 = h *^ f (y ^+^ k1 ^/ 2)
          k3 = h *^ f (y ^+^ k2)
      in y ^+^ (k0 ^+^ 2 *^ k1 ^+^ 2 *^ k2 ^+^ k3) ^/ 6

-- | Solve a first-order system of differential equations with 4th-order Runge-Kutta
integrateSystem :: (VectorSpace v, Fractional (Scalar v)) => (v -> v) -> Scalar v -> v -> [v]
integrateSystem systemDerivative dt
    = iterate (rungeKutta4 systemDerivative dt)
-}