From 1e002d595df89e4b5675ff911d118cbb3e3e01b9 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Fri, 7 Jun 2024 11:20:00 -0400 Subject: [PATCH] fULL support of sparse AD for NLS models --- src/ADNLPModels.jl | 4 ++++ src/ad_api.jl | 19 ++++++++++++++++ src/sparse_hessian.jl | 41 +++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + test/sparse_hessian_nls.jl | 45 +++++++++++++++++++++++++++++++++++++ test/sparse_jacobian.jl | 4 +++- test/sparse_jacobian_nls.jl | 4 +++- 7 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 test/sparse_hessian_nls.jl diff --git a/src/ADNLPModels.jl b/src/ADNLPModels.jl index 18650372..40ea4b00 100644 --- a/src/ADNLPModels.jl +++ b/src/ADNLPModels.jl @@ -4,7 +4,11 @@ module ADNLPModels using LinearAlgebra, SparseArrays # external +<<<<<<< HEAD using ADTypes: ADTypes, AbstractSparsityDetector +======= +using ADTypes: ADTypes +>>>>>>> 5cacfb0 (fULL support of sparse AD for NLS models) using SparseConnectivityTracer, ColPack, ForwardDiff, ReverseDiff # JSO diff --git a/src/ad_api.jl b/src/ad_api.jl index c9867b25..0c938b5b 100644 --- a/src/ad_api.jl +++ b/src/ad_api.jl @@ -289,6 +289,25 @@ function NLPModels.hess_coord!( return vals end +function NLPModels.hess_structure_residuals!( + b::ADBackend, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + nothing +end + +function NLPModels.hess_coord_residuals!( + b::ADBackend, + nls::AbstractADNLSModel, + x::AbstractVector, + v::AbstractVector, + vals::AbstractVector, +) + nothing +end + function NLPModels.hprod!( b::ADBackend, nlp::ADModel, diff --git a/src/sparse_hessian.jl b/src/sparse_hessian.jl index c28cdb1e..0e9e5e1b 100644 --- a/src/sparse_hessian.jl +++ b/src/sparse_hessian.jl @@ -294,3 +294,44 @@ function NLPModels.hess_coord!( sparse_hess_coord!(ℓ, b, x, obj_weight, b.y, vals) return vals end + +function NLPModels.hess_structure_residuals!( + b::Union{SparseADHessian, SparseReverseADHessian}, + nls::AbstractADNLSModel, + rows::AbstractVector{<:Integer}, + cols::AbstractVector{<:Integer}, +) + function objective(x) + F = get_F(nls, b) + Fx = F(x) + return dot(Fx, Fx) / 2 + end + + H = compute_hessian_sparsity(objective, nls.meta.nvar, nothing, 0) + trilH = tril(H) + rowval = trilH.rowval + colptr = trilH.colptr + rows .= rowval + for i = 1:(nls.meta.nvar) + for j = colptr[i]:(colptr[i + 1] - 1) + cols[j] = i + end + end + return rows, cols +end + +function NLPModels.hess_coord_residuals!( + b::Union{SparseADHessian, SparseReverseADHessian}, + nls::AbstractADNLSModel, + x::AbstractVector, + v::AbstractVector, + vals::AbstractVector, +) + function objective(x) + F = get_F(nls, b) + Fx = F(x) + return dot(Fx, Fx) / 2 + end + + sparse_hess_coord!(objective, b, x, 1.0, v, vals) +end diff --git a/test/runtests.jl b/test/runtests.jl index a69dc634..b98044b0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -29,6 +29,7 @@ end @testset "Basic Hessian derivative test" begin include("sparse_hessian.jl") + include("sparse_hessian_nls.jl") end for problem in NLPModelsTest.nlp_problems ∪ ["GENROSE"] diff --git a/test/sparse_hessian_nls.jl b/test/sparse_hessian_nls.jl new file mode 100644 index 00000000..69f791d3 --- /dev/null +++ b/test/sparse_hessian_nls.jl @@ -0,0 +1,45 @@ +list_sparse_hess_backend = ( + (ADNLPModels.SparseADHessian, Dict()), + (ADNLPModels.ForwardDiffADHessian, Dict()), +) + +dt = (Float32, Float64) + +@testset "Basic Hessian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, + (backend, kw) in list_sparse_hess_backend + + F!(Fx, x) = begin + Fx[1] = x[1] - 1 + Fx[2] = 10 * (x[2] - x[1]^2) + Fx[3] = x[2] + 1 + Fx + end + x0 = T[-1.2; 1.0] + nvar = 2 + nequ = 3 + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, hessian_residual_backend = backend; kw...) + + x = rand(T, nvar) + v = rand(T, nequ) + rows, cols = zeros(Int, nls.nls_meta.nnzh), zeros(Int, nls.nls_meta.nnzh) + vals = zeros(T, nls.nls_meta.nnzh) + hess_structure_residual!(nls, rows, cols) + hess_coord_residual!(nls, x, v, vals) + @test eltype(vals) == T + H = sparse(rows, cols, vals, nvar, nvar) + # @test H == [] + + # Test also the implementation of the backends + b = nls.adbackend.hessian_residual_backend + obj_weight = 0.5 + @test nls.nls_meta.nnzh == ADNLPModels.get_nln_nnzh(b, nvar) + ADNLPModels.hess_structure_residual!(b, nls, rows, cols) + ADNLPModels.hess_coord_residual!(b, nls, x, y, obj_weight, vals) + @test eltype(vals) == T + H = sparse(rows, cols, vals, nvar, nvar) + # @test H == [] + + nls = ADNLPModels.ADNLSModel!(F!, x0, 3, matrix_free = true; kw...) + @test nls.adbackend.hessian_backend isa ADNLPModels.EmptyADbackend + @test nls.adbackend.hessian_residual_backend isa ADNLPModels.EmptyADbackend +end diff --git a/test/sparse_jacobian.jl b/test/sparse_jacobian.jl index 20b4d05e..0a522eb1 100644 --- a/test/sparse_jacobian.jl +++ b/test/sparse_jacobian.jl @@ -1,8 +1,10 @@ list_sparse_jac_backend = ( - (ADNLPModels.SparseADJacobian, Dict()), # default + (ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()), ) + dt = (Float32, Float64) + @testset "Basic Jacobian derivative with backend=$(backend) and T=$(T)" for T in dt, (backend, kw) in list_sparse_jac_backend diff --git a/test/sparse_jacobian_nls.jl b/test/sparse_jacobian_nls.jl index 23df0dc3..0144c50f 100644 --- a/test/sparse_jacobian_nls.jl +++ b/test/sparse_jacobian_nls.jl @@ -1,8 +1,10 @@ list_sparse_jac_backend = ( - (ADNLPModels.SparseADJacobian, Dict()), # default + (ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()), ) + dt = (Float32, Float64) + @testset "Basic Jacobian of residual derivative with backend=$(backend) and T=$(T)" for T in dt, (backend, kw) in list_sparse_jac_backend