root finding with jacobian

4 files changed, 209 insertions, 18 deletions

diff --git a/examples/root.hs b/examples/root.hs
index 69db243..2a24f0f 100644
--- a/examples/root.hs
+++ b/examples/root.hs
@@ -11,6 +11,15 @@ test method = do
     print s -- solution
     disp p -- evolution of the algorithm
 
+jacobian a b [x,y] = [ [-a    , 0]
+                     , [-2*b*x, b] ]
+
+testJ method = do
+    print method
+    let (s,p) = rootJ method 1E-7 30 (rosenbrock 1 10) (jacobian 1 10) [-10,-5]
+    print s
+    disp p
+
 disp = putStrLn . format "  " (printf "%.3f")
 
 main = do
@@ -18,3 +27,8 @@ main = do
     test Hybrid
     test DNewton
     test Broyden
+
+    testJ HybridsJ
+    testJ HybridJ
+    testJ Newton
+    testJ GNewton
diff --git a/lib/Numeric/GSL/Root.hs b/lib/Numeric/GSL/Root.hs
index d674fad..6ce2c4c 100644
--- a/lib/Numeric/GSL/Root.hs
+++ b/lib/Numeric/GSL/Root.hs
@@ -45,7 +45,8 @@ main = do
 -----------------------------------------------------------------------------
 
 module Numeric.GSL.Root (
-    root, RootMethod(..)
+    root, RootMethod(..),
+    rootJ, RootMethodJ(..),
 ) where
 
 import Data.Packed.Internal
@@ -76,7 +77,7 @@ root method epsabs maxit fun xinit = rootGen (fi (fromEnum method)) fun xinit ep
 rootGen m f xi epsabs maxit = unsafePerformIO $ do
     let xiv = fromList xi
         n   = dim xiv
-    fp <- mkVecVecfun (aux_vTov (fromList . checkdim n f . toList))
+    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
     rawpath <- withVector xiv $ \xiv' ->
                    createMIO maxit (2*n+1)
                          (c_root m fp epsabs (fi maxit) // xiv')
@@ -89,22 +90,74 @@ rootGen m f xi epsabs maxit = unsafePerformIO $ do
 
 
 foreign import ccall "root"
-    c_root:: CInt -> FunPtr (CInt -> Ptr Double -> Ptr Double -> IO ()) -> Double -> CInt -> TVM
+    c_root:: CInt -> FunPtr TVV -> Double -> CInt -> TVM
+
+-------------------------------------------------------------------------
+
+data RootMethodJ = HybridsJ
+                 | HybridJ
+                 | Newton
+                 | GNewton
+                deriving (Enum,Eq,Show)
+
+-- | Nonlinear multidimensional root finding using both the function and its derivatives.
+rootJ :: RootMethodJ
+      -> Double                     -- ^ maximum residual
+      -> Int                        -- ^ maximum number of iterations allowed
+      -> ([Double] -> [Double])     -- ^ function to minimize
+      -> ([Double] -> [[Double]])   -- ^ Jacobian
+      -> [Double]                   -- ^ starting point
+      -> ([Double], Matrix Double)  -- ^ solution vector and optimization path
+
+rootJ method epsabs maxit fun jac xinit = rootJGen (fi (fromEnum method)) fun jac xinit epsabs maxit
+
+rootJGen m f jac xi epsabs maxit = unsafePerformIO $ do
+    let xiv = fromList xi
+        n   = dim xiv
+    fp <- mkVecVecfun (aux_vTov (checkdim1 n . fromList . f . toList))
+    jp <- mkVecMatfun (aux_vTom (checkdim2 n . fromLists . jac . toList))
+    rawpath <- withVector xiv $ \xiv' ->
+                   createMIO maxit (2*n+1)
+                         (c_rootj m fp jp epsabs (fi maxit) // xiv')
+                         "root"
+    let it = round (rawpath @@> (maxit-1,0))
+        path = takeRows it rawpath
+        [sol] = toLists $ dropRows (it-1) path
+    freeHaskellFunPtr fp
+    return (take n $ drop 1 sol, path)
+
+
+foreign import ccall "rootj"
+    c_rootj:: CInt -> FunPtr TVV -> FunPtr TVM -> Double -> CInt -> TVM
+
 
 ---------------------------------------------------------------------
 
 foreign import ccall "wrapper"
-    mkVecVecfun :: (CInt -> Ptr Double -> Ptr Double -> IO ())
-                -> IO (FunPtr (CInt -> Ptr Double -> Ptr Double->IO()))
+    mkVecVecfun :: TVV -> IO (FunPtr TVV)
 
-aux_vTov :: (Vector Double -> Vector Double) -> (CInt -> Ptr Double -> Ptr Double -> IO())
-aux_vTov f n p r = g where
+aux_vTov :: (Vector Double -> Vector Double) -> TVV
+aux_vTov f n p nr r = g where
     V {fptr = pr} = f x
     x = createV (fromIntegral n) copy "aux_vTov"
     copy n' q = do
         copyArray q p (fromIntegral n')
         return 0
-    g = withForeignPtr pr $ \p' -> copyArray r p' (fromIntegral n)
+    g = do withForeignPtr pr $ \p' -> copyArray r p' (fromIntegral nr)
+           return 0
+
+foreign import ccall "wrapper"
+    mkVecMatfun :: TVM -> IO (FunPtr TVM)
+
+aux_vTom :: (Vector Double -> Matrix Double) -> TVM
+aux_vTom f n p rr cr r = g where
+    V {fptr = pr} = flatten $ f x
+    x = createV (fromIntegral n) copy "aux_vTov"
+    copy n' q = do
+        copyArray q p (fromIntegral n')
+        return 0
+    g = do withForeignPtr pr $ \p' -> copyArray r p' (fromIntegral $ rr*cr)
+           return 0
 
 createV n fun msg = unsafePerformIO $ do
     r <- createVector n
@@ -116,8 +169,12 @@ createMIO r c fun msg = do
     app1 fun mat res msg
     return res
 
-checkdim n f x
-    | length y /= n = error $ "Error: "++ show n
-                              ++ " results expected in the function supplied to root"
-    | otherwise = y
-  where y = f x
+checkdim1 n v
+    | dim v == n = v
+    | otherwise = error $ "Error: "++ show n
+                        ++ " results expected in the function supplied to root"
+
+checkdim2 n m
+    | rows m == n && cols m == n = m
+    | otherwise = error $ "Error: "++ show n ++ "x" ++ show n
+                        ++ " Jacobian expected in root"
diff --git a/lib/Numeric/GSL/gsl-aux.c b/lib/Numeric/GSL/gsl-aux.c
index 80c23fc..c6b052f 100644
--- a/lib/Numeric/GSL/gsl-aux.c
+++ b/lib/Numeric/GSL/gsl-aux.c
@@ -511,17 +511,17 @@ int minimizeWithDeriv(int method, double f(int, double*), void df(int, double*,
 
 //---------------------------------------------------------------
 
-typedef void TrawfunV(int, double*, double*);
+typedef void TrawfunV(int, double*, int, double*);
 
 int only_f_aux_root(const gsl_vector*x, void *pars, gsl_vector*y) {
     TrawfunV * f = (TrawfunV*) pars;
     double* p = (double*)calloc(x->size,sizeof(double));
-    double* q = (double*)calloc(x->size,sizeof(double));
+    double* q = (double*)calloc(y->size,sizeof(double));
     int k;
     for(k=0;k<x->size;k++) {
         p[k] = gsl_vector_get(x,k);
     }
-    f(x->size,p,q);
+    f(x->size,p,y->size,q);
     for(k=0;k<y->size;k++) {
         gsl_vector_set(y,k,q[k]);
     }
@@ -588,3 +588,118 @@ int root(int method, void f(int, double*, int, double*),
     gsl_multiroot_fsolver_free(s);
     OK
 }
+
+// working with the jacobian
+
+typedef struct {int (*f)(int, double*, int, double *); int (*jf)(int, double*, int, int, double*);} Tfjf;
+
+int f_aux_root(const gsl_vector*x, void *pars, gsl_vector*y) {
+    Tfjf * fjf = ((Tfjf*) pars);
+    double* p = (double*)calloc(x->size,sizeof(double));
+    double* q = (double*)calloc(y->size,sizeof(double));
+    int k;
+    for(k=0;k<x->size;k++) {
+        p[k] = gsl_vector_get(x,k);
+    }
+    (fjf->f)(x->size,p,y->size,q);
+    for(k=0;k<y->size;k++) {
+        gsl_vector_set(y,k,q[k]);
+    }
+    free(p);
+    free(q);
+    return 0;
+}
+
+int jf_aux_root(const gsl_vector * x, void * pars, gsl_matrix * jac) {
+    Tfjf * fjf = ((Tfjf*) pars);
+    double* p = (double*)calloc(x->size,sizeof(double));
+    double* q = (double*)calloc((x->size)*(x->size),sizeof(double));
+    int i,j,k;
+    for(k=0;k<x->size;k++) {
+        p[k] = gsl_vector_get(x,k);
+    }
+
+    (fjf->jf)(x->size,p,x->size,x->size,q);
+
+    k=0;
+    for(i=0;i<x->size;i++) {
+        for(j=0;j<x->size;j++){
+            gsl_matrix_set(jac,i,j,q[k++]);
+        }
+    }
+    free(p);
+    free(q);
+    return 0;
+}
+
+int fjf_aux_root(const gsl_vector * x, void * pars, gsl_vector * f, gsl_matrix * g) {
+    f_aux_root(x,pars,f);
+    jf_aux_root(x,pars,g);
+    return 0;
+}
+
+int rootj(int method, int f(int, double*, int, double*),
+                      int jac(int, double*, int, int, double*),
+         double epsabs, int maxit,
+         KRVEC(xi), RMAT(sol)) {
+    REQUIRES(solr == maxit && solc == 1+2*xin,BAD_SIZE);
+    DEBUGMSG("root_fjf");
+    gsl_multiroot_function_fdf my_func;
+    // extract function from pars
+    my_func.f = f_aux_root;
+    my_func.df = jf_aux_root;
+    my_func.fdf = fjf_aux_root;
+    my_func.n = xin;
+    Tfjf stfjf;
+    stfjf.f = f;
+    stfjf.jf = jac;
+    my_func.params = &stfjf;
+    size_t iter = 0;
+    int status;
+    const gsl_multiroot_fdfsolver_type *T;
+    gsl_multiroot_fdfsolver *s;
+    // Starting point
+    KDVVIEW(xi);
+    switch(method) {
+        case 0 : {T = gsl_multiroot_fdfsolver_hybridsj;; break; }
+        case 1 : {T = gsl_multiroot_fdfsolver_hybridj; break; }
+        case 2 : {T = gsl_multiroot_fdfsolver_newton; break; }
+        case 3 : {T = gsl_multiroot_fdfsolver_gnewton; break; }
+        default: ERROR(BAD_CODE);
+    }
+    s = gsl_multiroot_fdfsolver_alloc (T, my_func.n);
+
+    gsl_multiroot_fdfsolver_set (s, &my_func, V(xi));
+
+    do {
+           status = gsl_multiroot_fdfsolver_iterate (s);
+
+           solp[iter*solc+0] = iter;
+
+           int k;
+           for(k=0;k<xin;k++) {
+               solp[iter*solc+k+1] = gsl_vector_get(s->x,k);
+           }
+           for(k=xin;k<2*xin;k++) {
+               solp[iter*solc+k+1] = gsl_vector_get(s->f,k-xin);
+           }
+
+           iter++;
+           if (status)   /* check if solver is stuck */
+             break;
+
+           status =
+             gsl_multiroot_test_residual (s->f, epsabs);
+        }
+        while (status == GSL_CONTINUE && iter < maxit);
+
+    int i,j;
+    for (i=iter; i<solr; i++) {
+        solp[i*solc+0] = iter;
+        for(j=1;j<solc;j++) {
+            solp[i*solc+j]=0.;
+        }
+    }
+    gsl_multiroot_fdfsolver_free(s);
+    OK
+}
\ No newline at end of file
diff --git a/lib/Numeric/LinearAlgebra/Tests.hs b/lib/Numeric/LinearAlgebra/Tests.hs
index 174e418..83f581f 100644
--- a/lib/Numeric/LinearAlgebra/Tests.hs
+++ b/lib/Numeric/LinearAlgebra/Tests.hs
@@ -119,9 +119,14 @@ minimizationTest = TestList [ utest "minimization conj grad" (minim1 f df [5,7]
 
 ---------------------------------------------------------------------
 
-rootFindingTest = utest "root Hybrids" (sol ~~ [1,1])
-    where sol = fst $ root Hybrids 1E-7 30 (rosenbrock 1 10) [-10,-5]
+rootFindingTest = TestList [ utest "root Hybrids" (fst sol1 ~~ [1,1])
+                           , utest "root Newton"  (rows (snd sol2) == 2)
+                           ]
+    where sol1 = root Hybrids 1E-7 30 (rosenbrock 1 10) [-10,-5]
+          sol2 = rootJ Newton 1E-7 30 (rosenbrock 1 10) (jacobian 1 10) [-10,-5]
           rosenbrock a b [x,y] = [ a*(1-x), b*(y-x^2) ]
+          jacobian a b [x,_y] = [ [-a    , 0]
+                                , [-2*b*x, b] ]
 
 ---------------------------------------------------------------------