Merge pull request #198 from baggepinnen/autoauto

RFC: add auto-auto hessian

Author: ChrisRackauckas

src/SparseDiffTools.jl

@@ -30,7 +30,10 @@ export  contract_color,
         ForwardColorJacCache,
         numauto_color_hessian!,
         numauto_color_hessian,
+        autoauto_color_hessian!,
+        autoauto_color_hessian,
         ForwardColorHesCache,
+        ForwardAutoColorHesCache,
         auto_jacvec,auto_jacvec!,
         num_jacvec,num_jacvec!,
         num_vecjac,num_vecjac!,

src/differentiation/compute_hessian_ad.jl

@@ -96,3 +96,76 @@ function numauto_color_hessian(f,
     numauto_color_hessian!(H, f, x, hes_cache)
     return H
 end
+
+
+
+## autoauto_color_hessian
+
+mutable struct ForwardAutoColorHesCache{TJC,TG,TS,TC}
+    jac_cache::TJC
+    grad!::TG
+    sparsity::TS
+    colorvec::TC
+end
+
+struct AutoAutoTag end
+
+function ForwardAutoColorHesCache(f,
+    x::AbstractVector{V},
+    colorvec::AbstractVector{<:Integer}=eachindex(x),
+    sparsity::Union{AbstractMatrix,Nothing}=nothing) where V
+
+    if sparsity === nothing
+        sparsity = sparse(ones(length(x), length(x)))
+    end
+
+    tag = ForwardDiff.Tag(AutoAutoTag(), V)
+    chunksize = ForwardDiff.pickchunksize(maximum(colorvec))
+    chunk = ForwardDiff.Chunk(chunksize)
+
+    jacobian_config = ForwardDiff.JacobianConfig(f, x, chunk, tag)
+    gradient_config = ForwardDiff.GradientConfig(f, jacobian_config.duals, chunk, tag)
+
+    outer_tag = get_tag(jacobian_config.duals)
+    g! = (G, x) -> ForwardDiff.gradient!(G, f, x, gradient_config, Val(false))
+
+    jac_cache = ForwardColorJacCache(g!, x; colorvec, sparsity, tag=outer_tag)
+    
+    return ForwardAutoColorHesCache(jac_cache, g!, sparsity, colorvec)
+end
+
+function autoauto_color_hessian!(H::AbstractMatrix{<:Number},
+    f,
+    x::AbstractArray{<:Number},
+    hes_cache::ForwardAutoColorHesCache)
+
+    forwarddiff_color_jacobian!(H, hes_cache.grad!, x, hes_cache.jac_cache)
+end
+
+function autoauto_color_hessian!(H::AbstractMatrix{<:Number},
+    f,
+    x::AbstractArray{<:Number},
+    colorvec::AbstractVector{<:Integer}=eachindex(x),
+    sparsity::Union{AbstractMatrix,Nothing}=nothing)
+    hes_cache = ForwardAutoColorHesCache(f, x, colorvec, sparsity)
+    autoauto_color_hessian!(H, f, x, hes_cache)
+    return H
+end
+
+function autoauto_color_hessian(f,
+    x::AbstractArray{<:Number},
+    hes_cache::ForwardAutoColorHesCache)
+    H = convert.(eltype(x), hes_cache.sparsity)
+    autoauto_color_hessian!(H, f, x, hes_cache)
+    return H
+end
+
+function autoauto_color_hessian(f,
+    x::AbstractArray{<:Number},
+    colorvec::AbstractVector{<:Integer}=eachindex(x),
+    sparsity::Union{AbstractMatrix,Nothing}=nothing)
+    hes_cache = ForwardAutoColorHesCache(f, x, colorvec, sparsity)
+    H = convert.(eltype(x), hes_cache.sparsity)
+    autoauto_color_hessian!(H, f, x, hes_cache)
+    return H
+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