module Numerics.DifferentialEquation (
      rungeKuttaConstantStep
    , rungeKuttaAdaptiveStep
) where

import Algebra.VectorSpace

-- | Single step of RK4 (»the standard Runge-Kutta«).
rungeKutta4Step
    :: VectorSpace vec
    => (Double -> vec -> vec) -- ^ \= dy/dt = f(t, y)
    -> vec                    -- ^ Current y
    -> Double                 -- ^ Time
    -> Double                 -- ^ Step size
    -> (Double, vec)          -- ^ time, y
rungeKutta4Step :: forall vec.
VectorSpace vec =>
(Double -> vec -> vec) -> vec -> Double -> Double -> (Double, vec)
rungeKutta4Step Double -> vec -> vec
f vec
y Double
t Double
dt
  = let k1 :: vec
k1 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f Double
t vec
y
        k2 :: vec
k2 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
tDouble -> Double -> Double
forall a. Num a => a -> a -> a
+Double
0.5Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt) (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
0.5Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1)
        k3 :: vec
k3 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
tDouble -> Double -> Double
forall a. Num a => a -> a -> a
+Double
0.5Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt) (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
0.5Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2)
        k4 :: vec
k4 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
tDouble -> Double -> Double
forall a. Num a => a -> a -> a
+Double
dt) (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. vec
k3)

        dy :: vec
dy = Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
6 Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. (vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
2Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
2Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. vec
k4)

    in (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dt, vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. vec
dy)

-- | Solve a system of first-order differential equations with RK4 (»the standard Runge-Kutta«).
--
-- Solution of a planetary model (with adjusted gravity and small friction for a
-- prettier non-elliptic trajectory),
--
-- \[
--     \begin{align}
--         \mathbf {\ddot x} &=
--             - \underbrace{ g \frac {\mathbf x} {\|\mathbf x\|^{3-\varepsilon}}}_{\text{Gravity} }
--             - \underbrace{ \mu \|\mathbf{\dot x}\| \mathbf{\dot x}}_{\text{Friction} }
--         \\
--         \mathbf{x}(0)      &= \begin{pmatrix}100\\0\end{pmatrix} \\
--         \mathbf{\dot x}(0) &= \begin{pmatrix}4\\4\end{pmatrix} \\
--         \varepsilon        &= 0.06 \\
--         g                  &= 2200 \\
--         \mu                &= 0.0001
--     \end{align}
-- \]
--
-- <<docs/differential_equations/1_two_body_problem.svg>>
rungeKuttaConstantStep
    :: VectorSpace vec
    => (Double -> vec -> vec) -- ^ \= dy/dt = f(t, y)
    -> vec                    -- ^ Initial y
    -> Double                 -- ^ Initial time
    -> Double                 -- ^ Step size
    -> [(Double, vec)]        -- ^ Infinite list of (time, y). Use e.g. 'takeWhile' to narrow it down if necessary.
rungeKuttaConstantStep :: forall vec.
VectorSpace vec =>
(Double -> vec -> vec)
-> vec -> Double -> Double -> [(Double, vec)]
rungeKuttaConstantStep Double -> vec -> vec
f vec
y0 Double
t0 Double
dt
  = ((Double, vec) -> (Double, vec))
-> (Double, vec) -> [(Double, vec)]
forall a. (a -> a) -> a -> [a]
iterate (\(Double
t, vec
y) -> (Double -> vec -> vec) -> vec -> Double -> Double -> (Double, vec)
forall vec.
VectorSpace vec =>
(Double -> vec -> vec) -> vec -> Double -> Double -> (Double, vec)
rungeKutta4Step Double -> vec -> vec
f vec
y Double
t Double
dt) (Double
t0, vec
y0)

rkf45step
    :: (VectorSpace vec)
    => (Double -> vec -> vec) -- ^ \= dy/dt = f(t, y)
    -> vec                    -- ^ Current y
    -> Double                 -- ^ Current time
    -> Double                 -- ^ Step size
    -> (vec -> Double)        -- ^ Norm function to calculate how good our estimate is
    -> Double                 -- ^ Error tolerance
    -> (Double, vec, Double)  -- ^ New time, new y, new step size
rkf45step :: forall vec.
VectorSpace vec =>
(Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> (Double, vec, Double)
rkf45step Double -> vec -> vec
f vec
y Double
t Double
dt vec -> Double
toleranceNorm Double
tolerance
  = let k1 :: vec
k1 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f Double
t vec
y
        k2 :: vec
k2 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt)   (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4)       Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1)
        k3 :: vec
k3 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
3Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
8Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt)   (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
3Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
32)      Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
9Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
32)     Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2)
        k4 :: vec
k4 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
12Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
13Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt) (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
1932Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2197) Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
7200Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2197)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
7296Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2197)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3)
        k5 :: vec
k5 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
dt)       (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
439Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
216)   Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. Double
8          Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
3680Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
513) Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
845Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4104) Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k4)
        k6 :: vec
k6 = Double
dt Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*. Double -> vec -> vec
f (Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
dt)   (vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
8Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
27)      Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
2          Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
3544Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2565)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
1859Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4104)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k4 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
11Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
40)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k5)

        y'rk4 :: vec
y'rk4 = vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
25Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
216)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
0Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
1408Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
2565) Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
2197Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4101)  Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k4 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
5) Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k5
        y'rk5 :: vec
y'rk5 = vec
y vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
16Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
135)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k1 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. Double
0Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k2 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
6656Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
12826)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k3 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
28561Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
56430)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k4 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. (Double
9Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
50)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k5 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
+. (Double
2Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
55)Double -> vec -> vec
forall v. VectorSpace v => Double -> v -> v
*.vec
k6

        -- How far are we from the tolerance threshold?
        deviation :: Double
deviation = vec -> Double
toleranceNorm (vec
y'rk4 vec -> vec -> vec
forall v. VectorSpace v => v -> v -> v
-. vec
y'rk5)

        -- Safety constant. Mathematically unnecessary, but some authors seem to
        -- use it to err on the safer side when it comes to reducing step size
        -- to avoid flip-flopping.
        safety :: Double
safety = Double
0.9

    in if Double
deviation Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
tolerance
        then let dtFactor :: Double
dtFactor = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
2 (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
safety Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
tolerance Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
deviation) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
4)
                 dt' :: Double
dt' = Double
dt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
dtFactor
             in (Double
tDouble -> Double -> Double
forall a. Num a => a -> a -> a
+Double
dt, vec
y'rk5, Double
dt')
        else let dtFactor :: Double
dtFactor = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0.1 (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
safety Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
tolerance Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
deviation) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (Double
1Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/Double
5)
                 dt' :: Double
dt' = Double
dt Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
dtFactor
             in (Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> (Double, vec, Double)
forall vec.
VectorSpace vec =>
(Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> (Double, vec, Double)
rkf45step Double -> vec -> vec
f vec
y Double
t Double
dt' vec -> Double
toleranceNorm Double
tolerance

-- | Solve a system of first-order differential equations with RKF45
-- (Runge-Kutta-Feinberg, adaptive step size using 4th-and-5th-order Runge-Kutta).
--
-- Solution for a double pendulum:
--
-- <<docs/differential_equations/2_double_pendulum.svg>>
rungeKuttaAdaptiveStep
    :: (VectorSpace vec)
    => (Double -> vec -> vec) -- ^ \= dy/dt = f(t, y)
    -> vec                    -- ^ Current y
    -> Double                 -- ^ Current time
    -> Double                 -- ^ Initial step size
    -> (vec -> Double)        -- ^ Norm function to calculate how good our estimate is
    -> Double                 -- ^ Error tolerance
    -> [(Double, vec)]        -- ^ Infinite list of (time, y). Use e.g. 'takeWhile' to narrow it down if necessary.
rungeKuttaAdaptiveStep :: forall vec.
VectorSpace vec =>
(Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> [(Double, vec)]
rungeKuttaAdaptiveStep Double -> vec -> vec
f vec
y0 Double
t0 Double
dt0 vec -> Double
toleranceNorm Double
tolerance
  = [ (Double
t, vec
y) | (Double
t,vec
y,Double
_dt) <- ((Double, vec, Double) -> (Double, vec, Double))
-> (Double, vec, Double) -> [(Double, vec, Double)]
forall a. (a -> a) -> a -> [a]
iterate (\(Double
t, vec
y, Double
dt) -> (Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> (Double, vec, Double)
forall vec.
VectorSpace vec =>
(Double -> vec -> vec)
-> vec
-> Double
-> Double
-> (vec -> Double)
-> Double
-> (Double, vec, Double)
rkf45step Double -> vec -> vec
f vec
y Double
t Double
dt vec -> Double
toleranceNorm Double
tolerance) (Double
t0, vec
y0, Double
dt0) ]