Generalize step control function

Extend API by odeSolveVWith and StepControl.

2 files changed, 103 insertions, 54 deletions

diff --git a/packages/gsl/src/Numeric/GSL/ODE.hs b/packages/gsl/src/Numeric/GSL/ODE.hs
index 3258b83..9e52873 100644
--- a/packages/gsl/src/Numeric/GSL/ODE.hs
+++ b/packages/gsl/src/Numeric/GSL/ODE.hs
@@ -32,7 +32,7 @@ main = mplot (ts : toColumns sol)
 -----------------------------------------------------------------------------
 
 module Numeric.GSL.ODE (
-    odeSolve, odeSolveV, ODEMethod(..), Jacobian
+    odeSolve, odeSolveV, odeSolveVWith, ODEMethod(..), Jacobian, StepControl(..)
 ) where
 
 import Numeric.LinearAlgebra.HMatrix
@@ -44,9 +44,10 @@ import System.IO.Unsafe(unsafePerformIO)
 
 -------------------------------------------------------------------------
 
-type TVV  = TV (TV Res)
-type TVM  = TV (TM Res)
-type TVVM = TV (TV (TM Res))
+type TVV   = TV (TV Res)
+type TVM   = TV (TM Res)
+type TVVM  = TV (TV (TM Res))
+type TVVVM = TV (TV (TV (TM Res)))
 
 type Jacobian = Double -> Vector Double -> Matrix Double
 
@@ -63,68 +64,100 @@ data ODEMethod = RK2 -- ^ Embedded Runge-Kutta (2, 3) method.
                | MSAdams -- ^ A variable-coefficient linear multistep Adams method in Nordsieck form. This stepper uses explicit Adams-Bashforth (predictor) and implicit Adams-Moulton (corrector) methods in P(EC)^m functional iteration mode. Method order varies dynamically between 1 and 12. 
                | MSBDF Jacobian -- ^ A variable-coefficient linear multistep backward differentiation formula (BDF) method in Nordsieck form. This stepper uses the explicit BDF formula as predictor and implicit BDF formula as corrector. A modified Newton iteration method is used to solve the system of non-linear equations. Method order varies dynamically between 1 and 5. The method is generally suitable for stiff problems.
 
+-- | Adaptive step-size control functions
+data StepControl = X     Double Double -- ^ abs. and rel. tolerance for x(t)
+                 | X'    Double Double -- ^ abs. and rel. tolerance for x'(t)
+                 | XX'   Double Double Double Double -- ^ include both via rel. tolerance scaling factors a_x, a_x'
+                 | ScXX' Double Double Double Double (Vector Double) -- ^ scale abs. tolerance of x(t) components
 
 -- | A version of 'odeSolveV' with reasonable default parameters and system of equations defined using lists.
 odeSolve
-    :: (Double -> [Double] -> [Double])        -- ^ xdot(t,x)
+    :: (Double -> [Double] -> [Double])        -- ^ x'(t,x)
     -> [Double]        -- ^ initial conditions
     -> Vector Double   -- ^ desired solution times
     -> Matrix Double   -- ^ solution
 odeSolve xdot xi ts = odeSolveV RKf45 hi epsAbs epsRel (l2v xdot) (fromList xi) ts
     where hi = (ts!1 - ts!0)/100
           epsAbs = 1.49012e-08
-          epsRel = 1.49012e-08
-          l2v f = \t -> fromList  . f t . toList
+          epsRel = epsAbs
+          l2v f  = \t -> fromList . f t . toList
 
--- | Evolution of the system with adaptive step-size control.
+-- | A version of 'odeSolveVWith' with reasonable default step control.
 odeSolveV
     :: ODEMethod
-    -> Double -- ^ initial step size
-    -> Double -- ^ absolute tolerance for the state vector
-    -> Double -- ^ relative tolerance for the state vector
-    -> (Double -> Vector Double -> Vector Double)   -- ^ xdot(t,x)
+    -> Double            -- ^ initial step size
+    -> Double            -- ^ absolute tolerance for the state vector
+    -> Double            -- ^ relative tolerance for the state vector
+    -> (Double -> Vector Double -> Vector Double)   -- ^ x'(t,x)
     -> Vector Double     -- ^ initial conditions
     -> Vector Double     -- ^ desired solution times
     -> Matrix Double     -- ^ solution
-odeSolveV RK2 = odeSolveV' 0 Nothing
-odeSolveV RK4 = odeSolveV' 1 Nothing
-odeSolveV RKf45 = odeSolveV' 2 Nothing
-odeSolveV RKck = odeSolveV' 3 Nothing
-odeSolveV RK8pd = odeSolveV' 4 Nothing
-odeSolveV (RK2imp jac) = odeSolveV' 5 (Just jac)
-odeSolveV (RK4imp jac) = odeSolveV' 6 (Just jac)
-odeSolveV (BSimp jac) = odeSolveV' 7 (Just jac)
-odeSolveV (RK1imp jac) = odeSolveV' 8 (Just jac)
-odeSolveV MSAdams = odeSolveV' 9 Nothing
-odeSolveV (MSBDF jac) = odeSolveV' 10 (Just jac)
-
-
-odeSolveV'
-    :: CInt
-    -> Maybe (Double -> Vector Double -> Matrix Double)   -- ^ optional jacobian
-    -> Double -- ^ initial step size
-    -> Double -- ^ absolute tolerance for the state vector
-    -> Double -- ^ relative tolerance for the state vector
-    -> (Double -> Vector Double -> Vector Double)   -- ^ xdot(t,x)
+odeSolveV meth hi epsAbs epsRel = odeSolveVWith meth (XX' epsAbs epsRel 1 1) hi
+
+-- | Evolution of the system with adaptive step-size control.
+odeSolveVWith
+    :: ODEMethod
+    -> StepControl
+    -> Double            -- ^ initial step size
+    -> (Double -> Vector Double -> Vector Double)   -- ^ x'(t,x)
     -> Vector Double     -- ^ initial conditions
     -> Vector Double     -- ^ desired solution times
     -> Matrix Double     -- ^ solution
-odeSolveV' method mbjac h epsAbs epsRel f  xiv ts = unsafePerformIO $ do
-    let n   = size xiv
-    fp <- mkDoubleVecVecfun (\t -> aux_vTov (checkdim1 n . f t))
-    jp <- case mbjac of
-        Just jac -> mkDoubleVecMatfun (\t -> aux_vTom (checkdim2 n . jac t))
-        Nothing  -> return nullFunPtr
-    sol <- vec xiv $ \xiv' ->
-            vec (checkTimes ts) $ \ts' ->
-             createMIO (size ts) n
-              (ode_c (method) h epsAbs epsRel fp jp // xiv' // ts' )
-              "ode"
-    freeHaskellFunPtr fp
-    return sol
+odeSolveVWith method control = odeSolveVWith' m mbj c epsAbs epsRel aX aX' mbsc
+    where (m, mbj) = case method of
+              RK2        -> (0 , Nothing )
+              RK4        -> (1 , Nothing )
+              RKf45      -> (2 , Nothing )
+              RKck       -> (3 , Nothing )
+              RK8pd      -> (4 , Nothing )
+              RK2imp jac -> (5 , Just jac)
+              RK4imp jac -> (6 , Just jac)
+              BSimp  jac -> (7 , Just jac)
+              RK1imp jac -> (8 , Just jac)
+              MSAdams    -> (9 , Nothing )
+              MSBDF  jac -> (10, Just jac)
+          (c, epsAbs, epsRel, aX, aX', mbsc) = case control of
+              X     ea er           -> (0, ea, er, 1 , 0  , Nothing)
+              X'    ea er           -> (0, ea, er, 0 , 1  , Nothing)
+              XX'   ea er ax ax'    -> (0, ea, er, ax, ax', Nothing)
+              ScXX' ea er ax ax' sc -> (1, ea, er, ax, ax', Just sc)
+
+odeSolveVWith'
+    :: CInt     -- ^ stepping function
+    -> Maybe (Double -> Vector Double -> Matrix Double)   -- ^ optional jacobian
+    -> CInt     -- ^ step-size control function
+    -> Double   -- ^ absolute tolerance for step-size control
+    -> Double   -- ^ relative tolerance for step-size control
+    -> Double   -- ^ scaling factor for relative tolerance of x(t)
+    -> Double   -- ^ scaling factor for relative tolerance of x'(t)
+    -> Maybe (Vector Double)    -- ^ optional scaling for absolute error
+    -> Double   -- ^ initial step size
+    -> (Double -> Vector Double -> Vector Double)        -- ^ x'(t,x)
+    -> Vector Double  -- ^ initial conditions
+    -> Vector Double  -- ^ desired solution times
+    -> Matrix Double  -- ^ solution
+odeSolveVWith' method mbjac control epsAbs epsRel aX aX' mbsc h f xiv ts =
+    unsafePerformIO $ do
+        let n  = size xiv
+            sc = case mbsc of
+                Just scv -> checkdim1 n scv
+                Nothing  -> xiv
+        fp <- mkDoubleVecVecfun (\t -> aux_vTov (checkdim1 n . f t))
+        jp <- case mbjac of
+            Just jac -> mkDoubleVecMatfun (\t -> aux_vTom (checkdim2 n . jac t))
+            Nothing  -> return nullFunPtr
+        sol <- vec sc $ \sc' -> vec xiv $ \xiv' ->
+            vec (checkTimes ts) $ \ts' -> createMIO (size ts) n
+                (ode_c method control h epsAbs epsRel aX aX' fp jp
+                // sc' // xiv' // ts' )
+                "ode"
+        freeHaskellFunPtr fp
+        return sol
 
 foreign import ccall safe "ode"
-    ode_c :: CInt -> Double -> Double -> Double -> FunPtr (Double -> TVV) -> FunPtr (Double -> TVM) -> TVVM
+    ode_c :: CInt -> CInt -> Double
+          -> Double -> Double -> Double -> Double
+          -> FunPtr (Double -> TVV) -> FunPtr (Double -> TVM) -> TVVVM
 
 -------------------------------------------------------
 
@@ -141,4 +174,3 @@ checkdim2 n m
 checkTimes ts | size ts > 1 && all (>0) (zipWith subtract ts' (tail ts')) = ts
               | otherwise = error "odeSolve requires increasing times"
     where ts' = toList ts
-
diff --git a/packages/gsl/src/Numeric/GSL/gsl-ode.c b/packages/gsl/src/Numeric/GSL/gsl-ode.c
index 3f2771b..a6bdb55 100644
--- a/packages/gsl/src/Numeric/GSL/gsl-ode.c
+++ b/packages/gsl/src/Numeric/GSL/gsl-ode.c
@@ -23,10 +23,11 @@ int odejac (double t, const double y[], double *dfdy, double dfdt[], void *param
 }
 
 
-int ode(int method, double h, double eps_abs, double eps_rel,
+int ode(int method, int control, double h,
+        double eps_abs, double eps_rel, double a_y, double a_dydt,
         int f(double, int, const double*, int, double*),
         int jac(double, int, const double*, int, int, double*),
-        KRVEC(xi), KRVEC(ts), RMAT(sol)) {
+        KRVEC(sc), KRVEC(xi), KRVEC(ts), RMAT(sol)) {
 
     const gsl_odeiv_step_type * T;
 
@@ -46,8 +47,16 @@ int ode(int method, double h, double eps_abs, double eps_rel,
     }
 
     gsl_odeiv_step * s = gsl_odeiv_step_alloc (T, xin);
-    gsl_odeiv_control * c = gsl_odeiv_control_y_new (eps_abs, eps_rel);
     gsl_odeiv_evolve * e = gsl_odeiv_evolve_alloc (xin);
+    gsl_odeiv_control * c;
+
+    switch(control) {
+        case 0: { c = gsl_odeiv_control_standard_new
+                      (eps_abs, eps_rel, a_y, a_dydt); break; }
+        case 1: { c = gsl_odeiv_control_scaled_new
+                      (eps_abs, eps_rel, a_y, a_dydt, scp, scn); break; }
+        default: ERROR(BAD_CODE);
+    }
 
     Tode P;
     P.f = f;
@@ -112,10 +121,11 @@ int odejac (double t, const double y[], double *dfdy, double dfdt[], void *param
 }
 
 
-int ode(int method, double h, double eps_abs, double eps_rel,
+int ode(int method, int control, double h,
+        double eps_abs, double eps_rel, double a_y, double a_dydt,
         int f(double, int, const double*, int, double*),
         int jac(double, int, const double*, int, int, double*),
-        KRVEC(xi), KRVEC(ts), RMAT(sol)) {
+        KRVEC(sc), KRVEC(xi), KRVEC(ts), RMAT(sol)) {
 
     const gsl_odeiv2_step_type * T;
 
@@ -141,8 +151,15 @@ int ode(int method, double h, double eps_abs, double eps_rel,
 
     gsl_odeiv2_system sys = {odefunc, odejac, xin, &P};
 
-    gsl_odeiv2_driver * d =
-         gsl_odeiv2_driver_alloc_y_new (&sys, T, h, eps_abs, eps_rel);
+    gsl_odeiv2_driver * d;
+
+    switch(control) {
+        case 0: { d = gsl_odeiv2_driver_alloc_standard_new
+                      (&sys, T, h, eps_abs, eps_rel, a_y, a_dydt); break; }
+        case 1: { d = gsl_odeiv2_driver_alloc_scaled_new
+                      (&sys, T, h, eps_abs, eps_rel, a_y, a_dydt, scp); break; }
+        default: ERROR(BAD_CODE);
+    }
 
     double t = tsp[0];