diff --git a/src/KKT/KKTsystem.jl b/src/KKT/KKTsystem.jl new file mode 100644 index 00000000..6916f3f6 --- /dev/null +++ b/src/KKT/KKTsystem.jl @@ -0,0 +1,247 @@ + +abstract type AbstractKKTSystem{T, MT<:AbstractMatrix{T}} end + +""" + AbstractUnreducedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} + +Augmented KKT system associated to the linearization of the KKT +conditions at the current primal-dual iterate ``(x, s, y, z, ν, w)``. + +The associated matrix is +``` +[Wₓₓ 0 Aₑ' Aᵢ' -I 0 ] [Δx] +[ 0 0 0 -I 0 -I ] [Δs] +[Aₑ 0 0 0 0 0 ] [Δy] +[Aᵢ -I 0 0 0 0 ] [Δz] +[V 0 0 0 X 0 ] [Δν] +[0 W 0 0 0 S ] [Δw] +``` +with +* ``Wₓₓ``: Hessian of the Lagrangian. +* ``Aₑ``: Jacobian of the equality constraints +* ``Aᵢ``: Jacobian of the inequality constraints +* ``X = diag(x)`` +* ``S = diag(s)`` +* ``V = diag(ν)`` +* ``W = diag(w)`` +""" +abstract type AbstractUnreducedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} end + +""" + AbstractReducedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} + +The reduced KKT system is a simplification of the original Augmented KKT system. +Comparing to [`AbstractUnreducedKKTSystem`](@ref)), `AbstractReducedKKTSystem` removes +the two last rows associated to the bounds' duals ``(ν, w)``. + +At a primal-dual iterate ``(x, s, y, z)``, the matrix writes +``` +[Wₓₓ + Σₓ 0 Aₑ' Aᵢ'] [Δx] +[ 0 Σₛ 0 -I ] [Δs] +[Aₑ 0 0 0 ] [Δy] +[Aᵢ -I 0 0 ] [Δz] +``` +with +* ``Wₓₓ``: Hessian of the Lagrangian. +* ``Aₑ``: Jacobian of the equality constraints +* ``Aᵢ``: Jacobian of the inequality constraints +* ``Σₓ = X⁻¹ V`` +* ``Σₛ = S⁻¹ W`` + +""" +abstract type AbstractReducedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} end + +""" + AbstractCondensedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} + +The condensed KKT system simplifies further the [``AbstractReducedKKTSystem`](@ref) +by removing the rows associated to the slack variables ``s`` and the inequalities. + +At the primal-dual iterate ``(x, y)``, the matrix writes +``` +[Wₓₓ + Σₓ + Aᵢ' Σₛ Aᵢ Aₑ'] [Δx] +[ Aₑ 0 ] [Δy] +``` +with +* ``Wₓₓ``: Hessian of the Lagrangian. +* ``Aₑ``: Jacobian of the equality constraints +* ``Aᵢ``: Jacobian of the inequality constraints +* ``Σₓ = X⁻¹ V`` +* ``Σₛ = S⁻¹ W`` +""" +abstract type AbstractCondensedKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} end + +#= + Templates +=# + +"Number of primal variables associated to the KKT system." +function num_variables end + +""" + get_kkt(kkt::AbstractKKTSystem)::AbstractMatrix + +Return a pointer to the KKT matrix implemented in `kkt`. +The pointer is passed afterward to a linear solver. +""" +function get_kkt end + +"Get Jacobian matrix" +function get_jacobian end + +"Get Hessian matrix" +function get_hessian end + +""" + initialize!(kkt::AbstractKKTSystem) + +Initialize KKT system with default values. +Called when we initialize the `InteriorPointSolver` storing the current KKT system `kkt`. +""" +function initialize! end + +""" + build_kkt!(kkt::AbstractKKTSystem) + +Assemble the KKT matrix before calling the factorization routine. + +""" +function build_kkt! end + +""" + compress_hessian!(kkt::AbstractKKTSystem) + +Compress the Hessian inside `kkt`'s internals. +This function is called every time a new Hessian is evaluated. + +Default implementation do nothing. + +""" +function compress_hessian! end + +""" + compress_jacobian!(kkt::AbstractKKTSystem) + +Compress the Jacobian inside `kkt`'s internals. +This function is called every time a new Jacobian is evaluated. + +By default, the function updates in the Jacobian the coefficients +associated to the slack variables. + +""" +function compress_jacobian! end + +""" + jtprod!(y::AbstractVector, kkt::AbstractKKTSystem, x::AbstractVector) + +Multiply with transpose of Jacobian and store the result +in `y`, such that ``y = A' x`` (with ``A`` current Jacobian). +""" +function jtprod! end + +""" + regularize_diagonal!(kkt::AbstractKKTSystem, primal_values::AbstractVector, dual_values::AbstractVector) + +Regularize the values in the diagonal of the KKT system. +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) + +Check if the inertia ``(n, m, p)`` returned by the linear solver is adapted +to the KKT system implemented in `kkt`. + +""" +function is_inertia_correct end + +# TODO: temporary +"Return true if KKT system is reduced." +function is_reduced end + +"Nonzero in Jacobian" +function nnz_jacobian end + +# TODO: we need these two templates as NLPModels does not implement +# a template for dense Jacobian and dense Hessian +"Dense Jacobian callback" +function jac_dense! end + +"Dense Hessian callback" +function hess_dense! end + +#= + Generic functions +=# +function initialize!(kkt::AbstractKKTSystem) + fill!(kkt.pr_diag, 1.0) + fill!(kkt.du_diag, 0.0) + fill!(kkt.hess, 0.0) +end + +function regularize_diagonal!(kkt::AbstractKKTSystem, primal, dual) + kkt.pr_diag .+= primal + kkt.du_diag .= .-dual +end + +Base.size(kkt::AbstractKKTSystem) = size(kkt.aug_com) +Base.size(kkt::AbstractKKTSystem, dim::Int) = size(kkt.aug_com, dim) + +# Getters +get_kkt(kkt::AbstractKKTSystem) = kkt.aug_com +get_jacobian(kkt::AbstractKKTSystem) = kkt.jac +get_hessian(kkt::AbstractKKTSystem) = kkt.hess +get_raw_jacobian(kkt::AbstractKKTSystem) = kkt.jac_raw + + +# Fix variable treatment +function treat_fixed_variable!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:SparseMatrixCSC{T, Int32}} + length(kkt.ind_fixed) == 0 && return + aug = kkt.aug_com + + fixed_aug_diag = view(aug.nzval, aug.colptr[kkt.ind_fixed]) + fixed_aug_diag .= 1.0 + fixed_aug = view(aug.nzval, kkt.ind_aug_fixed) + fixed_aug .= 0.0 + return +end +function treat_fixed_variable!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:Matrix{T}} + length(kkt.ind_fixed) == 0 && return + aug = kkt.aug_com + @inbounds for i in kkt.ind_fixed + aug[i, :] .= 0.0 + aug[:, i] .= 0.0 + aug[i, i] = 1.0 + end +end + +function is_inertia_correct(kkt::AbstractKKTSystem, num_pos, num_zero, num_neg) + return (num_zero == 0) && (num_pos == num_variables(kkt)) +end + +function build_kkt!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:Matrix{T}} + copyto!(kkt.aug_com, kkt.aug_raw) + treat_fixed_variable!(kkt) +end + +function build_kkt!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:SparseMatrixCSC{T, Int32}} + transfer!(kkt.aug_com, kkt.aug_raw, kkt.aug_csc_map) + treat_fixed_variable!(kkt) +end + +compress_hessian!(kkt::AbstractKKTSystem) = nothing + + +include("sparse.jl") +include("dense.jl") + diff --git a/src/KKT/dense.jl b/src/KKT/dense.jl new file mode 100644 index 00000000..2fb83657 --- /dev/null +++ b/src/KKT/dense.jl @@ -0,0 +1,155 @@ + +#= + DenseKKTSystem +=# + +""" + DenseKKTSystem{T, VT, MT} <: AbstractReducedKKTSystem{T, MT} + +Implement [`AbstractReducedKKTSystem`](@ref) with dense matrices. + +Requires a dense linear solver to be factorized (otherwise an error is returned). + +""" +struct DenseKKTSystem{T, VT, MT} <: AbstractReducedKKTSystem{T, MT} + hess::MT + jac::MT + pr_diag::VT + du_diag::VT + diag_hess::VT + # KKT system + aug_com::MT + # Info + ind_ineq::Vector{Int} + ind_fixed::Vector{Int} + jacobian_scaling::VT + # Buffers + etc::Dict{Symbol, Any} +end + +function DenseKKTSystem{T, VT, MT}(n, m, ind_ineq, ind_fixed) where {T, VT, MT} + ns = length(ind_ineq) + hess = MT(undef, n, n) + jac = MT(undef, m, n+ns) + pr_diag = VT(undef, n+ns) + du_diag = VT(undef, m) + diag_hess = VT(undef, n) + + # If the the problem is unconstrained, then KKT system is directly equal + # to the Hessian (+ some regularization terms) + aug_com = if (m == 0) + hess + else + MT(undef, n+ns+m, n+ns+m) + end + + jacobian_scaling = VT(undef, m) + + # Init! + fill!(aug_com, zero(T)) + fill!(hess, zero(T)) + fill!(jac, zero(T)) + fill!(pr_diag, zero(T)) + fill!(du_diag, zero(T)) + fill!(diag_hess, zero(T)) + fill!(jacobian_scaling, one(T)) + + return DenseKKTSystem{T, VT, MT}( + hess, jac, pr_diag, du_diag, diag_hess, aug_com, + ind_ineq, ind_fixed, jacobian_scaling, Dict{Symbol, Any}(), + ) +end + +function DenseKKTSystem{T, VT, MT}(nlp::AbstractNLPModel, info_constraints=get_index_constraints(nlp)) where {T, VT, MT} + return DenseKKTSystem{T, VT, MT}( + get_nvar(nlp), get_ncon(nlp), info_constraints.ind_ineq, info_constraints.ind_fixed + ) +end + +is_reduced(::DenseKKTSystem) = true +num_variables(kkt::DenseKKTSystem) = length(kkt.pr_diag) + +# Special getters for Jacobian +function get_jacobian(kkt::DenseKKTSystem) + n = size(kkt.hess, 1) + ns = length(kkt.ind_ineq) + return view(kkt.jac, :, 1:n) +end + +get_raw_jacobian(kkt::DenseKKTSystem) = kkt.jac + +nnz_jacobian(kkt::DenseKKTSystem) = length(kkt.jac) + +function diag_add!(dest::AbstractMatrix, d1::AbstractVector, d2::AbstractVector) + n = length(d1) + @inbounds for i in 1:n + dest[i, i] = d1[i] + d2[i] + end +end + +function _build_dense_kkt_system!(dest, hess, jac, pr_diag, du_diag, diag_hess, n, m, ns) + # Transfer Hessian + for i in 1:n, j in 1:i + if i == j + dest[i, i] = pr_diag[i] + diag_hess[i] + else + dest[i, j] = hess[i, j] + dest[j, i] = hess[j, i] + end + end + # Transfer slack diagonal + for i in 1:ns + dest[i+n, i+n] = pr_diag[i+n] + end + # Transfer Jacobian + for i in 1:m, j in 1:(n+ns) + dest[i + n + ns, j] = jac[i, j] + dest[j, i + n + ns] = jac[i, j] + end + # Transfer dual regularization + for i in 1:m + dest[i + n + ns, i + n + ns] = du_diag[i] + end +end + +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) + if m == 0 # If problem is unconstrained, just need to update the diagonal + diag_add!(kkt.aug_com, kkt.diag_hess, kkt.pr_diag) + else # otherwise, we update the full matrix + _build_dense_kkt_system!(kkt.aug_com, kkt.hess, kkt.jac, kkt.pr_diag, kkt.du_diag, kkt.diag_hess, n, m, ns) + end + treat_fixed_variable!(kkt) +end + +function compress_jacobian!(kkt::DenseKKTSystem{T, VT, MT}) where {T, VT, MT} + m = size(kkt.jac, 1) + n = size(kkt.hess, 1) + # Add slack indexes + for i in kkt.ind_ineq + kkt.jac[i, i+n] = -one(T) + end + # Scale + kkt.jac .*= kkt.jacobian_scaling + return +end + +function compress_hessian!(kkt::DenseKKTSystem) + # Transfer diagonal term for future regularization + diag!(kkt.diag_hess, kkt.hess) +end + +function mul!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector) + mul!(y, kkt.aug_com, x) +end + +function jtprod!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector) + mul!(y, kkt.jac', x) +end + +function set_jacobian_scaling!(kkt::DenseKKTSystem, constraint_scaling::AbstractVector) + copyto!(kkt.jacobian_scaling, constraint_scaling) +end + diff --git a/src/kktsystem.jl b/src/KKT/sparse.jl similarity index 54% rename from src/kktsystem.jl rename to src/KKT/sparse.jl index e73388d7..a62a73f2 100644 --- a/src/kktsystem.jl +++ b/src/KKT/sparse.jl @@ -1,134 +1,76 @@ -abstract type AbstractKKTSystem{T, MT<:AbstractMatrix{T}} end +""" + SparseKKTSystem{T, MT} <: AbstractReducedKKTSystem{T, MT} -"Sparse KKT system" -abstract type AbstractSparseKKTSystem{T, MT} <: AbstractKKTSystem{T, MT} end +Implement the [`AbstractReducedKKTSystem`](@ref) in sparse COO format. -#= - Templates -=# -"Initialize KKT system with default values." -function initialize! end - -"Assemble KKT matrix." -function build_kkt! end - -"Compress Hessian matrix." -function compress_hessian! end - -"Compress Jacobian matrix." -function compress_jacobian! end - -"Get KKT system" -function get_kkt end - -"Get Jacobian matrix" -function get_jacobian end - -"Get Hessian matrix" -function get_hessian end - -"Multiply with Jacobian" -function jtprod! end - -"Regularize values in the diagonal of the KKT system." -function regularize_diagonal! end - -"Set scaling of Jacobian" -function set_jacobian_scaling! end - -"Return true if KKT system is reduced." -function is_reduced end - -"Nonzero in Jacobian" -function nnz_jacobian end - -"Nonzero in Hessian" -function nnz_kkt end - -"Dense Jacobian callback" -function jac_dense! end - -"Dense Hessian callback" -function hess_dense! end - -"Number of variables associated to the KKT system." -function num_variables end - -"Check if inertia is correct." -function is_inertia_correct end - -#= - Generic functions -=# -function initialize!(kkt::AbstractKKTSystem) - fill!(kkt.pr_diag, 1.0) - fill!(kkt.du_diag, 0.0) - fill!(kkt.hess, 0.0) +""" +struct SparseKKTSystem{T, MT} <: AbstractReducedKKTSystem{T, MT} + hess::StrideOneVector{T} + jac::StrideOneVector{T} + pr_diag::StrideOneVector{T} + du_diag::StrideOneVector{T} + # Augmented system + aug_raw::SparseMatrixCOO{T,Int32} + aug_com::MT + aug_csc_map::Union{Nothing, Vector{Int}} + # Jacobian + jac_raw::SparseMatrixCOO{T,Int32} + jac_com::MT + jac_csc_map::Union{Nothing, Vector{Int}} + # Info + ind_ineq::Vector{Int} + ind_fixed::Vector{Int} + ind_aug_fixed::Vector{Int} + jacobian_scaling::Vector{T} end -function regularize_diagonal!(kkt::AbstractKKTSystem, primal, dual) - kkt.pr_diag .+= primal - kkt.du_diag .= .-dual -end +""" + SparseUnreducedKKTSystem{T, MT} <: AbstractUnreducedKKTSystem{T, MT} -Base.size(kkt::AbstractKKTSystem) = size(kkt.aug_com) -Base.size(kkt::AbstractKKTSystem, dim::Int) = size(kkt.aug_com, dim) +Implement the [`AbstractUnreducedKKTSystem`](@ref) in sparse COO format. -# Getters -get_kkt(kkt::AbstractKKTSystem) = kkt.aug_com -get_jacobian(kkt::AbstractKKTSystem) = kkt.jac -get_hessian(kkt::AbstractKKTSystem) = kkt.hess -get_raw_jacobian(kkt::AbstractKKTSystem) = kkt.jac_raw +""" +struct SparseUnreducedKKTSystem{T, MT} <: AbstractUnreducedKKTSystem{T, MT} + hess::StrideOneVector{T} + jac::StrideOneVector{T} + pr_diag::StrideOneVector{T} + du_diag::StrideOneVector{T} + l_diag::StrideOneVector{T} + u_diag::StrideOneVector{T} + l_lower::StrideOneVector{T} + u_lower::StrideOneVector{T} -# Fix variable treatment -function treat_fixed_variable!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:SparseMatrixCSC{T, Int32}} - length(kkt.ind_fixed) == 0 && return - aug = kkt.aug_com + aug_raw::SparseMatrixCOO{T,Int32} + aug_com::MT + aug_csc_map::Union{Nothing, Vector{Int}} - fixed_aug_diag = view(aug.nzval, aug.colptr[kkt.ind_fixed]) - fixed_aug_diag .= 1.0 - fixed_aug = view(aug.nzval, kkt.ind_aug_fixed) - fixed_aug .= 0.0 - return -end -function treat_fixed_variable!(kkt::AbstractKKTSystem{T, MT}) where {T, MT<:Matrix{T}} - length(kkt.ind_fixed) == 0 && return - aug = kkt.aug_com - @inbounds for i in kkt.ind_fixed - aug[i, :] .= 0.0 - aug[:, i] .= 0.0 - aug[i, i] = 1.0 - end + jac_raw::SparseMatrixCOO{T,Int32} + jac_com::MT + jac_csc_map::Union{Nothing, Vector{Int}} + ind_ineq::Vector{Int} + ind_fixed::Vector{Int} + ind_aug_fixed::Vector{Int} + jacobian_scaling::Vector{T} end -function is_inertia_correct(kkt::AbstractKKTSystem, num_pos, num_zero, num_neg) - return (num_zero == 0) && (num_pos == num_variables(kkt)) -end +# Template to dispatch on sparse representation +const AbstractSparseKKTSystem{T, MT} = Union{SparseKKTSystem{T, MT}, SparseUnreducedKKTSystem{T, MT}} #= - SparseKKTSystem + Generic sparse methods =# -# Generic function for sparse KKT systems -function build_kkt!(kkt::AbstractSparseKKTSystem{T, MT}) where {T, MT<:Matrix{T}} - copyto!(kkt.aug_com, kkt.aug_raw) - treat_fixed_variable!(kkt) +function mul!(y::AbstractVector, kkt::AbstractSparseKKTSystem, x::AbstractVector) + mul!(y, Symmetric(kkt.aug_com, :L), x) end -function build_kkt!(kkt::AbstractSparseKKTSystem{T, MT}) where {T, MT<:SparseMatrixCSC{T, Int32}} - transfer!(kkt.aug_com, kkt.aug_raw, kkt.aug_csc_map) - treat_fixed_variable!(kkt) +function jtprod!(y::AbstractVector, kkt::AbstractSparseKKTSystem, x::AbstractVector) + mul!(y, kkt.jac_com', x) end -function set_jacobian_scaling!(kkt::AbstractSparseKKTSystem{T, MT}, constraint_scaling::AbstractVector) where {T, MT} - nnzJ = length(kkt.jac)::Int - @inbounds for i in 1:nnzJ - index = kkt.jac_raw.I[i] - kkt.jacobian_scaling[i] = constraint_scaling[index] - end -end +nnz_jacobian(kkt::AbstractSparseKKTSystem) = nnz(kkt.jac_raw) function compress_jacobian!(kkt::AbstractSparseKKTSystem{T, MT}) where {T, MT<:SparseMatrixCSC{T, Int32}} ns = length(kkt.ind_ineq) @@ -144,45 +86,19 @@ function compress_jacobian!(kkt::AbstractSparseKKTSystem{T, MT}) where {T, MT<:M copyto!(kkt.jac_com, kkt.jac_raw) end -function compress_hessian!(kkt::AbstractSparseKKTSystem) - nothing -end - -function mul!(y::AbstractVector, kkt::AbstractSparseKKTSystem, x::AbstractVector) - mul!(y, Symmetric(kkt.aug_com, :L), x) -end - -function jtprod!(y::AbstractVector, kkt::AbstractSparseKKTSystem, x::AbstractVector) - mul!(y, kkt.jac_com', x) +function set_jacobian_scaling!(kkt::AbstractSparseKKTSystem{T, MT}, constraint_scaling::AbstractVector) where {T, MT} + nnzJ = length(kkt.jac)::Int + @inbounds for i in 1:nnzJ + index = kkt.jac_raw.I[i] + kkt.jacobian_scaling[i] = constraint_scaling[index] + end end -nnz_jacobian(kkt::AbstractSparseKKTSystem) = nnz(kkt.jac_raw) -nnz_kkt(kkt::AbstractSparseKKTSystem) = nnz(kkt.aug_raw) #= SparseKKTSystem =# -struct SparseKKTSystem{T, MT} <: AbstractSparseKKTSystem{T, MT} - hess::StrideOneVector{T} - jac::StrideOneVector{T} - pr_diag::StrideOneVector{T} - du_diag::StrideOneVector{T} - # Augmented system - aug_raw::SparseMatrixCOO{T,Int32} - aug_com::MT - aug_csc_map::Union{Nothing, Vector{Int}} - # Jacobian - jac_raw::SparseMatrixCOO{T,Int32} - jac_com::MT - jac_csc_map::Union{Nothing, Vector{Int}} - # Info - ind_ineq::Vector{Int} - ind_fixed::Vector{Int} - ind_aug_fixed::Vector{Int} - jacobian_scaling::Vector{T} -end - function SparseKKTSystem{T, MT}( n::Int, m::Int, ind_ineq::Vector{Int}, ind_fixed::Vector{Int}, hess_sparsity_I, hess_sparsity_J, @@ -270,34 +186,11 @@ end is_reduced(::SparseKKTSystem) = true num_variables(kkt::SparseKKTSystem) = length(kkt.pr_diag) + #= SparseUnreducedKKTSystem =# -struct SparseUnreducedKKTSystem{T, MT} <: AbstractSparseKKTSystem{T, MT} - hess::StrideOneVector{T} - jac::StrideOneVector{T} - pr_diag::StrideOneVector{T} - du_diag::StrideOneVector{T} - - l_diag::StrideOneVector{T} - u_diag::StrideOneVector{T} - l_lower::StrideOneVector{T} - u_lower::StrideOneVector{T} - - aug_raw::SparseMatrixCOO{T,Int32} - aug_com::MT - aug_csc_map::Union{Nothing, Vector{Int}} - - jac_raw::SparseMatrixCOO{T,Int32} - jac_com::MT - jac_csc_map::Union{Nothing, Vector{Int}} - ind_ineq::Vector{Int} - ind_fixed::Vector{Int} - ind_aug_fixed::Vector{Int} - jacobian_scaling::Vector{T} -end - function SparseUnreducedKKTSystem{T, MT}( n::Int, m::Int, nlb::Int, nub::Int, ind_ineq, ind_fixed, hess_sparsity_I, hess_sparsity_J, @@ -412,149 +305,3 @@ end is_reduced(::SparseUnreducedKKTSystem) = false num_variables(kkt::SparseUnreducedKKTSystem) = length(kkt.pr_diag) -#= - DenseKKTSystem -=# - -struct DenseKKTSystem{T, VT, MT} <: AbstractKKTSystem{T, MT} - hess::MT - jac::MT - pr_diag::VT - du_diag::VT - diag_hess::VT - # KKT system - aug_com::MT - # Info - ind_ineq::Vector{Int} - ind_fixed::Vector{Int} - jacobian_scaling::VT - # Buffers - etc::Dict{Symbol, Any} -end - -function DenseKKTSystem{T, VT, MT}(n, m, ind_ineq, ind_fixed) where {T, VT, MT} - ns = length(ind_ineq) - hess = MT(undef, n, n) - jac = MT(undef, m, n+ns) - pr_diag = VT(undef, n+ns) - du_diag = VT(undef, m) - diag_hess = VT(undef, n) - - # If the the problem is unconstrained, then KKT system is directly equal - # to the Hessian (+ some regularization terms) - aug_com = if (m == 0) - hess - else - MT(undef, n+ns+m, n+ns+m) - end - - jacobian_scaling = VT(undef, m) - - # Init! - fill!(aug_com, zero(T)) - fill!(hess, zero(T)) - fill!(jac, zero(T)) - fill!(pr_diag, zero(T)) - fill!(du_diag, zero(T)) - fill!(diag_hess, zero(T)) - fill!(jacobian_scaling, one(T)) - - return DenseKKTSystem{T, VT, MT}( - hess, jac, pr_diag, du_diag, diag_hess, aug_com, - ind_ineq, ind_fixed, jacobian_scaling, Dict{Symbol, Any}(), - ) -end - -function DenseKKTSystem{T, VT, MT}(nlp::AbstractNLPModel, info_constraints=get_index_constraints(nlp)) where {T, VT, MT} - return DenseKKTSystem{T, VT, MT}( - get_nvar(nlp), get_ncon(nlp), info_constraints.ind_ineq, info_constraints.ind_fixed - ) -end - -is_reduced(::DenseKKTSystem) = true -num_variables(kkt::DenseKKTSystem) = length(kkt.pr_diag) - -# Special getters for Jacobian -function get_jacobian(kkt::DenseKKTSystem) - n = size(kkt.hess, 1) - ns = length(kkt.ind_ineq) - return view(kkt.jac, :, 1:n) -end -get_raw_jacobian(kkt::DenseKKTSystem) = kkt.jac - -nnz_jacobian(kkt::DenseKKTSystem) = length(kkt.jac) -nnz_kkt(kkt::DenseKKTSystem) = length(kkt.aug_com) - -function diag_add!(dest::AbstractMatrix, d1::AbstractVector, d2::AbstractVector) - n = length(d1) - @inbounds for i in 1:n - dest[i, i] = d1[i] + d2[i] - end -end - -function _build_dense_kkt_system!(dest, hess, jac, pr_diag, du_diag, diag_hess, n, m, ns) - # Transfer Hessian - for i in 1:n, j in 1:i - if i == j - dest[i, i] = pr_diag[i] + diag_hess[i] - else - dest[i, j] = hess[i, j] - dest[j, i] = hess[j, i] - end - end - # Transfer slack diagonal - for i in 1:ns - dest[i+n, i+n] = pr_diag[i+n] - end - # Transfer Jacobian - for i in 1:m, j in 1:(n+ns) - dest[i + n + ns, j] = jac[i, j] - dest[j, i + n + ns] = jac[i, j] - end - # Transfer dual regularization - for i in 1:m - dest[i + n + ns, i + n + ns] = du_diag[i] - end -end - -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) - if m == 0 # If problem is unconstrained, just need to update the diagonal - diag_add!(kkt.aug_com, kkt.diag_hess, kkt.pr_diag) - else # otherwise, we update the full matrix - _build_dense_kkt_system!(kkt.aug_com, kkt.hess, kkt.jac, kkt.pr_diag, kkt.du_diag, kkt.diag_hess, n, m, ns) - end - treat_fixed_variable!(kkt) -end - -function compress_jacobian!(kkt::DenseKKTSystem{T, VT, MT}) where {T, VT, MT} - m = size(kkt.jac, 1) - n = size(kkt.hess, 1) - # Add slack indexes - for i in kkt.ind_ineq - kkt.jac[i, i+n] = -one(T) - end - # Scale - kkt.jac .*= kkt.jacobian_scaling - return -end - -function compress_hessian!(kkt::DenseKKTSystem) - # Transfer diagonal term for future regularization - diag!(kkt.diag_hess, kkt.hess) -end - -function mul!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector) - mul!(y, kkt.aug_com, x) -end - -function jtprod!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector) - mul!(y, kkt.jac', x) -end - -function set_jacobian_scaling!(kkt::DenseKKTSystem, constraint_scaling::AbstractVector) - copyto!(kkt.jacobian_scaling, constraint_scaling) -end - diff --git a/src/MadNLP.jl b/src/MadNLP.jl index fd44bd84..eaf99ac1 100644 --- a/src/MadNLP.jl +++ b/src/MadNLP.jl @@ -32,7 +32,7 @@ include("utils.jl") include("options.jl") include("matrixtools.jl") include(joinpath("LinearSolvers","linearsolvers.jl")) -include("kktsystem.jl") +include(joinpath("KKT", "KKTsystem.jl")) include("interiorpointsolver.jl") include(joinpath("Interfaces","interfaces.jl"))