    {-# 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 Double -> [Double] -> [Double]
    xdot [Double]
    xi 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
    !Int
    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
    !Int
    0)Double -> Double -> Double
    forall a. Fractional a => a -> a -> a
    /Double
    100
              epsAbs :: Double
    epsAbs = Double
    1.49012e-08
              epsRel :: Double
    epsRel = Double
    epsAbs
              l2v :: (t -> [a] -> [a]) -> t -> Vector a -> Vector a
    l2v t -> [a] -> [a]
    f  = \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 ODEMethod
    meth Double
    hi Double
    epsAbs 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 Double
    1 Double
    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 ODEMethod
    method 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 (CInt
    m, Maybe (Double -> Vector Double -> Matrix Double)
    mbj) = case ODEMethod
    method of
                  ODEMethod
    RK2        -> (CInt
    0 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  ODEMethod
    RK4        -> (CInt
    1 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  ODEMethod
    RKf45      -> (CInt
    2 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  ODEMethod
    RKck       -> (CInt
    3 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  ODEMethod
    RK8pd      -> (CInt
    4 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  RK2imp Double -> Vector Double -> Matrix Double
    jac -> (CInt
    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 Double -> Vector Double -> Matrix Double
    jac -> (CInt
    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  Double -> Vector Double -> Matrix Double
    jac -> (CInt
    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 Double -> Vector Double -> Matrix Double
    jac -> (CInt
    8 , (Double -> Vector Double -> Matrix Double)
    -> Maybe (Double -> Vector Double -> Matrix Double)
    forall a. a -> Maybe a
    Just Double -> Vector Double -> Matrix Double
    jac)
                  ODEMethod
    MSAdams    -> (CInt
    9 , Maybe (Double -> Vector Double -> Matrix Double)
    forall a. Maybe a
    Nothing )
                  MSBDF  Double -> Vector Double -> Matrix Double
    jac -> (CInt
    10, (Double -> Vector Double -> Matrix Double)
    -> Maybe (Double -> Vector Double -> Matrix Double)
    forall a. a -> Maybe a
    Just Double -> Vector Double -> Matrix Double
    jac)
              (CInt
    c, Double
    epsAbs, Double
    epsRel, Double
    aX, Double
    aX', Maybe (Vector Double)
    mbsc) = case StepControl
    control of
                  X     Double
    ea Double
    er           -> (CInt
    0, Double
    ea, Double
    er, Double
    1 , Double
    0  , Maybe (Vector Double)
    forall a. Maybe a
    Nothing)
                  X'    Double
    ea Double
    er           -> (CInt
    0, Double
    ea, Double
    er, Double
    0 , Double
    1  , Maybe (Vector Double)
    forall a. Maybe a
    Nothing)
                  XX'   Double
    ea Double
    er Double
    ax Double
    ax'    -> (CInt
    0, Double
    ea, Double
    er, Double
    ax, Double
    ax', Maybe (Vector Double)
    forall a. Maybe a
    Nothing)
                  ScXX' Double
    ea Double
    er Double
    ax Double
    ax' Vector Double
    sc -> (CInt
    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' CInt
    method Maybe (Double -> Vector Double -> Matrix Double)
    mbjac CInt
    control Double
    epsAbs Double
    epsRel Double
    aX Double
    aX' Maybe (Vector Double)
    mbsc Double
    h Double -> Vector Double -> Vector Double
    f Vector Double
    xiv 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 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
                    Maybe (Vector Double)
    Nothing  -> Vector Double
    xiv
            fp <- (Double -> TVV) -> IO (FunPtr (Double -> TVV))
    mkDoubleVecVecfun (\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))
            jp <- case mbjac of
                Just Double -> Vector Double -> Matrix Double
    jac -> (Double -> TVM) -> IO (FunPtr (Double -> TVM))
    mkDoubleVecMatfun (\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))
                Maybe (Double -> Vector Double -> Matrix Double)
    Nothing  -> FunPtr (Double -> TVM) -> IO (FunPtr (Double -> TVM))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return FunPtr (Double -> TVM)
    forall a. FunPtr a
    nullFunPtr
            sol <- vec 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
    $ \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
    $ \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' )
                    String
    "ode"
            freeHaskellFunPtr fp
            if (jp /= nullFunPtr) then freeHaskellFunPtr jp else pure ()
            return 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 IndexOf c
    n 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
    $ String
    "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]
    ++ String
    " components expected in the result of the function supplied to odeSolve"

    checkdim2 :: Int -> Matrix t -> Matrix t
    checkdim2 Int
    n 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
    $ String
    "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]
    ++ String
    "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]
    ++ String
    " Jacobian expected in odeSolve"

    checkTimes :: Vector a -> Vector a
    checkTimes Vector a
    ts | Vector a -> IndexOf Vector
    forall (c :: * -> *) t. Container c t => c t -> IndexOf c
    size Vector a
    ts Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    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
    >a
    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. HasCallStack => [a] -> [a]
    tail [a]
    ts')) = Vector a
    ts
                  | Bool
    otherwise = String -> Vector a
    forall a. HasCallStack => String -> a
    error String
    "odeSolve requires increasing times"
        where ts' :: [a]
    ts' = Vector a -> [a]
    forall a. Storable a => Vector a -> [a]
    toList Vector a
    ts