Add support for Jacobian-vector products

Author: devmotion

src/NLSolversBase.jl

@@ -14,15 +14,19 @@ export AbstractObjective,
        value!,
        value_gradient!,
        value_jacobian!,
+       value_jvp!,
        gradient,
        gradient!,
        jacobian,
        jacobian!,
+       jvp,
+       jvp!,
        hessian,
        hessian!,
        value!!,
        value_gradient!!,
        value_jacobian!!,
+       value_jvp!!,
        hessian!!,
        hv_product,
        hv_product!,
@@ -38,6 +42,7 @@ export AbstractObjective,
        clear!,
        f_calls,
        g_calls,
+       jvp_calls,
        h_calls,
        hv_calls
 

src/interface.jl

@@ -153,6 +153,126 @@ function jacobian(obj::AbstractObjective, x)
     return newdf
 end
 
+"""
+    jvp(obj::OnceDifferentiable)
+
+Return the most recently evaluated Jacobian-vector product of the objective function `obj`.
+"""
+jvp(obj::OnceDifferentiable) = obj.JVP
+
+"""
+    jvp(obj::OnceDifferentiable, x::AbstractArray, v::AbstractArray)
+
+Return the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`.
+
+!!! note
+    This function does neither use cached results nor cache the results.
+"""
+function jvp(obj::OnceDifferentiable, x::AbstractArray, v::AbstractArray)
+    obj.jvp_calls += 1
+    return obj.jvp(x, v)
+end
+
+"""
+    jvp!(obj::OnceDifferentiable, 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::OnceDifferentiable, x::AbstractArray, v::AbstractArray)
+    if x != obj.x_jvp || v != obj.v_jvp
+        jvp!!(obj, x, v)
+    end
+    return jvp(obj)
+end
+
+"""
+    jvp!!(obj::OnceDifferentiable, 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::OnceDifferentiable, 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 jvp(obj)
+end
+
+"""
+    value_jvp(obj::OnceDifferentiable, x::AbstractArray, v::AbstractArray)
+
+Return the value and the Jacobian-vector product of the objective function `obj` at point `x` with tangents `v`.
+
+!!! note
+    This function does neither use cached results nor cache the results.
+"""
+function value_jvp(obj::OnceDifferentiable, x::AbstractArray, v::AbstractArray)
+    obj.f_calls += 1
+    obj.jvp_calls += 1
+    return obj.fjvp(x, v)
+end
+
+"""
+    value_jvp!(obj::OnceDifferentiable, 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::OnceDifferentiable, 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::OnceDifferentiable, 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::OnceDifferentiable, 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)
@@ -188,6 +308,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,6 +340,7 @@ clear!(d::NonDifferentiable)  = _clear_f!(d)
 function clear!(d::OnceDifferentiable)
     _clear_f!(d)
     _clear_df!(d)
+    _clear_jvp!(d)
     nothing
 end
 
@@ -229,6 +362,8 @@ g_calls(d::NonDifferentiable) = 0
 h_calls(d::Union{NonDifferentiable, OnceDifferentiable}) = 0
 f_calls(d) = d.f_calls
 g_calls(d) = d.df_calls
+jvp_calls(d) = 0
+jvp_calls(d::OnceDifferentiable) = d.jvp_calls
 h_calls(d) = d.h_calls
 hv_calls(d) = 0
 h_calls(d::TwiceDifferentiableHV) = 0

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)'