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

@@ -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,14 +37,13 @@ using NLSolversBase:
     NonDifferentiable,
     OnceDifferentiable,
     TwiceDifferentiable,
-    TwiceDifferentiableHV,
     AbstractConstraints,
     ConstraintBounds,
     TwiceDifferentiableConstraints,
     nconstraints,
     nconstraints_x,
     hessian!,
-    hv_product!
+    hvp!
 
 # var for NelderMead
 using Statistics: var

src/api.jl

@@ -99,16 +99,16 @@ g_norm_trace(r::OptimizationResults) =
     error("g_norm_trace is not implemented for $(summary(r)).")
 g_norm_trace(r::MultivariateOptimizationResults) = [state.g_norm for state in trace(r)]
 
+# TODO: Overload `NLSolversBase.xxx` instead of defining separate `Optim.xxx` methods?
 f_calls(r::OptimizationResults) = r.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::OptimizationResults) = error(LazyString("`g_calls` is not implemented for ", summary(r), "."))
 g_calls(r::MultivariateOptimizationResults) = r.g_calls
-g_calls(d::AbstractObjective) = NLSolversBase.g_calls(d)
-
-h_calls(r::OptimizationResults) = error("h_calls is not implemented for $(summary(r)).")
+jvp_calls(r::OptimizationResults) = error(LazyString("`jvp_calls` is not implemented for ", summary(r), "."))
+jvp_calls(r::MultivariateOptimizationResults) = r.jvp_calls
+h_calls(r::OptimizationResults) = error(LazyString("`h_calls` is not implemented for ", summary(r), "."))
 h_calls(r::MultivariateOptimizationResults) = r.h_calls
-h_calls(d::AbstractObjective) = NLSolversBase.h_calls(d) + NLSolversBase.hv_calls(d)
+hvp_calls(r::OptimizationResults) = error(LazyString("`hvp_calls` is not implemented for ", summary(r), "."))
+hvp_calls(r::MultivariateOptimizationResults) = r.hvp_calls
 
 converged(r::UnivariateOptimizationResults) = r.stopped_by.converged
 function converged(r::MultivariateOptimizationResults)

src/maximize.jl

@@ -109,7 +109,9 @@ for api_method in (
     :iteration_limit_reached,
     :f_calls,
     :g_calls,
+    :jvp_calls,
     :h_calls,
+    :hvp_calls,
 )
     @eval $api_method(r::MaximizationWrapper) = $api_method(res(r))
 end