Autodiff support for 2nd order constraints

src/objective_types/constraints.jl

@@ -1,17 +1,17 @@
 ### Constraints
-#
+# 
 # Constraints are specified by the user as
 #    lx_i ≤   x[i]  ≤ ux_i  # variable (box) constraints
 #    lc_i ≤ c(x)[i] ≤ uc_i  # linear/nonlinear constraints
 # and become equality constraints with l_i = u_i. ±∞ are allowed for l
 # and u, in which case the relevant side(s) are unbounded.
-#
+# 
 # The user supplies functions to calculate c(x) and its derivatives.
-#
+# 
 # Of course we could unify the box-constraints into the
 # linear/nonlinear constraints, but that would force the user to
 # provide the variable-derivatives manually, which would be silly.
-#
+# 
 # This parametrization of the constraints gets "parsed" into a form
 # that speeds and simplifies the IPNewton algorithm, at the cost of many
 # additional variables. See `parse_constraints` for details.
@@ -35,7 +35,7 @@ function ConstraintBounds(lx, ux, lc, uc)
     _cb(symmetrize(lx, ux)..., symmetrize(lc, uc)...)
 end
 function _cb(lx::AbstractArray{Tx}, ux::AbstractArray{Tx}, lc::AbstractVector{Tc}, uc::AbstractVector{Tc}) where {Tx,Tc}
-    T = promote_type(Tx,Tc)
+    T = promote_type(Tx, Tc)
     ConstraintBounds{T}(length(lc), parse_constraints(T, lx, ux)..., parse_constraints(T, lc, uc)...)
 end
 
@@ -114,8 +114,8 @@ function OnceDifferentiableConstraints(lx::AbstractArray, ux::AbstractArray)
 end
 
 function OnceDifferentiableConstraints(bounds::ConstraintBounds)
-    c! = (c,x)->nothing
-    J! = (J,x)->nothing
+    c! = (c, x)->nothing
+    J! = (J, x)->nothing
     OnceDifferentiableConstraints(c!, J!, bounds)
 end
 
@@ -136,19 +136,19 @@ function OnceDifferentiableConstraints(c!, lx::AbstractVector, ux::AbstractVecto
     sizex = size(lx)
     sizec = size(lc)
 
-    xcache = zeros(T,sizex)
-    ccache = zeros(T,sizec)
+    xcache = zeros(T, sizex)
+    ccache = zeros(T, sizec)
 
-    if any(autodiff .== (:finite, :central))
+    if is_finitediff(autodiff)
         ccache2 = similar(ccache)
-        central_cache = FiniteDiff.JacobianCache(xcache, ccache,
-                                                      ccache2)
+        fdtype = finitediff_fdtype(autodiff)
+        jacobian_cache = FiniteDiff.JacobianCache(xcache, ccache,ccache2,fdtype)
         function jfinite!(J, x)
-            FiniteDiff.finite_difference_jacobian!(J, c!, x, central_cache)
+            FiniteDiff.finite_difference_jacobian!(J, c!, x, jacobian_cache)
             J
         end
         return OnceDifferentiableConstraints(c!, jfinite!, bounds)
-    elseif autodiff == :forward
+    elseif is_forwarddiff(autodiff)
         jac_cfg = ForwardDiff.JacobianConfig(c!, ccache, xcache, chunk)
         ForwardDiff.checktag(jac_cfg, c!, xcache)
 
@@ -169,20 +169,167 @@ struct TwiceDifferentiableConstraints{F,J,H,T} <: AbstractConstraints
     h!::H   # h!(storage, x) stores the hessian of the constraint functions
     bounds::ConstraintBounds{T}
 end
+
 function TwiceDifferentiableConstraints(c!, jacobian!, h!, lx, ux, lc, uc)
     b = ConstraintBounds(lx, ux, lc, uc)
     TwiceDifferentiableConstraints(c!, jacobian!, h!, b)
 end
 
+function TwiceDifferentiableConstraints(c!, lx::AbstractVector, ux::AbstractVector,
+    lc::AbstractVector, uc::AbstractVector,
+    autodiff::Symbol = :central,
+    chunk::ForwardDiff.Chunk = checked_chunk(lx))
+       if is_finitediff(autodiff)
+        fdtype = finitediff_fdtype(autodiff)
+        return twicediff_constraints_finite(c!,lx,ux,lc,uc,fdtype,nothing)
+    elseif is_forwarddiff(autodiff)
+        return twicediff_constraints_forward(c!,lx,ux,lc,uc,chunk,nothing)
+    else
+        error("The autodiff value $autodiff is not support. Use :finite or :forward.")
+    end
+end
+
+function TwiceDifferentiableConstraints(c!, con_jac!,lx::AbstractVector, ux::AbstractVector,
+    lc::AbstractVector, uc::AbstractVector,
+    autodiff::Symbol = :central,
+    chunk::ForwardDiff.Chunk = checked_chunk(lx))
+       if is_finitediff(autodiff)
+        fdtype = finitediff_fdtype(autodiff)
+        return twicediff_constraints_finite(c!,lx,ux,lc,uc,fdtype,con_jac!)
+    elseif is_forwarddiff(autodiff)
+        return  twicediff_constraints_forward(c!,lx,ux,lc,uc,chunk,con_jac!)
+    else
+        error("The autodiff value $autodiff is not support. Use :finite or :forward.")
+    end
+end
+
+
+
 function TwiceDifferentiableConstraints(lx::AbstractArray, ux::AbstractArray)
     bounds = ConstraintBounds(lx, ux, [], [])
     TwiceDifferentiableConstraints(bounds)
 end
 
+
+function twicediff_constraints_forward(c!, lx, ux, lc, uc,chunk,con_jac! = nothing)
+    bounds = ConstraintBounds(lx, ux, lc, uc)
+    T = eltype(bounds)
+    nc = length(lc)
+    nx = length(lx)
+    ccache = zeros(T, nc)
+    xcache = zeros(T, nx)
+    cache_check =  Ref{DataType}(Missing) #the datatype Missing, not the singleton
+    ref_f= Ref{Any}() #cache for intermediate jacobian used in the hessian
+    cxxcache = zeros(T, nx * nc, nx) #output cache for hessian
+    h = reshape(cxxcache, (nc, nx, nx)) #reshaped output 
+    hi = [@view h[i,:,:] for i in 1:nc]
+        #ref_f caches the closure function with its caches. other aproaches include using a Dict, but the 
+        #cost of switching happens just once per optimize call. 
+
+    if isnothing(con_jac!) #if the jacobian is not provided, generate one
+        jac_cfg = ForwardDiff.JacobianConfig(c!, ccache, xcache, chunk)
+        ForwardDiff.checktag(jac_cfg, c!, xcache)
+
+        jac! = (J, x) -> begin
+            ForwardDiff.jacobian!(J, c!, ccache, x, jac_cfg, Val{false}())
+            J
+        end
+
+        con_jac_cached = x -> begin
+            exists_cache = (cache_check[] == eltype(x))
+            if exists_cache
+                f = ref_f[]
+                return f(x)
+            else
+                jcache = zeros(eltype(x), nc)
+                out_cache = zeros(eltype(x), nc, nx)
+                cfg_cache = ForwardDiff.JacobianConfig(c!,jcache,x)
+                f = z->ForwardDiff.jacobian!(out_cache, c!, jcache, z,cfg_cache,Val{false}())
+                ref_f[] = f
+                cache_check[]= eltype(x)
+                return f(x)
+            end
+        end
+
+    else
+        jac! = (J,x) -> con_jac!(J,x)
+
+        #here, the cache should also include a JacobianConfig
+         con_jac_cached = x -> begin
+            exists_cache = (cache_check[] == eltype(x))
+            if exists_cache
+                f = ref_f[]
+                return f(x)
+            else
+                out_cache = zeros(eltype(x), nc, nx)
+                f = z->jac!(out_cache,x)
+                ref_f[] = f
+                cache_check[]= eltype(x)
+                return f(x)
+            end
+        end
+    end
+
+    hess_config_cache = ForwardDiff.JacobianConfig(typeof(con_jac_cached),lx)
+    function con_hess!(hess, x, λ)
+        ForwardDiff.jacobian!(cxxcache, con_jac_cached, x,hess_config_cache,Val{false}()) 
+        for i = 1:nc  #hot hessian loop
+            hess+=λ[i].*hi[i]
+        end
+        return hess
+    end
+
+    return TwiceDifferentiableConstraints(c!, jac!, con_hess!, bounds) 
+end
+
+
+function twicediff_constraints_finite(c!,lx,ux,lc,uc,fdtype,con_jac! = nothing)
+    bounds = ConstraintBounds(lx, ux, lc, uc)
+    T = eltype(bounds)
+    nx = length(lx)
+    nc = length(lc)
+    xcache = zeros(T, nx)
+    ccache = zeros(T, nc)
+
+    if isnothing(con_jac!)   
+        jac_ccache = similar(ccache)
+        jacobian_cache = FiniteDiff.JacobianCache(xcache, ccache,jac_ccache,fdtype)
+        function jac!(J, x)
+            FiniteDiff.finite_difference_jacobian!(J, c!, x, jacobian_cache)
+            J
+        end
+    else
+        jac! = (J,x) -> con_jac!(J,x)
+    end
+        cxxcache = zeros(T,nc*nx,nx) # to create cached jacobian
+        h = reshape(cxxcache, (nc, nx, nx)) #reshaped output 
+        hi = [@view h[i,:,:] for i in 1:nc]
+
+        function jac_vec!(J,x) #to evaluate the jacobian of a jacobian, FiniteDiff needs a vector version of that
+            j_mat = reshape(J,nc,nx)
+            return jac!(j_mat,x)
+            return J
+        end
+        hess_xcache =similar(xcache)
+        hess_cxcache =zeros(T,nc*nx) #output of jacobian, as a vector
+        hess_cxxcache =similar(hess_cxcache)
+        hess_config_cache = FiniteDiff.JacobianCache(hess_xcache,hess_cxcache,hess_cxxcache,fdtype)
+        function con_hess!(hess, x, λ)
+            FiniteDiff.finite_difference_jacobian!(cxxcache, jac_vec!, x,hess_config_cache) 
+            for i = 1:nc  
+                hi = @view h[i,:,:]
+                hess+=λ[i].*hi
+            end
+            return hess
+        end 
+        return TwiceDifferentiableConstraints(c!, jac!, con_hess!, bounds)
+end
+
+
 function TwiceDifferentiableConstraints(bounds::ConstraintBounds)
-    c! = (x,c)->nothing
-    J! = (x,J)->nothing
-    h! = (x,λ,h)->nothing
+    c! = (x, c)->nothing
+    J! = (x, J)->nothing
+    h! = (x, λ, h)->nothing
     TwiceDifferentiableConstraints(c!, J!, h!, bounds)
 end
 
@@ -278,11 +425,11 @@ function showeq(io, indent, eq, val, chr, style)
     if !isempty(eq)
         print(io, '\n', indent)
         if style == :bracket
-            eqstrs = map((i,v) -> "$chr[$i]=$v", eq, val)
+            eqstrs = map((i, v)->"$chr[$i]=$v", eq, val)
         else
-            eqstrs = map((i,v) -> "$(chr)_$i=$v", eq, val)
+            eqstrs = map((i, v)->"$(chr)_$i=$v", eq, val)
         end
-        foreach(s->print(io, s*", "), eqstrs[1:end-1])
+        foreach(s->print(io, s * ", "), eqstrs[1:end - 1])
         print(io, eqstrs[end])
     end
 end
@@ -291,12 +438,12 @@ function showineq(io, indent, ineqs, σs, bs, chr, style)
     if !isempty(ineqs)
         print(io, '\n', indent)
         if style == :bracket
-            ineqstrs = map((i,σ,b) -> "$chr[$i]"*ineqstr(σ,b), ineqs, σs, bs)
+            ineqstrs = map((i, σ, b)->"$chr[$i]" * ineqstr(σ, b), ineqs, σs, bs)
         else
-            ineqstrs = map((i,σ,b) -> "$(chr)_$i"*ineqstr(σ,b), ineqs, σs, bs)
+            ineqstrs = map((i, σ, b)->"$(chr)_$i" * ineqstr(σ, b), ineqs, σs, bs)
         end
-        foreach(s->print(io, s*", "), ineqstrs[1:end-1])
+        foreach(s->print(io, s * ", "), ineqstrs[1:end - 1])
         print(io, ineqstrs[end])
     end
 end
-ineqstr(σ,b) = σ>0 ? "≥$b" : "≤$b"
+ineqstr(σ, b) = σ > 0 ? "≥$b" : "≤$b"

test/constraints.jl

@@ -5,7 +5,7 @@
         @test NLSolversBase.nconstraints_x(cb) == 0
         @test cb.valc == [0.0]
         @test eltype(cb) == Float64
-
+        @test eltype(cb) == eltype(typeof(cb))
         cb = ConstraintBounds([0], [0], [], [])
         @test NLSolversBase.nconstraints(cb) == 0
         @test NLSolversBase.nconstraints_x(cb) == 1
@@ -69,7 +69,8 @@
                 lx = fill(-Inf, nx)
                 ux = fill(Inf, nx)
             end
-
+            @test_throws ErrorException OnceDifferentiableConstraints(cbd.c!, lx, ux,
+            lc, uc, :wuoah)
             for autodiff in (:finite, :forward)
                 odca = OnceDifferentiableConstraints(cbd.c!, lx, ux,
                                                      lc, uc, autodiff)
@@ -111,5 +112,56 @@
         @test isempty(odc.bounds.bc)
         @test odc.bounds.valc == [0.0]
 
+        @testset "second order autodiff" begin
+
+            prob = MVP.ConstrainedProblems.examples["HS9"]
+            cbd = prob.constraintdata
+            nc = length(cbd.lc)
+            nx = length(prob.initial_x)
+            lx = fill(-Inf, nx)
+            ux = fill(Inf, nx)
+            cb = ConstraintBounds(lx, ux, cbd.lc, cbd.uc)
+            odc = TwiceDifferentiableConstraints(cbd.c!, cbd.jacobian!, cbd.h!,lx, ux, cbd.lc, cbd.uc)
+
+            T = eltype(odc.bounds)
+            jac_result = zeros(T, nc,nx)
+            jac_result_autodiff = zeros(T, nc,nx)
+
+            hess_result = zeros(T, nx,nx)
+            hess_result_autodiff = zeros(T, nx,nx)
+            λ = rand(T,nc)
+            λ0 = zeros(T,nc)
+            @test_throws ErrorException TwiceDifferentiableConstraints(cbd.c!,lx, ux, cbd.lc, cbd.uc,:campanario)
+            for autodiff in (:finite,:forward) #testing double differentiation
+                odca2 = TwiceDifferentiableConstraints(cbd.c!,lx, ux, cbd.lc, cbd.uc,autodiff)
+
+                odca2.jacobian!(jac_result_autodiff,prob.initial_x) 
+                odc.jacobian!(jac_result,prob.initial_x)
+                odca2.h!(hess_result_autodiff,prob.initial_x,λ0) #warmup,λ0 means no modification 
+                odca2.h!(hess_result_autodiff,prob.initial_x,λ) 
+
+                odc.h!(hess_result,prob.initial_x,λ)
+
+
+                @test isapprox(jac_result, jac_result_autodiff, atol=1e-10) 
+                @test isapprox(hess_result, hess_result_autodiff, atol=1e-10) 
+
+                fill!(hess_result_autodiff,zero(T))
+                fill!(hess_result,zero(T))
+                fill!(jac_result_autodiff,zero(T))
+                fill!(jac_result,zero(T))
+            end
+            @test_throws ErrorException TwiceDifferentiableConstraints(cbd.c!, cbd.jacobian!,lx, ux, cbd.lc, cbd.uc,:qloctm)
+            for autodiff in (:finite,:forward) #testing autodiff hessian from constraint jacobian
+                odca2 = TwiceDifferentiableConstraints(cbd.c!, cbd.jacobian!,lx, ux, cbd.lc, cbd.uc,autodiff)
+                odca2.h!(hess_result_autodiff,prob.initial_x,λ0) #warmup,λ0 means no modification 
+                odca2.h!(hess_result_autodiff,prob.initial_x,λ) 
+                odc.h!(hess_result,prob.initial_x,λ)
+                @test isapprox(hess_result, hess_result_autodiff, atol=1e-10)
+                fill!(hess_result_autodiff,zero(T))
+                fill!(hess_result,zero(T))
+            end
+        end
     end
 end
+