diff --git a/benchmark/benchmark.hs b/benchmark/benchmark.hs new file mode 100644 --- /dev/null +++ b/benchmark/benchmark.hs @@ -0,0 +1,61 @@ +import Prelude hiding ((<>)) +import Numeric.LinearAlgebra hiding (size, step) +import qualified System.Clock as Clock +import Data.Csv (FromNamedRecord, ToNamedRecord, DefaultOrdered) +import Data.Csv.Incremental +import GHC.Generics (Generic) +import qualified Data.ByteString.Lazy as LBS +import Numeric.Sundials.CVode.ODE +import Control.Exception +import System.IO.Unsafe + +type Time = Double + +getTime :: IO Time +getTime = do + t <- Clock.getTime Clock.Monotonic + let ns = realToFrac $ Clock.toNanoSecs t + return $ ns / 10 ^ (9 :: Int) + +timed :: IO a -> IO Time +timed t = do + start <- getTime + !_ <- t + end <- getTime + return (end-start) + +data Measurement t = Measurement + { nts :: !Int -- the number of time steps in the output + , step :: !Double -- the size of a single time step + , time :: !t -- the total solving time (if known) + } + deriving Generic + +instance FromNamedRecord (Measurement Time) +instance ToNamedRecord (Measurement Time) +instance DefaultOrdered (Measurement Time) + +measure :: Measurement () -> IO (Measurement Time) +measure m@Measurement{..} = do + let + rhs _t y = + let + mu = 1000 + y1 = y ! 0 + y2 = y ! 1 + in + fromList [y2, mu * (1 - y1*y1) * y2 - y1] + y0 = fromList [2,0] + !times = fromList . take (nts+1) $ iterate (+ step) 0 + time <- (timed . evaluate) + (odeSolveV BDF Nothing 1e-2 1e-2 rhs y0 times :: Matrix Double) + return m{time} + +main :: IO () +main = do + measurements <- sequence $ do + nts <- [2] + step <- [1e5, 2e5 .. 2e9] + let m0 = Measurement { time = (), .. } + return . unsafeInterleaveIO $ measure m0 + LBS.writeFile "timings.csv" . encodeDefaultOrderedByName . mconcat $ encodeNamedRecord <$> measurements diff --git a/diagrams/boundedSine.png b/diagrams/boundedSine.png new file mode 100644 Binary files /dev/null and b/diagrams/boundedSine.png differ diff --git a/diagrams/brussRoot.png b/diagrams/brussRoot.png new file mode 100644 Binary files /dev/null and b/diagrams/brussRoot.png differ diff --git a/diagrams/lorenz.png b/diagrams/lorenz.png Binary files a/diagrams/lorenz.png and b/diagrams/lorenz.png differ diff --git a/diagrams/lorenz1.png b/diagrams/lorenz1.png Binary files a/diagrams/lorenz1.png and b/diagrams/lorenz1.png differ diff --git a/diagrams/lorenz2.png b/diagrams/lorenz2.png Binary files a/diagrams/lorenz2.png and b/diagrams/lorenz2.png differ diff --git a/diagrams/predatorPrey1.png b/diagrams/predatorPrey1.png Binary files a/diagrams/predatorPrey1.png and b/diagrams/predatorPrey1.png differ diff --git a/diagrams/transport.png b/diagrams/transport.png new file mode 100644 Binary files /dev/null and b/diagrams/transport.png differ diff --git a/hmatrix-sundials.cabal b/hmatrix-sundials.cabal --- a/hmatrix-sundials.cabal +++ b/hmatrix-sundials.cabal @@ -1,11 +1,11 @@ name: hmatrix-sundials -version: 0.19.1.0 +version: 0.20.1.0 synopsis: hmatrix interface to sundials description: An interface to the solving suite SUNDIALS. Currently, it mimics the solving interface in hmstrix-gsl but provides more diagnostic information and the Butcher Tableaux (for Runge-Kutta methods). -homepage: https://github.com/idontgetoutmuch/hmatrix/tree/sundials +homepage: https://github.com/haskell-numerics/hmatrix-sundials license: BSD3 license-file: LICENSE author: Dominic Steinitz @@ -19,43 +19,71 @@ library - build-depends: base >=4.10 && <4.11, - inline-c >=0.6 && <0.7, - vector >=0.12 && <0.13, - template-haskell >=2.12 && <2.13, - containers >=0.5 && <0.6, - hmatrix>=0.18 + build-depends: base >=4.12 && <4.13, + inline-c, + vector, + template-haskell >=2.14 && <2.15, + containers, + split, + hmatrix, + deepseq extra-libraries: sundials_arkode, sundials_cvode other-extensions: QuasiQuotes hs-source-dirs: src - exposed-modules: Numeric.Sundials.ODEOpts, + exposed-modules: Numeric.Sundials.Types, + Numeric.Sundials.Arkode, Numeric.Sundials.ARKode.ODE, Numeric.Sundials.CVode.ODE - other-modules: Numeric.Sundials.Arkode - c-sources: src/helpers.c src/helpers.h + c-sources: src/helpers.c default-language: Haskell2010 test-suite hmatrix-sundials-testsuite type: exitcode-stdio-1.0 main-is: Main.hs - other-modules: Numeric.Sundials.ODEOpts, - Numeric.Sundials.ARKode.ODE, - Numeric.Sundials.CVode.ODE, - Numeric.Sundials.Arkode - build-depends: base >=4.10 && <4.11, - inline-c >=0.6 && <0.7, - vector >=0.12 && <0.13, - template-haskell >=2.12 && <2.13, - containers >=0.5 && <0.6, - hmatrix>=0.18, + build-depends: base >=4.12 && <4.13, + inline-c, + vector, + template-haskell >=2.14 && <2.15, + containers, + split, + hmatrix, + hmatrix-sundials, plots, diagrams-lib, diagrams-rasterific, lens, - hspec - hs-source-dirs: src + hspec, + hmatrix-sundials + hs-source-dirs: test extra-libraries: sundials_arkode, sundials_cvode - c-sources: src/helpers.c src/helpers.h + c-sources: src/helpers.c default-language: Haskell2010 + +benchmark benchmark + type: + exitcode-stdio-1.0 + hs-source-dirs: + benchmark + main-is: + benchmark.hs + default-language: + Haskell2010 + build-depends: + base, + hmatrix, + hmatrix-sundials, + clock >= 0.7.1, + optparse-applicative, + cassava, + bytestring + ghc-options: + -Wall -Wno-name-shadowing + default-extensions: + BangPatterns + DeriveGeneric + FlexibleInstances + RecordWildCards + NamedFieldPuns + diff --git a/src/Main.hs b/src/Main.hs deleted file mode 100644 --- a/src/Main.hs +++ /dev/null @@ -1,186 +0,0 @@ -{-# OPTIONS_GHC -Wall #-} - -import qualified Numeric.Sundials.ARKode.ODE as ARK -import qualified Numeric.Sundials.CVode.ODE as CV -import Numeric.LinearAlgebra - -import Plots as P -import qualified Diagrams.Prelude as D -import Diagrams.Backend.Rasterific - -import Control.Lens - -import Test.Hspec - - -lorenz :: Double -> [Double] -> [Double] -lorenz _t u = [ sigma * (y - x) - , x * (rho - z) - y - , x * y - beta * z - ] - where - rho = 28.0 - sigma = 10.0 - beta = 8.0 / 3.0 - x = u !! 0 - y = u !! 1 - z = u !! 2 - -_lorenzJac :: Double -> Vector Double -> Matrix Double -_lorenzJac _t u = (3><3) [ (-sigma), rho - z, y - , sigma , -1.0 , x - , 0.0 , (-x) , (-beta) - ] - where - rho = 28.0 - sigma = 10.0 - beta = 8.0 / 3.0 - x = u ! 0 - y = u ! 1 - z = u ! 2 - -brusselator :: Double -> [Double] -> [Double] -brusselator _t x = [ a - (w + 1) * u + v * u * u - , w * u - v * u * u - , (b - w) / eps - w * u - ] - where - a = 1.0 - b = 3.5 - eps = 5.0e-6 - u = x !! 0 - v = x !! 1 - w = x !! 2 - -_brussJac :: Double -> Vector Double -> Matrix Double -_brussJac _t x = (3><3) [ (-(w + 1.0)) + 2.0 * u * v, w - 2.0 * u * v, (-w) - , u * u , (-(u * u)) , 0.0 - , (-u) , u , (-1.0) / eps - u - ] - where - y = toList x - u = y !! 0 - v = y !! 1 - w = y !! 2 - eps = 5.0e-6 - -stiffish :: Double -> [Double] -> [Double] -stiffish t v = [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ] - where - lamda = -100.0 - u = v !! 0 - -stiffishV :: Double -> Vector Double -> Vector Double -stiffishV t v = fromList [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ] - where - lamda = -100.0 - u = v ! 0 - -_stiffJac :: Double -> Vector Double -> Matrix Double -_stiffJac _t _v = (1><1) [ lamda ] - where - lamda = -100.0 - -predatorPrey :: Double -> [Double] -> [Double] -predatorPrey _t v = [ x * a - b * x * y - , d * x * y - c * y - e * y * z - , (-f) * z + g * y * z - ] - where - x = v!!0 - y = v!!1 - z = v!!2 - a = 1.0 - b = 1.0 - c = 1.0 - d = 1.0 - e = 1.0 - f = 1.0 - g = 1.0 - -lSaxis :: [[Double]] -> P.Axis B D.V2 Double -lSaxis xs = P.r2Axis &~ do - let ts = xs!!0 - us = xs!!1 - vs = xs!!2 - ws = xs!!3 - P.linePlot' $ zip ts us - P.linePlot' $ zip ts vs - P.linePlot' $ zip ts ws - -kSaxis :: [(Double, Double)] -> P.Axis B D.V2 Double -kSaxis xs = P.r2Axis &~ do - P.linePlot' xs - -main :: IO () -main = do - - let res1 = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0]) - renderRasterific "diagrams/brusselator.png" - (D.dims2D 500.0 500.0) - (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1)) - - let res1a = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0]) - renderRasterific "diagrams/brusselatorA.png" - (D.dims2D 500.0 500.0) - (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1a)) - - let res2 = ARK.odeSolve stiffish [0.0] (fromList [0.0, 0.1 .. 10.0]) - renderRasterific "diagrams/stiffish.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip [0.0, 0.1 .. 10.0] (concat $ toLists res2)) - - let res2a = ARK.odeSolveV (ARK.SDIRK_5_3_4') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) - - let res2b = ARK.odeSolveV (ARK.TRBDF2_3_3_2') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) - - let maxDiffA = maximum $ map abs $ - zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2b)!!0) - - let res2c = CV.odeSolveV (CV.BDF) Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) - - let maxDiffB = maximum $ map abs $ - zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2c)!!0) - - let maxDiffC = maximum $ map abs $ - zipWith (-) ((toLists $ tr res2b)!!0) ((toLists $ tr res2c)!!0) - - let res3 = ARK.odeSolve lorenz [-5.0, -5.0, 1.0] (fromList [0.0, 0.01 .. 10.0]) - - renderRasterific "diagrams/lorenz.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!1)) - - renderRasterific "diagrams/lorenz1.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!2)) - - renderRasterific "diagrams/lorenz2.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!1) ((toLists $ tr res3)!!2)) - - let res4 = CV.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0]) - - renderRasterific "diagrams/predatorPrey.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!1)) - - renderRasterific "diagrams/predatorPrey1.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!2)) - - renderRasterific "diagrams/predatorPrey2.png" - (D.dims2D 500.0 500.0) - (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!1) ((toLists $ tr res4)!!2)) - - let res4a = ARK.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0]) - - let maxDiffPpA = maximum $ map abs $ - zipWith (-) ((toLists $ tr res4)!!0) ((toLists $ tr res4a)!!0) - - hspec $ describe "Compare results" $ do - it "for SDIRK_5_3_4' and TRBDF2_3_3_2'" $ maxDiffA < 1.0e-6 - it "for SDIRK_5_3_4' and BDF" $ maxDiffB < 1.0e-6 - it "for TRBDF2_3_3_2' and BDF" $ maxDiffC < 1.0e-6 - it "for CV and ARK for the Predator Prey model" $ maxDiffPpA < 1.0e-3 - diff --git a/src/Numeric/Sundials/ARKode/ODE.hs b/src/Numeric/Sundials/ARKode/ODE.hs --- a/src/Numeric/Sundials/ARKode/ODE.hs +++ b/src/Numeric/Sundials/ARKode/ODE.hs @@ -1,8 +1,11 @@ +{-# OPTIONS_GHC -Wall -Wno-partial-type-signatures #-} + {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE MultiWayIf #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PartialTypeSignatures #-} {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE KindSignatures #-} @@ -64,7 +67,9 @@ -- (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1)) -- @ -- --- With Sundials ARKode, it is possible to retrieve the Butcher tableau for the solver. +-- With Sundials ARKode, it is possible to retrieve the Butcher +-- tableau for the solver. FIXME: Not available just now and hopefully +-- normal service will be resumed soon. -- -- @ -- import Numeric.Sundials.ARKode.ODE @@ -157,8 +162,7 @@ , odeSolveV , odeSolveVWith , odeSolveVWith' - , ButcherTable(..) - , butcherTable + , odeSolveWithEvents , ODEMethod(..) , StepControl(..) ) where @@ -169,9 +173,9 @@ import Data.Monoid ((<>)) import Data.Maybe (isJust) -import Foreign.C.Types (CDouble, CInt, CLong) +import Foreign.C.Types (CDouble, CInt) import Foreign.Ptr (Ptr) -import Foreign.Storable (poke) +import Foreign.Storable (poke, peek) import qualified Data.Vector.Storable as V @@ -184,12 +188,12 @@ import Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, rows, cols, toLists, size, reshape, - subVector, subMatrix, (><)) + (><)) -import Numeric.Sundials.ODEOpts (ODEOpts(..), Jacobian, SundialsDiagnostics(..)) +import Numeric.Sundials.Types import qualified Numeric.Sundials.Arkode as T -import Numeric.Sundials.Arkode (getDataFromContents, putDataInContents, arkSMax, - sDIRK_2_1_2, +import Numeric.Sundials.Arkode (SunIndexType) +import Numeric.Sundials.Arkode (sDIRK_2_1_2, bILLINGTON_3_3_2, tRBDF2_3_3_2, kVAERNO_4_2_3, @@ -215,16 +219,16 @@ fEHLBERG_13_7_8) -C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> T.sunCtx) +C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> sunCtx) C.include "" C.include "" C.include "" -C.include "" -- prototypes for ARKODE fcts., consts. +C.include "" -- prototypes for ARKODE fcts., consts. +C.include "" C.include "" -- serial N_Vector types, fcts., macros C.include "" -- access to dense SUNMatrix C.include "" -- access to dense SUNLinearSolver -C.include "" -- access to ARKDls interface C.include "" -- definition of type realtype C.include "" C.include "../../../helpers.h" @@ -433,14 +437,14 @@ where opts = ODEOpts { maxNumSteps = 10000 , minStep = 1.0e-12 - , relTol = error "relTol" - , absTols = error "absTol" - , initStep = error "initStep" , maxFail = 10 + , odeMethod = error "ARKode: unexpected use of ODEOpts.odeMethod" + , stepControl = error "ARKode: unexpected use of ODEOpts.stepControl" + , initStep = error "ARKode: unexpected use of ODEOpts.initStep" } odeSolveVWith' :: - ODEOpts + ODEOpts ODEMethod -> ODEMethod -> StepControl -> Maybe Double -- ^ initial step size - by default, ARKode @@ -459,7 +463,9 @@ (fromIntegral $ getMethod method) (coerce initStepSize) jacH (scise control) (coerce f) (coerce y0) (coerce tt) of Left (v, c) -> Left (reshape l (coerce v), fromIntegral c) - Right (v, d) -> Right (reshape l (coerce v), d) + Right (v, d) + | V.null y0 -> Right ((V.length tt >< 0) [], emptyDiagnostics) + | otherwise -> Right (reshape l (coerce v), d) where l = size y0 scise (X aTol rTol) = coerce (V.replicate l aTol, rTol) @@ -476,9 +482,61 @@ -- FIXME: efficiency vs = V.fromList $ map coerce $ concat $ toLists m +-- | This function implements the same interface as +-- 'Numeric.Sundials.CVode.ODE.odeSolveWithEvents', although it does not +-- currently support events. +odeSolveWithEvents + :: ODEOpts ODEMethod + -> [EventSpec] + -- ^ Event specifications + -> Int + -- ^ Maximum number of events + -> (Double -> V.Vector Double -> V.Vector Double) + -- ^ The RHS of the system \(\dot{y} = f(t,y)\) + -> Maybe (Double -> Vector Double -> Matrix Double) + -- ^ The Jacobian (optional) + -> V.Vector Double + -- ^ Initial conditions + -> V.Vector Double + -- ^ Desired solution times + -> Either Int SundialsSolution + -- ^ Either an error code or a solution +odeSolveWithEvents opts events _ rhs _mb_jac y0 times + | (not . null) events = + -- Call error rather than return a Left because this is a programming + -- error, not just a runtime issue. + error $ "ARKode called with a non-empty list of events (" ++ show (length events) ++ + " in total).\ + \ ARKode does not support events at this point and should not be passed any." + | otherwise = + let + result :: Either (Matrix Double, Int) + (Matrix Double, SundialsDiagnostics) + result = + odeSolveVWith' opts + (odeMethod opts) + (stepControl opts) + (initStep opts) + rhs y0 times + in + case result of + Left (_, code) -> Left code + Right (mx, diagn) -> + Right $ SundialsSolution + { actualTimeGrid = times + , solutionMatrix = + -- Note: at this time, ARKode's output matrix does not + -- include the time column, so we're not dropping it + -- here unlike in CVode. If/when we add event support + -- to ARKode, this is going to change. + mx + , eventInfo = [] + , diagnostics = diagn + } + solveOdeC :: CInt -> - CLong -> + SunIndexType -> CDouble -> CInt -> Maybe CDouble -> @@ -490,8 +548,11 @@ -> Either (V.Vector CDouble, CInt) (V.Vector CDouble, SundialsDiagnostics) -- ^ Partial solution and error code or -- solution and diagnostics solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize - jacH (aTols, rTol) fun f0 ts = unsafePerformIO $ do - + jacH (aTols, rTol) fun f0 ts + | V.null f0 = -- 0-dimensional (empty) system + Right (V.empty, emptyDiagnostics) + | otherwise = + unsafePerformIO $ do let isInitStepSize :: CInt isInitStepSize = fromIntegral $ fromEnum $ isJust initStepSize ss :: CDouble @@ -503,25 +564,21 @@ Just x -> x let dim = V.length f0 - nEq :: CLong + nEq :: SunIndexType nEq = fromIntegral dim nTs :: CInt nTs = fromIntegral $ V.length ts quasiMatrixRes <- createVector ((fromIntegral dim) * (fromIntegral nTs)) qMatMut <- V.thaw quasiMatrixRes - diagnostics :: V.Vector CLong <- createVector 10 -- FIXME - diagMut <- V.thaw diagnostics + diagMut :: V.MVector _ SunIndexType <- V.thaw =<< createVector 10 -- FIXME -- We need the types that sundials expects. These are tied together -- in 'CLangToHaskellTypes'. FIXME: The Haskell type is currently empty! let funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt - funIO x y f _ptr = do - -- Convert the pointer we get from C (y) to a vector, and then - -- apply the user-supplied function. - fImm <- fun x <$> getDataFromContents dim y - -- Fill in the provided pointer with the resulting vector. - putDataInContents fImm dim f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. + funIO t y f _ptr = do + sv <- peek y + poke f $ T.SunVector { T.sunVecN = T.sunVecN sv + , T.sunVecVals = fun t (T.sunVecVals sv) + } [CU.exp| int{ 0 } |] let isJac :: CInt isJac = fromIntegral $ fromEnum $ isJust jacH @@ -531,7 +588,7 @@ jacIO t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do case jacH of Nothing -> error "Numeric.Sundials.ARKode.ODE: Jacobian not defined" - Just jacI -> do j <- jacI t <$> getDataFromContents dim y + Just jacI -> do j <- jacI t <$> (T.sunVecVals <$> peek y) poke jacS j -- FIXME: I don't understand what this comment means -- Unsafe since the function will be called many times. @@ -548,7 +605,7 @@ SUNLinearSolver LS = NULL; /* empty linear solver object */ void *arkode_mem = NULL; /* empty ARKode memory structure */ realtype t; - long int nst, nst_a, nfe, nfi, nsetups, nje, nfeLS, nni, ncfn, netf; + long nst, nst_a, nfe, nfi, nsetups, nje, nfeLS, nni, ncfn, netf; /* general problem parameters */ @@ -571,34 +628,31 @@ NV_Ith_S(tv,i) = ($vec-ptr:(double *aTols))[i]; }; - arkode_mem = ARKodeCreate(); /* Create the solver memory */ - if (check_flag((void *)arkode_mem, "ARKodeCreate", 0)) return 1; - - /* Call ARKodeInit to initialize the integrator memory and specify the */ - /* right-hand side function in y'=f(t,y), the inital time T0, and */ - /* the initial dependent variable vector y. Note: we treat the */ - /* problem as fully implicit and set f_E to NULL and f_I to f. */ + /* Call ARKStepCreate to initialize the ARK timestepper module and */ + /* specify the right-hand side function in y'=f(t,y), the inital time */ + /* T0, and the initial dependent variable vector y. Note: since this */ + /* problem is fully implicit, we set f_E to NULL and f_I to f. */ /* Here we use the C types defined in helpers.h which tie up with */ - /* the Haskell types defined in CLangToHaskellTypes */ + /* the Haskell types defined in CLangToHaskellTypes */ if ($(int method) < MIN_DIRK_NUM) { - flag = ARKodeInit(arkode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), NULL, T0, y); - if (check_flag(&flag, "ARKodeInit", 1)) return 1; - } else { - flag = ARKodeInit(arkode_mem, NULL, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y); - if (check_flag(&flag, "ARKodeInit", 1)) return 1; + arkode_mem = ARKStepCreate($fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), NULL, T0, y); + if (check_flag((void *)arkode_mem, "ARKStepCreate", 0)) return 1; + } else { + arkode_mem = ARKStepCreate(NULL, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y); + if (check_flag(&flag, "ARKStepCreate", 0)) return 1; } - flag = ARKodeSetMinStep(arkode_mem, $(double minStep_)); - if (check_flag(&flag, "ARKodeSetMinStep", 1)) return 1; - flag = ARKodeSetMaxNumSteps(arkode_mem, $(long int maxNumSteps_)); - if (check_flag(&flag, "ARKodeSetMaxNumSteps", 1)) return 1; - flag = ARKodeSetMaxErrTestFails(arkode_mem, $(int maxErrTestFails)); - if (check_flag(&flag, "ARKodeSetMaxErrTestFails", 1)) return 1; + flag = ARKStepSetMinStep(arkode_mem, $(double minStep_)); + if (check_flag(&flag, "ARKStepSetMinStep", 1)) return 1; + flag = ARKStepSetMaxNumSteps(arkode_mem, $(sunindextype maxNumSteps_)); + if (check_flag(&flag, "ARKStepSetMaxNumSteps", 1)) return 1; + flag = ARKStepSetMaxErrTestFails(arkode_mem, $(int maxErrTestFails)); + if (check_flag(&flag, "ARKStepSetMaxErrTestFails", 1)) return 1; /* Set routines */ - flag = ARKodeSVtolerances(arkode_mem, $(double rTol), tv); - if (check_flag(&flag, "ARKodeSVtolerances", 1)) return 1; + flag = ARKStepSVtolerances(arkode_mem, $(double rTol), tv); + if (check_flag(&flag, "ARKStepSVtolerances", 1)) return 1; /* Initialize dense matrix data structure and solver */ A = SUNDenseMatrix(NEQ, NEQ); @@ -607,21 +661,21 @@ if (check_flag((void *)LS, "SUNDenseLinearSolver", 0)) return 1; /* Attach matrix and linear solver */ - flag = ARKDlsSetLinearSolver(arkode_mem, LS, A); - if (check_flag(&flag, "ARKDlsSetLinearSolver", 1)) return 1; + flag = ARKStepSetLinearSolver(arkode_mem, LS, A); + if (check_flag(&flag, "ARKStepSetLinearSolver", 1)) return 1; /* Set the initial step size if there is one */ if ($(int isInitStepSize)) { /* FIXME: We could check if the initial step size is 0 */ /* or even NaN and then throw an error */ - flag = ARKodeSetInitStep(arkode_mem, $(double ss)); - if (check_flag(&flag, "ARKodeSetInitStep", 1)) return 1; + flag = ARKStepSetInitStep(arkode_mem, $(double ss)); + if (check_flag(&flag, "ARKStepSetInitStep", 1)) return 1; } /* Set the Jacobian if there is one */ if ($(int isJac)) { - flag = ARKDlsSetJacFn(arkode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[]))); - if (check_flag(&flag, "ARKDlsSetJacFn", 1)) return 1; + flag = ARKStepSetJacFn(arkode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[]))); + if (check_flag(&flag, "ARKStepSetJacFn", 1)) return 1; } /* Store initial conditions */ @@ -631,19 +685,19 @@ /* Explicitly set the method */ if ($(int method) >= MIN_DIRK_NUM) { - flag = ARKodeSetIRKTableNum(arkode_mem, $(int method)); - if (check_flag(&flag, "ARKodeSetIRKTableNum", 1)) return 1; + flag = ARKStepSetTableNum(arkode_mem, $(int method), -1); + if (check_flag(&flag, "ARKStepSetTableNum", 1)) return 1; } else { - flag = ARKodeSetERKTableNum(arkode_mem, $(int method)); - if (check_flag(&flag, "ARKodeSetERKTableNum", 1)) return 1; + flag = ARKStepSetTableNum(arkode_mem, -1, $(int method)); + if (check_flag(&flag, "ERKStepSetTableNum", 1)) return 1; } - /* Main time-stepping loop: calls ARKode to perform the integration */ + /* Main time-stepping loop: calls ARKStep to perform the integration */ /* Stops when the final time has been reached */ for (i = 1; i < $(int nTs); i++) { - flag = ARKode(arkode_mem, ($vec-ptr:(double *ts))[i], y, &t, ARK_NORMAL); /* call integrator */ - if (check_flag(&flag, "ARKode solver failure, stopping integration", 1)) return 1; + flag = ARKStepEvolve(arkode_mem, ($vec-ptr:(double *ts))[i], y, &t, ARK_NORMAL); /* call integrator */ + if (check_flag(&flag, "ARKStep solver failure, stopping integration", 1)) return 1; /* Store the results for Haskell */ for (j = 0; j < NEQ; j++) { @@ -653,47 +707,47 @@ /* Get some final statistics on how the solve progressed */ - flag = ARKodeGetNumSteps(arkode_mem, &nst); - check_flag(&flag, "ARKodeGetNumSteps", 1); - ($vec-ptr:(long int *diagMut))[0] = nst; + flag = ARKStepGetNumSteps(arkode_mem, &nst); + check_flag(&flag, "ARKStepGetNumSteps", 1); + ($vec-ptr:(sunindextype *diagMut))[0] = nst; - flag = ARKodeGetNumStepAttempts(arkode_mem, &nst_a); - check_flag(&flag, "ARKodeGetNumStepAttempts", 1); - ($vec-ptr:(long int *diagMut))[1] = nst_a; + flag = ARKStepGetNumStepAttempts(arkode_mem, &nst_a); + check_flag(&flag, "ARKStepGetNumStepAttempts", 1); + ($vec-ptr:(sunindextype *diagMut))[1] = nst_a; - flag = ARKodeGetNumRhsEvals(arkode_mem, &nfe, &nfi); - check_flag(&flag, "ARKodeGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[2] = nfe; - ($vec-ptr:(long int *diagMut))[3] = nfi; + flag = ARKStepGetNumRhsEvals(arkode_mem, &nfe, &nfi); + check_flag(&flag, "ARKStepGetNumRhsEvals", 1); + ($vec-ptr:(sunindextype *diagMut))[2] = nfe; + ($vec-ptr:(sunindextype *diagMut))[3] = nfi; - flag = ARKodeGetNumLinSolvSetups(arkode_mem, &nsetups); - check_flag(&flag, "ARKodeGetNumLinSolvSetups", 1); - ($vec-ptr:(long int *diagMut))[4] = nsetups; + flag = ARKStepGetNumLinSolvSetups(arkode_mem, &nsetups); + check_flag(&flag, "ARKStepGetNumLinSolvSetups", 1); + ($vec-ptr:(sunindextype *diagMut))[4] = nsetups; - flag = ARKodeGetNumErrTestFails(arkode_mem, &netf); - check_flag(&flag, "ARKodeGetNumErrTestFails", 1); - ($vec-ptr:(long int *diagMut))[5] = netf; + flag = ARKStepGetNumErrTestFails(arkode_mem, &netf); + check_flag(&flag, "ARKStepGetNumErrTestFails", 1); + ($vec-ptr:(sunindextype *diagMut))[5] = netf; - flag = ARKodeGetNumNonlinSolvIters(arkode_mem, &nni); - check_flag(&flag, "ARKodeGetNumNonlinSolvIters", 1); - ($vec-ptr:(long int *diagMut))[6] = nni; + flag = ARKStepGetNumNonlinSolvIters(arkode_mem, &nni); + check_flag(&flag, "ARKStepGetNumNonlinSolvIters", 1); + ($vec-ptr:(sunindextype *diagMut))[6] = nni; - flag = ARKodeGetNumNonlinSolvConvFails(arkode_mem, &ncfn); - check_flag(&flag, "ARKodeGetNumNonlinSolvConvFails", 1); - ($vec-ptr:(long int *diagMut))[7] = ncfn; + flag = ARKStepGetNumNonlinSolvConvFails(arkode_mem, &ncfn); + check_flag(&flag, "ARKStepGetNumNonlinSolvConvFails", 1); + ($vec-ptr:(sunindextype *diagMut))[7] = ncfn; - flag = ARKDlsGetNumJacEvals(arkode_mem, &nje); - check_flag(&flag, "ARKDlsGetNumJacEvals", 1); - ($vec-ptr:(long int *diagMut))[8] = ncfn; + flag = ARKStepGetNumJacEvals(arkode_mem, &nje); + check_flag(&flag, "ARKStepGetNumJacEvals", 1); + ($vec-ptr:(sunindextype *diagMut))[8] = nje; - flag = ARKDlsGetNumRhsEvals(arkode_mem, &nfeLS); - check_flag(&flag, "ARKDlsGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[9] = ncfn; + flag = ARKStepGetNumLinRhsEvals(arkode_mem, &nfeLS); + check_flag(&flag, "ARKStepGetNumLinRhsEvals", 1); + ($vec-ptr:(sunindextype *diagMut))[9] = nfeLS; /* Clean up and return */ N_VDestroy(y); /* Free y vector */ N_VDestroy(tv); /* Free tv vector */ - ARKodeFree(&arkode_mem); /* Free integrator memory */ + ARKStepFree(&arkode_mem); /* Free integrator memory */ SUNLinSolFree(LS); /* Free linear solver */ SUNMatDestroy(A); /* Free A matrix */ @@ -716,181 +770,3 @@ return $ Right (m, d) else do return $ Left (m, res) - -data ButcherTable = ButcherTable { am :: Matrix Double - , cv :: Vector Double - , bv :: Vector Double - , b2v :: Vector Double - } - deriving Show - -data ButcherTable' a = ButcherTable' { am' :: V.Vector a - , cv' :: V.Vector a - , bv' :: V.Vector a - , b2v' :: V.Vector a - } - deriving Show - -butcherTable :: ODEMethod -> ButcherTable -butcherTable method = - case getBT method of - Left c -> error $ show c -- FIXME - Right (ButcherTable' v w x y, sqp) -> - ButcherTable { am = subMatrix (0, 0) (s, s) $ (arkSMax >< arkSMax) (V.toList v) - , cv = subVector 0 s w - , bv = subVector 0 s x - , b2v = subVector 0 s y - } - where - s = fromIntegral $ sqp V.! 0 - -getBT :: ODEMethod -> Either Int (ButcherTable' Double, V.Vector Int) -getBT method = case getButcherTable method of - Left c -> - Left $ fromIntegral c - Right (ButcherTable' a b c d, sqp) -> - Right $ ( ButcherTable' (coerce a) (coerce b) (coerce c) (coerce d) - , V.map fromIntegral sqp ) - -getButcherTable :: ODEMethod - -> Either CInt (ButcherTable' CDouble, V.Vector CInt) -getButcherTable method = unsafePerformIO $ do - -- ARKode seems to want an ODE in order to set and then get the - -- Butcher tableau so here's one to keep it happy - let funI :: CDouble -> V.Vector CDouble -> V.Vector CDouble - funI _t ys = V.fromList [ ys V.! 0 ] - let funE :: CDouble -> V.Vector CDouble -> V.Vector CDouble - funE _t ys = V.fromList [ ys V.! 0 ] - f0 = V.fromList [ 1.0 ] - ts = V.fromList [ 0.0 ] - dim = V.length f0 - nEq :: CLong - nEq = fromIntegral dim - mN :: CInt - mN = fromIntegral $ getMethod method - - btSQP :: V.Vector CInt <- createVector 3 - btSQPMut <- V.thaw btSQP - btAs :: V.Vector CDouble <- createVector (arkSMax * arkSMax) - btAsMut <- V.thaw btAs - btCs :: V.Vector CDouble <- createVector arkSMax - btBs :: V.Vector CDouble <- createVector arkSMax - btB2s :: V.Vector CDouble <- createVector arkSMax - btCsMut <- V.thaw btCs - btBsMut <- V.thaw btBs - btB2sMut <- V.thaw btB2s - let funIOI :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt - funIOI x y f _ptr = do - fImm <- funI x <$> getDataFromContents dim y - putDataInContents fImm dim f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] - let funIOE :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt - funIOE x y f _ptr = do - fImm <- funE x <$> getDataFromContents dim y - putDataInContents fImm dim f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] - res <- [C.block| int { - /* general problem variables */ - - int flag; /* reusable error-checking flag */ - N_Vector y = NULL; /* empty vector for storing solution */ - void *arkode_mem = NULL; /* empty ARKode memory structure */ - int i, j; /* reusable loop indices */ - - /* general problem parameters */ - - realtype T0 = RCONST(($vec-ptr:(double *ts))[0]); /* initial time */ - sunindextype NEQ = $(sunindextype nEq); /* number of dependent vars */ - - /* Initialize data structures */ - - y = N_VNew_Serial(NEQ); /* Create serial vector for solution */ - if (check_flag((void *)y, "N_VNew_Serial", 0)) return 1; - /* Specify initial condition */ - for (i = 0; i < NEQ; i++) { - NV_Ith_S(y,i) = ($vec-ptr:(double *f0))[i]; - }; - arkode_mem = ARKodeCreate(); /* Create the solver memory */ - if (check_flag((void *)arkode_mem, "ARKodeCreate", 0)) return 1; - - flag = ARKodeInit(arkode_mem, $fun:(int (* funIOE) (double t, SunVector y[], SunVector dydt[], void * params)), $fun:(int (* funIOI) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y); - if (check_flag(&flag, "ARKodeInit", 1)) return 1; - - if ($(int mN) >= MIN_DIRK_NUM) { - flag = ARKodeSetIRKTableNum(arkode_mem, $(int mN)); - if (check_flag(&flag, "ARKodeSetIRKTableNum", 1)) return 1; - } else { - flag = ARKodeSetERKTableNum(arkode_mem, $(int mN)); - if (check_flag(&flag, "ARKodeSetERKTableNum", 1)) return 1; - } - - int s, q, p; - realtype *ai = (realtype *)malloc(ARK_S_MAX * ARK_S_MAX * sizeof(realtype)); - realtype *ae = (realtype *)malloc(ARK_S_MAX * ARK_S_MAX * sizeof(realtype)); - realtype *ci = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - realtype *ce = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - realtype *bi = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - realtype *be = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - realtype *b2i = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - realtype *b2e = (realtype *)malloc(ARK_S_MAX * sizeof(realtype)); - flag = ARKodeGetCurrentButcherTables(arkode_mem, &s, &q, &p, ai, ae, ci, ce, bi, be, b2i, b2e); - if (check_flag(&flag, "ARKode", 1)) return 1; - $vec-ptr:(int *btSQPMut)[0] = s; - $vec-ptr:(int *btSQPMut)[1] = q; - $vec-ptr:(int *btSQPMut)[2] = p; - for (i = 0; i < s; i++) { - for (j = 0; j < s; j++) { - /* FIXME: double should be realtype */ - ($vec-ptr:(double *btAsMut))[i * ARK_S_MAX + j] = ai[i * ARK_S_MAX + j]; - } - } - - for (i = 0; i < s; i++) { - ($vec-ptr:(double *btCsMut))[i] = ci[i]; - ($vec-ptr:(double *btBsMut))[i] = bi[i]; - ($vec-ptr:(double *btB2sMut))[i] = b2i[i]; - } - - /* Clean up and return */ - N_VDestroy(y); /* Free y vector */ - ARKodeFree(&arkode_mem); /* Free integrator memory */ - - return flag; - } |] - if res == 0 - then do - x <- V.freeze btAsMut - y <- V.freeze btSQPMut - z <- V.freeze btCsMut - u <- V.freeze btBsMut - v <- V.freeze btB2sMut - return $ Right (ButcherTable' { am' = x, cv' = z, bv' = u, b2v' = v }, y) - else do - return $ Left res - --- | Adaptive step-size control --- functions. --- --- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control) --- allows the user to control the step size adjustment using --- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where --- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\) --- is the required relative error, \(s_i\) is a vector of scaling --- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and --- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\). --- --- [ARKode](https://computation.llnl.gov/projects/sundials/arkode) --- allows the user to control the step size adjustment using --- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with --- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl), --- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no --- effect. -data StepControl = X Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical - | X' Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem - | XX' Double Double Double Double -- ^ include both via relative tolerance - -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\) - | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\) diff --git a/src/Numeric/Sundials/Arkode.hsc b/src/Numeric/Sundials/Arkode.hsc --- a/src/Numeric/Sundials/Arkode.hsc +++ b/src/Numeric/Sundials/Arkode.hsc @@ -3,7 +3,44 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE EmptyDataDecls #-} -module Numeric.Sundials.Arkode where +module Numeric.Sundials.Arkode ( getDataFromContents + , putDataInContents + , cV_ADAMS + , cV_BDF + , vectorToC + , cV_SUCCESS + , cV_ROOT_RETURN + , SunIndexType + , SunRealType + , SunMatrix(..) + , SunVector(..) + , sunContentLengthOffset + , sunContentDataOffset + , hEUN_EULER_2_1_2 + , bOGACKI_SHAMPINE_4_2_3 + , aRK324L2SA_ERK_4_2_3 + , zONNEVELD_5_3_4 + , aRK436L2SA_ERK_6_3_4 + , sAYFY_ABURUB_6_3_4 + , cASH_KARP_6_4_5 + , fEHLBERG_6_4_5 + , dORMAND_PRINCE_7_4_5 + , aRK548L2SA_ERK_8_4_5 + , vERNER_8_5_6 + , fEHLBERG_13_7_8 + , sDIRK_2_1_2 + , bILLINGTON_3_3_2 + , tRBDF2_3_3_2 + , kVAERNO_4_2_3 + , aRK324L2SA_DIRK_4_2_3 + , cASH_5_2_4 + , cASH_5_3_4 + , sDIRK_5_3_4 + , kVAERNO_5_3_4 + , aRK436L2SA_DIRK_6_3_4 + , kVAERNO_7_4_5 + , aRK548L2SA_DIRK_8_4_5 + ) where import Foreign import Foreign.C.Types @@ -26,28 +63,22 @@ #include #include #include +#include +#include #include -data SunVector +data SunVector = SunVector { sunVecN :: SunIndexType + , sunVecVals :: V.Vector CDouble + } + data SunMatrix = SunMatrix { rows :: CInt , cols :: CInt , vals :: V.Vector CDouble } --- | This is true only if configured/ built as 64 bits -type SunIndexType = CLong - -sunTypesTable :: Map.Map TypeSpecifier TH.TypeQ -sunTypesTable = Map.fromList - [ - (TypeName "sunindextype", [t| SunIndexType |] ) - , (TypeName "SunVector", [t| SunVector |] ) - , (TypeName "SunMatrix", [t| SunMatrix |] ) - ] - -sunCtx :: Context -sunCtx = mempty {ctxTypesTable = sunTypesTable} +type SunIndexType = #type sunindextype +type SunRealType = #type realtype getMatrixDataFromContents :: Ptr SunMatrix -> IO SunMatrix getMatrixDataFromContents ptr = do @@ -69,6 +100,15 @@ rtr <- getMatrixData qtr vectorToC vs (fromIntegral $ rs * cs) rtr +instance Storable SunVector where + poke p v = putDataInContents (sunVecVals v) (fromIntegral $ sunVecN v) p + peek p = do (l, v) <- getDataFromContents p + return $ SunVector { sunVecN = fromIntegral l + , sunVecVals = v + } + sizeOf _ = error "sizeOf not supported for SunVector" + alignment _ = error "alignment not supported for SunVector" + instance Storable SunMatrix where poke = flip putMatrixDataFromContents peek = getMatrixDataFromContents @@ -85,16 +125,19 @@ ptr' <- newForeignPtr_ ptr VS.copy (VM.unsafeFromForeignPtr0 ptr' len) vec -getDataFromContents :: Int -> Ptr SunVector -> IO (VS.Vector CDouble) -getDataFromContents len ptr = do +getDataFromContents :: Ptr SunVector -> IO (SunIndexType, VS.Vector CDouble) +getDataFromContents ptr = do qtr <- getContentPtr ptr rtr <- getData qtr - vectorFromC len rtr + len' <- getLength qtr + v <- vectorFromC (fromIntegral len') rtr + return (len', v) -putDataInContents :: Storable a => VS.Vector a -> Int -> Ptr b -> IO () +putDataInContents :: VS.Vector CDouble -> Int -> Ptr SunVector -> IO () putDataInContents vec len ptr = do qtr <- getContentPtr ptr rtr <- getData qtr + putLength (fromIntegral len) qtr vectorToC vec len rtr #def typedef struct _generic_N_Vector SunVector; @@ -103,6 +146,12 @@ #def typedef struct _generic_SUNMatrix SunMatrix; #def typedef struct _SUNMatrixContent_Dense SunMatrixContent; +sunContentLengthOffset :: Int +sunContentLengthOffset = #offset SunContent, length + +sunContentDataOffset :: Int +sunContentDataOffset = #offset SunContent, data + getContentMatrixPtr :: Storable a => Ptr b -> IO a getContentMatrixPtr ptr = (#peek SunMatrix, content) ptr @@ -125,6 +174,12 @@ getData :: Storable a => Ptr b -> IO a getData ptr = (#peek SunContent, data) ptr +getLength :: Ptr b -> IO SunIndexType +getLength ptr = (#peek SunContent, length) ptr + +putLength :: SunIndexType -> Ptr b -> IO () +putLength l ptr = (#poke SunContent, length) ptr l + cV_SUCCESS :: Int cV_SUCCESS = #const CV_SUCCESS cV_ROOT_RETURN :: Int @@ -134,9 +189,6 @@ cV_ADAMS = #const CV_ADAMS cV_BDF :: Int cV_BDF = #const CV_BDF - -arkSMax :: Int -arkSMax = #const ARK_S_MAX mIN_DIRK_NUM, mAX_DIRK_NUM :: Int mIN_DIRK_NUM = #const MIN_DIRK_NUM diff --git a/src/Numeric/Sundials/CVode/ODE.hs b/src/Numeric/Sundials/CVode/ODE.hs --- a/src/Numeric/Sundials/CVode/ODE.hs +++ b/src/Numeric/Sundials/CVode/ODE.hs @@ -1,10 +1,12 @@ -{-# OPTIONS_GHC -Wall #-} +{-# OPTIONS_GHC -Wall -Wno-partial-type-signatures #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE MultiWayIf #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE ForeignFunctionInterface #-} ----------------------------------------------------------------------------- -- | @@ -67,6 +69,7 @@ , odeSolveVWith , odeSolveVWith' , odeSolveRootVWith' + , odeSolveWithEvents , ODEMethod(..) , StepControl(..) , SolverResult(..) @@ -76,11 +79,12 @@ import qualified Language.C.Inline.Unsafe as CU import Data.Monoid ((<>)) -import Data.Maybe (isJust) +import Data.Maybe (isJust, fromJust) +import Data.List (genericLength) -import Foreign.C.Types (CDouble, CInt, CLong) -import Foreign.Ptr (Ptr) -import Foreign.Storable (poke) +import Foreign.C.Types (CDouble, CInt) +import Foreign.Ptr +import Foreign.Storable (peek, poke) import qualified Data.Vector.Storable as V @@ -90,16 +94,17 @@ import Numeric.LinearAlgebra.Devel (createVector) import Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, rows, - cols, toLists, size, reshape) + cols, toLists, size, reshape, + subVector, subMatrix, toColumns, fromColumns, asColumn) import Numeric.Sundials.Arkode (cV_ADAMS, cV_BDF, - getDataFromContents, putDataInContents, - vectorToC, cV_SUCCESS, cV_ROOT_RETURN) + vectorToC, cV_SUCCESS, + SunVector(..), SunIndexType) import qualified Numeric.Sundials.Arkode as T -import Numeric.Sundials.ODEOpts (ODEOpts(..), Jacobian, SundialsDiagnostics(..)) +import Numeric.Sundials.Types -C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> T.sunCtx) +C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> sunCtx) C.include "" C.include "" @@ -108,17 +113,37 @@ C.include "" -- serial N_Vector types, fcts., macros C.include "" -- access to dense SUNMatrix C.include "" -- access to dense SUNLinearSolver +C.include "" C.include "" -- access to CVDls interface C.include "" -- definition of type realtype C.include "" C.include "../../../helpers.h" C.include "Numeric/Sundials/Arkode_hsc.h" - -- | Stepping functions data ODEMethod = ADAMS | BDF + deriving (Eq, Ord, Show, Read) +-- Contrary to the documentation, it appears that CVodeGetRootInfo +-- may use both 1 and -1 to indicate a root, depending on the +-- direction of the sign change. See near the end of cvRootfind. +intToDirection :: Integral d => d -> Maybe CrossingDirection +intToDirection d = + case d of + 1 -> Just Upwards + -1 -> Just Downwards + _ -> Nothing + +-- | Almost inverse of 'intToDirection'. Map 'Upwards' to 1, 'Downwards' to +-- -1, and 'AnyDirection' to 0. +directionToInt :: Integral d => CrossingDirection -> d +directionToInt d = + case d of + Upwards -> 1 + Downwards -> -1 + AnyDirection -> 0 + getMethod :: ODEMethod -> Int getMethod (ADAMS) = cV_ADAMS getMethod (BDF) = cV_BDF @@ -160,6 +185,8 @@ where g t x0 = V.fromList $ f t (V.toList x0) +-- | A version of 'odeSolveVWith'' with reasonable default solver +-- options. odeSolveVWith :: ODEMethod -> StepControl @@ -174,39 +201,39 @@ -> V.Vector Double -- ^ Desired solution times -> Matrix Double -- ^ Error code or solution odeSolveVWith method control initStepSize f y0 tt = - case odeSolveVWith' opts method control initStepSize f y0 tt of + case odeSolveVWith' opts f y0 tt of Left (c, _v) -> error $ show c -- FIXME Right (v, _d) -> v where opts = ODEOpts { maxNumSteps = 10000 , minStep = 1.0e-12 - , relTol = error "relTol" - , absTols = error "absTol" - , initStep = error "initStep" , maxFail = 10 + , odeMethod = method + , stepControl = control + , initStep = initStepSize } odeSolveVWith' :: - ODEOpts - -> ODEMethod - -> StepControl - -> Maybe Double -- ^ initial step size - by default, CVode - -- estimates the initial step size to be the - -- solution \(h\) of the equation - -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where - -- \(\ddot{y}\) is an estimated value of the second - -- derivative of the solution at \(t_0\) + ODEOpts ODEMethod -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\) -> V.Vector Double -- ^ Initial conditions -> V.Vector Double -- ^ Desired solution times -> Either (Matrix Double, Int) (Matrix Double, SundialsDiagnostics) -- ^ Error code or solution -odeSolveVWith' opts method control initStepSize f y0 tt = +odeSolveVWith' opts f y0 tt = case solveOdeC (fromIntegral $ maxFail opts) - (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts) - (fromIntegral $ getMethod method) (coerce initStepSize) jacH (scise control) - (coerce f) (coerce y0) (coerce tt) of - Left (v, c) -> Left (reshape l (coerce v), fromIntegral c) - Right (v, d) -> Right (reshape l (coerce v), d) + (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts) + (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts) + (OdeRhsHaskell $ coerce f) (coerce y0) + 0 (\_ x -> x) [] 0 (\_ _ y -> y) (coerce tt) of + -- Remove the time column for backwards compatibility + SolverError m c -> Left + ( subMatrix (0, 1) (V.length tt, l) m + , fromIntegral c + ) + SolverSuccess _ m d -> Right + ( subMatrix (0, 1) (V.length tt, l) m + , d + ) where l = size y0 scise (X aTol rTol) = coerce (V.replicate l aTol, rTol) @@ -215,257 +242,47 @@ -- FIXME; Should we check that the length of ss is correct? scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol) jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $ - getJacobian method - matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs } - where - nr = fromIntegral $ rows m - nc = fromIntegral $ cols m - -- FIXME: efficiency - vs = V.fromList $ map coerce $ concat $ toLists m - -solveOdeC :: - CInt -> - CLong -> - CDouble -> - CInt -> - Maybe CDouble -> - (Maybe (CDouble -> V.Vector CDouble -> T.SunMatrix)) -> - (V.Vector CDouble, CDouble) -> - (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The RHS of the system \(\dot{y} = f(t,y)\) - -> V.Vector CDouble -- ^ Initial conditions - -> V.Vector CDouble -- ^ Desired solution times - -> Either (V.Vector CDouble, CInt) (V.Vector CDouble, SundialsDiagnostics) -- ^ Partial solution and error code or - -- solution and diagnostics -solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize - jacH (aTols, rTol) fun f0 ts = - unsafePerformIO $ do - - let isInitStepSize :: CInt - isInitStepSize = fromIntegral $ fromEnum $ isJust initStepSize - ss :: CDouble - ss = case initStepSize of - -- It would be better to put an error message here but - -- inline-c seems to evaluate this even if it is never - -- used :( - Nothing -> 0.0 - Just x -> x - - let dim = V.length f0 - nEq :: CLong - nEq = fromIntegral dim - nTs :: CInt - nTs = fromIntegral $ V.length ts - quasiMatrixRes <- createVector ((fromIntegral dim) * (fromIntegral nTs)) - qMatMut <- V.thaw quasiMatrixRes - diagnostics :: V.Vector CLong <- createVector 10 -- FIXME - diagMut <- V.thaw diagnostics - -- We need the types that sundials expects. These are tied together - -- in 'CLangToHaskellTypes'. FIXME: The Haskell type is currently empty! - let funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt - funIO x y f _ptr = do - -- Convert the pointer we get from C (y) to a vector, and then - -- apply the user-supplied function. - fImm <- fun x <$> getDataFromContents dim y - -- Fill in the provided pointer with the resulting vector. - putDataInContents fImm dim f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] - let isJac :: CInt - isJac = fromIntegral $ fromEnum $ isJust jacH - jacIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunMatrix -> - Ptr () -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunVector -> - IO CInt - jacIO t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do - case jacH of - Nothing -> error "Numeric.Sundials.CVode.ODE: Jacobian not defined" - Just jacI -> do j <- jacI t <$> getDataFromContents dim y - poke jacS j - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] - - res <- [C.block| int { - /* general problem variables */ - - int flag; /* reusable error-checking flag */ - int i, j; /* reusable loop indices */ - N_Vector y = NULL; /* empty vector for storing solution */ - N_Vector tv = NULL; /* empty vector for storing absolute tolerances */ - - SUNMatrix A = NULL; /* empty matrix for linear solver */ - SUNLinearSolver LS = NULL; /* empty linear solver object */ - void *cvode_mem = NULL; /* empty CVODE memory structure */ - realtype t; - long int nst, nfe, nsetups, nje, nfeLS, nni, ncfn, netf, nge; - - /* general problem parameters */ - - realtype T0 = RCONST(($vec-ptr:(double *ts))[0]); /* initial time */ - sunindextype NEQ = $(sunindextype nEq); /* number of dependent vars. */ - - /* Initialize data structures */ - - y = N_VNew_Serial(NEQ); /* Create serial vector for solution */ - if (check_flag((void *)y, "N_VNew_Serial", 0)) return 1; - /* Specify initial condition */ - for (i = 0; i < NEQ; i++) { - NV_Ith_S(y,i) = ($vec-ptr:(double *f0))[i]; - }; - - cvode_mem = CVodeCreate($(int method), CV_NEWTON); - if (check_flag((void *)cvode_mem, "CVodeCreate", 0)) return(1); - - /* Call CVodeInit to initialize the integrator memory and specify the - * user's right hand side function in y'=f(t,y), the inital time T0, and - * the initial dependent variable vector y. */ - flag = CVodeInit(cvode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y); - if (check_flag(&flag, "CVodeInit", 1)) return(1); - - tv = N_VNew_Serial(NEQ); /* Create serial vector for absolute tolerances */ - if (check_flag((void *)tv, "N_VNew_Serial", 0)) return 1; - /* Specify tolerances */ - for (i = 0; i < NEQ; i++) { - NV_Ith_S(tv,i) = ($vec-ptr:(double *aTols))[i]; - }; - - flag = CVodeSetMinStep(cvode_mem, $(double minStep_)); - if (check_flag(&flag, "CVodeSetMinStep", 1)) return 1; - flag = CVodeSetMaxNumSteps(cvode_mem, $(long int maxNumSteps_)); - if (check_flag(&flag, "CVodeSetMaxNumSteps", 1)) return 1; - flag = CVodeSetMaxErrTestFails(cvode_mem, $(int maxErrTestFails)); - if (check_flag(&flag, "CVodeSetMaxErrTestFails", 1)) return 1; - - /* Call CVodeSVtolerances to specify the scalar relative tolerance - * and vector absolute tolerances */ - flag = CVodeSVtolerances(cvode_mem, $(double rTol), tv); - if (check_flag(&flag, "CVodeSVtolerances", 1)) return(1); - - /* Initialize dense matrix data structure and solver */ - A = SUNDenseMatrix(NEQ, NEQ); - if (check_flag((void *)A, "SUNDenseMatrix", 0)) return 1; - LS = SUNDenseLinearSolver(y, A); - if (check_flag((void *)LS, "SUNDenseLinearSolver", 0)) return 1; - - /* Attach matrix and linear solver */ - flag = CVDlsSetLinearSolver(cvode_mem, LS, A); - if (check_flag(&flag, "CVDlsSetLinearSolver", 1)) return 1; - - /* Set the initial step size if there is one */ - if ($(int isInitStepSize)) { - /* FIXME: We could check if the initial step size is 0 */ - /* or even NaN and then throw an error */ - flag = CVodeSetInitStep(cvode_mem, $(double ss)); - if (check_flag(&flag, "CVodeSetInitStep", 1)) return 1; - } - - /* Set the Jacobian if there is one */ - if ($(int isJac)) { - flag = CVDlsSetJacFn(cvode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[]))); - if (check_flag(&flag, "CVDlsSetJacFn", 1)) return 1; - } - - /* Store initial conditions */ - for (j = 0; j < NEQ; j++) { - ($vec-ptr:(double *qMatMut))[0 * $(int nTs) + j] = NV_Ith_S(y,j); - } - - /* Main time-stepping loop: calls CVode to perform the integration */ - /* Stops when the final time has been reached */ - for (i = 1; i < $(int nTs); i++) { - - flag = CVode(cvode_mem, ($vec-ptr:(double *ts))[i], y, &t, CV_NORMAL); /* call integrator */ - if (check_flag(&flag, "CVode solver failure, stopping integration", 1)) return 1; - - /* Store the results for Haskell */ - for (j = 0; j < NEQ; j++) { - ($vec-ptr:(double *qMatMut))[i * NEQ + j] = NV_Ith_S(y,j); - } - } - - /* Get some final statistics on how the solve progressed */ - - flag = CVodeGetNumSteps(cvode_mem, &nst); - check_flag(&flag, "CVodeGetNumSteps", 1); - ($vec-ptr:(long int *diagMut))[0] = nst; - - /* FIXME */ - ($vec-ptr:(long int *diagMut))[1] = 0; - - flag = CVodeGetNumRhsEvals(cvode_mem, &nfe); - check_flag(&flag, "CVodeGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[2] = nfe; - /* FIXME */ - ($vec-ptr:(long int *diagMut))[3] = 0; - - flag = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups); - check_flag(&flag, "CVodeGetNumLinSolvSetups", 1); - ($vec-ptr:(long int *diagMut))[4] = nsetups; - - flag = CVodeGetNumErrTestFails(cvode_mem, &netf); - check_flag(&flag, "CVodeGetNumErrTestFails", 1); - ($vec-ptr:(long int *diagMut))[5] = netf; - - flag = CVodeGetNumNonlinSolvIters(cvode_mem, &nni); - check_flag(&flag, "CVodeGetNumNonlinSolvIters", 1); - ($vec-ptr:(long int *diagMut))[6] = nni; - - flag = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn); - check_flag(&flag, "CVodeGetNumNonlinSolvConvFails", 1); - ($vec-ptr:(long int *diagMut))[7] = ncfn; - - flag = CVDlsGetNumJacEvals(cvode_mem, &nje); - check_flag(&flag, "CVDlsGetNumJacEvals", 1); - ($vec-ptr:(long int *diagMut))[8] = ncfn; - - flag = CVDlsGetNumRhsEvals(cvode_mem, &nfeLS); - check_flag(&flag, "CVDlsGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[9] = ncfn; - - /* Clean up and return */ + getJacobian $ odeMethod opts - N_VDestroy(y); /* Free y vector */ - N_VDestroy(tv); /* Free tv vector */ - CVodeFree(&cvode_mem); /* Free integrator memory */ - SUNLinSolFree(LS); /* Free linear solver */ - SUNMatDestroy(A); /* Free A matrix */ +matrixToSunMatrix :: Matrix Double -> T.SunMatrix +matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs } + where + nr = fromIntegral $ rows m + nc = fromIntegral $ cols m + -- FIXME: efficiency + vs = V.fromList $ map coerce $ concat $ toLists m - return flag; - } |] - preD <- V.freeze diagMut - let d = SundialsDiagnostics (fromIntegral $ preD V.!0) - (fromIntegral $ preD V.!1) - (fromIntegral $ preD V.!2) - (fromIntegral $ preD V.!3) - (fromIntegral $ preD V.!4) - (fromIntegral $ preD V.!5) - (fromIntegral $ preD V.!6) - (fromIntegral $ preD V.!7) - (fromIntegral $ preD V.!8) - (fromIntegral $ preD V.!9) - m <- V.freeze qMatMut - if res == 0 - then do - return $ Right (m, d) - else do - return $ Left (m, res) +foreign import ccall "wrapper" + mkOdeRhsC :: OdeRhsCType -> IO (FunPtr OdeRhsCType) -solveOdeC' :: +solveOdeC :: CInt -> - CLong -> + SunIndexType -> CDouble -> CInt -> Maybe CDouble -> (Maybe (CDouble -> V.Vector CDouble -> T.SunMatrix)) -> (V.Vector CDouble, CDouble) -> - (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The RHS of the system \(\dot{y} = f(t,y)\) + OdeRhs -- ^ The RHS of the system \(\dot{y} = f(t,y)\) -> V.Vector CDouble -- ^ Initial conditions - -> CInt -- ^ FIXME - -> (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ FIXME + -> CInt -- ^ Number of event equations + -> (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The event equations themselves + -> [CrossingDirection] -- ^ The required crossing direction for each event + -> CInt -- ^ Maximum number of events + -> (Int -> CDouble -> V.Vector CDouble -> V.Vector CDouble) + -- ^ Function to reset/update the state when an event occurs. The + -- 'Int' argument is the 0-based number of the event that has + -- occurred. If multiple events have occurred at the same time, they + -- are handled in the increasing order of the event index. The other + -- arguments are the time and the point in the state space. Return + -- the updated point in the state space. -> V.Vector CDouble -- ^ Desired solution times - -> SolverResult V.Vector V.Vector CInt CDouble -solveOdeC' maxErrTestFails maxNumSteps_ minStep_ method initStepSize - jacH (aTols, rTol) fun f0 nr g ts = + -> SolverResult +solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize + jacH (aTols, rTol) rhs f0 nr event_fn directions max_events apply_event ts + | V.null f0 = -- 0-dimensional (empty) system + SolverSuccess [] (asColumn (coerce ts)) emptyDiagnostics + | otherwise = unsafePerformIO $ do let isInitStepSize :: CInt @@ -479,44 +296,61 @@ Just x -> x let dim = V.length f0 - nEq :: CLong + nEq :: SunIndexType nEq = fromIntegral dim nTs :: CInt nTs = fromIntegral $ V.length ts - quasiMatrixRes <- createVector ((fromIntegral dim) * (fromIntegral nTs)) - qMatMut <- V.thaw quasiMatrixRes - diagnostics :: V.Vector CLong <- createVector 10 -- FIXME - diagMut <- V.thaw diagnostics - -- We need the types that sundials expects. - -- FIXME: The Haskell type is currently empty! - let funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt - funIO t y f _ptr = do - -- Convert the pointer we get from C (y) to a vector, and then - -- apply the user-supplied function. - fImm <- fun t <$> getDataFromContents dim y - -- Fill in the provided pointer with the resulting vector. - putDataInContents fImm dim f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] + output_mat_mut :: V.MVector _ CDouble <- V.thaw =<< createVector ((1 + fromIntegral dim) * (fromIntegral (2 * max_events) + fromIntegral nTs)) + diagMut :: V.MVector _ SunIndexType <- V.thaw =<< createVector 10 -- FIXME + (rhs_funptr :: FunPtr OdeRhsCType, userdata :: Ptr UserData) <- + case rhs of + OdeRhsC ptr u -> return (ptr, u) + OdeRhsHaskell fun -> do + let + funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr UserData -> IO CInt + funIO t y f _ptr = do + sv <- peek y + poke f $ SunVector { sunVecN = sunVecN sv + , sunVecVals = fun t (sunVecVals sv) + } + return 0 + funptr <- mkOdeRhsC funIO + return (funptr, nullPtr) let nrPre = fromIntegral nr gResults :: V.Vector CInt <- createVector nrPre + -- FIXME: Do we need to do this here? Maybe as it will get GC'd and + -- we'd have to do a malloc in C otherwise :( gResMut <- V.thaw gResults - tRoot :: V.Vector CDouble <- createVector 1 - tRootMut <- V.thaw tRoot + event_index_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events) + event_time_mut :: V.MVector _ CDouble <- V.thaw =<< createVector (fromIntegral max_events) + -- Total number of events. This is *not* directly re + n_events_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1 + -- Total number of rows in the output_mat_mut matrix. It *cannot* be + -- inferred from n_events_mut because when an event occurs k times, it + -- contributes k to n_events_mut but only 2 to n_rows_mut. + n_rows_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1 + actual_event_direction_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events) - let gIO :: CDouble -> Ptr T.SunVector -> Ptr CDouble -> Ptr () -> IO CInt - gIO x y f _ptr = do - -- Convert the pointer we get from C (y) to a vector, and then - -- apply the user-supplied function. - gImm <- g x <$> getDataFromContents dim y - -- Fill in the provided pointer with the resulting vector. - vectorToC gImm nrPre f - -- FIXME: I don't understand what this comment means - -- Unsafe since the function will be called many times. - [CU.exp| int{ 0 } |] + let event_fn_c :: CDouble -> Ptr T.SunVector -> Ptr CDouble -> Ptr () -> IO CInt + event_fn_c x y f _ptr = do + vals <- event_fn x <$> (sunVecVals <$> peek y) + -- FIXME: We should be able to use poke somehow + vectorToC vals nrPre f + return 0 + apply_event_c :: CInt -> CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> IO CInt + apply_event_c event_index t y y' = do + sv <- peek y + poke y' $ SunVector + { sunVecN = sunVecN sv + , sunVecVals = apply_event (fromIntegral event_index) t (sunVecVals sv) + } + return 0 + + requested_event_directions :: V.Vector CInt + requested_event_directions = V.fromList $ map directionToInt directions + let isJac :: CInt isJac = fromIntegral $ fromEnum $ isJust jacH jacIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunMatrix -> @@ -525,17 +359,16 @@ jacIO t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do case jacH of Nothing -> error "Numeric.Sundials.CVode.ODE: Jacobian not defined" - Just jacI -> do j <- jacI t <$> getDataFromContents dim y + Just jacI -> do j <- jacI t <$> (sunVecVals <$> peek y) poke jacS j -- FIXME: I don't understand what this comment means -- Unsafe since the function will be called many times. [CU.exp| int{ 0 } |] - res <- [C.block| int { + res :: Int <- fromIntegral <$> [C.block| int { /* general problem variables */ int flag; /* reusable error-checking flag */ - int flagr; /* root finding flag */ int i, j; /* reusable loop indices */ N_Vector y = NULL; /* empty vector for storing solution */ @@ -545,10 +378,25 @@ SUNLinearSolver LS = NULL; /* empty linear solver object */ void *cvode_mem = NULL; /* empty CVODE memory structure */ realtype t; - long int nst, nfe, nsetups, nje, nfeLS, nni, ncfn, netf, nge; + long nst, nfe, nsetups, nje, nfeLS, nni, ncfn, netf, nge; realtype tout; + /* input_ind tracks the current index into the ts array */ + int input_ind = 1; + /* output_ind tracks the current row into the output_mat_mut matrix. + If differs from input_ind because of the extra rows corresponding to events. */ + int output_ind = 1; + /* We need to update n_rows_mut every time we update output_ind because + of the possibility of early return (in which case we still need to assemble + the partial results matrix). We could even work with n_rows_mut only and ditch + output_ind, but the inline-c expression is quite verbose, and output_ind is + more convenient to use in index calculations. + */ + ($vec-ptr:(int *n_rows_mut))[0] = output_ind; + /* event_ind tracks the current event number */ + int event_ind = 0; + /* general problem parameters */ realtype T0 = RCONST(($vec-ptr:(double *ts))[0]); /* initial time */ @@ -563,14 +411,16 @@ NV_Ith_S(y,i) = ($vec-ptr:(double *f0))[i]; }; - cvode_mem = CVodeCreate($(int method), CV_NEWTON); + cvode_mem = CVodeCreate($(int method)); if (check_flag((void *)cvode_mem, "CVodeCreate", 0)) return(1); /* Call CVodeInit to initialize the integrator memory and specify the * user's right hand side function in y'=f(t,y), the inital time T0, and * the initial dependent variable vector y. */ - flag = CVodeInit(cvode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y); + flag = CVodeInit(cvode_mem, $(int (* rhs_funptr) (double t, SunVector y[], SunVector dydt[], UserData* params)), T0, y); if (check_flag(&flag, "CVodeInit", 1)) return(1); + flag = CVodeSetUserData(cvode_mem, $(UserData* userdata)); + if (check_flag(&flag, "CVodeSetUserData", 1)) return(1); tv = N_VNew_Serial(NEQ); /* Create serial vector for absolute tolerances */ if (check_flag((void *)tv, "N_VNew_Serial", 0)) return 1; @@ -581,7 +431,7 @@ flag = CVodeSetMinStep(cvode_mem, $(double minStep_)); if (check_flag(&flag, "CVodeSetMinStep", 1)) return 1; - flag = CVodeSetMaxNumSteps(cvode_mem, $(long int maxNumSteps_)); + flag = CVodeSetMaxNumSteps(cvode_mem, $(sunindextype maxNumSteps_)); if (check_flag(&flag, "CVodeSetMaxNumSteps", 1)) return 1; flag = CVodeSetMaxErrTestFails(cvode_mem, $(int maxErrTestFails)); if (check_flag(&flag, "CVodeSetMaxErrTestFails", 1)) return 1; @@ -591,8 +441,8 @@ flag = CVodeSVtolerances(cvode_mem, $(double rTol), tv); if (check_flag(&flag, "CVodeSVtolerances", 1)) return(1); - /* Call CVodeRootInit to specify the root function g with nr components */ - flag = CVodeRootInit(cvode_mem, $(int nr), $fun:(int (* gIO) (double t, SunVector y[], double gout[], void * params))); + /* Call CVodeRootInit to specify the root function event_fn_c with nr components */ + flag = CVodeRootInit(cvode_mem, $(int nr), $fun:(int (* event_fn_c) (double t, SunVector y[], double gout[], void * params))); if (check_flag(&flag, "CVodeRootInit", 1)) return(1); @@ -621,69 +471,120 @@ } /* Store initial conditions */ + ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + 0] = ($vec-ptr:(double *ts))[0]; for (j = 0; j < NEQ; j++) { - ($vec-ptr:(double *qMatMut))[0 * $(int nTs) + j] = NV_Ith_S(y,j); + ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j); } - /* Main time-stepping loop: calls CVode to perform the integration */ - /* Stops when the final time has been reached */ - for (i = 1; i < $(int nTs); i++) { - - flag = CVode(cvode_mem, ($vec-ptr:(double *ts))[i], y, &t, CV_NORMAL); /* call integrator */ + while (1) { + flag = CVode(cvode_mem, ($vec-ptr:(double *ts))[input_ind], y, &t, CV_NORMAL); /* call integrator */ if (check_flag(&flag, "CVode solver failure, stopping integration", 1)) return 1; /* Store the results for Haskell */ + ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t; for (j = 0; j < NEQ; j++) { - ($vec-ptr:(double *qMatMut))[i * NEQ + j] = NV_Ith_S(y,j); + ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j); } + output_ind++; + ($vec-ptr:(int *n_rows_mut))[0] = output_ind; if (flag == CV_ROOT_RETURN) { - flagr = CVodeGetRootInfo(cvode_mem, ($vec-ptr:(int *gResMut))); - if (check_flag(&flagr, "CVodeGetRootInfo", 1)) return(1); - ($vec-ptr:(double *tRootMut))[0] = t; - flagr = flag; - break; + if (event_ind >= $(int max_events)) { + /* We reached the maximum number of events. + Either the maximum number of events is set to 0, + or there's a bug in our code below. In any case return an error. + */ + return 1; + } + + /* Are we interested in this event? + If not, continue without any observable side-effects. + */ + int good_event = 0; + flag = CVodeGetRootInfo(cvode_mem, ($vec-ptr:(int *gResMut))); + if (check_flag(&flag, "CVodeGetRootInfo", 1)) return 1; + for (i = 0; i < $(int nr); i++) { + int ev = ($vec-ptr:(int *gResMut))[i]; + int req_dir = ($vec-ptr:(const int *requested_event_directions))[i]; + if (ev != 0 && ev * req_dir >= 0) { + good_event = 1; + + ($vec-ptr:(int *actual_event_direction_mut))[event_ind] = ev; + ($vec-ptr:(int *event_index_mut))[event_ind] = i; + ($vec-ptr:(double *event_time_mut))[event_ind] = t; + event_ind++; + + /* Update the state with the supplied function */ + $fun:(int (* apply_event_c) (int, double, SunVector y[], SunVector z[]))(i, t, y, y); + } + } + + if (good_event) { + ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t; + for (j = 0; j < NEQ; j++) { + ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j); + } + output_ind++; + ($vec-ptr:(int *n_rows_mut))[0] = output_ind; + + if (event_ind >= $(int max_events)) { + /* We collected the requested number of events. Stop the solver. */ + break; + } + flag = CVodeReInit(cvode_mem, t, y); + if (check_flag(&flag, "CVodeReInit", 1)) return(1); + } else { + /* Since this is not a wanted event, it shouldn't get a row */ + output_ind--; + ($vec-ptr:(int *n_rows_mut))[0] = output_ind; + } } + else { + if (++input_ind >= $(int nTs)) + break; + } } - /* Get some final statistics on how the solve progressed */ + /* The number of actual roots we found */ + ($vec-ptr:(int *n_events_mut))[0] = event_ind; + /* Get some final statistics on how the solve progressed */ flag = CVodeGetNumSteps(cvode_mem, &nst); check_flag(&flag, "CVodeGetNumSteps", 1); - ($vec-ptr:(long int *diagMut))[0] = nst; + ($vec-ptr:(sunindextype *diagMut))[0] = nst; /* FIXME */ - ($vec-ptr:(long int *diagMut))[1] = 0; + ($vec-ptr:(sunindextype *diagMut))[1] = 0; flag = CVodeGetNumRhsEvals(cvode_mem, &nfe); check_flag(&flag, "CVodeGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[2] = nfe; + ($vec-ptr:(sunindextype *diagMut))[2] = nfe; /* FIXME */ - ($vec-ptr:(long int *diagMut))[3] = 0; + ($vec-ptr:(sunindextype *diagMut))[3] = 0; flag = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups); check_flag(&flag, "CVodeGetNumLinSolvSetups", 1); - ($vec-ptr:(long int *diagMut))[4] = nsetups; + ($vec-ptr:(sunindextype *diagMut))[4] = nsetups; flag = CVodeGetNumErrTestFails(cvode_mem, &netf); check_flag(&flag, "CVodeGetNumErrTestFails", 1); - ($vec-ptr:(long int *diagMut))[5] = netf; + ($vec-ptr:(sunindextype *diagMut))[5] = netf; flag = CVodeGetNumNonlinSolvIters(cvode_mem, &nni); check_flag(&flag, "CVodeGetNumNonlinSolvIters", 1); - ($vec-ptr:(long int *diagMut))[6] = nni; + ($vec-ptr:(sunindextype *diagMut))[6] = nni; flag = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn); check_flag(&flag, "CVodeGetNumNonlinSolvConvFails", 1); - ($vec-ptr:(long int *diagMut))[7] = ncfn; + ($vec-ptr:(sunindextype *diagMut))[7] = ncfn; flag = CVDlsGetNumJacEvals(cvode_mem, &nje); check_flag(&flag, "CVDlsGetNumJacEvals", 1); - ($vec-ptr:(long int *diagMut))[8] = ncfn; + ($vec-ptr:(sunindextype *diagMut))[8] = ncfn; flag = CVDlsGetNumRhsEvals(cvode_mem, &nfeLS); check_flag(&flag, "CVDlsGetNumRhsEvals", 1); - ($vec-ptr:(long int *diagMut))[9] = ncfn; + ($vec-ptr:(sunindextype *diagMut))[9] = ncfn; /* Clean up and return */ @@ -693,13 +594,15 @@ SUNLinSolFree(LS); /* Free linear solver */ SUNMatDestroy(A); /* Free A matrix */ - if (flag == CV_SUCCESS && flagr == CV_ROOT_RETURN) { - return CV_ROOT_RETURN; - } - else { - return flag; - } + return CV_SUCCESS; } |] + + -- Free the allocated FunPtr. Ideally this should be done within + -- a bracket... + case rhs of + OdeRhsHaskell {} -> freeHaskellFunPtr rhs_funptr + OdeRhsC {} -> return () -- we didn't allocate this + preD <- V.freeze diagMut let d = SundialsDiagnostics (fromIntegral $ preD V.!0) (fromIntegral $ preD V.!1) @@ -711,48 +614,59 @@ (fromIntegral $ preD V.!7) (fromIntegral $ preD V.!8) (fromIntegral $ preD V.!9) - m <- V.freeze qMatMut - t <- V.freeze tRootMut - rs <- V.freeze gResMut - putStrLn $ show rs - let f r | r == cV_SUCCESS = SolverSuccess m d - | r == cV_ROOT_RETURN = SolverRoot (t V.!0) rs m d - | otherwise = SolverError m res - return $ f $ fromIntegral res + n_rows <- fromIntegral . V.head <$> V.freeze n_rows_mut + output_mat <- coerce . reshape (dim + 1) . subVector 0 ((dim + 1) * n_rows) <$> + V.freeze output_mat_mut + n_events <- fromIntegral . V.head <$> V.freeze n_events_mut + event_time :: V.Vector Double + <- coerce . V.take n_events <$> V.freeze event_time_mut + event_index :: V.Vector Int + <- V.map fromIntegral . V.take n_events <$> V.freeze event_index_mut + actual_event_direction :: V.Vector CInt + <- V.take n_events <$> V.freeze actual_event_direction_mut + let + events :: [EventInfo] + events = zipWith3 EventInfo + (V.toList event_time) + (V.toList event_index) + (map (fromJust . intToDirection) $ V.toList actual_event_direction) + return $ + if res == cV_SUCCESS + then + SolverSuccess events output_mat d + else + SolverError output_mat res -data SolverResult f g a b = - SolverError (f b) a -- ^ Partial results and error code - | SolverSuccess (f b) SundialsDiagnostics -- ^ Results and diagnostics - | SolverRoot b (g a) (f b) SundialsDiagnostics -- ^ Time at which the root was found, the root itself and the - -- results and diagnostics. NB the final result will be at the time - -- at which the root was found not as specified by the times given - -- to the solver. +data SolverResult + = SolverError !(Matrix Double) !Int + -- ^ Partial results and error code + | SolverSuccess + [EventInfo] + !(Matrix Double) + !SundialsDiagnostics + -- ^ Times at which the event was triggered, information about which root and the + -- results and diagnostics. deriving Show odeSolveRootVWith' :: - ODEOpts - -> ODEMethod - -> StepControl - -> Maybe Double -- ^ initial step size - by default, CVode - -- estimates the initial step size to be the - -- solution \(h\) of the equation - -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where - -- \(\ddot{y}\) is an estimated value of the second - -- derivative of the solution at \(t_0\) - -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\) - -> V.Vector Double -- ^ Initial conditions - -> Int -- ^ Dimension of the range of the roots function - -> (Double -> V.Vector Double -> V.Vector Double) -- ^ Roots function - -> V.Vector Double -- ^ Desired solution times - -> SolverResult Matrix Vector Int Double -odeSolveRootVWith' opts method control initStepSize f y0 is gg tt = - case solveOdeC' (fromIntegral $ maxFail opts) + ODEOpts ODEMethod + -> OdeRhs + -- ^ The RHS of the system \(\dot{y} = f(t,y)\) + -> Maybe (Double -> Vector Double -> Matrix Double) + -- ^ The Jacobian (optional) + -> V.Vector Double -- ^ Initial conditions + -> [EventSpec] -- ^ Event specifications + -> Int -- ^ Maximum number of events + -> V.Vector Double -- ^ Desired solution times + -> SolverResult +odeSolveRootVWith' opts rhs mb_jacobian y0 event_specs nRootEvs tt = + solveOdeC (fromIntegral $ maxFail opts) (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts) - (fromIntegral $ getMethod method) (coerce initStepSize) jacH (scise control) - (coerce f) (coerce y0) (fromIntegral is) (coerce gg) (coerce tt) of - SolverError v c -> SolverError (reshape l (coerce v)) (fromIntegral c) - SolverSuccess v d -> SolverSuccess (reshape l (coerce v)) d - SolverRoot t rs v d -> SolverRoot (coerce t) (V.map fromIntegral rs) (reshape l (coerce v)) d + (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts) + rhs (coerce y0) + (genericLength event_specs) event_equations event_directions + (fromIntegral nRootEvs) reset_state + (coerce tt) where l = size y0 scise (X aTol rTol) = coerce (V.replicate l aTol, rTol) @@ -760,34 +674,50 @@ scise (XX' aTol rTol yScale _yDotScale) = coerce (V.replicate l aTol, yScale * rTol) -- FIXME; Should we check that the length of ss is correct? scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol) - jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $ - getJacobian method - matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs } - where - nr = fromIntegral $ rows m - nc = fromIntegral $ cols m - -- FIXME: efficiency - vs = V.fromList $ map coerce $ concat $ toLists m + jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $ mb_jacobian + event_equations :: CDouble -> Vector CDouble -> Vector CDouble + event_equations t y = V.fromList $ + map (\ev -> coerce (eventCondition ev) t y) event_specs + event_directions :: [CrossingDirection] + event_directions = map eventDirection event_specs + reset_state :: Int -> CDouble -> Vector CDouble -> Vector CDouble + reset_state n_event = coerce $ eventUpdate (event_specs !! n_event) --- | Adaptive step-size control --- functions. --- --- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control) --- allows the user to control the step size adjustment using --- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where --- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\) --- is the required relative error, \(s_i\) is a vector of scaling --- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and --- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\). --- --- [ARKode](https://computation.llnl.gov/projects/sundials/arkode) --- allows the user to control the step size adjustment using --- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with --- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl), --- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no --- effect. -data StepControl = X Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical - | X' Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem - | XX' Double Double Double Double -- ^ include both via relative tolerance - -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\) - | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\) +odeSolveWithEvents + :: ODEOpts ODEMethod + -> [EventSpec] + -- ^ Event specifications + -> Int + -- ^ Maximum number of events + -> OdeRhs + -- ^ The RHS of the system \(\dot{y} = f(t,y)\) + -> Maybe (Double -> Vector Double -> Matrix Double) + -- ^ The Jacobian (optional) + -> V.Vector Double + -- ^ Initial conditions + -> V.Vector Double + -- ^ Desired solution times + -> Either Int SundialsSolution + -- ^ Either an error code or a solution +odeSolveWithEvents opts event_specs max_events rhs mb_jacobian initial sol_times = + let + result :: SolverResult + result = + odeSolveRootVWith' opts rhs mb_jacobian initial event_specs + max_events sol_times + in + case result of + SolverError _ code -> Left code + SolverSuccess events mx diagn -> + Right $ SundialsSolution + { actualTimeGrid = extractTimeGrid mx + , solutionMatrix = dropTimeGrid mx + , eventInfo = events + , diagnostics = diagn + } + where + -- The time grid is the first column of the result matrix + extractTimeGrid :: Matrix Double -> Vector Double + extractTimeGrid = head . toColumns + dropTimeGrid :: Matrix Double -> Matrix Double + dropTimeGrid = fromColumns . tail . toColumns diff --git a/src/Numeric/Sundials/ODEOpts.hs b/src/Numeric/Sundials/ODEOpts.hs deleted file mode 100644 --- a/src/Numeric/Sundials/ODEOpts.hs +++ /dev/null @@ -1,32 +0,0 @@ -module Numeric.Sundials.ODEOpts where - -import Data.Word (Word32) -import qualified Data.Vector.Storable as VS - -import Numeric.LinearAlgebra.HMatrix (Vector, Matrix) - - -type Jacobian = Double -> Vector Double -> Matrix Double - -data ODEOpts = ODEOpts { - maxNumSteps :: Word32 - , minStep :: Double - , relTol :: Double - , absTols :: VS.Vector Double - , initStep :: Maybe Double - , maxFail :: Word32 - } deriving (Read, Show, Eq, Ord) - -data SundialsDiagnostics = SundialsDiagnostics { - aRKodeGetNumSteps :: Int - , aRKodeGetNumStepAttempts :: Int - , aRKodeGetNumRhsEvals_fe :: Int - , aRKodeGetNumRhsEvals_fi :: Int - , aRKodeGetNumLinSolvSetups :: Int - , aRKodeGetNumErrTestFails :: Int - , aRKodeGetNumNonlinSolvIters :: Int - , aRKodeGetNumNonlinSolvConvFails :: Int - , aRKDlsGetNumJacEvals :: Int - , aRKDlsGetNumRhsEvals :: Int - } deriving Show - diff --git a/src/Numeric/Sundials/Types.hs b/src/Numeric/Sundials/Types.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Sundials/Types.hs @@ -0,0 +1,148 @@ +{-# LANGUAGE DeriveAnyClass, DeriveGeneric, TemplateHaskell, OverloadedStrings #-} +module Numeric.Sundials.Types + ( OdeRhsCType + , OdeRhs(..) + , UserData + , Jacobian + , StepControl(..) + , ODEOpts(..) + , SundialsDiagnostics(..) + , emptyDiagnostics + , SundialsSolution(..) + , EventInfo(..) + , CrossingDirection(..) + , EventSpec(..) + , SunVector(..) + , SunIndexType + , SunRealType + , sunContentLengthOffset + , sunContentDataOffset + , sunCtx + ) + where + +import Data.Int (Int32) +import qualified Data.Vector.Storable as VS +import qualified Data.Map.Strict as Map +import qualified Language.Haskell.TH as TH + +import Numeric.LinearAlgebra.HMatrix (Vector, Matrix) +import Control.DeepSeq (NFData) +import GHC.Generics (Generic) +import Foreign.C.Types +import Foreign.Ptr +import Language.C.Types as CT +import Language.C.Inline.Context +import Numeric.Sundials.Arkode (SunVector(..), SunMatrix(..), + SunIndexType, SunRealType, + sunContentLengthOffset, + sunContentDataOffset) + +-- | The type of the C ODE RHS function. +type OdeRhsCType = CDouble -> Ptr SunVector -> Ptr SunVector -> Ptr UserData -> IO CInt + +data UserData + +-- | The right-hand side of the ODE system. +-- +-- Can be either a Haskell function or a pointer to a C function. +data OdeRhs + = OdeRhsHaskell (CDouble -> VS.Vector CDouble -> VS.Vector CDouble) + | OdeRhsC (FunPtr OdeRhsCType) (Ptr UserData) + +type Jacobian = Double -> Vector Double -> Matrix Double + +-- | Adaptive step-size control +-- functions. +-- +-- [GSL](https://www.gnu.org/software/gsl/doc/html/ode-initval.html#adaptive-step-size-control) +-- allows the user to control the step size adjustment using +-- \(D_i = \epsilon^{abs}s_i + \epsilon^{rel}(a_{y} |y_i| + a_{dy/dt} h |\dot{y}_i|)\) where +-- \(\epsilon^{abs}\) is the required absolute error, \(\epsilon^{rel}\) +-- is the required relative error, \(s_i\) is a vector of scaling +-- factors, \(a_{y}\) is a scaling factor for the solution \(y\) and +-- \(a_{dydt}\) is a scaling factor for the derivative of the solution \(dy/dt\). +-- +-- [CVode](https://computation.llnl.gov/projects/sundials/cvode) +-- and [ARKode](https://computation.llnl.gov/projects/sundials/arkode) +-- allow the user to control the step size adjustment using +-- \(\eta^{rel}|y_i| + \eta^{abs}_i\). For compatibility with +-- [hmatrix-gsl](https://hackage.haskell.org/package/hmatrix-gsl), +-- tolerances for \(y\) and \(\dot{y}\) can be specified but the latter have no +-- effect. +data StepControl = X Double Double -- ^ absolute and relative tolerance for \(y\); in GSL terms, \(a_{y} = 1\) and \(a_{dy/dt} = 0\); in ARKode terms, the \(\eta^{abs}_i\) are identical + | X' Double Double -- ^ absolute and relative tolerance for \(\dot{y}\); in GSL terms, \(a_{y} = 0\) and \(a_{dy/dt} = 1\); in ARKode terms, the latter is treated as the relative tolerance for \(y\) so this is the same as specifying 'X' which may be entirely incorrect for the given problem + | XX' Double Double Double Double -- ^ include both via relative tolerance + -- scaling factors \(a_y\), \(a_{{dy}/{dt}}\); in ARKode terms, the latter is ignored and \(\eta^{rel} = a_{y}\epsilon^{rel}\) + | ScXX' Double Double Double Double (Vector Double) -- ^ scale absolute tolerance of \(y_i\); in ARKode terms, \(a_{{dy}/{dt}}\) is ignored, \(\eta^{abs}_i = s_i \epsilon^{abs}\) and \(\eta^{rel} = a_{y}\epsilon^{rel}\) + deriving (Eq, Ord, Show, Read) + +data ODEOpts method = ODEOpts { + maxNumSteps :: Int32 + , minStep :: Double + , maxFail :: Int32 + , odeMethod :: method + , stepControl :: StepControl + , initStep :: Maybe Double + -- ^ initial step size - by default, CVode + -- estimates the initial step size to be the + -- solution \(h\) of the equation + -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where + -- \(\ddot{y}\) is an estimated value of the second + -- derivative of the solution at \(t_0\) + } deriving (Read, Show, Eq, Ord) + +data SundialsDiagnostics = SundialsDiagnostics { + odeGetNumSteps :: Int + , odeGetNumStepAttempts :: Int + , odeGetNumRhsEvals_fe :: Int + , odeGetNumRhsEvals_fi :: Int + , odeGetNumLinSolvSetups :: Int + , odeGetNumErrTestFails :: Int + , odeGetNumNonlinSolvIters :: Int + , odeGetNumNonlinSolvConvFails :: Int + , dlsGetNumJacEvals :: Int + , dlsGetNumRhsEvals :: Int + } deriving Show + +emptyDiagnostics :: SundialsDiagnostics +emptyDiagnostics = SundialsDiagnostics 0 0 0 0 0 0 0 0 0 0 + +data SundialsSolution = + SundialsSolution + { actualTimeGrid :: VS.Vector Double -- ^ actual time grid returned by the solver (with duplicated event times) + , solutionMatrix :: Matrix Double -- ^ matrix of solutions: each column is an unknwown + , eventInfo :: [EventInfo] -- ^ event infos, as many items as triggered events during the simulation + , diagnostics :: SundialsDiagnostics -- ^ usual Sundials diagnostics + } + +data EventInfo = + EventInfo + { eventTime :: !Double -- ^ time at which event was triggered + , eventIndex :: !Int -- ^ which index was triggered + , rootDirection :: !CrossingDirection -- ^ in which direction ((+)->(-) or (-)->(+)) the root is crossed + } + deriving (Generic, Show, NFData) + +-- | The direction in which a function should cross the x axis +data CrossingDirection = Upwards | Downwards | AnyDirection + deriving (Generic, Eq, Show, NFData) + +data EventSpec = EventSpec + { eventCondition :: Double -> VS.Vector Double -> Double + , eventDirection :: CrossingDirection + , eventUpdate :: Double -> VS.Vector Double -> VS.Vector Double + } + +sunTypesTable :: Map.Map TypeSpecifier TH.TypeQ +sunTypesTable = Map.fromList + [ + (TypeName "sunindextype", [t| SunIndexType |] ) + , (TypeName "SunVector", [t| SunVector |] ) + , (TypeName "SunMatrix", [t| SunMatrix |] ) + , (TypeName "UserData", [t| UserData |] ) + ] + +-- | Allows to map between Haskell and C types +sunCtx :: Context +sunCtx = mempty {ctxTypesTable = sunTypesTable} diff --git a/src/helpers.c b/src/helpers.c --- a/src/helpers.c +++ b/src/helpers.c @@ -4,7 +4,6 @@ #include /* serial N_Vector types, fcts., macros */ #include /* access to dense SUNMatrix */ #include /* access to dense SUNLinearSolver */ -#include /* access to ARKDls interface */ #include /* definition of type realtype */ #include diff --git a/src/helpers.h b/src/helpers.h deleted file mode 100644 --- a/src/helpers.h +++ /dev/null @@ -1,9 +0,0 @@ -/* Check function return value... - opt == 0 means SUNDIALS function allocates memory so check if - returned NULL pointer - opt == 1 means SUNDIALS function returns a flag so check if - flag >= 0 - opt == 2 means function allocates memory so check if returned - NULL pointer -*/ -int check_flag(void *flagvalue, const char *funcname, int opt); diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,539 @@ +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# OPTIONS_GHC -Wall #-} + +import qualified Numeric.Sundials.ARKode.ODE as ARK +import qualified Numeric.Sundials.CVode.ODE as CV +import Numeric.LinearAlgebra as L +import Numeric.Sundials.Types + +import Plots as P +import qualified Diagrams.Prelude as D +import Diagrams.Backend.Rasterific +import qualified Data.Vector.Storable as V + +import Control.Lens +import Control.Monad +import Data.Coerce +import Foreign.C.Types + +import Test.Hspec + + +lorenz :: Double -> [Double] -> [Double] +lorenz _t u = [ sigma * (y - x) + , x * (rho - z) - y + , x * y - beta * z + ] + where + rho = 28.0 + sigma = 10.0 + beta = 8.0 / 3.0 + x = u !! 0 + y = u !! 1 + z = u !! 2 + +_lorenzJac :: Double -> Vector Double -> Matrix Double +_lorenzJac _t u = (3><3) [ (-sigma), rho - z, y + , sigma , -1.0 , x + , 0.0 , (-x) , (-beta) + ] + where + rho = 28.0 + sigma = 10.0 + beta = 8.0 / 3.0 + x = u ! 0 + y = u ! 1 + z = u ! 2 + +brusselator :: Double -> [Double] -> [Double] +brusselator _t x = [ a - (w + 1) * u + v * u * u + , w * u - v * u * u + , (b - w) / eps - w * u + ] + where + a = 1.0 + b = 3.5 + eps = 5.0e-6 + u = x !! 0 + v = x !! 1 + w = x !! 2 + +brussJac :: Double -> Vector Double -> Matrix Double +brussJac _t x = tr $ + (3><3) [ (-(w + 1.0)) + 2.0 * u * v, w - 2.0 * u * v, (-w) + , u * u , (-(u * u)) , 0.0 + , (-u) , u , (-1.0) / eps - u + ] + where + y = toList x + u = y !! 0 + v = y !! 1 + w = y !! 2 + eps = 5.0e-6 + +brusselatorWithJacobian :: Vector Double -> Bool -> CV.SolverResult +brusselatorWithJacobian ts usejac = CV.odeSolveRootVWith' opts + (OdeRhsHaskell . coerce $ \t v -> vector $ brusselator t (toList v)) + (if usejac then Just brussJac else Nothing) + (vector [1.2, 3.1, 3.0]) + [] 0 + ts + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1 + , initStep = Nothing + } + +brussRoot :: CV.SolverResult +brussRoot = CV.odeSolveRootVWith' opts + (OdeRhsHaskell . coerce $ \t v -> vector $ brusselator t (toList v)) + Nothing + (vector [1.2, 3.1, 3.0]) + events 100 + (vector [0.0, 0.1 .. 10.0]) + where + events = + [ EventSpec { eventCondition = brussRootFn + , eventUpdate = + \_ev x -> let y = toList x in vector [(y!!0) + 0.5 , (y!!1), (y!!2)] + , eventDirection = AnyDirection + } + ] + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1 + , initStep = Nothing + } + +brussRootFn :: Double -> Vector Double -> Double +brussRootFn _ v = case xs of + [y1, _y2, y3] -> y1 - y3 + _ -> error "brusselator root function RHS not defined" + where + xs = toList v + +exponential :: CV.SolverResult +exponential = CV.odeSolveRootVWith' opts + (OdeRhsHaskell . coerce $ \(t :: Double) y -> vector [y ! 0]) + Nothing + (vector [1]) + events 100 + (vector [ fromIntegral k / 100 | k <- [0..(22::Int)]]) + where + events = + [ EventSpec { eventCondition = \t y -> y ! 0 - 1.1 + , eventUpdate = \ev y -> vector [ 2 ] + , eventDirection = Upwards + } + ] + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1 + , initStep = Nothing + } + +-- A sine wave that only changes direction once it reaches ±0.9. +-- Illustrates event-specific reset function +boundedSine :: CV.SolverResult +boundedSine = CV.odeSolveRootVWith' + opts + (OdeRhsHaskell . coerce $ \(_t :: Double) y -> vector [y ! 1, - y ! 0]) -- ODE RHS + Nothing + (vector [0, 1]) -- initial conditions + events + 100 -- maximum number of events + (vector [ 2 * pi * k / 360 | k <- [0..360]]) -- solution times + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.ADAMS + , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1 + , initStep = Nothing + } + events = + [ EventSpec { eventCondition = \_t y -> y ! 0 - 0.9 + , eventUpdate = \_ y -> vector [ y ! 0, - abs (y ! 1) ] + , eventDirection = Upwards + } + , EventSpec { eventCondition = \_t y -> y ! 0 + 0.9 + , eventUpdate = \_ y -> vector [ y ! 0, abs (y ! 1) ] + , eventDirection = Downwards + } + ] + +stiffish :: Double -> [Double] -> [Double] +stiffish t v = [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ] + where + lamda = -100.0 + u = v !! 0 + +stiffishV :: Double -> Vector Double -> Vector Double +stiffishV t v = fromList [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ] + where + lamda = -100.0 + u = v ! 0 + +_stiffJac :: Double -> Vector Double -> Matrix Double +_stiffJac _t _v = (1><1) [ lamda ] + where + lamda = -100.0 + +predatorPrey :: Double -> [Double] -> [Double] +predatorPrey _t v = [ x * a - b * x * y + , d * x * y - c * y - e * y * z + , (-f) * z + g * y * z + ] + where + x = v!!0 + y = v!!1 + z = v!!2 + a = 1.0 + b = 1.0 + c = 1.0 + d = 1.0 + e = 1.0 + f = 1.0 + g = 1.0 + +roberts :: OdeRhs +roberts = OdeRhsHaskell . coerce $ \(t :: Double) v -> vector $ robertsAux t (toList v) + where + robertsAux _ [y1, y2, y3] = + [ -0.04 * y1 + 1.0e4 * y2 * y3 + , 0.04 * y1 - 1.0e4 * y2 * y3 - 3.0e7 * (y2)^(2 :: Int) + , 3.0e7 * (y2)^(2 :: Int) + ] + robertsAux _ _ = error "roberts RHS not defined" + +robertsJac :: Double -> Vector Double -> Matrix Double +robertsJac _t (toList -> [y1, y2, y3]) = (3 >< 3) + [ -0.04, 1.0e4 * y3, 1.0e4 * y2 + , 0.04, -1.0e4*y3 - 3.0e7*2*y2, -1.0e4*y2 + , 0, 3.0e7*2*y2, 0 + ] + +ts :: [Double] +ts = take 12 $ map (* 10.0) (0.04 : ts) + +solve :: CV.SolverResult +solve = CV.odeSolveRootVWith' opts + roberts Nothing (vector [1.0, 0.0, 0.0]) + events 100 + (vector (0.0 : ts)) + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6]) + , initStep = Nothing + } + events = + [ EventSpec { eventCondition = \_t y -> y ! 0 - 0.0001 + , eventUpdate = const id + , eventDirection = AnyDirection + } + , EventSpec { eventCondition = \_t y -> y ! 2 - 0.01 + , eventUpdate = const id + , eventDirection = AnyDirection + } + ] + +solve2 :: CV.SolverResult +solve2 = CV.odeSolveRootVWith' opts + roberts Nothing (vector [1.0, 0.0, 0.0]) + events 100 + (vector (0.0 : ts)) + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6]) + , initStep = Nothing + } + events = + [ EventSpec { eventCondition = \_t y -> y ! 0 - 0.0001 + , eventUpdate = upd + , eventDirection = AnyDirection + } + , EventSpec { eventCondition = \_t y -> y ! 2 - 0.01 + , eventUpdate = upd + , eventDirection = AnyDirection + } + ] + upd _ _ = vector [1.0, 0.0, 0.0] + +solve1 :: CV.SolverResult +solve1 = CV.odeSolveRootVWith' opts + roberts Nothing (vector [1.0, 0.0, 0.0]) + events 100 + (vector (0.0 : ts)) + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6]) + , initStep = Nothing + } + events = + [ EventSpec { eventCondition = \t _y -> t - 1.0 + , eventUpdate = \t y -> vector [2.0, y!1, y!2] + , eventDirection = AnyDirection + } + ] + +robertsonWithJacobian :: Vector Double -> Bool -> CV.SolverResult +robertsonWithJacobian ts usejac = CV.odeSolveRootVWith' opts + roberts (if usejac then Just robertsJac else Nothing) (vector [1.0, 0.0, 0.0]) + [] 0 + ts + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6]) + , initStep = Nothing + } + +lSaxis :: [[Double]] -> P.Axis B D.V2 Double +lSaxis xs = P.r2Axis &~ do + let zs = xs!!0 + us = xs!!1 + vs = xs!!2 + ws = xs!!3 + P.linePlot' $ zip zs us + P.linePlot' $ zip zs vs + P.linePlot' $ zip zs ws + +lSaxis2 :: [[Double]] -> P.Axis B D.V2 Double +lSaxis2 xs = P.r2Axis &~ do + let zs = xs!!0 + us = xs!!1 + vs = xs!!2 + P.linePlot' $ zip zs us + P.linePlot' $ zip zs vs + +kSaxis :: [(Double, Double)] -> P.Axis B D.V2 Double +kSaxis xs = P.r2Axis &~ do + P.linePlot' xs + + +main :: IO () +main = do + let res1 = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0]) + renderRasterific "diagrams/brusselator.png" + (D.dims2D 500.0 500.0) + (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1)) + + let res1a = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0]) + renderRasterific "diagrams/brusselatorA.png" + (D.dims2D 500.0 500.0) + (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1a)) + + let res2 = ARK.odeSolve stiffish [0.0] (fromList [0.0, 0.1 .. 10.0]) + renderRasterific "diagrams/stiffish.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip [0.0, 0.1 .. 10.0] (concat $ toLists res2)) + + let res2a = ARK.odeSolveV (ARK.SDIRK_5_3_4') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) + + let res2b = ARK.odeSolveV (ARK.TRBDF2_3_3_2') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) + + let maxDiffA = maximum $ map abs $ + zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2b)!!0) + + let res2c = CV.odeSolveV (CV.BDF) Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0]) + + let maxDiffB = maximum $ map abs $ + zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2c)!!0) + + let maxDiffC = maximum $ map abs $ + zipWith (-) ((toLists $ tr res2b)!!0) ((toLists $ tr res2c)!!0) + + let res3 = ARK.odeSolve lorenz [-5.0, -5.0, 1.0] (fromList [0.0, 0.01 .. 20.0]) + + renderRasterific "diagrams/lorenz.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!1)) + + renderRasterific "diagrams/lorenz1.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!2)) + + renderRasterific "diagrams/lorenz2.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!1) ((toLists $ tr res3)!!2)) + + let res4 = CV.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0]) + + renderRasterific "diagrams/predatorPrey.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!1)) + + renderRasterific "diagrams/predatorPrey1.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!2)) + + renderRasterific "diagrams/predatorPrey2.png" + (D.dims2D 500.0 500.0) + (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!1) ((toLists $ tr res4)!!2)) + + let res4a = ARK.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0]) + + let maxDiffPpA = maximum $ map abs $ + zipWith (-) ((toLists $ tr res4)!!0) ((toLists $ tr res4a)!!0) + + let cond5 = + case solve of + CV.SolverSuccess events _ _ -> do + length events `shouldBe` 2 + (abs (eventTime (events!!0) - 0.2640208751331032) / 0.2640208751331032 < 1.0e-8) `shouldBe` True + (abs (eventTime (events!!1) - 2.0786731062254436e7) / 2.0786731062254436e7 < 1.0e-8) `shouldBe` True + CV.SolverError _ _ -> + error "Root finding error!" + + let cond6 = + case solve1 of + CV.SolverSuccess events _ _ -> do + length events `shouldBe` 1 + (abs (eventTime (events!!0) - 1.0) / 1.0 < 1.0e-10) `shouldBe` True + CV.SolverError _ _ -> + error "Root finding error!" + + let cond7 = + case solve2 of + CV.SolverSuccess events _ _ -> length events `shouldBe` 100 + CV.SolverError _ _ -> error "solver failed" + True + + case brussRoot of + CV.SolverSuccess events m _diagn -> do + renderRasterific + "diagrams/brussRoot.png" + (D.dims2D 500.0 500.0) + (renderAxis $ lSaxis $ toLists $ tr m) + CV.SolverError m n -> + expectationFailure $ show n + + let boundedSineSpec = do + case boundedSine of + CV.SolverSuccess events m _ -> do + renderRasterific + "diagrams/boundedSine.png" + (D.dims2D 500.0 500.0) + (renderAxis $ lSaxis2 $ toLists $ tr m) + length events `shouldBe` 3 + map rootDirection events `shouldBe` [Upwards, Downwards, Upwards] + map eventIndex events `shouldBe` [0, 1, 0] + forM_ (zip (map eventTime events) [1.1197660081724263,3.3592952656818404,5.5988203973243]) $ \(et_got, et_exp) -> + et_got `shouldSatisfy` ((< 1e-8) . abs . subtract et_exp) + CV.SolverError m n -> + expectationFailure "Solver error" + let exponentialSpec = do + case exponential of + CV.SolverSuccess events _m _diagn -> do + length events `shouldBe` 1 + (abs (eventTime (events!!0) - log 1.1) < 1e-4) `shouldBe` True + rootDirection (events!!0) `shouldBe` Upwards + eventIndex (events!!0) `shouldBe` 0 + CV.SolverError m n -> + expectationFailure $ show n + + robertsonJac = do + let ts = vector [0, 1 .. 10] + CV.SolverSuccess _ m1 _ = robertsonWithJacobian ts True + CV.SolverSuccess _ m2 _ = robertsonWithJacobian ts False + norm_2 (m1-m2) `shouldSatisfy` (< 1e-4) + + brusselatorJac = do + let ts = [0.0, 0.1 .. 10.0] + CV.SolverSuccess _ m1 _ = brusselatorWithJacobian (vector ts) True + CV.SolverSuccess _ m2 _ = brusselatorWithJacobian (vector ts) False + norm_2 (m1-m2) `shouldSatisfy` (< 1e-3) + + hspec $ do + describe "Compare results" $ do + it "Robertson should stop early" cond7 + it "Robertson time only" $ cond6 + it "Robertson from SUNDIALS manual" $ cond5 + it "Robertson with explicit Jacobian up to t=10" robertsonJac + it "Brusselator with explicit Jacobian" brusselatorJac + it "for SDIRK_5_3_4' and TRBDF2_3_3_2'" $ maxDiffA < 1.0e-6 + it "for SDIRK_5_3_4' and BDF" $ maxDiffB < 1.0e-6 + it "for TRBDF2_3_3_2' and BDF" $ maxDiffC < 1.0e-6 + it "for CV and ARK for the Predator Prey model" $ maxDiffPpA < 1.0e-3 + describe "Handling empty systems" $ + forM_ [("CVOde",CV.odeSolve),("ARKOde",ARK.odeSolve)] $ \(name, solveFn) -> + it name $ + solveFn (\_ _ -> []) [] (V.enumFromTo 0 10) `shouldSatisfy` \sol -> + L.size sol == (11,0) + describe "Events" $ do + it "Bounded sine events" $ boundedSineSpec + it "Exponential events" $ exponentialSpec + describe "Discontinuous zero crossings" $ do + let + eq :: OdeRhs + eq = OdeRhsHaskell $ \_ _ -> V.singleton 1 + + cond + :: (Double -> Double -> Bool) + -> (Double -> Vector Double -> Double) + cond op _t y = + if V.head y `op` 0 + then 1 + else -1 + + solve op = CV.odeSolveWithEvents + opts + [EventSpec + { eventCondition = cond op + , eventDirection = AnyDirection + , eventUpdate = \_t y -> V.map (+1) y + } + ] + 5 -- max # of events + eq + Nothing + (V.singleton (-1)) + (V.fromList [0, 2]) + + ops :: [(String, Double -> Double -> Bool)] + ops = + [ (">=", (>=)) + , (">", (>)) + , ("<=", (<=)) + , ("<", (<)) + ] + forM_ ops $ \(op_name, op) -> it ("Event condition expressed as " ++ op_name) $ do + let + Right soln = solve op + evs = eventInfo soln + [ev] = evs + length evs `shouldBe` 1 + eventTime ev `shouldSatisfy` (\t -> abs (t-1) < 1e-3) + -- row 0 is time 0, rows 1 and 2 are right before and right after + -- the event (time 1), row 3 is the end point (time 2) + solutionMatrix soln ! 1 ! 0 `shouldSatisfy` (\y -> abs y < 1e-3) + solutionMatrix soln ! 2 ! 0 `shouldSatisfy` (\y -> abs (y-1) < 1e-3) + solutionMatrix soln ! 3 ! 0 `shouldSatisfy` (\y -> abs (y-2) < 1e-3) + + where + opts = ODEOpts { maxNumSteps = 10000 + , minStep = 1.0e-12 + , maxFail = 10 + , odeMethod = CV.BDF + , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6]) + , initStep = Nothing + }