Compute JVP in line searches

Author: devmotion

.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

@@ -34,7 +34,7 @@ LineSearches = "7.4.0"
 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"
@@ -65,3 +65,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "Aqua", "Distributions", "ExplicitImports", "ForwardDiff", "JET", "MathOptInterface", "Measurements", "OptimTestProblems", "Random", "RecursiveArrayTools", "StableRNGs", "ReverseDiff"]
+
+[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" }

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,9 @@ 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} <:
+mutable struct ManifoldObjective{M<:Manifold,T<:NLSolversBase.AbstractObjective} <:
                NLSolversBase.AbstractObjective
-    manifold::Manifold
+    manifold::M
     inner_obj::T
 end
 # TODO: is it safe here to call retract! and change x?
@@ -52,6 +52,20 @@ function NLSolversBase.value_gradient!(obj::ManifoldObjective, x)
     return value(obj.inner_obj)
 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)
+    gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, gradient(obj.inner_obj), xin)
+    return dot(gradient(obj.inner_obj), v)
+end
+function NLSolversBase.value_jvp!(obj::ManifoldObjective, x, v)
+    xin = retract(obj.manifold, x)
+    value_gradient!(obj.inner_obj, xin)
+    project_tangent!(obj.manifold, gradient(obj.inner_obj), xin)
+    return value(obj.inner_obj), dot(gradient(obj.inner_obj), 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
@@ -62,6 +76,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) = jvp!(obj.inner_obj, x, v)
+NLSolversBase.value_jvp!(obj::ManifoldObjective{Flat}, x, v) = 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
@@ -41,6 +43,15 @@ using NLSolversBase:
     TwiceDifferentiableConstraints,
     nconstraints,
     nconstraints_x,
+    value,
+    value!,
+    value!!,
+    gradient,
+    gradient!,
+    value_gradient!,
+    value_gradient!!,
+    jvp!,
+    value_jvp!,
     hessian,
     hessian!,
     hessian!!,

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/solvers/constrained/fminbox.jl

@@ -1,5 +1,3 @@
-import NLSolversBase:
-    value, value!, value!!, gradient, gradient!, value_gradient!, value_gradient!!
 ####### FIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIX THE MIDDLE OF BOX CASE THAT WAS THERE
 mutable struct BarrierWrapper{TO,TB,Tm,TF,TDF} <: AbstractObjective
     obj::TO
@@ -64,24 +62,24 @@ function _barrier_term_gradient(x::T, l, u) where {T}
     end
     return g
 end
-function value_gradient!(bb::BoxBarrier, g, x)
+function NLSolversBase.value_gradient!(bb::BoxBarrier, g, x)
     g .= _barrier_term_gradient.(x, bb.lower, bb.upper)
     value(bb, x)
 end
-function gradient(bb::BoxBarrier, g, x)
+function NLSolversBase.gradient(bb::BoxBarrier, g, x)
     g = copy(g)
     g .= _barrier_term_gradient.(x, bb.lower, bb.upper)
 end
 # Wrappers
-function value!!(bw::BarrierWrapper, x)
+function NLSolversBase.value!!(bw::BarrierWrapper, x)
     bw.Fb = value(bw.b, x)
     bw.Ftotal = bw.mu * bw.Fb
     if in_box(bw, x)
         value!!(bw.obj, x)
         bw.Ftotal += value(bw.obj)
     end
 end
-function value_gradient!!(bw::BarrierWrapper, x)
+function NLSolversBase.value_gradient!!(bw::BarrierWrapper, x)
     bw.Fb = value(bw.b, x)
     bw.Ftotal = bw.mu * bw.Fb
     bw.DFb .= _barrier_term_gradient.(x, bw.b.lower, bw.b.upper)
@@ -93,7 +91,7 @@ function value_gradient!!(bw::BarrierWrapper, x)
     end
 
 end
-function value_gradient!(bb::BarrierWrapper, x)
+function NLSolversBase.value_gradient!(bb::BarrierWrapper, x)
     bb.DFb .= _barrier_term_gradient.(x, bb.b.lower, bb.b.upper)
     bb.Fb = value(bb.b, x)
     bb.DFtotal .= bb.mu .* bb.DFb
@@ -105,9 +103,9 @@ function value_gradient!(bb::BarrierWrapper, x)
         bb.Ftotal += value(bb.obj)
     end
 end
-value(bb::BoxBarrier, x) =
+NLSolversBase.value(bb::BoxBarrier, x) =
     mapreduce(x -> _barrier_term_value(x...), +, zip(x, bb.lower, bb.upper))
-function value!(obj::BarrierWrapper, x)
+function NLSolversBase.value!(obj::BarrierWrapper, x)
     obj.Fb = value(obj.b, x)
     obj.Ftotal = obj.mu * obj.Fb
     if in_box(obj, x)
@@ -116,20 +114,20 @@ function value!(obj::BarrierWrapper, x)
     end
     obj.Ftotal
 end
-value(obj::BarrierWrapper) = obj.Ftotal
-function value(obj::BarrierWrapper, x)
+NLSolversBase.value(obj::BarrierWrapper) = obj.Ftotal
+function NLSolversBase.value(obj::BarrierWrapper, x)
     F = obj.mu * value(obj.b, x)
     if in_box(obj, x)
         F += value(obj.obj, x)
     end
     F
 end
-function gradient!(obj::BarrierWrapper, x)
+function NLSolversBase.gradient!(obj::BarrierWrapper, x)
     gradient!(obj.obj, x)
     obj.DFb .= gradient(obj.b, obj.DFb, x) # this should just be inplace?
     obj.DFtotal .= gradient(obj.obj) .+ obj.mu * obj.Fb
 end
-gradient(obj::BarrierWrapper) = obj.DFtotal
+NLSolversBase.gradient(obj::BarrierWrapper) = obj.DFtotal
 
 # this mutates mu but not the gradients
 # Super unsafe in that it depends on x_df being correct!
@@ -299,7 +297,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))
@@ -536,12 +534,13 @@ function optimize(
     end
     results = optimize(dfbox, x, _optimizer, options, state)
     stopped_by_callback = results.stopped_by.callback
-    dfbox.obj.f_calls[1] = 0
+    # TODO: Define an API (e.g. `reset_calls!`?) in NLSolversBase
+    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
     copyto!(x, minimizer(results))
     boxdist = Base.minimum(((xi, li, ui),) -> min(xi - li, ui - xi), zip(x, l, u)) # Base.minimum !== minimum
@@ -613,12 +612,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: Define an API (e.g. `reset_calls!`?) in NLSolversBase
+            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
             copyto!(x, minimizer(results))
             boxdist = Base.minimum(((xi, li, ui),) -> min(xi - li, ui - xi), zip(x, l, u)) # Base.minimum !== minimum