{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wno-pattern-namespace-specifier #-}

-- | Ordinary differential equation integrators as streaming 'Process' machines.
--
-- A vector field is represented as a 'Diff' so the same field can be reused
-- in differentiable contexts; the integrators themselves use only the forward
-- pass.
--
-- The step size is kept separate from the state so that vector-valued states
-- (for example 'NumHask.Algebra.Metric.EuclideanPair') can be stepped with a
-- scalar step size.
--
-- The scalar self-actions for 'Double' (and 'Float') that this module used to
-- supply as orphan instances now live in "NumHask.Algebra.Action".
module Circuit.Stats.ODE
  ( -- * Vector fields
    vectorField,

    -- * Single steps
    eulerStep,
    rk4Step,

    -- * Trajectory generators
    euler,
    rk4,

    -- * Process machines
    eulerProcess,
    rk4Process,
  )
where

import Circuit.Diff (Diff, Diff', runDiff, pattern Diff)
import Circuit.Stats (Process (..))
import Data.List (scanl')
import NumHask.Prelude

-- $setup
--
-- >>> :m -Prelude
-- >>> :set -XRebindableSyntax
-- >>> import NumHask.Prelude
-- >>> import Circuit.Stats.ODE
-- >>> import NumHask.Algebra.Metric (EuclideanPair (..))
-- >>> import Circuit.Diff (Diff, Diff', runDiff)

-- | Lift a pure vector field into a 'Diff' with zero pullback.
--
-- >>> let f = vectorField (\y -> y) :: Diff' Double Double
-- >>> fst (runDiff f 2.0)
-- 2.0
vectorField :: (Additive s) => (s -> s) -> Diff' s s
vectorField :: forall s. Additive s => (s -> s) -> Diff' s s
vectorField s -> s
f = (s -> (s, s -> s)) -> Diff () s s
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((s -> (s, s -> s)) -> Diff () s s)
-> (s -> (s, s -> s)) -> Diff () s s
forall a b. (a -> b) -> a -> b
$ \s
x -> (s -> s
f s
x, s -> s -> s
forall a b. a -> b -> a
const s
forall a. Additive a => a
zero)

-- | Evaluate a differentiable vector field at a point, keeping only the
-- forward value.
evalField :: Diff' s s -> s -> s
evalField :: forall s. Diff' s s -> s -> s
evalField Diff' s s
f s
x = (s, s -> s) -> s
forall a b. (a, b) -> a
fst (Diff' s s -> s -> (s, s -> s)
forall k (p :: k) a b. Diff p a b -> a -> (b, b -> a)
runDiff Diff' s s
f s
x)

-- | One Euler step: @y' = y + h · f(y)@.
--
-- >>> let f = vectorField (\y -> y) :: Diff' Double Double
-- >>> eulerStep f 1.0 0.1
-- 1.1
eulerStep ::
  (Additive s, MultiplicativeAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  h ->
  s
eulerStep :: forall s h.
(Additive s, MultiplicativeAction s, Scalar s ~ h) =>
Diff' s s -> s -> h -> s
eulerStep Diff' s s
f s
y h
h = s
y s -> s -> s
forall a. Additive a => a -> a -> a
+ (h
Scalar s
h Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| Diff' s s -> s -> s
forall s. Diff' s s -> s -> s
evalField Diff' s s
f s
y)

-- | One RK4 step.
--
-- >>> let f = vectorField (\y -> y) :: Diff' Double Double
-- >>> rk4Step f 1.0 0.1
-- 1.1051708333333334
rk4Step ::
  (Additive s, Additive (Scalar s), DivisiveAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  h ->
  s
rk4Step :: forall s h.
(Additive s, Additive (Scalar s), DivisiveAction s,
 Scalar s ~ h) =>
Diff' s s -> s -> h -> s
rk4Step Diff' s s
f s
y h
h =
  let tw :: h
tw = h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one
      sx :: h
sx = h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one h -> h -> h
forall a. Additive a => a -> a -> a
+ h
forall a. Multiplicative a => a
one
      k1 :: s
k1 = h
Scalar s
h Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| Diff' s s -> s -> s
forall s. Diff' s s -> s -> s
evalField Diff' s s
f s
y
      k2 :: s
k2 = h
Scalar s
h Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| Diff' s s -> s -> s
forall s. Diff' s s -> s -> s
evalField Diff' s s
f (s
y s -> s -> s
forall a. Additive a => a -> a -> a
+ s
k1 s -> Scalar s -> s
forall m. DivisiveAction m => m -> Scalar m -> m
|/ h
Scalar s
tw)
      k3 :: s
k3 = h
Scalar s
h Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| Diff' s s -> s -> s
forall s. Diff' s s -> s -> s
evalField Diff' s s
f (s
y s -> s -> s
forall a. Additive a => a -> a -> a
+ s
k2 s -> Scalar s -> s
forall m. DivisiveAction m => m -> Scalar m -> m
|/ h
Scalar s
tw)
      k4 :: s
k4 = h
Scalar s
h Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| Diff' s s -> s -> s
forall s. Diff' s s -> s -> s
evalField Diff' s s
f (s
y s -> s -> s
forall a. Additive a => a -> a -> a
+ s
k3)
   in s
y s -> s -> s
forall a. Additive a => a -> a -> a
+ (s
k1 s -> s -> s
forall a. Additive a => a -> a -> a
+ h
Scalar s
tw Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| s
k2 s -> s -> s
forall a. Additive a => a -> a -> a
+ h
Scalar s
tw Scalar s -> s -> s
forall m. MultiplicativeAction m => Scalar m -> m -> m
*| s
k3 s -> s -> s
forall a. Additive a => a -> a -> a
+ s
k4) s -> Scalar s -> s
forall m. DivisiveAction m => m -> Scalar m -> m
|/ h
Scalar s
sx

-- | Integrate a 'Diff' vector field over a list of step sizes using Euler.
--
-- The result includes the initial state as the first element.
--
-- >>> let f = vectorField (\y -> y) :: Diff' Double Double
-- >>> euler f 1.0 [0.1, 0.1, 0.1]
-- [1.0,1.1,1.2100000000000002,1.3310000000000002]
euler ::
  (Additive s, MultiplicativeAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  [h] ->
  [s]
euler :: forall s h.
(Additive s, MultiplicativeAction s, Scalar s ~ h) =>
Diff' s s -> s -> [h] -> [s]
euler Diff' s s
f = (s -> h -> s) -> s -> [h] -> [s]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' (Diff' s s -> s -> h -> s
forall s h.
(Additive s, MultiplicativeAction s, Scalar s ~ h) =>
Diff' s s -> s -> h -> s
eulerStep Diff' s s
f)

-- | Integrate a 'Diff' vector field over a list of step sizes using RK4.
--
-- The result includes the initial state as the first element.
--
-- Harmonic oscillator @x'' = −x@ written as @x' = v, v' = −x@:
--
-- >>> let f = vectorField (\(EuclideanPair (x, v)) -> EuclideanPair (v, -x)) :: Diff' (EuclideanPair Double) (EuclideanPair Double)
-- >>> take 5 (rk4 f (EuclideanPair (1.0, 0.0)) (replicate 40 ((pi :: Double) / 20)))
-- [EuclideanPair {euclidPair = (1.0,0.0)},EuclideanPair {euclidPair = (0.9876883614494284,-0.1564336685819834)},EuclideanPair {euclidPair = (0.9510568066766389,-0.3090154275945243)},EuclideanPair {euclidPair = (0.8910073220447337,-0.45398824664172294)},EuclideanPair {euclidPair = (0.8090185150345391,-0.5877824515637287)}]
rk4 ::
  (Additive s, Additive (Scalar s), DivisiveAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  [h] ->
  [s]
rk4 :: forall s h.
(Additive s, Additive (Scalar s), DivisiveAction s,
 Scalar s ~ h) =>
Diff' s s -> s -> [h] -> [s]
rk4 Diff' s s
f = (s -> h -> s) -> s -> [h] -> [s]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl' (Diff' s s -> s -> h -> s
forall s h.
(Additive s, Additive (Scalar s), DivisiveAction s,
 Scalar s ~ h) =>
Diff' s s -> s -> h -> s
rk4Step Diff' s s
f)

-- | A 'Process' machine that performs Euler integration.
--
-- Input is the step size @h@; output is the current state.  The first input
-- is used only to kick off the machine, so its value is ignored.
eulerProcess ::
  (Additive s, MultiplicativeAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  Process h s
eulerProcess :: forall s h.
(Additive s, MultiplicativeAction s, Scalar s ~ h) =>
Diff' s s -> s -> Process h s
eulerProcess Diff' s s
f s
y0 = (h -> s) -> (s -> h -> s) -> (s -> s) -> Process h s
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process (s -> h -> s
forall a b. a -> b -> a
const s
y0) (Diff' s s -> s -> h -> s
forall s h.
(Additive s, MultiplicativeAction s, Scalar s ~ h) =>
Diff' s s -> s -> h -> s
eulerStep Diff' s s
f) s -> s
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id

-- | A 'Process' machine that performs RK4 integration.
--
-- Input is the step size @h@; output is the current state.  The first input
-- is used only to kick off the machine, so its value is ignored.
rk4Process ::
  (Additive s, Additive (Scalar s), DivisiveAction s, Scalar s ~ h) =>
  Diff' s s ->
  s ->
  Process h s
rk4Process :: forall s h.
(Additive s, Additive (Scalar s), DivisiveAction s,
 Scalar s ~ h) =>
Diff' s s -> s -> Process h s
rk4Process Diff' s s
f s
y0 = (h -> s) -> (s -> h -> s) -> (s -> s) -> Process h s
forall s a b. (a -> s) -> (s -> a -> s) -> (s -> b) -> Process a b
Process (s -> h -> s
forall a b. a -> b -> a
const s
y0) (Diff' s s -> s -> h -> s
forall s h.
(Additive s, Additive (Scalar s), DivisiveAction s,
 Scalar s ~ h) =>
Diff' s s -> s -> h -> s
rk4Step Diff' s s
f) s -> s
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id