Compute JVP in line searches

.github/workflows/CI.yml

@@ -11,8 +11,8 @@ jobs:
     strategy:
       matrix:
         version:
-          - "min"
-          - "lts"
+          # - "min"
+          # - "lts"
           - "1"
         os:
           - ubuntu-latest

.github/workflows/Docs.yml

@@ -19,7 +19,7 @@ jobs:
         with:
           version: '1'
       - name: Install dependencies
-        run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: julia --project=docs/ -e 'using Pkg; Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # If authenticating with GitHub Actions token

Project.toml

@@ -18,6 +18,10 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [weakdeps]
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 
+[sources]
+LineSearches = { url = "https://github.com/devmotion/LineSearches.jl.git", rev = "dmw/jvp" }
+NLSolversBase = { url = "https://github.com/devmotion/NLSolversBase.jl.git", rev = "dmw/jvp" }
+
 [extensions]
 OptimMOIExt = "MathOptInterface"
 
@@ -30,11 +34,11 @@ ExplicitImports = "1.13.2"
 FillArrays = "0.6.2, 0.7, 0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
 ForwardDiff = "0.10, 1"
 JET = "0.9, 0.10"
-LineSearches = "7.5.1"
+LineSearches = "7.6"
 LinearAlgebra = "<0.0.1, 1.6"
 MathOptInterface = "1.17"
 Measurements = "2.14.1"
-NLSolversBase = "7.9.0"
+NLSolversBase = "8"
 NaNMath = "0.3.2, 1"
 OptimTestProblems = "2.0.3"
 PositiveFactorizations = "0.2.2"

docs/Project.toml

@@ -12,5 +12,6 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Documenter = "1"
 Literate = "2"
 
-[sources.Optim]
-path = ".."
+[sources]
+Optim = { path = ".." }
+NLSolversBase = { url = "https://github.com/devmotion/NLSolversBase.jl.git", rev = "dmw/jvp" }

docs/src/examples/ipnewton_basics.jl

@@ -22,6 +22,7 @@
 # constraint is unbounded from below or above respectively.
 
 using Optim, NLSolversBase #hide
+import ADTypes #hide
 import NLSolversBase: clear! #hide
 
 # # Constrained optimization with `IPNewton`

ext/OptimMOIExt.jl

@@ -1,7 +1,7 @@
 module OptimMOIExt
 
 using Optim
-using Optim.LinearAlgebra: rmul! 
+using Optim.LinearAlgebra: rmul!
 import MathOptInterface as MOI
 
 function __init__()

src/Manifolds.jl

@@ -20,9 +20,8 @@ project_tangent(M::Manifold, x) = project_tangent!(M, similar(x), x)
 retract(M::Manifold, x) = retract!(M, copy(x))
 
 # Fake objective function implementing a retraction
-mutable struct ManifoldObjective{T<:NLSolversBase.AbstractObjective} <:
-               NLSolversBase.AbstractObjective
-    manifold::Manifold
+struct ManifoldObjective{M<:Manifold,T<:AbstractObjective} <: AbstractObjective
+    manifold::M
     inner_obj::T
 end
 # TODO: is it safe here to call retract! and change x?
@@ -43,6 +42,20 @@ function NLSolversBase.value_gradient!(obj::ManifoldObjective, x)
     return f_x, g_x
 end
 
+# In general, we have to compute the gradient/Jacobian separately as it has to be projected
+function NLSolversBase.jvp!(obj::ManifoldObjective, x, v)
+    xin = retract(obj.manifold, x)
+    g_x = gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, g_x, xin)
+    return dot(g_x, v)
+end
+function NLSolversBase.value_jvp!(obj::ManifoldObjective, x, v)
+    xin = retract(obj.manifold, x)
+    f_x, g_x = value_gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, g_x, xin)
+    return f_x, dot(g_x, v)
+end
+
 """Flat Euclidean space {R,C}^N, with projections equal to the identity."""
 struct Flat <: Manifold end
 # all the functions below are no-ops, and therefore the generated code
@@ -53,6 +66,10 @@ retract!(M::Flat, x) = x
 project_tangent(M::Flat, g, x) = g
 project_tangent!(M::Flat, g, x) = g
 
+# Optimizations for `Flat` manifold
+NLSolversBase.jvp!(obj::ManifoldObjective{Flat}, x, v) = NLSolversBase.jvp!(obj.inner_obj, x, v)
+NLSolversBase.value_jvp!(obj::ManifoldObjective{Flat}, x, v) = NLSolversBase.value_jvp!(obj.inner_obj, x, v)
+
 """Spherical manifold {|x| = 1}."""
 struct Sphere <: Manifold end
 retract!(S::Sphere, x) = (x ./= norm(x))

src/Optim.jl

@@ -16,6 +16,8 @@ documentation online at http://julianlsolvers.github.io/Optim.jl/stable/ .
 """
 module Optim
 
+import ADTypes
+
 using PositiveFactorizations: Positive # for globalization strategy in Newton
 
 using LineSearches: LineSearches # for globalization strategy in Quasi-Newton algs
@@ -35,7 +37,6 @@ using NLSolversBase:
     NonDifferentiable,
     OnceDifferentiable,
     TwiceDifferentiable,
-    TwiceDifferentiableHV,
     AbstractConstraints,
     ConstraintBounds,
     TwiceDifferentiableConstraints,

src/api.jl

@@ -104,7 +104,7 @@ f_calls(d::AbstractObjective) = NLSolversBase.f_calls(d)
 
 g_calls(r::OptimizationResults) = error("g_calls is not implemented for $(summary(r)).")
 g_calls(r::MultivariateOptimizationResults) = r.g_calls
-g_calls(d::AbstractObjective) = NLSolversBase.g_calls(d)
+g_calls(d::AbstractObjective) = NLSolversBase.g_calls(d) + NLSolversBase.jvp_calls(d)
 
 h_calls(r::OptimizationResults) = error("h_calls is not implemented for $(summary(r)).")
 h_calls(r::MultivariateOptimizationResults) = r.h_calls

src/multivariate/optimize/interface.jl

@@ -66,20 +66,6 @@ promote_objtype(
     inplace::Bool,
     f::InplaceObjective,
 ) = TwiceDifferentiable(f, x, real(zero(eltype(x))))
-promote_objtype(
-    method::SecondOrderOptimizer,
-    x,
-    autodiff::ADTypes.AbstractADType,
-    inplace::Bool,
-    f::NLSolversBase.InPlaceObjectiveFGHv,
-) = TwiceDifferentiableHV(f, x)
-promote_objtype(
-    method::SecondOrderOptimizer,
-    x,
-    autodiff::ADTypes.AbstractADType,
-    inplace::Bool,
-    f::NLSolversBase.InPlaceObjectiveFG_Hv,
-) = TwiceDifferentiableHV(f, x)
 promote_objtype(method::SecondOrderOptimizer, x, autodiff::ADTypes.AbstractADType, inplace::Bool, f, g) =
     TwiceDifferentiable(
         f,

src/multivariate/optimize/optimize.jl

@@ -1,25 +1,23 @@
 # Update function value, gradient and Hessian
 function update_fgh!(d, state, ::ZerothOrderOptimizer)
-    f_x = value!(d, state.x)
+    f_x = NLSolversBase.value!(d, state.x)
     state.f_x = f_x
     return nothing
 end
 function update_fgh!(d, state, method::FirstOrderOptimizer)
-    f_x, g_x = value_gradient!(d, state.x)
+    f_x, g_x = NLSolversBase.value_gradient!(d, state.x)
+    copyto!(state.g_x, g_x)
     if hasproperty(method, :manifold)
-        project_tangent!(method.manifold, g_x, state.x)
+        project_tangent!(method.manifold, state.g_x, state.x)
     end
     state.f_x = f_x
-    copyto!(state.g_x, g_x)
     return nothing
 end
 function update_fgh!(d, state, method::SecondOrderOptimizer)
     # Manifold optimization is currently not supported for second order optimization algorithms
     @assert !hasproperty(method, :manifold)
 
-    # TODO: Switch to `value_gradient_hessian!` when it becomes available
-    f_x, g_x = value_gradient!(d, state.x)
-    H_x = hessian!(d, state.x)
+    f_x, g_x, H_x = NLSolversBase.value_gradient_hessian!(d, state.x)
     state.f_x = f_x
     copyto!(state.g_x, g_x)
     copyto!(state.H_x, H_x)

src/multivariate/solvers/constrained/fminbox.jl

@@ -13,6 +13,7 @@ end
 
 NLSolversBase.f_calls(obj::BarrierWrapper) = NLSolversBase.f_calls(obj.obj)
 NLSolversBase.g_calls(obj::BarrierWrapper) = NLSolversBase.g_calls(obj.obj)
+NLSolversBase.jvp_calls(obj::BarrierWrapper) = NLSolversBase.jvp_calls(obj.obj)
 NLSolversBase.h_calls(obj::BarrierWrapper) = NLSolversBase.h_calls(obj.obj)
 NLSolversBase.hv_calls(obj::BarrierWrapper) = NLSolversBase.hv_calls(obj.obj)
 
@@ -40,9 +41,9 @@ struct BoxBarrier{L,U}
     upper::U
 end
 function in_box(bb::BoxBarrier, x)
-    all(x -> x[1] >= x[2] && x[1] <= x[3], zip(x, bb.lower, bb.upper))
+    all(x -> x[1] <= x[2] <= x[3], zip(bb.lower, x, bb.upper))
 end
-in_box(bw::BarrierWrapper, x) = in_box(bw.b, x)
+
 # evaluates the value and gradient components comming from the log barrier
 function _barrier_term_value(x::T, l, u) where {T}
     dxl = x - l
@@ -71,11 +72,15 @@ function _barrier_term_gradient(x::T, l, u) where {T}
     return g
 end
 
+function _barrier_jvp(bb::BoxBarrier, x, v)
+    return sum(Broadcast.instantiate(Broadcast.broadcasted((xi, li, ui, vi) -> dot(_barrier_term_gradient(xi, li, ui), vi), x, bb.lower, bb.upper, v)))
+end
+
 # Wrappers
 function NLSolversBase.value_gradient!(bb::BarrierWrapper, x)
     bb.DFb .= _barrier_term_gradient.(x, bb.b.lower, bb.b.upper)
     bb.Fb = _barrier_value(bb.b, x)
-    if in_box(bb, x)
+    if in_box(bb.b, x)
         F, DF = value_gradient!(bb.obj, x)
         bb.DFtotal .= muladd.(bb.mu, bb.DFb, DF)
         bb.Ftotal = muladd(bb.mu, bb.Fb, F)
@@ -87,7 +92,7 @@ function NLSolversBase.value_gradient!(bb::BarrierWrapper, x)
 end
 function NLSolversBase.value!(obj::BarrierWrapper, x)
     obj.Fb = _barrier_value(obj.b, x)
-    if in_box(obj, x)
+    if in_box(obj.b, x)
         F = value!(obj.obj, x)
         obj.Ftotal = muladd(obj.mu, obj.Fb, F)
     else
@@ -97,15 +102,15 @@ function NLSolversBase.value!(obj::BarrierWrapper, x)
 end
 function NLSolversBase.value(obj::BarrierWrapper, x)
     Fb = _barrier_value(obj.b, x)
-    if in_box(obj, x)
+    if in_box(obj.b, x)
         return muladd(obj.mu, Fb, value(obj.obj, x))
     else
         return obj.mu * Fb
     end
 end
 function NLSolversBase.gradient!(obj::BarrierWrapper, x)
     obj.DFb .= _barrier_term_gradient.(x, obj.b.lower, obj.b.upper)
-    if in_box(obj, x)
+    if in_box(obj.b, x)
         DF = gradient!(obj.obj, x)
         obj.DFtotal .= muladd.(obj.mu, obj.DFb, DF)
     else
@@ -114,6 +119,29 @@ function NLSolversBase.gradient!(obj::BarrierWrapper, x)
     return obj.DFtotal
 end
 
+function NLSolversBase.jvp!(obj::BarrierWrapper, x, v)
+    JVPb = _barrier_jvp(obj.b, x, v)
+    if in_box(obj.b, x)
+        JVP = NLSolversBase.jvp!(obj.obj, x, v)
+        return muladd(obj.mu, JVPb, JVP)
+    else
+        return obj.mu * JVPb
+    end
+end
+function NLSolversBase.value_jvp!(obj::BarrierWrapper, x, v)
+    JVPb = _barrier_jvp(obj.b, x, v)
+    obj.Fb = _barrier_value(obj.b, x)
+    if in_box(obj.b, x)
+        F, JVP = NLSolversBase.value_jvp!(obj.obj, x, v)
+        JVPtotal = muladd(obj.mu, JVPb, JVP)
+        obj.Ftotal = muladd(obj.mu, obj.Fb, F)
+    else
+        JVPtotal = obj.mu * JVPb
+        obj.Ftotal = obj.mu * obj.Fb
+    end
+    return obj.Ftotal, JVPtotal
+end
+
 function limits_box(
     x::AbstractArray{T},
     d::AbstractArray{T},
@@ -274,7 +302,7 @@ function optimize(
     initial_x::AbstractArray,
     F::Fminbox = Fminbox(),
     options::Options = Options();
-    inplace = true,
+    inplace::Bool = true,
 )
 
     g! = inplace ? g : (G, x) -> copyto!(G, g(x))
@@ -499,12 +527,13 @@ function optimize(
 
     results = optimize(dfbox, x, _optimizer, options, state)
     stopped_by_callback = results.stopped_by.callback
-    dfbox.obj.f_calls[1] = 0
+    # TODO: Reset everything? Add upstream API for resetting call counters?
+    dfbox.obj.f_calls = 0
     if hasfield(typeof(dfbox.obj), :df_calls)
-        dfbox.obj.df_calls[1] = 0
+        dfbox.obj.df_calls = 0
     end
     if hasfield(typeof(dfbox.obj), :h_calls)
-        dfbox.obj.h_calls[1] = 0
+        dfbox.obj.h_calls = 0
     end
 
     # Compute function value and gradient (without barrier term)
@@ -579,12 +608,13 @@ function optimize(
             resultsnew = optimize(dfbox, x, _optimizer, options, state)
             stopped_by_callback = resultsnew.stopped_by.callback
             append!(results, resultsnew)
-            dfbox.obj.f_calls[1] = 0
+            # TODO: Reset everything? Add upstream API for resetting call counters?
+            dfbox.obj.f_calls = 0
             if hasfield(typeof(dfbox.obj), :df_calls)
-                dfbox.obj.df_calls[1] = 0
+                dfbox.obj.df_calls = 0
             end
             if hasfield(typeof(dfbox.obj), :h_calls)
-                dfbox.obj.h_calls[1] = 0
+                dfbox.obj.h_calls = 0
             end
 
             # Compute function value and gradient (without barrier term)

src/multivariate/solvers/constrained/ipnewton/ipnewton.jl

@@ -107,8 +107,7 @@ function initial_state(
     # Allocate fields for the objective function
     n = length(initial_x)
 
-    # TODO: Switch to `value_gradient_hessian!`
-    f_x, g_x, H_x = NLSolversBase.value_gradient_hessian!!(d, initial_x)
+    f_x, g_x, H_x = NLSolversBase.value_gradient_hessian!(d, initial_x)
 
     # More constraints
     constr_J = fill!(Matrix{T}(undef, mc, n), NaN)
@@ -157,8 +156,7 @@ end
 
 function update_fgh!(d, constraints::TwiceDifferentiableConstraints, state, method::IPNewton)
     # Compute objective function, gradient and Hessian matrix
-    # TODO: Switch to `value_gradient_hessian!!`
-    f_x, g_x, H_x = NLSolversBase.value_gradient_hessian!!(d, state.x)
+    f_x, g_x, H_x = NLSolversBase.value_gradient_hessian!(d, state.x)
     copyto!(state.g_x, g_x)
     copyto!(state.H_x, H_x)
     state.f_x = f_x

src/multivariate/solvers/first_order/accelerated_gradient_descent.jl

@@ -43,12 +43,13 @@ function initial_state(
 )
     initial_x = copy(initial_x)
     retract!(method.manifold, initial_x)
-    f_x, g_x = value_gradient!(d, initial_x)
+    f_x, g_x = NLSolversBase.value_gradient!(d, initial_x)
+    g_x = copy(g_x)
     project_tangent!(method.manifold, g_x, initial_x)
 
     AcceleratedGradientDescentState(
         initial_x, # Maintain current state in state.x
-        copy(g_x), # Maintain current gradient in state.g_x
+        g_x, # Maintain current gradient in state.g_x
         f_x, # Maintain current f in state.f_x
         fill!(similar(initial_x), NaN), # Maintain previous state in state.x_previous
         oftype(f_x, NaN), # Store previous f in state.f_x_previous

src/multivariate/solvers/first_order/adam.jl

@@ -62,7 +62,7 @@ function reset!(method::Adam, state::AdamState, obj, x)
     # Update function value and gradient
     copyto!(state.x, x)
     retract!(method.manifold, state.x)
-    f_x, g_x = value_gradient!(obj, state.x)
+    f_x, g_x = NLSolversBase.value_gradient!(obj, state.x)
     copyto!(state.g_x, g_x)
     project_tangent!(method.manifold, state.g_x, state.x)
     state.f_x = f_x
@@ -88,7 +88,7 @@ function initial_state(method::Adam, ::Options, d, initial_x::AbstractArray)
     # Compute function value and gradient
     initial_x = copy(initial_x)
     retract!(method.manifold, initial_x)
-    f_x, g_x = value_gradient!(d, initial_x)
+    f_x, g_x = NLSolversBase.value_gradient!(d, initial_x)
     g_x = copy(g_x)
     project_tangent!(method.manifold, g_x, initial_x)