-- SPDX-License-Identifier: MPL-2.0

{- | Explicit optimizer state for deterministic gradients, SGD and minibatches.
Defaults: beta1=.9,beta2=.999,epsilon=1e-8; RMSProp/AdaDelta rho=.9;
AdamW decays weights by .01*learningRate. Bias corrections use one-based steps.
-}
module HMath.Optimization.FirstOrder (
    GradientMethod (..),
    OptimizerState,
    initializeOptimizer,
    optimizerStep,
    minimizeGradient,
    minimizeGradientWithCallback,
    stochasticGradient,
    miniBatchGradient,
    projectedGradient,
    proximalGradient,
    fista,
    coordinateDescent,
) where

import Control.Monad (foldM)
import qualified Data.Foldable as Foldable (foldl')
import HMath.Error
import HMath.Optimization.Common

{- | First-order optimizer selector, including momentum and adaptive-moment variants.
Stateful parameters are validated by the selected optimizer.
-}
data GradientMethod
    = GradientDescent
    | Momentum
    | Nesterov
    | AdaGrad
    | RMSProp
    | AdaDelta
    | Adam
    | AdamW
    | AMSGrad
    | AdaMax
    | Nadam
    | RAdam
    deriving (GradientMethod -> GradientMethod -> Bool
(GradientMethod -> GradientMethod -> Bool)
-> (GradientMethod -> GradientMethod -> Bool) -> Eq GradientMethod
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: GradientMethod -> GradientMethod -> Bool
== :: GradientMethod -> GradientMethod -> Bool
$c/= :: GradientMethod -> GradientMethod -> Bool
/= :: GradientMethod -> GradientMethod -> Bool
Eq, Int -> GradientMethod -> ShowS
[GradientMethod] -> ShowS
GradientMethod -> String
(Int -> GradientMethod -> ShowS)
-> (GradientMethod -> String)
-> ([GradientMethod] -> ShowS)
-> Show GradientMethod
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GradientMethod -> ShowS
showsPrec :: Int -> GradientMethod -> ShowS
$cshow :: GradientMethod -> String
show :: GradientMethod -> String
$cshowList :: [GradientMethod] -> ShowS
showList :: [GradientMethod] -> ShowS
Show, Int -> GradientMethod
GradientMethod -> Int
GradientMethod -> [GradientMethod]
GradientMethod -> GradientMethod
GradientMethod -> GradientMethod -> [GradientMethod]
GradientMethod
-> GradientMethod -> GradientMethod -> [GradientMethod]
(GradientMethod -> GradientMethod)
-> (GradientMethod -> GradientMethod)
-> (Int -> GradientMethod)
-> (GradientMethod -> Int)
-> (GradientMethod -> [GradientMethod])
-> (GradientMethod -> GradientMethod -> [GradientMethod])
-> (GradientMethod -> GradientMethod -> [GradientMethod])
-> (GradientMethod
    -> GradientMethod -> GradientMethod -> [GradientMethod])
-> Enum GradientMethod
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: GradientMethod -> GradientMethod
succ :: GradientMethod -> GradientMethod
$cpred :: GradientMethod -> GradientMethod
pred :: GradientMethod -> GradientMethod
$ctoEnum :: Int -> GradientMethod
toEnum :: Int -> GradientMethod
$cfromEnum :: GradientMethod -> Int
fromEnum :: GradientMethod -> Int
$cenumFrom :: GradientMethod -> [GradientMethod]
enumFrom :: GradientMethod -> [GradientMethod]
$cenumFromThen :: GradientMethod -> GradientMethod -> [GradientMethod]
enumFromThen :: GradientMethod -> GradientMethod -> [GradientMethod]
$cenumFromTo :: GradientMethod -> GradientMethod -> [GradientMethod]
enumFromTo :: GradientMethod -> GradientMethod -> [GradientMethod]
$cenumFromThenTo :: GradientMethod
-> GradientMethod -> GradientMethod -> [GradientMethod]
enumFromThenTo :: GradientMethod
-> GradientMethod -> GradientMethod -> [GradientMethod]
Enum, GradientMethod
GradientMethod -> GradientMethod -> Bounded GradientMethod
forall a. a -> a -> Bounded a
$cminBound :: GradientMethod
minBound :: GradientMethod
$cmaxBound :: GradientMethod
maxBound :: GradientMethod
Bounded)

{- | Opaque accumulated first-order optimizer moments/history. Initialize and advance
through the stateful optimizer API.
-}
data OptimizerState = OptimizerState Int [Double] [Double] [Double] deriving (OptimizerState -> OptimizerState -> Bool
(OptimizerState -> OptimizerState -> Bool)
-> (OptimizerState -> OptimizerState -> Bool) -> Eq OptimizerState
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OptimizerState -> OptimizerState -> Bool
== :: OptimizerState -> OptimizerState -> Bool
$c/= :: OptimizerState -> OptimizerState -> Bool
/= :: OptimizerState -> OptimizerState -> Bool
Eq, Int -> OptimizerState -> ShowS
[OptimizerState] -> ShowS
OptimizerState -> String
(Int -> OptimizerState -> ShowS)
-> (OptimizerState -> String)
-> ([OptimizerState] -> ShowS)
-> Show OptimizerState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> OptimizerState -> ShowS
showsPrec :: Int -> OptimizerState -> ShowS
$cshow :: OptimizerState -> String
show :: OptimizerState -> String
$cshowList :: [OptimizerState] -> ShowS
showList :: [OptimizerState] -> ShowS
Show)

-- | Allocate zero momentum/moment buffers for a nonnegative parameter count.
initializeOptimizer :: Int -> Either MathError OptimizerState
initializeOptimizer :: Int -> Either MathError OptimizerState
initializeOptimizer Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError OptimizerState
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative optimizer dimension")
    | Bool
otherwise = OptimizerState -> Either MathError OptimizerState
forall a b. b -> Either a b
Right (Int -> [Double] -> [Double] -> [Double] -> OptimizerState
OptimizerState Int
0 (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0) (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
n Double
0))

-- | Apply one checked stateful update using the supplied gradient and positive rate. Nesterov here applies a corrected momentum direction; minimizeGradient evaluates the actual look-ahead gradient.
optimizerStep ::
    GradientMethod ->
    Double ->
    OptimizerState ->
    [Double] ->
    [Double] ->
    Either MathError ([Double], OptimizerState)
optimizerStep :: GradientMethod
-> Double
-> OptimizerState
-> [Double]
-> [Double]
-> Either MathError ([Double], OptimizerState)
optimizerStep GradientMethod
method Double
rate (OptimizerState Int
t [Double]
m [Double]
v [Double]
u) [Double]
x [Double]
g = do
    [Double] -> Either MathError ()
checkedPoint (Double
rate Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
x [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
g [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
m [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
v [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
u)
    if Double
rate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| ([Double] -> Bool) -> [[Double]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Double]
g, [Double]
m, [Double]
v, [Double]
u]
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"optimizer rate or state shape")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        b1 :: Double
b1 = Double
0.9
        b2 :: Double
b2 = Double
0.999
        eps :: Double
eps = Double
1e-8
        m' :: [Double]
m' = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
b1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b1) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
m [Double]
g
        v' :: [Double]
v' = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
b2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b2) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
v [Double]
g
        mh :: [Double]
mh = Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b1 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t')) [Double]
m'
        vh :: [Double]
vh = Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b2 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t')) [Double]
v'
        plain :: Double -> Double -> Double
plain Double
a Double
b = Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double -> Double
forall a. Floating a => a -> a
sqrt Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)
        ([Double]
direction, [Double]
newM, [Double]
newV, [Double]
newU) = case GradientMethod
method of
            GradientMethod
GradientDescent -> ([Double]
g, [Double]
m, [Double]
v, [Double]
u)
            GradientMethod
Momentum -> let vel :: [Double]
vel = [Double] -> [Double] -> [Double]
add (Double -> [Double] -> [Double]
scale Double
b1 [Double]
m) [Double]
g in ([Double]
vel, [Double]
vel, [Double]
v, [Double]
u)
            GradientMethod
Nesterov -> let vel :: [Double]
vel = [Double] -> [Double] -> [Double]
add (Double -> [Double] -> [Double]
scale Double
b1 [Double]
m) [Double]
g in ([Double] -> [Double] -> [Double]
add [Double]
g (Double -> [Double] -> [Double]
scale Double
b1 [Double]
vel), [Double]
vel, [Double]
v, [Double]
u)
            GradientMethod
AdaGrad -> let vs :: [Double]
vs = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
v [Double]
g in ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
g [Double]
vs, [Double]
m, [Double]
vs, [Double]
u)
            GradientMethod
RMSProp -> let vs :: [Double]
vs = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
0.9 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
v [Double]
g in ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
g [Double]
vs, [Double]
m, [Double]
vs, [Double]
u)
            GradientMethod
AdaDelta ->
                let vs :: [Double]
vs = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
0.9 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
v [Double]
g
                    ds :: [Double]
ds = (Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
a Double
b Double
c -> Double -> Double
forall a. Floating a => a -> a
sqrt (Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
c) [Double]
u [Double]
vs [Double]
g
                    us :: [Double]
us = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
0.9 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.1 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b) [Double]
u [Double]
ds
                 in ([Double]
ds, [Double]
m, [Double]
vs, [Double]
us)
            GradientMethod
Adam -> ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
mh [Double]
vh, [Double]
m', [Double]
v', [Double]
u)
            GradientMethod
AdamW -> ([Double] -> [Double] -> [Double]
add ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
mh [Double]
vh) (Double -> [Double] -> [Double]
scale Double
0.01 [Double]
x), [Double]
m', [Double]
v', [Double]
u)
            GradientMethod
AMSGrad -> let us :: [Double]
us = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Ord a => a -> a -> a
max [Double]
u [Double]
v' in ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
mh (Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b2 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t')) [Double]
us), [Double]
m', [Double]
v', [Double]
us)
            GradientMethod
AdaMax ->
                let us :: [Double]
us = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Double
b2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a) (Double -> Double
forall a. Num a => a -> a
abs Double
b)) [Double]
u [Double]
g
                 in ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
a Double
b -> Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
b Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)) [Double]
mh [Double]
us, [Double]
m', [Double]
v', [Double]
us)
            GradientMethod
Nadam ->
                let nm :: [Double]
nm = [Double] -> [Double] -> [Double]
add (Double -> [Double] -> [Double]
scale Double
b1 [Double]
mh) (Double -> [Double] -> [Double]
scale ((Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b1 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t')) [Double]
g) in ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
nm [Double]
vh, [Double]
m', [Double]
v', [Double]
u)
            GradientMethod
RAdam ->
                let rhoInf :: Double
rhoInf = Double
2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b2) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1
                    rho :: Double
rho = Double
rhoInf Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
t' Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
b2 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t' Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
b2 Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
t')
                    rect :: Double
rect = Double -> Double
forall a. Floating a => a -> a
sqrt ((Double
rho Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
rho Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rhoInf Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ ((Double
rhoInf Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
4) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
rhoInf Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
2) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rho))
                 in (if Double
rho Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
4 then Double -> [Double] -> [Double]
scale Double
rect ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
plain [Double]
mh [Double]
vh) else [Double]
mh, [Double]
m', [Double]
v', [Double]
u)
        next :: [Double]
next = [Double] -> [Double] -> [Double]
sub [Double]
x (Double -> [Double] -> [Double]
scale Double
rate [Double]
direction)
    [Double] -> Either MathError ()
checkedPoint ([Double]
next [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
direction [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
newM [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
newV [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
newU)
    ([Double], OptimizerState)
-> Either MathError ([Double], OptimizerState)
forall a b. b -> Either a b
Right ([Double]
next, Int -> [Double] -> [Double] -> [Double] -> OptimizerState
OptimizerState Int
t' [Double]
newM [Double]
newV [Double]
newU)

-- | Minimize with explicit gradient/state and fixed learning rate. Only finite gradient residuals below tolerance certify convergence.
minimizeGradient ::
    GradientMethod ->
    OptimizeOptions ->
    Objective ->
    Gradient ->
    [Double] ->
    Either MathError OptimizeResult
minimizeGradient :: GradientMethod
-> OptimizeOptions
-> Objective
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
minimizeGradient = (OptimizeResult -> Either MathError Bool)
-> GradientMethod
-> OptimizeOptions
-> Objective
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
minimizeGradientWithCallback (Either MathError Bool -> OptimizeResult -> Either MathError Bool
forall a b. a -> b -> a
const (Bool -> Either MathError Bool
forall a b. b -> Either a b
Right Bool
True))

{- | Observe each finite iterate after its objective/gradient evaluation. False
stops explicitly with convergence=False; callback failures propagate. Callback
observations do not themselves count as objective or derivative evaluations.
-}
minimizeGradientWithCallback ::
    (OptimizeResult -> Either MathError Bool) ->
    GradientMethod ->
    OptimizeOptions ->
    Objective ->
    Gradient ->
    [Double] ->
    Either MathError OptimizeResult
minimizeGradientWithCallback :: (OptimizeResult -> Either MathError Bool)
-> GradientMethod
-> OptimizeOptions
-> Objective
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
minimizeGradientWithCallback OptimizeResult -> Either MathError Bool
callback GradientMethod
method OptimizeOptions
o Objective
f Gradient
g [Double]
x0 = do
    OptimizeOptions -> Either MathError ()
checkedOptions OptimizeOptions
o
    state <- Int -> Either MathError OptimizerState
initializeOptimizer ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x0)
    loop 0 0 state x0
  where
    loop :: Int
-> Int
-> OptimizerState
-> [Double]
-> Either MathError OptimizeResult
loop Int
k Int
n OptimizerState
state [Double]
x = do
        OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
        y <- Objective -> Objective
checkedValue Objective
f [Double]
x
        gs <- checkedGradient g x
        let current = [Double]
-> Double -> Double -> Int -> Int -> Bool -> OptimizeResult
OptimizeResult [Double]
x Double
y ([Double] -> Double
norm [Double]
gs) Int
k (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) ([Double] -> Double
norm [Double]
gs Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= OptimizeOptions -> Double
optTolerance OptimizeOptions
o)
        continue <- callback current
        if not continue
            then Right current{optConverged = False}
            else
                if norm gs <= optTolerance o
                    then Right current
                    else
                        if k >= optMaxIterations o
                            then Left (NonConvergence (show method) k)
                            else do
                                -- True Nesterov evaluates the gradient at the look-ahead point.
                                actual <- case (method, state) of
                                    (GradientMethod
Nesterov, OptimizerState Int
_ [Double]
velocity [Double]
_ [Double]
_) -> do OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3); Gradient -> Gradient
checkedGradient Gradient
g ([Double] -> [Double] -> [Double]
sub [Double]
x (Double -> [Double] -> [Double]
scale (Double
0.9 Double -> Double -> Double
forall a. Num a => a -> a -> a
* OptimizeOptions -> Double
optLearningRate OptimizeOptions
o) [Double]
velocity))
                                    (GradientMethod, OptimizerState)
_ -> Gradient
forall a b. b -> Either a b
Right [Double]
gs
                                (next, state') <-
                                    optimizerStep (if method == Nesterov then Momentum else method) (optLearningRate o) state x actual
                                loop (k + 1) (n + if method == Nesterov then 3 else 2) state' next

{- | One pass over explicit sample/batch gradients; ordering and shuffling are
caller-owned, allowing exact replay. All gradients are checked before update.
-}
stochasticGradient ::
    GradientMethod -> Double -> [Gradient] -> [Double] -> Either MathError ([Double], OptimizerState)
stochasticGradient :: GradientMethod
-> Double
-> [Gradient]
-> [Double]
-> Either MathError ([Double], OptimizerState)
stochasticGradient GradientMethod
method Double
rate [Gradient]
grads [Double]
x0 = do
    [Double] -> Either MathError ()
checkedPoint (Double
rate Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: [Double]
x0)
    if Double
rate Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"positive stochastic learning rate required") else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    s <- Int -> Either MathError OptimizerState
initializeOptimizer ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x0)
    foldM
        (\([Double]
x, OptimizerState
state) Gradient
g -> do gs <- Gradient -> Gradient
checkedGradient Gradient
g [Double]
x; optimizerStep method rate state x gs)
        (x0, s)
        grads

-- | Average a nonempty explicit set of same-length checked sample gradients at the current parameters.
miniBatchGradient :: [Gradient] -> Gradient
miniBatchGradient :: [Gradient] -> Gradient
miniBatchGradient [] [Double]
_ = MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"empty mini batch")
miniBatchGradient [Gradient]
gradients [Double]
x = do
    [Double] -> Either MathError ()
checkedPoint [Double]
x
    gs <- (Gradient -> Either MathError [Double])
-> [Gradient] -> Either MathError [[Double]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Gradient -> Gradient
`checkedGradient` [Double]
x) [Gradient]
gradients
    let average = ([Double] -> [Double] -> [Double])
-> [Double] -> [[Double]] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' [Double] -> [Double] -> [Double]
add (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x) Double
0) (([Double] -> [Double]) -> [[Double]] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
gs))) [[Double]]
gs)
    checkedPoint average
    Right average

-- | Project every update and stop on the projected-gradient mapping, including constrained boundary optima.
projectedGradient ::
    OptimizeOptions ->
    Objective ->
    Gradient ->
    ([Double] -> Either MathError [Double]) ->
    [Double] ->
    Either MathError OptimizeResult
projectedGradient :: OptimizeOptions
-> Objective
-> Gradient
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
projectedGradient OptimizeOptions
o Objective
f Gradient
g Gradient
project = OptimizeOptions
-> Objective
-> Gradient
-> (Double -> Gradient)
-> [Double]
-> Either MathError OptimizeResult
proximalGradient OptimizeOptions
o Objective
f Gradient
g (\Double
_ -> Gradient
project)

-- | Apply prox(rate,point) after each gradient step; residual is the proximal-gradient mapping norm.
proximalGradient ::
    OptimizeOptions ->
    Objective ->
    Gradient ->
    (Double -> [Double] -> Either MathError [Double]) ->
    [Double] ->
    Either MathError OptimizeResult
proximalGradient :: OptimizeOptions
-> Objective
-> Gradient
-> (Double -> Gradient)
-> [Double]
-> Either MathError OptimizeResult
proximalGradient OptimizeOptions
o Objective
f Gradient
g Double -> Gradient
prox [Double]
x0 = do OptimizeOptions -> Either MathError ()
checkedOptions OptimizeOptions
o; Int -> Int -> [Double] -> Either MathError OptimizeResult
loop Int
0 Int
0 [Double]
x0
  where
    loop :: Int -> Int -> [Double] -> Either MathError OptimizeResult
loop Int
k Int
n [Double]
x = do
        OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
        gs <- Gradient -> Gradient
checkedGradient Gradient
g [Double]
x
        y <- prox (optLearningRate o) (sub x (scale (optLearningRate o) gs))
        if length y /= length x then Left (InvalidArgument "proximal shape") else checkedPoint y
        val <- checkedValue f y
        let residual = [Double] -> Double
norm (Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ OptimizeOptions -> Double
optLearningRate OptimizeOptions
o) ([Double] -> [Double] -> [Double]
sub [Double]
x [Double]
y))
        if residual <= optTolerance o
            then Right (OptimizeResult y val residual k (n + 3) True)
            else
                if k >= optMaxIterations o
                    then Left (NonConvergence "proximal gradient" k)
                    else loop (k + 1) (n + 3) y

-- | Accelerated proximal gradient using the Beck-Teboulle t recurrence and a fresh mapping residual at accepted points.
fista ::
    OptimizeOptions ->
    Objective ->
    Gradient ->
    (Double -> [Double] -> Either MathError [Double]) ->
    [Double] ->
    Either MathError OptimizeResult
fista :: OptimizeOptions
-> Objective
-> Gradient
-> (Double -> Gradient)
-> [Double]
-> Either MathError OptimizeResult
fista OptimizeOptions
o Objective
f Gradient
g Double -> Gradient
prox [Double]
x0 = do OptimizeOptions -> Either MathError ()
checkedOptions OptimizeOptions
o; Int
-> Int
-> Double
-> [Double]
-> [Double]
-> Either MathError OptimizeResult
loop Int
0 Int
0 Double
1 [Double]
x0 [Double]
x0
  where
    loop :: Int
-> Int
-> Double
-> [Double]
-> [Double]
-> Either MathError OptimizeResult
loop Int
k Int
n Double
t [Double]
x [Double]
y = do
        OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
        gy <- Gradient -> Gradient
checkedGradient Gradient
g [Double]
y
        z <- prox (optLearningRate o) (sub y (scale (optLearningRate o) gy))
        if length z /= length x then Left (InvalidArgument "FISTA proximal shape") else checkedPoint z
        gz <- checkedGradient g z
        p <- prox (optLearningRate o) (sub z (scale (optLearningRate o) gz))
        if length p /= length z then Left (InvalidArgument "FISTA proximal shape") else checkedPoint p
        val <- checkedValue f z
        let residual = [Double] -> Double
norm (Double -> [Double] -> [Double]
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ OptimizeOptions -> Double
optLearningRate OptimizeOptions
o) ([Double] -> [Double] -> [Double]
sub [Double]
z [Double]
p)); t' = (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
4 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
        if residual <= optTolerance o
            then Right (OptimizeResult z val residual k (n + 3) True)
            else
                if k >= optMaxIterations o
                    then Left (NonConvergence "FISTA" k)
                    else loop (k + 1) (n + 3) t' z (add z (scale ((t - 1) / t') (sub z x)))

-- | Cyclic coordinates or blocks; every coordinate must occur exactly once.
coordinateDescent ::
    [[Int]] -> OptimizeOptions -> Objective -> Gradient -> [Double] -> Either MathError OptimizeResult
coordinateDescent :: [[Int]]
-> OptimizeOptions
-> Objective
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
coordinateDescent [[Int]]
blocks OptimizeOptions
o Objective
f Gradient
g [Double]
x0 = do
    OptimizeOptions -> Either MathError ()
checkedOptions OptimizeOptions
o
    let ids :: [Int]
ids = [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Int]]
blocks; n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x0
    if ([Int] -> Bool) -> [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any [Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Int]]
blocks Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ids Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i) [Int]
ids) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        then MathError -> Either MathError OptimizeResult
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"coordinate blocks must partition variables")
        else Int -> Int -> [Double] -> Either MathError OptimizeResult
loop Int
0 Int
0 [Double]
x0
  where
    loop :: Int -> Int -> [Double] -> Either MathError OptimizeResult
loop Int
k Int
n [Double]
x = do
        OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [[Int]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
blocks)
        gs <- Gradient -> Gradient
checkedGradient Gradient
g [Double]
x
        val <- checkedValue f x
        if norm gs <= optTolerance o
            then Right (OptimizeResult x val (norm gs) k (n + 2) True)
            else
                if k >= optMaxIterations o
                    then Left (NonConvergence "coordinate descent" k)
                    else do
                        y <-
                            foldM
                                ( \[Double]
z [Int]
block -> do
                                    dz <- Gradient -> Gradient
checkedGradient Gradient
g [Double]
z
                                    Right [if i `elem` block then a - optLearningRate o * (dz !! i) else a | (i, a) <- zip [0 ..] z]
                                )
                                x
                                blocks
                        loop (k + 1) (n + 2 + length blocks) y