Skip to content

Commit

Permalink
add proper testing for KKT systems (#278)
Browse files Browse the repository at this point in the history
  • Loading branch information
@frapac
frapac authored Nov 22, 2023
1 parent 6d694cd commit 6c4cecb
Show file tree
Hide file tree
Showing 10 changed files with 301 additions and 213 deletions.
104 changes: 80 additions & 24 deletions lib/MadNLPTests/src/MadNLPTests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,37 +21,93 @@ function solcmp(x,sol;atol=1e-4,rtol=1e-4)
return (aerr < atol || rerr < rtol)
end

function test_linear_solver(solver,T; kwargs...)

function test_linear_solver(solver, T; kwargs...)
m = 2
n = 2
row = Int32[1,2,2]
col = Int32[1,1,2]
val = T[1.,.1,2.]
b = T[1.0,3.0]
x = similar(b)
sol= [0.8542713567839195, 1.4572864321608041]

@testset "Linear solver $solver" begin

csc = sparse(row,col,val,m,n)
sol= [0.8542713567839195, 1.4572864321608041]
if MadNLP.input_type(solver) == :csc
opt = MadNLP.default_options(solver)
M = solver(csc; opt=opt)
elseif MadNLP.input_type(solver) == :dense
dense = Array(csc)
opt = MadNLP.default_options(solver)
M = solver(dense; opt=opt)
end
MadNLP.introduce(M)
MadNLP.improve!(M)
MadNLP.factorize!(M)
if MadNLP.is_inertia(M)
@test MadNLP.inertia(M) == (2, 0, 0)
end
x = MadNLP.solve!(M,copy(b))
@test solcmp(x,sol)
csc = sparse(row,col,val,m,n)
if MadNLP.input_type(solver) == :csc
opt = MadNLP.default_options(solver)
M = solver(csc; opt=opt)
elseif MadNLP.input_type(solver) == :dense
dense = Array(csc)
opt = MadNLP.default_options(solver)
M = solver(dense; opt=opt)
end
MadNLP.introduce(M)
MadNLP.improve!(M)
MadNLP.factorize!(M)
if MadNLP.is_inertia(M)
@test MadNLP.inertia(M) == (2, 0, 0)
end
x = MadNLP.solve!(M,copy(b))
@test solcmp(x,sol)
end

function test_kkt_system(kkt, cb)
# Getters
n = MadNLP.num_variables(kkt)
(m, p) = size(kkt)
# system should be square
@test m == p
@test isa(MadNLP.is_reduced(kkt), Bool)

# Interface
MadNLP.initialize!(kkt)

# Update internal structure
x0 = NLPModels.get_x0(cb.nlp)
y0 = NLPModels.get_y0(cb.nlp)
# Update Jacobian manually
jac = MadNLP.get_jacobian(kkt)
MadNLP._eval_jac_wrapper!(cb, x0, jac)
MadNLP.compress_jacobian!(kkt)
# Update Hessian manually
hess = MadNLP.get_hessian(kkt)
MadNLP._eval_lag_hess_wrapper!(cb, x0, y0, hess)
MadNLP.compress_hessian!(kkt)

# N.B.: set non-trivial dual's bounds to ensure
# l_lower and u_lower are positive. If not we run into
# an issue inside SparseUnreducedKKTSystem, which symmetrize
# the system using the values in l_lower and u_lower.
fill!(kkt.l_lower, 1e-3)
fill!(kkt.u_lower, 1e-3)

# Update diagonal terms manually.
MadNLP._set_aug_diagonal!(kkt)

# Factorization
MadNLP.build_kkt!(kkt)
MadNLP.factorize!(kkt.linear_solver)

# Backsolve
x = MadNLP.UnreducedKKTVector(kkt)
fill!(MadNLP.full(x), 1.0) # fill RHS with 1
out1 = MadNLP.solve!(kkt, x)
@test out1 === x

y = copy(x)
fill!(MadNLP.full(y), 0.0)
out2 = mul!(y, kkt, x)
@test out2 === y
@test MadNLP.full(y) ≈ ones(length(x))

if MadNLP.is_inertia(kkt.linear_solver)
ni, mi, pi = MadNLP.inertia(kkt.linear_solver)
@test MadNLP.is_inertia_correct(kkt, ni, mi, pi)
end

prim_reg, dual_reg = 1.0, 1.0
MadNLP.regularize_diagonal!(kkt, prim_reg, dual_reg)

return
end

function test_madnlp(name,optimizer_constructor::Function,exclude; Arr = Array)
Expand All @@ -75,7 +131,7 @@ function infeasible(optimizer_constructor::Function; Arr = Array)
)
optimizer = optimizer_constructor()
result = MadNLP.madnlp(nlp; optimizer.options...)

@test result.status == MadNLP.INFEASIBLE_PROBLEM_DETECTED
end
end
Expand All @@ -92,7 +148,7 @@ function unbounded(optimizer_constructor::Function; Arr = Array)
)
optimizer = optimizer_constructor()
result = MadNLP.madnlp(nlp; optimizer.options...)

@test result.status == MadNLP.DIVERGING_ITERATES
end
end
Expand Down
17 changes: 12 additions & 5 deletions src/IPM/factorization.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,14 @@ function solve!(kkt::SparseUnreducedKKTSystem, w::AbstractKKTVector)
solve!(kkt.linear_solver, full(w))
wzl .*= .-kkt.l_lower_aug
wzu .*= kkt.u_lower_aug
return w
end

function solve!(kkt::AbstractReducedKKTSystem, w::AbstractKKTVector)
reduce_rhs!(w.xp_lr, dual_lb(w), kkt.l_diag, w.xp_ur, dual_ub(w), kkt.u_diag)
solve!(kkt.linear_solver, primal_dual(w))
finish_aug_solve!(kkt, w)
return w
end

function solve!(
Expand Down Expand Up @@ -86,6 +88,7 @@ function solve!(
end

finish_aug_solve!(kkt, w)
return w
end


Expand All @@ -112,7 +115,7 @@ function solve!(kkt::SparseCondensedKKTSystem{T}, w::AbstractKKTVector) where T
ws .= (ws .+ wz) ./ Σs

finish_aug_solve!(kkt, w)

return w
end

function solve!(
Expand Down Expand Up @@ -153,13 +156,15 @@ function solve!(
ws .= (ws .+ wz) ./ Σs

finish_aug_solve!(kkt, w)
return w
end

function mul!(w::AbstractKKTVector{T}, kkt::Union{SparseKKTSystem{T,VT,MT,QN},SparseUnreducedKKTSystem{T,VT,MT,QN}}, x::AbstractKKTVector, alpha = one(T), beta = zero(T)) where {T, VT, MT, QN<:ExactHessian}
mul!(primal(w), Symmetric(kkt.hess_com, :L), primal(x), alpha, beta)
mul!(primal(w), kkt.jac_com', dual(x), alpha, one(T))
mul!(dual(w), kkt.jac_com, primal(x), alpha, beta)
_kktmul!(w,x,kkt.reg,kkt.du_diag,kkt.l_lower,kkt.u_lower,kkt.l_diag,kkt.u_diag, alpha, beta)
return w
end

function mul!(w::AbstractKKTVector{T}, kkt::Union{SparseKKTSystem{T,VT,MT,QN},SparseUnreducedKKTSystem{T,VT,MT,QN}}, x::AbstractKKTVector, alpha = one(T), beta = zero(T)) where {T, VT, MT, QN<:CompactLBFGS}
Expand Down Expand Up @@ -199,12 +204,13 @@ function mul!(w::AbstractKKTVector{T}, kkt::SparseCondensedKKTSystem, x::Abstrac

mul!(wx, Symmetric(kkt.hess_com, :L), xx, alpha, beta) # TODO: make this symmetric

mul!(wx, kkt.jt_csc, xz, alpha, beta)
mul!(wz, kkt.jt_csc', xx, alpha, one(T))
axpy!(-alpha, xz, ws)
mul!(wx, kkt.jt_csc, xz, alpha, one(T))
mul!(wz, kkt.jt_csc', xx, alpha, beta)
axpy!(-alpha, xs, wz)
ws .= beta.*ws .- alpha.* xz

_kktmul!(w,x,kkt.reg,kkt.du_diag,kkt.l_lower,kkt.u_lower,kkt.l_diag,kkt.u_diag, alpha, beta)
return w
end

function mul!(w::AbstractKKTVector{T}, kkt::AbstractDenseKKTSystem, x::AbstractKKTVector, alpha = one(T), beta = zero(T)) where T
Expand All @@ -225,8 +231,9 @@ function mul!(w::AbstractKKTVector{T}, kkt::AbstractDenseKKTSystem, x::AbstractK
mul!(wy, kkt.jac, xx, alpha, beta)
end
ws .= beta.*ws .- alpha.* xz
wz .= beta.*wz .- alpha.* xs
wz .-= alpha.* xs
_kktmul!(w,x,kkt.reg,kkt.du_diag,kkt.l_lower,kkt.u_lower,kkt.l_diag,kkt.u_diag, alpha, beta)
return w
end

function mul_hess_blk!(wx, kkt::Union{DenseKKTSystem,DenseCondensedKKTSystem}, t)
Expand Down
14 changes: 3 additions & 11 deletions src/KKT/KKTsystem.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ abstract type AbstractCondensedKKTSystem{T, VT, MT, QN} <: AbstractKKTSystem{T,
Templates
=#

"Number of primal variables associated to the KKT system."
"Number of primal variables (including slacks) associated to the KKT system."
function num_variables end

"""
Expand Down Expand Up @@ -152,15 +152,6 @@ Called internally inside the interior-point routine.
"""
function regularize_diagonal! end

"""
set_jacobian_scaling!(kkt::AbstractKKTSystem, scaling::AbstractVector)
Set the scaling of the Jacobian with the vector `scaling` storing
the scaling for all the constraints in the problem.
"""
function set_jacobian_scaling! end

"""
is_inertia_correct(kkt::AbstractKKTSystem, n::Int, m::Int, p::Int)
Expand Down Expand Up @@ -193,6 +184,7 @@ function initialize!(kkt::AbstractKKTSystem)
fill!(kkt.pr_diag, 1.0)
fill!(kkt.du_diag, 0.0)
fill!(kkt.hess, 0.0)
return
end

function regularize_diagonal!(kkt::AbstractKKTSystem, primal, dual)
Expand All @@ -215,7 +207,7 @@ function is_inertia_correct(kkt::AbstractKKTSystem, num_pos, num_zero, num_neg)
return (num_zero == 0) && (num_pos == num_variables(kkt))
end

compress_hessian!(kkt::AbstractKKTSystem) = nothing
compress_hessian!(kkt::AbstractKKTSystem) = nothing


include("rhs.jl")
Expand Down
39 changes: 20 additions & 19 deletions src/KKT/dense.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ struct DenseKKTSystem{
LS,
VI <: AbstractVector{Int},
} <: AbstractReducedKKTSystem{T, VT, MT, QN}

hess::MT
jac::MT
quasi_newton::QN
Expand All @@ -41,27 +41,27 @@ end

function create_kkt_system(
::Type{DenseKKTSystem},
cb::AbstractCallback{T,VT}, opt,
cb::AbstractCallback{T,VT}, opt,
opt_linear_solver, cnt, ind_cons) where {T,VT}

ind_ineq = ind_cons.ind_ineq
ind_lb = ind_cons.ind_lb
ind_ub = ind_cons.ind_ub

n = cb.nvar
m = cb.ncon
ns = length(ind_ineq)
nlb = length(ind_cons.ind_lb)
nub = length(ind_cons.ind_ub)

hess = create_array(cb, n, n)
jac = create_array(cb, m, n)
aug_com = create_array(cb, n+ns+m, n+ns+m)
reg = create_array(cb, n+ns)
pr_diag = create_array(cb, n+ns)
du_diag = create_array(cb, m)
diag_hess = create_array(cb, n)

l_diag = fill!(VT(undef, nlb), one(T))
u_diag = fill!(VT(undef, nub), one(T))
l_lower = fill!(VT(undef, nlb), zero(T))
Expand All @@ -71,20 +71,21 @@ function create_kkt_system(
fill!(aug_com, zero(T))
fill!(hess, zero(T))
fill!(jac, zero(T))
fill!(reg, zero(T))
fill!(pr_diag, zero(T))
fill!(du_diag, zero(T))
fill!(diag_hess, zero(T))

quasi_newton = create_quasi_newton(opt.hessian_approximation, cb, n)
cnt.linear_solver_time += @elapsed linear_solver = opt.linear_solver(aug_com; opt = opt_linear_solver)

return DenseKKTSystem(
hess, jac, quasi_newton,
reg, pr_diag, du_diag, l_diag, u_diag, l_lower, u_lower,
diag_hess, aug_com,
ind_ineq, ind_cons.ind_lb, ind_cons.ind_ub,
linear_solver,
Dict{Symbol, Any}(),
Dict{Symbol, Any}(),
)
end

Expand All @@ -104,7 +105,7 @@ struct DenseCondensedKKTSystem{
LS,
VI <: AbstractVector{Int}
} <: AbstractCondensedKKTSystem{T, VT, MT, QN}

hess::MT
jac::MT
quasi_newton::QN
Expand Down Expand Up @@ -140,9 +141,9 @@ end

function create_kkt_system(
::Type{DenseCondensedKKTSystem},
cb::AbstractCallback{T,VT}, opt,
cb::AbstractCallback{T,VT}, opt,
opt_linear_solver, cnt, ind_cons) where {T,VT}

n = cb.nvar
m = cb.ncon
ns = length(ind_cons.ind_ineq)
Expand All @@ -162,11 +163,11 @@ function create_kkt_system(
u_diag = fill!(VT(undef, nub), one(T))
l_lower = fill!(VT(undef, nlb), zero(T))
u_lower = fill!(VT(undef, nub), zero(T))

pd_buffer = VT(undef, n + n_eq)
diag_buffer = VT(undef, ns)
buffer = VT(undef, m)

# Init!
fill!(aug_com, zero(T))
fill!(hess, zero(T))
Expand All @@ -180,12 +181,12 @@ function create_kkt_system(

quasi_newton = create_quasi_newton(opt.hessian_approximation, cb, n)
cnt.linear_solver_time += @elapsed linear_solver = opt.linear_solver(aug_com; opt = opt_linear_solver)

return DenseCondensedKKTSystem(
hess, jac, quasi_newton, jac_ineq,
reg, pr_diag, du_diag, l_diag, u_diag, l_lower, u_lower,
pd_buffer, diag_buffer, buffer,
aug_com,
aug_com,
n_eq, ind_cons.ind_eq, ind_eq_shifted,
ns,
ind_cons.ind_ineq, ind_cons.ind_lb, ind_cons.ind_ub,
Expand Down Expand Up @@ -285,10 +286,10 @@ function build_kkt!(kkt::DenseKKTSystem{T, VT, MT}) where {T, VT, MT}
n = size(kkt.hess, 1)
m = size(kkt.jac, 1)
ns = length(kkt.ind_ineq)

_build_dense_kkt_system!(kkt.aug_com, kkt.hess, kkt.jac,
kkt.pr_diag, kkt.du_diag, kkt.diag_hess,
kkt.ind_ineq,
kkt.ind_ineq,
n, m, ns)
end

Expand Down Expand Up @@ -337,12 +338,12 @@ end

function _build_ineq_jac!(
dest::AbstractMatrix, jac::AbstractMatrix, diag_buffer::AbstractVector,
ind_ineq::AbstractVector,
ind_ineq::AbstractVector,
n, m_ineq,
)
@inbounds for i in 1:m_ineq, j in 1:n
is = ind_ineq[i]
dest[i, j] = jac[is, j] * sqrt(diag_buffer[i])
dest[i, j] = jac[is, j] * sqrt(diag_buffer[i])
end
end

Expand Down
Loading

0 comments on commit 6c4cecb

Please sign in to comment.