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
@@ -35,12 +37,21 @@ using NLSolversBase:
     NonDifferentiable,
     OnceDifferentiable,
     TwiceDifferentiable,
-    TwiceDifferentiableHV,
     AbstractConstraints,
     ConstraintBounds,
     TwiceDifferentiableConstraints,
     nconstraints,
     nconstraints_x,
+    value,
+    value!,
+    value!!,
+    gradient,
+    gradient!,
+    value_gradient!,
+    value_gradient!!,
+    jvp,
+    jvp!,
+    value_jvp!,
     hessian,
     hessian!,
     hessian!!,

src/api.jl

@@ -100,18 +100,15 @@ g_norm_trace(r::OptimizationResults) =
 g_norm_trace(r::MultivariateOptimizationResults) = [state.g_norm for state in trace(r)]
 
 f_calls(r::OptimizationResults) = r.f_calls
-f_calls(d) = first(d.f_calls)
+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::NonDifferentiable) = 0
-g_calls(d) = first(d.df_calls)
+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
-h_calls(d::Union{NonDifferentiable,OnceDifferentiable}) = 0
-h_calls(d) = first(d.h_calls)
-h_calls(d::TwiceDifferentiableHV) = first(d.hv_calls)
+h_calls(d::AbstractObjective) = NLSolversBase.h_calls(d) + NLSolversBase.hv_calls(d)
 
 converged(r::UnivariateOptimizationResults) = r.stopped_by.converged
 function converged(r::MultivariateOptimizationResults)

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,