Add support for Jacobian-vector products

Author: devmotion

README.md

@@ -19,7 +19,7 @@ NLSolversBase.jl is the core, common dependency of several packages in the [Juli
 The package aims at establishing common ground for [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl), [LineSearches.jl](https://github.com/JuliaNLSolvers/LineSearches.jl), and [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl). The common ground is mainly the types used to hold objective related callables, information about the objectives, and an interface to interact with these types.
 
 ## NDifferentiable
-There are currently three main types: `NonDifferentiable`, `OnceDifferentiable`, and `TwiceDifferentiable`. There's also a more experimental `TwiceDifferentiableHV` for optimization algorithms that use Hessian-vector products. An `NDifferentiable` instance can be used to hold relevant functions for
+There are currently three main types: `NonDifferentiable`, `OnceDifferentiable`, and `TwiceDifferentiable`. An `NDifferentiable` instance can be used to hold relevant functions for
 
  - Optimization: $F : \mathbb{R}^N \to \mathbb{R}$
  - Solving systems of equations: $F : \mathbb{R}^N \to \mathbb{R}^N$

src/NLSolversBase.jl

@@ -9,15 +9,17 @@ export AbstractObjective,
        NonDifferentiable,
        OnceDifferentiable,
        TwiceDifferentiable,
-       TwiceDifferentiableHV,
        value,
        value!,
        value_gradient!,
        value_jacobian!,
+       value_jvp!,
        gradient,
        gradient!,
        jacobian,
        jacobian!,
+       jvp,
+       jvp!,
        hessian,
        hessian!,
        value!!,
@@ -38,6 +40,7 @@ export AbstractObjective,
        clear!,
        f_calls,
        g_calls,
+       jvp_calls,
        h_calls,
        hv_calls
 
@@ -51,12 +54,11 @@ include("objective_types/abstract.jl")
 include("objective_types/nondifferentiable.jl")
 include("objective_types/oncedifferentiable.jl")
 include("objective_types/twicedifferentiable.jl")
-include("objective_types/twicedifferentiablehv.jl")
 include("objective_types/incomplete.jl")
 include("objective_types/constraints.jl")
 include("interface.jl")
 
 NonDifferentiable(f::OnceDifferentiable, x::AbstractArray) = NonDifferentiable(f.f, x, copy(f.F))
 NonDifferentiable(f::TwiceDifferentiable, x::AbstractArray) = NonDifferentiable(f.f, x, copy(f.F))
-NonDifferentiable(f::TwiceDifferentiableHV, x::AbstractArray) = NonDifferentiable(f.f, x, copy(f.F))
+
 end # module

src/interface.jl

@@ -103,10 +103,36 @@ value(obj::AbstractObjective) = obj.F
 gradient(obj::AbstractObjective) = obj.DF
 "Get the most recently evaluated Jacobian of `obj`."
 jacobian(obj::AbstractObjective) = obj.DF
+"""
+    jvp(obj::Union{OnceDifferentiable,TwiceDifferentiable})
+
+Return the most recently evaluated Jacobian-vector product of `obj`.
+
+!!! warn
+    Generally, it is unsafe to rely on the state of `obj`.
+    This function should only be used to, e.g., check whether the
+    most recently evaluated Jacobian-vector product is finite.
+    In most cases one should use [`jvp!(obj, x, v)`](@ref)
+    instead of this function.
+"""
+jvp(obj::Union{OnceDifferentiable, TwiceDifferentiable}) = obj.JVP
 "Get the `i`th element of the most recently evaluated gradient of `obj`."
 gradient(obj::AbstractObjective, i::Integer) = obj.DF[i]
 "Get the most recently evaluated Hessian of `obj`"
 hessian(obj::AbstractObjective) = obj.H
+"""
+    hv_product(obj::TwiceDifferentiable)
+
+Return the most recently evaluated Hessian-vector product of `obj`.
+
+!!! warn
+    Generally, it is unsafe to rely on the state of `obj`.
+    This function should only be used to, e.g., check whether the
+    most recently evaluated Hessian-vector product is finite.
+    In most cases one should use [`hv_product!(obj, x, v)`](@ref)
+    instead of this function.
+"""
+hv_product(obj::TwiceDifferentiable) = obj.Hv
 
 value_jacobian!(obj, x) = value_jacobian!(obj, obj.F, obj.DF, x)
 function value_jacobian!(obj, F, J, x)
@@ -153,6 +179,92 @@ function jacobian(obj::AbstractObjective, x)
     return newdf
 end
 
+"""
+    jvp!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+
+Return the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`,
+and cache the results in `obj`.
+
+!!! note
+    This function does use cached results if available.
+"""
+function jvp!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+    if x != obj.x_jvp || v != obj.v_jvp
+        jvp!!(obj, x, v)
+    end
+    return obj.JVP
+end
+
+"""
+    jvp!!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+
+Return the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`,
+and cache the results in `obj`.
+
+!!! note
+    This function does not use cached results but forces reevaluation of the Jacobian-vector product.
+"""
+function jvp!!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+    if obj.JVP isa Real
+        obj.JVP = obj.jvp(x, v)
+    else
+        obj.jvp(obj.JVP, x, v)
+    end
+    copyto!(obj.x_jvp, x)
+    copyto!(obj.v_jvp, v)
+    obj.jvp_calls += 1
+    return obj.JVP
+end
+
+"""
+    value_jvp!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+
+Return the value and the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`,
+and cache the results in `obj`.
+
+!!! note
+    This function does use cached results if available.
+"""
+function value_jvp!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+    if x != obj.x_f
+        if x != obj.x_jvp || v != obj.v_jvp
+            # Both value and Jacobian-vector product have to be recomputed
+            value_jvp!!(obj, x, v)
+        else
+            # Only value has to be recomputed
+            value!!(obj, x)
+        end
+    elseif x != obj.x_jvp || v != obj.v_jvp
+        jvp!!(obj, x, v)
+    end
+    return obj.F, obj.JVP
+end
+
+"""
+    value_jvp!!(obj::Union{OnceDifferentiable, TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+
+Return the value and the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`,
+and cache the results in `obj`.
+
+!!! note
+    This function does not use cached results but forces reevaluation of the Jacobian-vector product.
+"""
+function value_jvp!!(obj::Union{OnceDifferentiable,TwiceDifferentiable}, x::AbstractArray, v::AbstractArray)
+    if obj.F isa Real
+        y, ty = obj.fjvp(x, v)
+        obj.F = y
+        obj.JVP = ty
+    else
+        obj.fjvp(obj.F, obj.JVP, x, v)
+    end
+    copyto!(obj.x_f, x)
+    copyto!(obj.x_jvp, x)
+    copyto!(obj.v_jvp, v)
+    obj.f_calls += 1
+    obj.jvp_calls += 1
+    return obj.F, obj.JVP
+end
+
 value(obj::NonDifferentiable{TF, TX}, x) where {TF<:AbstractArray, TX} = value(obj, copy(obj.F), x)
 value(obj::OnceDifferentiable{TF, TDF, TX}, x) where {TF<:AbstractArray, TDF, TX} = value(obj, copy(obj.F), x)
 function value(obj::AbstractObjective, F, x)
@@ -170,6 +282,25 @@ function value!!(obj::AbstractObjective, F, x)
     F
 end
 
+function hv_product!(obj::TwiceDifferentiable, x, v)
+    if x != obj.x_hv || v != obj.v_hv
+        hv_product!!(obj, x, v)
+    end
+    obj.Hv
+end
+function hv_product!!(obj::TwiceDifferentiable, x, v)
+    if obj.hv === nothing
+        H = hessian!(obj, x)
+        LinearAlgebra.mul!(obj.Hv, H, v)
+    else
+        obj.hv_calls += 1
+        obj.hv(obj.Hv, x, v)
+    end
+    copyto!(obj.x_hv, x)
+    copyto!(obj.v_hv, v)
+    return obj.Hv
+end
+
 function _clear_f!(d::AbstractObjective)
     d.f_calls = 0
     if d.F isa AbstractArray
@@ -188,6 +319,18 @@ function _clear_df!(d::AbstractObjective)
     nothing
 end
 
+function _clear_jvp!(d::AbstractObjective)
+    d.jvp_calls = 0
+    if d.JVP isa AbstractArray
+        fill!(d.JVP, NaN)
+    else
+        d.JVP = NaN
+    end
+    fill!(d.x_jvp, NaN)
+    fill!(d.v_jvp, NaN)
+    nothing
+end
+
 function _clear_h!(d::AbstractObjective)
     d.h_calls = 0
     fill!(d.H, NaN)
@@ -208,28 +351,25 @@ clear!(d::NonDifferentiable)  = _clear_f!(d)
 function clear!(d::OnceDifferentiable)
     _clear_f!(d)
     _clear_df!(d)
+    _clear_jvp!(d)
     nothing
 end
 
 function clear!(d::TwiceDifferentiable)
     _clear_f!(d)
     _clear_df!(d)
+    _clear_jvp!(d)
     _clear_h!(d)
-    nothing
-end
-
-function clear!(d::TwiceDifferentiableHV)
-    _clear_f!(d)
-    _clear_df!(d)
     _clear_hv!(d)
     nothing
 end
 
+f_calls(d::Union{NonDifferentiable, OnceDifferentiable, TwiceDifferentiable}) = d.f_calls
 g_calls(d::NonDifferentiable) = 0
+g_calls(d::Union{OnceDifferentiable, TwiceDifferentiable}) = d.df_calls
+jvp_calls(d::NonDifferentiable) = 0
+jvp_calls(d::Union{OnceDifferentiable, TwiceDifferentiable}) = d.jvp_calls
 h_calls(d::Union{NonDifferentiable, OnceDifferentiable}) = 0
-f_calls(d) = d.f_calls
-g_calls(d) = d.df_calls
-h_calls(d) = d.h_calls
-hv_calls(d) = 0
-h_calls(d::TwiceDifferentiableHV) = 0
-hv_calls(d::TwiceDifferentiableHV) = d.hv_calls
+h_calls(d::TwiceDifferentiable) = d.h_calls
+hv_calls(d::Union{NonDifferentiable, OnceDifferentiable}) = 0
+hv_calls(d::TwiceDifferentiable) = d.hv_calls

src/objective_types/abstract.jl

@@ -15,6 +15,45 @@ function make_fdf(x, F::Number, f, g!)
     end
 end
 
+# Given Julia functions of the gradient/Jacobian, create a function that calculates the Jacobian-vector product.
+function make_jvp(x::AbstractArray, F::AbstractArray, j!)
+    let j! = j!, jx = alloc_DF(x, F)
+        function jvp!(jvp, x, v)
+            j!(jx, x)
+            LinearAlgebra.mul!(jvp, jx, v)
+            return jvp
+        end
+    end
+end
+function make_jvp(x::AbstractArray, F::Number, g!)
+    let g! = g!, gx = alloc_DF(x, F)
+        function jvp(x, v)
+            g!(gx, x)
+            return LinearAlgebra.dot(gx, v)
+        end
+    end
+end
+
+# Given a Julia function of the objective function and its gradient/Jacobian, create a function that calculates the function value and the Jacobian-vector product.
+function make_fjvp(x::AbstractArray, F::AbstractArray, fj!)
+    let fj! = fj!, jx = alloc_DF(x, F)
+        function fjvp!(fx, jvp, x, v)
+            fj!(fx, jx, x)
+            LinearAlgebra.mul!(jvp, jx, v)
+            return fx, jvp
+        end
+    end
+end
+function make_fjvp(x::AbstractArray, F::Number, fg!)
+    let fg! = fg!, gx = alloc_DF(x, F)
+        function jvp(x, v)
+            fx = fg!(gx, x)
+            return fx, LinearAlgebra.dot(gx, v)
+        end
+    end
+end
+
+
 # Initialize an n-by-n Jacobian
 alloc_DF(x, F) = eltype(x)(NaN) .* vec(F) .* vec(x)'