----------------------------------------------------
-- |
-- Module     :  Math.LevenbergMarquardt
-- License    :  GPL
--
-- Maintainer :  Kiet Lam <ktklam9@gmail.com>
--
--
-- The Levenberg-Marquardt algorithm is a minimization
-- algorithm for functions expressed as a sum of squared
-- errors
--
-- This can be used for curve-fitting, multidimensional
-- function minimization, and neural networks training
--
--
----------------------------------------------------

module Math.LevenbergMarquardt (
  Function,
  Jacobian,
  lmMinimize
  ) where


import Data.Packed.Vector
import Data.Packed.Matrix
import Numeric.LinearAlgebra.Algorithms
import Numeric.Container


-- | Type that represents the function that can calculate the residues
type Function = Vector Double -> Vector Double


-- | Type that represents the function that can calculate the jacobian matrix
--   of the residue with respect to each parameter
type Jacobian = Vector Double -> Matrix Double


-- | Name says it all
sumSquared :: Vector Double -> Double
sumSquared = foldVector (+) 0 . mapVector (**2)


-- | Evolves the parameter x for f(x) = sum-square(e(x)) so that f(x)
--   will be minimized, where:
--
--   f     = real-valued error function,
--   e(x)  = {e1(x),e2(x),..,eN(x)}, where
--   e1(x) = the residue at the vector x
--
--   NOTE: eN(x) is usually represented as (sample - hypothesis(x))
--
--         e.g.: In training neural networks, hypothesis(x) would be the
--               network's output for a training set, and sample would be the
--               expected output for that training set
--
--   NOTE: The dampening constant(lambda) should be set to 0.01 and the dampening
--         update value (beta) should be set to be 10
lmMinimize :: Function         -- ^ Multi-dimensional function that will
                               --   return a vector of residues
              -> Jacobian      -- ^ The function that calculate the Jacobian
                               --   matrix of each residue with respect to
                               --   each parameter
              -> Vector Double -- ^ The initial guess for the parameter
              -> Double        -- ^ Dampening constant (usually lambda in most literature)
              -> Double        -- ^ Dampening update value (usually beta in most literature)
              -> Double        -- ^ The precision desired
              -> Int           -- ^ The maximum iteration
              -> (Vector Double, Matrix Double) -- ^ Returns the optimal parameter
                                                --   and the matrix path
lmMinimize func jacob param lambda beta prec iter =
  let errValue = sumSquared $ func param
      (v, m) = lmMinimizeInternal func jacob param prec errValue lambda beta iter [param]
  in
   (v, (fromLists . reverse . map toList) m)


lmMinimizeInternal :: Function           -- ^ Multi-dimensional function that will
                                         --   return a vector of residues
                      -> Jacobian        -- ^ The Jacobian matrix of each residue
                                         --   with respect to each parameter
                      -> Vector Double   -- ^ The current guess for the parameter
                      -> Double          -- ^ The precision desired
                      -> Double          -- ^ The current value for the error function
                                         --   (which is the sum of the residues squared)
                      -> Double          -- ^ The dampening constant (lambda)
                      -> Double          -- ^ The beta constant to update lambda
                      -> Int             -- ^ The maximum iteration
                      -> [Vector Double] -- ^ The current matrix path
                      -> (Vector Double, [Vector Double]) -- ^ Returns the optimal parameter
                                                          --   and the matrix path
lmMinimizeInternal _ _ v _ _ _ _ 0 m = (v,m)
lmMinimizeInternal func jacob param prec errValue lambda beta iter mat =
  let jacobian = jacob param
      residues = func param

      gradients = (trans jacobian) `mXv` residues
      hessian = (trans jacobian) `multiply` jacobian

      diagonalHessian = diagRect 0 (takeDiag hessian) (cols hessian) (cols hessian)
      diagonalLambda = mapMatrix (*lambda) (ident (cols hessian))
      hessianModded = hessian `add` diagonalLambda

      newParam = param `sub` (flatten (linearSolveSVD hessianModded (reshape 1 gradients)))

      newErrValue = (sumSquared . func) newParam
  in
   case norm2(gradients) < prec of
     True -> (newParam, newParam:mat)
     False ->
       case newErrValue > errValue of
         True  -> lmMinimizeInternal func jacob param prec errValue (lambda * beta) beta iter mat
         False -> lmMinimizeInternal func jacob newParam prec newErrValue (lambda / beta) beta (iter - 1) (newParam:mat)