add tests for auto-auto hessians

Author: baggepinnen

src/differentiation/compute_hessian_ad.jl

@@ -1,4 +1,4 @@
-struct ForwardColorHesCache{THS,THC,TI<:Integer,TD,TGF,TGC,TG}
+struct ForwardColorHesCache{THS, THC, TI<:Integer, TD, TGF, TGC, TG}
     sparsity::THS
     colors::THC
     ncolors::TI
@@ -16,40 +16,40 @@ function make_hessian_buffers(colorvec, x)
     buffer = similar(D)
     G1 = similar(x)
     G2 = similar(x)
-    return (; ncolors, D, buffer, G1, G2)
+    return (;ncolors, D, buffer, G1, G2)
 end
 
-function ForwardColorHesCache(f,
-    x::AbstractVector{<:Number},
-    colorvec::AbstractVector{<:Integer}=eachindex(x),
-    sparsity::Union{AbstractMatrix,Nothing}=nothing,
-    (g!)=(G, x, grad_config) -> ForwardDiff.gradient!(G, f, x, grad_config))
+function ForwardColorHesCache(f, 
+                              x::AbstractVector{<:Number}, 
+                              colorvec::AbstractVector{<:Integer}=eachindex(x), 
+                              sparsity::Union{AbstractMatrix, Nothing}=nothing,
+                              g! = (G, x, grad_config) -> ForwardDiff.gradient!(G, f, x, grad_config))
     ncolors, D, buffer, G, G2 = make_hessian_buffers(colorvec, x)
     grad_config = ForwardDiff.GradientConfig(f, x)
-
+    
     # If user supplied their own gradient function, make sure it has the right
     # signature (i.e. g!(G, x) or g!(G, x, grad_config::ForwardDiff.GradientConfig))
-    if !hasmethod(g!, (typeof(G), typeof(G), typeof(grad_config)))
-        if !hasmethod(g!, (typeof(G), typeof(G)))
+    if ! hasmethod(g!, (typeof(G), typeof(G), typeof(grad_config)))
+        if ! hasmethod(g!, (typeof(G), typeof(G)))
             throw(ArgumentError("Signature of `g!` must be either `g!(G, x)` or `g!(G, x, grad_config::ForwardDiff.GradientConfig)`"))
         end
         # define new method that takes a GradientConfig but doesn't use it
         g1!(G, x, grad_config) = g!(G, x)
     else
         g1! = g!
     end
-
+    
     if sparsity === nothing
         sparsity = sparse(ones(length(x), length(x)))
     end
     return ForwardColorHesCache(sparsity, colorvec, ncolors, D, buffer, g1!, grad_config, G, G2)
 end
 
-function numauto_color_hessian!(H::AbstractMatrix{<:Number},
-    f,
-    x::AbstractArray{<:Number},
-    hes_cache::ForwardColorHesCache;
-    safe=true)
+function numauto_color_hessian!(H::AbstractMatrix{<:Number}, 
+                                    f, 
+                                    x::AbstractArray{<:Number}, 
+                                    hes_cache::ForwardColorHesCache;
+                                    safe = true)
     ϵ = cbrt(eps(eltype(x)))
     for j in 1:hes_cache.ncolors
         x .+= ϵ .* @view hes_cache.D[:, j]
@@ -69,28 +69,28 @@ function numauto_color_hessian!(H::AbstractMatrix{<:Number},
     return H
 end
 
-function numauto_color_hessian!(H::AbstractMatrix{<:Number},
-    f,
-    x::AbstractArray{<:Number},
-    colorvec::AbstractVector{<:Integer}=eachindex(x),
-    sparsity::Union{AbstractMatrix,Nothing}=nothing)
+function numauto_color_hessian!(H::AbstractMatrix{<:Number}, 
+                                    f, 
+                                    x::AbstractArray{<:Number},
+                                    colorvec::AbstractVector{<:Integer}=eachindex(x), 
+                                    sparsity::Union{AbstractMatrix, Nothing}=nothing)
     hes_cache = ForwardColorHesCache(f, x, colorvec, sparsity)
     numauto_color_hessian!(H, f, x, hes_cache)
     return H
 end
 
-function numauto_color_hessian(f,
-    x::AbstractArray{<:Number},
-    hes_cache::ForwardColorHesCache)
+function numauto_color_hessian(f, 
+                                   x::AbstractArray{<:Number}, 
+                                   hes_cache::ForwardColorHesCache)
     H = convert.(eltype(x), hes_cache.sparsity)
     numauto_color_hessian!(H, f, x, hes_cache)
     return H
 end
 
 function numauto_color_hessian(f,
-    x::AbstractArray{<:Number},
-    colorvec::AbstractVector{<:Integer}=eachindex(x),
-    sparsity::Union{AbstractMatrix,Nothing}=nothing)
+                                   x::AbstractArray{<:Number},
+                                   colorvec::AbstractVector{<:Integer}=eachindex(x), 
+                                   sparsity::Union{AbstractMatrix, Nothing}=nothing)
     hes_cache = ForwardColorHesCache(f, x, colorvec, sparsity)
     H = convert.(eltype(x), hes_cache.sparsity)
     numauto_color_hessian!(H, f, x, hes_cache)
@@ -103,7 +103,7 @@ end
 
 mutable struct ForwardAutoColorHesCache{TJC,TG,TS,TC}
     jac_cache::TJC
-    grad!TG
+    grad!::TG
     sparsity::TS
     colorvec::TC
 end

test/test_sparse_hessian.jl

@@ -98,3 +98,30 @@ for (i, hescache) in enumerate([hescache1, hescache2, hescache3, hescache4, hesc
     #     for _ in 1:100)
     # @test t_unsafe <= t_safe
 end
+
+
+hescache1 = ForwardAutoColorHesCache(fscalar, x, colors, sparsity)
+hescache2 = ForwardAutoColorHesCache(fscalar, x)
+
+
+for (i, hescache) in enumerate([hescache1, hescache2])
+
+    H = SparseDiffTools.autoauto_color_hessian(fscalar, x, colors, sparsity)
+    H1 = SparseDiffTools.autoauto_color_hessian(fscalar, x, hescache)
+    H2 = SparseDiffTools.autoauto_color_hessian(fscalar, x)
+    @test all(isapprox.(Hforward, H, rtol=1e-6))
+    @test all(isapprox.(H, H1, rtol=1e-6))
+    @test all(isapprox.(H2, H1, rtol=1e-6))
+
+    H1 = similar(H)
+
+    SparseDiffTools.autoauto_color_hessian!(H1, fscalar, x, collect(hescache.colorvec), hescache.sparsity)
+    @test all(isapprox.(H1, H))
+
+    SparseDiffTools.autoauto_color_hessian!(H2, fscalar, x)
+    @test all(isapprox.(H2, H))
+
+    SparseDiffTools.autoauto_color_hessian!(H1, fscalar, x, hescache)
+    @test all(isapprox.(H1, H))
+
+end