With a specific implicit RK method

3 files changed, 84 insertions, 20 deletions

diff --git a/packages/sundials/hmatrix-sundials.cabal b/packages/sundials/hmatrix-sundials.cabal
index 331e6c4..762537e 100644
--- a/packages/sundials/hmatrix-sundials.cabal
+++ b/packages/sundials/hmatrix-sundials.cabal
@@ -33,14 +33,17 @@ library
 
 executable sundials
   main-is:             Main.hs
-  other-modules:       Types
+  other-modules:       Types, Numeric.Sundials.Arkode.ODE
   other-extensions:    QuasiQuotes, TemplateHaskell, MultiWayIf, OverloadedStrings
   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
+                       hmatrix>=0.18,
+                       plots,
+                       diagrams-lib,
+                       diagrams-rasterific
   hs-source-dirs:      src
   default-language:    Haskell2010
   extra-libraries:     sundials_arkode
diff --git a/packages/sundials/src/Main.hs b/packages/sundials/src/Main.hs
index 2a561c4..5e51372 100644
--- a/packages/sundials/src/Main.hs
+++ b/packages/sundials/src/Main.hs
@@ -2,28 +2,58 @@
 
 import qualified Data.Vector.Storable as V
 import           Numeric.Sundials.Arkode.ODE
+import           Numeric.LinearAlgebra
 
-brusselator :: Double -> V.Vector Double -> V.Vector Double
-brusselator _t x = V.fromList [ a - (w + 1) * u + v * u^2
-                              , w * u - v * u^2
-                              , (b - w) / eps - w * u
-                              ]
+import           Plots as P
+import qualified Diagrams.Prelude as D
+import           Diagrams.Backend.Rasterific
+
+import           Control.Lens
+import           Data.List (zip4)
+
+
+brusselator _t x = [ a - (w + 1) * u + v * u^2
+                   , w * u - v * u^2
+                   , (b - w) / eps - w * u
+                   ]
   where
     a = 1.0
     b = 3.5
     eps = 5.0e-6
-    u = x V.! 0
-    v = x V.! 1
-    w = x V.! 2
+    u = x !! 0
+    v = x !! 1
+    w = x !! 2
 
-stiffish t v = V.fromList [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
+stiffish t v = [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
   where
     lamda = -100.0
-    u = v V.! 0
+    u = v !! 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 res = solveOde brusselator (V.fromList [1.2, 3.1, 3.0]) (V.fromList [0.0, 1.0 .. 10.0])
+  let res = odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
   putStrLn $ show res
-  let res = solveOde stiffish (V.fromList [1.0]) (V.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 res))
+
+  let res = odeSolve stiffish [0.0] (fromList [0.0, 0.1 .. 10.0])
   putStrLn $ show res
+  renderRasterific "diagrams/stiffish.png"
+                   (D.dims2D 500.0 500.0)
+                   (renderAxis $ kSaxis $ zip [0.0, 0.1 .. 10.0] (concat $ toLists res))
+  
diff --git a/packages/sundials/src/Numeric/Sundials/ARKode/ODE.hs b/packages/sundials/src/Numeric/Sundials/ARKode/ODE.hs
index 630827c..f432951 100644
--- a/packages/sundials/src/Numeric/Sundials/ARKode/ODE.hs
+++ b/packages/sundials/src/Numeric/Sundials/ARKode/ODE.hs
@@ -6,7 +6,9 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Numeric.Sundials.Arkode.ODE ( solveOde ) where
+module Numeric.Sundials.Arkode.ODE ( solveOde
+                                   , odeSolve
+                                   ) where
 
 import qualified Language.C.Inline as C
 import qualified Language.C.Inline.Unsafe as CU
@@ -26,13 +28,14 @@ import           System.IO.Unsafe (unsafePerformIO)
 
 import           Numeric.LinearAlgebra.Devel (createVector)
 
-import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix)
+import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, (><))
 
 import qualified Types as T
 
 C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> T.sunCtx)
 
 -- C includes
+C.include "<stdlib.h>"
 C.include "<stdio.h>"
 C.include "<math.h>"
 C.include "<arkode/arkode.h>"                 -- prototypes for ARKODE fcts., consts.
@@ -96,7 +99,14 @@ odeSolve :: (Double -> [Double] -> [Double]) -- ^ The RHS of the system \(\dot{y
          -> [Double]                         -- ^ initial conditions
          -> Vector Double                    -- ^ desired solution times
          -> Matrix Double                    -- ^ solution
-odeSolve = undefined
+odeSolve f y0 ts = case solveOde g (V.fromList y0) (V.fromList $ toList ts) of
+                     Left c -> error $ show c -- FIXME
+                     Right (v, _) -> (nR >< nC) (V.toList v)
+  where
+    us = toList ts
+    nR = length us
+    nC = length y0
+    g t x0 = V.fromList $ f t (V.toList x0)
 
 solveOde :: 
   (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
@@ -146,13 +156,12 @@ solveOdeC fun f0 ts = unsafePerformIO $ do
                          SUNMatrix A = NULL;             /* empty matrix for linear solver */
                          SUNLinearSolver LS = NULL;      /* empty linear solver object */
                          void *arkode_mem = NULL;        /* empty ARKode memory structure */
-                         FILE *UFID;
                          realtype t;
                          long int nst, nst_a, nfe, nfi, nsetups, nje, nfeLS, nni, ncfn, netf;
 
                          /* general problem parameters */
                          realtype T0 = RCONST(($vec-ptr:(double *tMut))[0]);    /* initial time */
-                         realtype Tf = RCONST(($vec-ptr:(double *tMut))[$(int nTs) - 1]);   /* final time */
+
                          sunindextype NEQ = $(sunindextype nEq); /* number of dependent vars. */
                          realtype reltol = 1.0e-6;     /* tolerances */
                          realtype abstol = 1.0e-10;
@@ -198,7 +207,29 @@ solveOdeC fun f0 ts = unsafePerformIO $ do
                          for (j = 0; j < NEQ; j++) {
                            ($vec-ptr:(double *qMatMut))[0 * $(int nTs) + j] = NV_Ith_S(y,j);
                          }
-                         
+
+                         flag = ARKodeSetIRKTableNum(arkode_mem, KVAERNO_4_2_3);
+                         if (check_flag(&flag, "ARKode", 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;
+                         fprintf(stderr, "s = %d, q = %d, p = %d\n", s, q, p);
+                         for (i = 0; i < s; i++) {
+                           for (j = 0; j < s; j++) {
+                             fprintf(stderr, "ai[%d,%d] = %f", i, j, ai[i * ARK_S_MAX + j]);
+                           }
+                           fprintf(stderr, "\n");
+                         }
+
                          /* Main time-stepping loop: calls ARKode to perform the integration */
                          /* Stops when the final time has been reached */
                          for (i = 1; i < $(int nTs); i++) {