-- SPDX-License-Identifier: MPL-2.0

{- | Shared checked Double objective contracts. Finite deterministic callbacks
return structured failures. Budgets count objective/gradient/Hessian calls,
including line-search trials, rather than only accepted steps.
-}
module HMath.Optimization.Common (
    Objective,
    Gradient,
    Hessian,
    OptimizeOptions (..),
    defaultOptimizeOptions,
    OptimizeResult (..),
    checkedOptions,
    checkedPoint,
    checkedValue,
    checkedGradient,
    dot,
    norm,
    add,
    sub,
    scale,
    identity,
    matVec,
    transpose,
    matMul,
    outer,
    solveDense,
    result,
    checkBudget,
    LearningRate (..),
    learningRate,
    EarlyStopping (..),
    earlyStopping,
    LineSearch (..),
    lineSearch,
) where

import Data.List (transpose)
import HMath.Error
import qualified HMath.Matrix as M
import qualified HMath.Matrix.LU as LU

-- | Scalar objective function over a parameter vector, supplied by the optimization caller.
type Objective = [Double] -> Either MathError Double

{- | Caller-supplied gradient components in parameter order; dimension/finite checks occur
in the consuming optimizer.
-}
type Gradient = [Double] -> Either MathError [Double]

{- | Caller-supplied second-derivative matrix in parameter order, consumed by checked Newton
methods.
-}
type Hessian = [Double] -> Either MathError [[Double]]

{- | Finite optimization work budgets,tolerance,and method step rate. Start from
defaultOptimizeOptions.
-}
data OptimizeOptions = OptimizeOptions
    { OptimizeOptions -> Int
optMaxIterations :: Int
    -- ^ Maximum optimization iterations.
    , OptimizeOptions -> Int
optMaxEvaluations :: Int
    -- ^ Maximum objective/evaluation work budget.
    , OptimizeOptions -> Double
optTolerance :: Double
    -- ^ Positive finite method-specific stopping tolerance.
    , OptimizeOptions -> Double
optLearningRate :: Double
    -- ^ Positive finite step-size/rate where used by the method.
    }
    deriving (OptimizeOptions -> OptimizeOptions -> Bool
(OptimizeOptions -> OptimizeOptions -> Bool)
-> (OptimizeOptions -> OptimizeOptions -> Bool)
-> Eq OptimizeOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OptimizeOptions -> OptimizeOptions -> Bool
== :: OptimizeOptions -> OptimizeOptions -> Bool
$c/= :: OptimizeOptions -> OptimizeOptions -> Bool
/= :: OptimizeOptions -> OptimizeOptions -> Bool
Eq, Int -> OptimizeOptions -> ShowS
[OptimizeOptions] -> ShowS
OptimizeOptions -> String
(Int -> OptimizeOptions -> ShowS)
-> (OptimizeOptions -> String)
-> ([OptimizeOptions] -> ShowS)
-> Show OptimizeOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> OptimizeOptions -> ShowS
showsPrec :: Int -> OptimizeOptions -> ShowS
$cshow :: OptimizeOptions -> String
show :: OptimizeOptions -> String
$cshowList :: [OptimizeOptions] -> ShowS
showList :: [OptimizeOptions] -> ShowS
Show)

{- | Default optimization options: 1000 iterations, 20000 evaluations, tolerance 1e-8 and
learning rate 0.01. Customize with record updates; algorithms validate the selected
settings.
-}
defaultOptimizeOptions :: OptimizeOptions
defaultOptimizeOptions :: OptimizeOptions
defaultOptimizeOptions = Int -> Int -> Double -> Double -> OptimizeOptions
OptimizeOptions Int
1000 Int
20000 Double
1e-8 Double
0.01

{- | Accepted parameters and diagnostics. optGradientNorm is the Euclidean
gradient norm for smooth methods, the projected/proximal mapping for proximal
methods, and a method-specific geometric stopping residual for derivative-free
methods. Stochastic searches cannot certify global optimality.
-}
data OptimizeResult = OptimizeResult
    { OptimizeResult -> [Double]
optPoint :: [Double]
    -- ^ Returned parameter vector at the accepted estimate.
    , OptimizeResult -> Double
optValue :: Double
    -- ^ Objective value at optPoint.
    , OptimizeResult -> Double
optGradientNorm :: Double
    -- ^ Gradient,projected/proximal,or geometric stopping residual as defined by the selected method.
    , OptimizeResult -> Int
optIterations :: Int
    -- ^ Completed optimization iterations.
    , OptimizeResult -> Int
optEvaluations :: Int
    -- ^ Recorded objective/evaluation count.
    , OptimizeResult -> Bool
optConverged :: Bool
    -- ^ Whether the selected method met its stopping criterion; this does not certify a global optimum.
    }
    deriving (OptimizeResult -> OptimizeResult -> Bool
(OptimizeResult -> OptimizeResult -> Bool)
-> (OptimizeResult -> OptimizeResult -> Bool) -> Eq OptimizeResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: OptimizeResult -> OptimizeResult -> Bool
== :: OptimizeResult -> OptimizeResult -> Bool
$c/= :: OptimizeResult -> OptimizeResult -> Bool
/= :: OptimizeResult -> OptimizeResult -> Bool
Eq, Int -> OptimizeResult -> ShowS
[OptimizeResult] -> ShowS
OptimizeResult -> String
(Int -> OptimizeResult -> ShowS)
-> (OptimizeResult -> String)
-> ([OptimizeResult] -> ShowS)
-> Show OptimizeResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> OptimizeResult -> ShowS
showsPrec :: Int -> OptimizeResult -> ShowS
$cshow :: OptimizeResult -> String
show :: OptimizeResult -> String
$cshowList :: [OptimizeResult] -> ShowS
showList :: [OptimizeResult] -> ShowS
Show)

-- | Require finite positive tolerance/rate and positive iteration/evaluation budgets.
checkedOptions :: OptimizeOptions -> Either MathError ()
checkedOptions :: OptimizeOptions -> Either MathError ()
checkedOptions OptimizeOptions
o
    | OptimizeOptions -> Int
optMaxIterations OptimizeOptions
o Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| OptimizeOptions -> Int
optMaxEvaluations OptimizeOptions
o Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 =
        MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"positive optimization budgets required")
    | (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x Bool -> Bool -> Bool
|| Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0) [OptimizeOptions -> Double
optTolerance OptimizeOptions
o, OptimizeOptions -> Double
optLearningRate OptimizeOptions
o] =
        MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
InvalidTolerance
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

-- | Reject NaN/infinity, including finite-input arithmetic overflow caught at solver boundaries.
checkedPoint :: [Double] -> Either MathError ()
checkedPoint :: [Double] -> Either MathError ()
checkedPoint [Double]
xs
    | (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Double
x -> Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x) [Double]
xs = MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
NonFiniteResult
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

-- | Evaluate a deterministic objective at finite coordinates and reject a nonfinite value.
checkedValue :: Objective -> [Double] -> Either MathError Double
checkedValue :: Objective -> Objective
checkedValue Objective
f [Double]
x = do [Double] -> Either MathError ()
checkedPoint [Double]
x; y <- Objective
f [Double]
x; checkedPoint [y]; Right y

-- | Evaluate a finite gradient of exactly the input dimension.
checkedGradient :: Gradient -> [Double] -> Either MathError [Double]
checkedGradient :: Gradient -> Gradient
checkedGradient Gradient
g [Double]
x = do
    [Double] -> Either MathError ()
checkedPoint [Double]
x
    ys <- Gradient
g [Double]
x
    if length x /= length ys
        then Left (VectorLengthMismatch (length x) (length ys))
        else checkedPoint ys
    Right ys

-- | Reject a proposed evaluation total that exceeds the supplied limit; equality is permitted.
checkBudget :: OptimizeOptions -> Int -> Either MathError ()
checkBudget :: OptimizeOptions -> Int -> Either MathError ()
checkBudget OptimizeOptions
o Int
n =
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> OptimizeOptions -> Int
optMaxEvaluations OptimizeOptions
o
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"optimization evaluation budget" Int
n)
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()

-- | Internal same-length dot kernel; public algorithms validate shapes and finite arithmetic at their boundaries.
dot :: [Double] -> [Double] -> Double
dot :: [Double] -> [Double] -> Double
dot [Double]
a [Double]
b = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
a [Double]
b)

-- | Scaled Euclidean norm. A mathematical norm beyond Double range remains positive infinity.
norm :: [Double] -> Double
norm :: [Double] -> Double
norm [Double]
xs =
    let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
xs)
     in if Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) [Double]
xs))

-- | Internal same-length vector kernels; shape validation belongs to the public checked caller.
add, sub :: [Double] -> [Double] -> [Double]
add :: [Double] -> [Double] -> [Double]
add = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)
sub :: [Double] -> [Double] -> [Double]
sub = (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-)

-- | Multiply each vector component by a scalar.
scale :: Double -> [Double] -> [Double]
scale :: Double -> [Double] -> [Double]
scale Double
a = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
a)

-- | Dense identity matrix for a checked nonnegative dimension.
identity :: Int -> [[Double]]
identity :: Int -> [[Double]]
identity Int
n = [[if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j then Double
1 else Double
0 | Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

-- | Dense row-major matrix-vector kernel for shapes already validated by the caller.
matVec :: [[Double]] -> [Double] -> [Double]
matVec :: [[Double]] -> [Double] -> [Double]
matVec [[Double]]
m [Double]
v = ([Double] -> Double) -> [[Double]] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map ([Double] -> [Double] -> Double
`dot` [Double]
v) [[Double]]
m

-- | Dense row-major matrix multiplication for compatible checked inner dimensions.
matMul :: [[Double]] -> [[Double]] -> [[Double]]
matMul :: [[Double]] -> [[Double]] -> [[Double]]
matMul [[Double]]
a [[Double]]
b = [[[Double] -> [Double] -> Double
dot [Double]
row [Double]
col | [Double]
col <- [[Double]] -> [[Double]]
forall a. [[a]] -> [[a]]
transpose [[Double]]
b] | [Double]
row <- [[Double]]
a]

-- | Outer product preserving input coordinate order.
outer :: [Double] -> [Double] -> [[Double]]
outer :: [Double] -> [Double] -> [[Double]]
outer [Double]
a [Double]
b = (Double -> [Double]) -> [Double] -> [[Double]]
forall a b. (a -> b) -> [a] -> [b]
map (\Double
x -> Double -> [Double] -> [Double]
scale Double
x [Double]
b) [Double]
a

-- | Checked square dense LU solve; singular and nonfinite systems return MathError.
solveDense :: [[Double]] -> [Double] -> Either MathError [Double]
solveDense :: [[Double]] -> Gradient
solveDense [[Double]]
a [Double]
b
    | [[Double]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Double]]
a 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]
b 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]
b) (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]]
a =
        MathError -> Either MathError [Double]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"square system shape")
    | Bool
otherwise = do
        [Double] -> Either MathError ()
checkedPoint ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
a [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
b); m <- Int -> Int -> [Double] -> Either MathError (Matrix Double)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
M.fromListChecked ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b) ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
b) ([[Double]] -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Double]]
a); LU.solve m b

-- | Build diagnostics from a final objective and gradient, charging both callbacks.
result :: Int -> Int -> Bool -> Objective -> Gradient -> [Double] -> Either MathError OptimizeResult
result :: Int
-> Int
-> Bool
-> Objective
-> Gradient
-> [Double]
-> Either MathError OptimizeResult
result Int
k Int
n Bool
converged Objective
f Gradient
g [Double]
x = do
    y <- Objective -> Objective
checkedValue Objective
f [Double]
x
    gs <- checkedGradient g x
    checkedPoint [norm gs]
    Right (OptimizeResult x y (norm gs) k (n + 2) converged)

{- | Selected constant or scheduled learning-rate policy, with method-specific parameters
validated by the schedule function.
-}
data LearningRate
    = ConstantRate Double
    | StepDecay Double Int Double
    | ExponentialDecay Double Double
    | CosineAnnealing Double Double Int
    | Warmup Int LearningRate
    deriving (LearningRate -> LearningRate -> Bool
(LearningRate -> LearningRate -> Bool)
-> (LearningRate -> LearningRate -> Bool) -> Eq LearningRate
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LearningRate -> LearningRate -> Bool
== :: LearningRate -> LearningRate -> Bool
$c/= :: LearningRate -> LearningRate -> Bool
/= :: LearningRate -> LearningRate -> Bool
Eq, Int -> LearningRate -> ShowS
[LearningRate] -> ShowS
LearningRate -> String
(Int -> LearningRate -> ShowS)
-> (LearningRate -> String)
-> ([LearningRate] -> ShowS)
-> Show LearningRate
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LearningRate -> ShowS
showsPrec :: Int -> LearningRate -> ShowS
$cshow :: LearningRate -> String
show :: LearningRate -> String
$cshowList :: [LearningRate] -> ShowS
showList :: [LearningRate] -> ShowS
Show)

-- | Evaluate a nonnegative explicit schedule at a zero-based step; warmup scales the initial steps linearly.
learningRate :: LearningRate -> Int -> Either MathError Double
learningRate :: LearningRate -> Int -> Either MathError Double
learningRate LearningRate
schedule Int
step
    | Int
step Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative schedule step")
    | Bool
otherwise = do
        value <- case LearningRate
schedule of
            ConstantRate Double
r -> Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
r
            StepDecay Double
r Int
period Double
decay
                | Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Double
decay Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
decay Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
step Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
period))
                | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"step schedule parameters")
            ExponentialDecay Double
r Double
decay
                | Double
decay Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 -> Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
decay Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
step)
                | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"exponential schedule decay")
            CosineAnnealing Double
hi Double
lo Int
period
                | Int
period Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Double
hi Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
lo Bool -> Bool -> Bool
&& Double
lo Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 ->
                    Double -> Either MathError Double
forall a b. b -> Either a b
Right (Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double -> Double
forall a. Floating a => a -> a
cos (Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
step Int
period) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
period)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
                | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"cosine schedule parameters")
            Warmup Int
steps LearningRate
base
                | Int
steps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> do
                    r <- LearningRate -> Int -> Either MathError Double
learningRate LearningRate
base (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
step Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
steps))
                    Right (r * min 1 (fromIntegral (step + 1) / fromIntegral steps))
                | Bool
otherwise -> MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"warmup steps")
        checkedPoint [value]
        if value < 0 then Left (InvalidArgument "negative learning rate") else Right value

{- | Explicit monitoring state retaining the best score,nonimproving count,and stop flag; it
does not mutate an optimizer automatically.
-}
data EarlyStopping = EarlyStopping
    { EarlyStopping -> Maybe Double
bestScore :: Maybe Double
    -- ^ Best monitored score so far, or Nothing before the first observation.
    , EarlyStopping -> Int
staleSteps :: Int
    -- ^ Consecutive monitoring updates without sufficient improvement.
    , EarlyStopping -> Bool
stopped :: Bool
    -- ^ Whether early-stopping patience has been exhausted.
    }
    deriving (EarlyStopping -> EarlyStopping -> Bool
(EarlyStopping -> EarlyStopping -> Bool)
-> (EarlyStopping -> EarlyStopping -> Bool) -> Eq EarlyStopping
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EarlyStopping -> EarlyStopping -> Bool
== :: EarlyStopping -> EarlyStopping -> Bool
$c/= :: EarlyStopping -> EarlyStopping -> Bool
/= :: EarlyStopping -> EarlyStopping -> Bool
Eq, Int -> EarlyStopping -> ShowS
[EarlyStopping] -> ShowS
EarlyStopping -> String
(Int -> EarlyStopping -> ShowS)
-> (EarlyStopping -> String)
-> ([EarlyStopping] -> ShowS)
-> Show EarlyStopping
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EarlyStopping -> ShowS
showsPrec :: Int -> EarlyStopping -> ShowS
$cshow :: EarlyStopping -> String
show :: EarlyStopping -> String
$cshowList :: [EarlyStopping] -> ShowS
showList :: [EarlyStopping] -> ShowS
Show)

-- | Track strict improvement beyond minDelta; stop after patience stale observations.
earlyStopping :: Int -> Double -> Double -> EarlyStopping -> Either MathError EarlyStopping
earlyStopping :: Int
-> Double
-> Double
-> EarlyStopping
-> Either MathError EarlyStopping
earlyStopping Int
patience Double
minDelta Double
value EarlyStopping
state = do
    [Double] -> Either MathError ()
checkedPoint [Double
minDelta, Double
value]
    if Int
patience Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Double
minDelta 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
"early stopping parameters")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let improved :: Bool
improved = Bool -> (Double -> Bool) -> Maybe Double -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (\Double
old -> Double
value Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
old Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
minDelta) (EarlyStopping -> Maybe Double
bestScore EarlyStopping
state)
        stale :: Int
stale = if Bool
improved then Int
0 else EarlyStopping -> Int
staleSteps EarlyStopping
state Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    EarlyStopping -> Either MathError EarlyStopping
forall a b. b -> Either a b
Right (Maybe Double -> Int -> Bool -> EarlyStopping
EarlyStopping (if Bool
improved then Double -> Maybe Double
forall a. a -> Maybe a
Just Double
value else EarlyStopping -> Maybe Double
bestScore EarlyStopping
state) Int
stale (Int
stale Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
patience))

-- | Step selection policy: Armijo backtracking or strong-Wolfe conditions.
data LineSearch = Armijo | StrongWolfe deriving (LineSearch -> LineSearch -> Bool
(LineSearch -> LineSearch -> Bool)
-> (LineSearch -> LineSearch -> Bool) -> Eq LineSearch
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LineSearch -> LineSearch -> Bool
== :: LineSearch -> LineSearch -> Bool
$c/= :: LineSearch -> LineSearch -> Bool
/= :: LineSearch -> LineSearch -> Bool
Eq, Int -> LineSearch -> ShowS
[LineSearch] -> ShowS
LineSearch -> String
(Int -> LineSearch -> ShowS)
-> (LineSearch -> String)
-> ([LineSearch] -> ShowS)
-> Show LineSearch
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LineSearch -> ShowS
showsPrec :: Int -> LineSearch -> ShowS
$cshow :: LineSearch -> String
show :: LineSearch -> String
$cshowList :: [LineSearch] -> ShowS
showList :: [LineSearch] -> ShowS
Show)

{- | Backtracking Armijo or bracket-and-zoom strong Wolfe (c1=1e-4,c2=.9).
The return includes all objective/gradient trials; descent is required.
-}
lineSearch ::
    LineSearch ->
    Int ->
    Objective ->
    Gradient ->
    [Double] ->
    [Double] ->
    Double ->
    [Double] ->
    Either MathError (Double, Double, [Double], Int)
lineSearch :: LineSearch
-> Int
-> Objective
-> Gradient
-> [Double]
-> [Double]
-> Double
-> [Double]
-> Either MathError (Double, Double, [Double], Int)
lineSearch LineSearch
method Int
budget Objective
f Gradient
g [Double]
x [Double]
direction Double
fx [Double]
gx
    | [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x 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]
direction Bool -> Bool -> Bool
|| [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
x 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]
gx =
        MathError -> Either MathError (Double, Double, [Double], Int)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"line search vector shape")
    | Int
budget Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = MathError -> Either MathError (Double, Double, [Double], Int)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"line search budget" Int
0)
    | [Double] -> [Double] -> Double
dot [Double]
gx [Double]
direction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 = MathError -> Either MathError (Double, Double, [Double], Int)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"line search requires descent")
    | Bool
otherwise = Int
-> Double
-> Double
-> Double
-> Either MathError (Double, Double, [Double], Int)
search Int
0 Double
0 Double
1 Double
fx
  where
    slope :: Double
slope = [Double] -> [Double] -> Double
dot [Double]
gx [Double]
direction
    trial :: Double -> Either MathError (Double, [Double], Double)
trial Double
a = do
        let point :: [Double]
point = [Double] -> [Double] -> [Double]
add [Double]
x (Double -> [Double] -> [Double]
scale Double
a [Double]
direction)
        y <- Objective -> Objective
checkedValue Objective
f [Double]
point
        gy <- checkedGradient g point
        Right (y, gy, dot gy direction)
    search :: Int
-> Double
-> Double
-> Double
-> Either MathError (Double, Double, [Double], Int)
search Int
k Double
prev Double
a Double
prevY
        | Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
budget = MathError -> Either MathError (Double, Double, [Double], Int)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"line search" Int
k)
        | Bool
otherwise = do
            (y, gy, d) <- Double -> Either MathError (Double, [Double], Double)
trial Double
a
            if y > fx + 1e-4 * a * slope || (k > 0 && y >= prevY)
                then zoom (k + 1) prev a prevY
                else
                    if method == Armijo || abs d <= -0.9 * slope
                        then Right (a, y, gy, 2 * k + 2)
                        else
                            if d >= 0
                                then zoom (k + 1) a prev y
                                else search (k + 1) a (2 * a) y
    zoom :: Int
-> Double
-> Double
-> Double
-> Either MathError (Double, Double, [Double], Int)
zoom Int
k Double
lo Double
hi Double
ylo
        | Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
budget Bool -> Bool -> Bool
|| Double -> Double
forall a. Num a => a -> a
abs (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-16 = MathError -> Either MathError (Double, Double, [Double], Int)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"Wolfe zoom" Int
k)
        | Bool
otherwise = do
            let a :: Double
a = (Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
hi) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
            (y, gy, d) <- Double -> Either MathError (Double, [Double], Double)
trial Double
a
            if y > fx + 1e-4 * a * slope || y >= ylo
                then zoom (k + 1) lo a ylo
                else
                    if abs d <= -0.9 * slope
                        then Right (a, y, gy, 2 * k + 2)
                        else if d * (hi - lo) >= 0 then zoom (k + 1) a lo y else zoom (k + 1) a hi y