{-# LANGUAGE FlexibleContexts #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}

{- |
Module      :  Numeric.GSL.ODE
Copyright   :  (c) Alberto Ruiz 2010
License     :  GPL
Maintainer  :  Alberto Ruiz
Stability   :  provisional

Solution of ordinary differential equation (ODE) initial value problems.

<http://www.gnu.org/software/gsl/manual/html_node/Ordinary-Differential-Equations.html>

A simple example:

@
import Numeric.GSL.ODE
import Numeric.LinearAlgebra
import Graphics.Plot(mplot)

xdot t [x,v] = [v, -0.95*x - 0.1*v]

ts = linspace 100 (0,20 :: Double)

sol = odeSolve xdot [10,0] ts

main = mplot (ts : toColumns sol)
@

-}
-----------------------------------------------------------------------------

module Numeric.GSL.ODE (
    odeSolve, odeSolveV, odeSolveVWith, ODEMethod(..), Jacobian, StepControl(..)
) where

import Numeric.LinearAlgebra.HMatrix
import Numeric.GSL.Internal

import Foreign.Ptr(FunPtr, nullFunPtr, freeHaskellFunPtr)
import Foreign.C.Types
import System.IO.Unsafe(unsafePerformIO)

-------------------------------------------------------------------------

type TVV   = TV (TV Res)
type TVM   = TV (TM Res)
type TVVM  = TV (TV (TM Res))
type TVVVM = TV (TV (TV (TM Res)))

type Jacobian = Double -> Vector Double -> Matrix Double

-- | Stepping functions
data ODEMethod = RK2 -- ^ Embedded Runge-Kutta (2, 3) method.
               | RK4 -- ^ 4th order (classical) Runge-Kutta. The error estimate is obtained by halving the step-size. For more efficient estimate of the error, use the embedded methods.
               | RKf45 -- ^ Embedded Runge-Kutta-Fehlberg (4, 5) method. This method is a good general-purpose integrator.
               | RKck -- ^ Embedded Runge-Kutta Cash-Karp (4, 5) method.
               | RK8pd -- ^ Embedded Runge-Kutta Prince-Dormand (8,9) method.
               | RK2imp Jacobian -- ^ Implicit 2nd order Runge-Kutta at Gaussian points.
               | RK4imp Jacobian -- ^ Implicit 4th order Runge-Kutta at Gaussian points.
               | BSimp Jacobian -- ^ Implicit Bulirsch-Stoer method of Bader and Deuflhard. The method is generally suitable for stiff problems.
               | RK1imp Jacobian -- ^ Implicit Gaussian first order Runge-Kutta. Also known as implicit Euler or backward Euler method. Error estimation is carried out by the step doubling method.
               | MSAdams -- ^ A variable-coefficient linear multistep Adams method in Nordsieck form. This stepper uses explicit Adams-Bashforth (predictor) and implicit Adams-Moulton (corrector) methods in P(EC)^m functional iteration mode. Method order varies dynamically between 1 and 12. 
               | MSBDF Jacobian -- ^ A variable-coefficient linear multistep backward differentiation formula (BDF) method in Nordsieck form. This stepper uses the explicit BDF formula as predictor and implicit BDF formula as corrector. A modified Newton iteration method is used to solve the system of non-linear equations. Method order varies dynamically between 1 and 5. The method is generally suitable for stiff problems.

-- | Adaptive step-size control functions
data StepControl = X     Double Double -- ^ abs. and rel. tolerance for x(t)
                 | X'    Double Double -- ^ abs. and rel. tolerance for x'(t)
                 | XX'   Double Double Double Double -- ^ include both via rel. tolerance scaling factors a_x, a_x'
                 | ScXX' Double Double Double Double (Vector Double) -- ^ scale abs. tolerance of x(t) components

-- | A version of 'odeSolveV' with reasonable default parameters and system of equations defined using lists.
odeSolve
    :: (Double -> [Double] -> [Double])        -- ^ x'(t,x)
    -> [Double]        -- ^ initial conditions
    -> Vector Double   -- ^ desired solution times
    -> Matrix Double   -- ^ solution
odeSolve :: (Double -> [Double] -> [Double])
-> [Double] -> Vector Double -> Matrix Double
odeSolve xdot :: Double -> [Double] -> [Double]
xdot xi :: [Double]
xi ts :: Vector Double
ts = ODEMethod
-> Double
-> Double
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveV ODEMethod
RKf45 Double
hi Double
epsAbs Double
epsRel ((Double -> [Double] -> [Double])
-> Double -> Vector Double -> Vector Double
forall a a t.
(Storable a, Storable a) =>
(t -> [a] -> [a]) -> t -> Vector a -> Vector a
l2v Double -> [Double] -> [Double]
xdot) ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Double]
xi) Vector Double
ts
    where hi :: Double
hi = (Vector Double
tsVector Double -> Int -> Double
forall c t. Indexable c t => c -> Int -> t
!1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Vector Double
tsVector Double -> Int -> Double
forall c t. Indexable c t => c -> Int -> t
!0)Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/100
          epsAbs :: Double
epsAbs = 1.49012e-08
          epsRel :: Double
epsRel = Double
epsAbs
          l2v :: (t -> [a] -> [a]) -> t -> Vector a -> Vector a
l2v f :: t -> [a] -> [a]
f  = \t :: t
t -> [a] -> Vector a
forall a. Storable a => [a] -> Vector a
fromList ([a] -> Vector a) -> (Vector a -> [a]) -> Vector a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> [a] -> [a]
f t
t ([a] -> [a]) -> (Vector a -> [a]) -> Vector a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList

-- | A version of 'odeSolveVWith' with reasonable default step control.
odeSolveV
    :: ODEMethod
    -> Double            -- ^ initial step size
    -> Double            -- ^ absolute tolerance for the state vector
    -> Double            -- ^ relative tolerance for the state vector
    -> (Double -> Vector Double -> Vector Double)   -- ^ x'(t,x)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveV :: ODEMethod
-> Double
-> Double
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveV meth :: ODEMethod
meth hi :: Double
hi epsAbs :: Double
epsAbs epsRel :: Double
epsRel = ODEMethod
-> StepControl
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveVWith ODEMethod
meth (Double -> Double -> Double -> Double -> StepControl
XX' Double
epsAbs Double
epsRel 1 1) Double
hi

-- | Evolution of the system with adaptive step-size control.
odeSolveVWith
    :: ODEMethod
    -> StepControl
    -> Double            -- ^ initial step size
    -> (Double -> Vector Double -> Vector Double)   -- ^ x'(t,x)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveVWith :: ODEMethod
-> StepControl
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveVWith method :: ODEMethod
method control :: StepControl
control = CInt
-> Maybe (Double -> Vector Double -> Matrix Double)
-> CInt
-> Double
-> Double
-> Double
-> Double
-> Maybe (Vector Double)
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveVWith' CInt
m Maybe (Double -> Vector Double -> Matrix Double)
mbj CInt
c Double
epsAbs Double
epsRel Double
aX Double
aX' Maybe (Vector Double)
mbsc
    where (m :: CInt
m, mbj :: Maybe (Double -> Vector Double -> Matrix Double)
mbj) = case ODEMethod
method of
              RK2        -> (0 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              RK4        -> (1 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              RKf45      -> (2 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              RKck       -> (3 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              RK8pd      -> (4 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              RK2imp jac :: Double -> Vector Double -> Matrix Double
jac -> (5 , (Double -> Vector Double -> Matrix Double)
-> Maybe (Double -> Vector Double -> Matrix Double)
forall a. a -> Maybe a
Just Double -> Vector Double -> Matrix Double
jac)
              RK4imp jac :: Double -> Vector Double -> Matrix Double
jac -> (6 , (Double -> Vector Double -> Matrix Double)
-> Maybe (Double -> Vector Double -> Matrix Double)
forall a. a -> Maybe a
Just Double -> Vector Double -> Matrix Double
jac)
              BSimp  jac :: Double -> Vector Double -> Matrix Double
jac -> (7 , (Double -> Vector Double -> Matrix Double)
-> Maybe (Double -> Vector Double -> Matrix Double)
forall a. a -> Maybe a
Just Double -> Vector Double -> Matrix Double
jac)
              RK1imp jac :: Double -> Vector Double -> Matrix Double
jac -> (8 , (Double -> Vector Double -> Matrix Double)
-> Maybe (Double -> Vector Double -> Matrix Double)
forall a. a -> Maybe a
Just Double -> Vector Double -> Matrix Double
jac)
              MSAdams    -> (9 , Maybe (Double -> Vector Double -> Matrix Double)
forall a. Maybe a
Nothing )
              MSBDF  jac :: Double -> Vector Double -> Matrix Double
jac -> (10, (Double -> Vector Double -> Matrix Double)
-> Maybe (Double -> Vector Double -> Matrix Double)
forall a. a -> Maybe a
Just Double -> Vector Double -> Matrix Double
jac)
          (c :: CInt
c, epsAbs :: Double
epsAbs, epsRel :: Double
epsRel, aX :: Double
aX, aX' :: Double
aX', mbsc :: Maybe (Vector Double)
mbsc) = case StepControl
control of
              X     ea :: Double
ea er :: Double
er           -> (0, Double
ea, Double
er, 1 , 0  , Maybe (Vector Double)
forall a. Maybe a
Nothing)
              X'    ea :: Double
ea er :: Double
er           -> (0, Double
ea, Double
er, 0 , 1  , Maybe (Vector Double)
forall a. Maybe a
Nothing)
              XX'   ea :: Double
ea er :: Double
er ax :: Double
ax ax' :: Double
ax'    -> (0, Double
ea, Double
er, Double
ax, Double
ax', Maybe (Vector Double)
forall a. Maybe a
Nothing)
              ScXX' ea :: Double
ea er :: Double
er ax :: Double
ax ax' :: Double
ax' sc :: Vector Double
sc -> (1, Double
ea, Double
er, Double
ax, Double
ax', Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just Vector Double
sc)

odeSolveVWith'
    :: CInt     -- ^ stepping function
    -> Maybe (Double -> Vector Double -> Matrix Double)   -- ^ optional jacobian
    -> CInt     -- ^ step-size control function
    -> Double   -- ^ absolute tolerance for step-size control
    -> Double   -- ^ relative tolerance for step-size control
    -> Double   -- ^ scaling factor for relative tolerance of x(t)
    -> Double   -- ^ scaling factor for relative tolerance of x'(t)
    -> Maybe (Vector Double)    -- ^ optional scaling for absolute error
    -> Double   -- ^ initial step size
    -> (Double -> Vector Double -> Vector Double)        -- ^ x'(t,x)
    -> Vector Double  -- ^ initial conditions
    -> Vector Double  -- ^ desired solution times
    -> Matrix Double  -- ^ solution
odeSolveVWith' :: CInt
-> Maybe (Double -> Vector Double -> Matrix Double)
-> CInt
-> Double
-> Double
-> Double
-> Double
-> Maybe (Vector Double)
-> Double
-> (Double -> Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
-> Matrix Double
odeSolveVWith' method :: CInt
method mbjac :: Maybe (Double -> Vector Double -> Matrix Double)
mbjac control :: CInt
control epsAbs :: Double
epsAbs epsRel :: Double
epsRel aX :: Double
aX aX' :: Double
aX' mbsc :: Maybe (Vector Double)
mbsc h :: Double
h f :: Double -> Vector Double -> Vector Double
f xiv :: Vector Double
xiv ts :: Vector Double
ts =
    IO (Matrix Double) -> Matrix Double
forall a. IO a -> a
unsafePerformIO (IO (Matrix Double) -> Matrix Double)
-> IO (Matrix Double) -> Matrix Double
forall a b. (a -> b) -> a -> b
$ do
        let n :: IndexOf Vector
n  = Vector Double -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Vector Double
xiv
            sc :: Vector Double
sc = case Maybe (Vector Double)
mbsc of
                Just scv :: Vector Double
scv -> IndexOf Vector -> Vector Double -> Vector Double
forall (c :: * -> *) t.
(Eq (IndexOf c), Container c t, Show (IndexOf c)) =>
IndexOf c -> c t -> c t
checkdim1 Int
IndexOf Vector
n Vector Double
scv
                Nothing  -> Vector Double
xiv
        FunPtr (Double -> TVV)
fp <- (Double -> TVV) -> IO (FunPtr (Double -> TVV))
mkDoubleVecVecfun (\t :: Double
t -> (Vector Double -> Vector Double) -> TVV
aux_vTov (IndexOf Vector -> Vector Double -> Vector Double
forall (c :: * -> *) t.
(Eq (IndexOf c), Container c t, Show (IndexOf c)) =>
IndexOf c -> c t -> c t
checkdim1 Int
IndexOf Vector
n (Vector Double -> Vector Double)
-> (Vector Double -> Vector Double)
-> Vector Double
-> Vector Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Vector Double -> Vector Double
f Double
t))
        FunPtr (Double -> TVM)
jp <- case Maybe (Double -> Vector Double -> Matrix Double)
mbjac of
            Just jac :: Double -> Vector Double -> Matrix Double
jac -> (Double -> TVM) -> IO (FunPtr (Double -> TVM))
mkDoubleVecMatfun (\t :: Double
t -> (Vector Double -> Matrix Double) -> TVM
aux_vTom (Int -> Matrix Double -> Matrix Double
forall t. Int -> Matrix t -> Matrix t
checkdim2 Int
n (Matrix Double -> Matrix Double)
-> (Vector Double -> Matrix Double)
-> Vector Double
-> Matrix Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Vector Double -> Matrix Double
jac Double
t))
            Nothing  -> FunPtr (Double -> TVM) -> IO (FunPtr (Double -> TVM))
forall (m :: * -> *) a. Monad m => a -> m a
return FunPtr (Double -> TVM)
forall a. FunPtr a
nullFunPtr
        Matrix Double
sol <- Vector Double
-> (((CInt -> Ptr Double -> TV TVM) -> TV TVM)
    -> IO (Matrix Double))
-> IO (Matrix Double)
forall a t b.
Storable a =>
Vector a -> (((CInt -> Ptr a -> t) -> t) -> IO b) -> IO b
vec Vector Double
sc ((((CInt -> Ptr Double -> TV TVM) -> TV TVM) -> IO (Matrix Double))
 -> IO (Matrix Double))
-> (((CInt -> Ptr Double -> TV TVM) -> TV TVM)
    -> IO (Matrix Double))
-> IO (Matrix Double)
forall a b. (a -> b) -> a -> b
$ \sc' :: (CInt -> Ptr Double -> TV TVM) -> TV TVM
sc' -> Vector Double
-> ((TV TVM -> TVM) -> IO (Matrix Double)) -> IO (Matrix Double)
forall a t b.
Storable a =>
Vector a -> (((CInt -> Ptr a -> t) -> t) -> IO b) -> IO b
vec Vector Double
xiv (((TV TVM -> TVM) -> IO (Matrix Double)) -> IO (Matrix Double))
-> ((TV TVM -> TVM) -> IO (Matrix Double)) -> IO (Matrix Double)
forall a b. (a -> b) -> a -> b
$ \xiv' :: TV TVM -> TVM
xiv' ->
            Vector Double
-> ((TVM -> TM Res) -> IO (Matrix Double)) -> IO (Matrix Double)
forall a t b.
Storable a =>
Vector a -> (((CInt -> Ptr a -> t) -> t) -> IO b) -> IO b
vec (Vector Double -> Vector Double
forall a.
(Container Vector a, Ord a, Num a) =>
Vector a -> Vector a
checkTimes Vector Double
ts) (((TVM -> TM Res) -> IO (Matrix Double)) -> IO (Matrix Double))
-> ((TVM -> TM Res) -> IO (Matrix Double)) -> IO (Matrix Double)
forall a b. (a -> b) -> a -> b
$ \ts' :: TVM -> TM Res
ts' -> Int -> Int -> TM Res -> String -> IO (Matrix Double)
forall a.
Storable a =>
Int
-> Int -> (CInt -> CInt -> Ptr a -> Res) -> String -> IO (Matrix a)
createMIO (Vector Double -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Vector Double
ts) Int
n
                (CInt
-> CInt
-> Double
-> Double
-> Double
-> Double
-> Double
-> FunPtr (Double -> TVV)
-> FunPtr (Double -> TVM)
-> CInt
-> Ptr Double
-> TV TVM
ode_c CInt
method CInt
control Double
h Double
epsAbs Double
epsRel Double
aX Double
aX' FunPtr (Double -> TVV)
fp FunPtr (Double -> TVM)
jp
                (CInt -> Ptr Double -> TV TVM)
-> ((CInt -> Ptr Double -> TV TVM) -> TV TVM) -> TV TVM
forall x y. x -> (x -> y) -> y
// (CInt -> Ptr Double -> TV TVM) -> TV TVM
sc' TV TVM -> (TV TVM -> TVM) -> TVM
forall x y. x -> (x -> y) -> y
// TV TVM -> TVM
xiv' TVM -> (TVM -> TM Res) -> TM Res
forall x y. x -> (x -> y) -> y
// TVM -> TM Res
ts' )
                "ode"
        FunPtr (Double -> TVV) -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr FunPtr (Double -> TVV)
fp
        if (FunPtr (Double -> TVM)
jp FunPtr (Double -> TVM) -> FunPtr (Double -> TVM) -> Bool
forall a. Eq a => a -> a -> Bool
/= FunPtr (Double -> TVM)
forall a. FunPtr a
nullFunPtr) then FunPtr (Double -> TVM) -> IO ()
forall a. FunPtr a -> IO ()
freeHaskellFunPtr FunPtr (Double -> TVM)
jp else () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Matrix Double -> IO (Matrix Double)
forall (m :: * -> *) a. Monad m => a -> m a
return Matrix Double
sol

foreign import ccall safe "ode"
    ode_c :: CInt -> CInt -> Double
          -> Double -> Double -> Double -> Double
          -> FunPtr (Double -> TVV) -> FunPtr (Double -> TVM) -> TVVVM

-------------------------------------------------------

checkdim1 :: IndexOf c -> c t -> c t
checkdim1 n :: IndexOf c
n v :: c t
v
    | c t -> IndexOf c
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size c t
v IndexOf c -> IndexOf c -> Bool
forall a. Eq a => a -> a -> Bool
== IndexOf c
n = c t
v
    | Bool
otherwise = String -> c t
forall a. HasCallStack => String -> a
error (String -> c t) -> String -> c t
forall a b. (a -> b) -> a -> b
$ "Error: "String -> String -> String
forall a. [a] -> [a] -> [a]
++ IndexOf c -> String
forall a. Show a => a -> String
show IndexOf c
n
                        String -> String -> String
forall a. [a] -> [a] -> [a]
++ " components expected in the result of the function supplied to odeSolve"

checkdim2 :: Int -> Matrix t -> Matrix t
checkdim2 n :: Int
n m :: Matrix t
m
    | Matrix t -> Int
forall t. Matrix t -> Int
rows Matrix t
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Bool -> Bool -> Bool
&& Matrix t -> Int
forall t. Matrix t -> Int
cols Matrix t
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Matrix t
m
    | Bool
otherwise = String -> Matrix t
forall a. HasCallStack => String -> a
error (String -> Matrix t) -> String -> Matrix t
forall a b. (a -> b) -> a -> b
$ "Error: "String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> String -> String
forall a. [a] -> [a] -> [a]
++ "x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
                        String -> String -> String
forall a. [a] -> [a] -> [a]
++ " Jacobian expected in odeSolve"

checkTimes :: Vector a -> Vector a
checkTimes ts :: Vector a
ts | Vector a -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Vector a
ts IndexOf Vector -> IndexOf Vector -> Bool
forall a. Ord a => a -> a -> Bool
> 1 Bool -> Bool -> Bool
&& (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>0) ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
subtract [a]
ts' ([a] -> [a]
forall a. [a] -> [a]
tail [a]
ts')) = Vector a
ts
              | Bool
otherwise = String -> Vector a
forall a. HasCallStack => String -> a
error "odeSolve requires increasing times"
    where ts' :: [a]
ts' = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
ts