Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[KKT] add two folders Sparse and Dense for KKT formulations #347

Merged
merged 1 commit into from
Jun 13, 2024
Merged

[KKT] add two folders Sparse and Dense for KKT formulations #347

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 162 additions & 0 deletions src/KKT/Dense/augmented.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@

"""
DenseKKTSystem{T, VT, MT, QN, VI} <: AbstractReducedKKTSystem{T, VT, MT, QN}

Implement [`AbstractReducedKKTSystem`](@ref) with dense matrices.

Requires a dense linear solver to be factorized (otherwise an error is returned).

"""
struct DenseKKTSystem{
T,
VT <: AbstractVector{T},
MT <: AbstractMatrix{T},
QN,
LS,
VI <: AbstractVector{Int},
} <: AbstractReducedKKTSystem{T, VT, MT, QN}

hess::MT
jac::MT
quasi_newton::QN
reg::VT
pr_diag::VT
du_diag::VT
l_diag::VT
u_diag::VT
l_lower::VT
u_lower::VT
diag_hess::VT
# KKT system
aug_com::MT
# Info
ind_ineq::VI
ind_lb::VI
ind_ub::VI
# Linear Solver
linear_solver::LS
# Buffers
etc::Dict{Symbol, Any}
end

function create_kkt_system(
::Type{DenseKKTSystem},
cb::AbstractCallback{T,VT},
ind_cons,
linear_solver::Type;
opt_linear_solver=default_options(linear_solver),
hessian_approximation=ExactHessian,
) 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))
u_lower = fill!(VT(undef, nub), zero(T))

# Init!
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(hessian_approximation, cb, n)
_linear_solver = 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}(),
)
end

num_variables(kkt::DenseKKTSystem) = length(kkt.pr_diag)

function mul!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector)
symul!(y, kkt.aug_com, x)
end

# 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

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::VT, hess, jac, pr_diag, du_diag, diag_hess, ind_ineq, n, m, ns) where {T, VT <: AbstractMatrix{T}}
# 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 / variables
for i in 1:m, j in 1:n
dest[i + n + ns, j] = jac[i, j]
dest[j, i + n + ns] = jac[i, j]
end
# Transfer Jacobian / slacks
for j in 1:ns
is = ind_ineq[j]
dest[is + n + ns, j + n] = - one(T)
dest[j + n, is + n + ns] = - one(T)
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)

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

function compress_hessian!(kkt::DenseKKTSystem)
# Transfer diagonal term for future regularization
diag!(kkt.diag_hess, kkt.hess)
end

198 changes: 0 additions & 198 deletions src/KKT/dense.jl → src/KKT/Dense/condensed.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,98 +1,4 @@

"""
DenseKKTSystem{T, VT, MT, QN, VI} <: AbstractReducedKKTSystem{T, VT, MT, QN}

Implement [`AbstractReducedKKTSystem`](@ref) with dense matrices.

Requires a dense linear solver to be factorized (otherwise an error is returned).

"""
struct DenseKKTSystem{
T,
VT <: AbstractVector{T},
MT <: AbstractMatrix{T},
QN,
LS,
VI <: AbstractVector{Int},
} <: AbstractReducedKKTSystem{T, VT, MT, QN}

hess::MT
jac::MT
quasi_newton::QN
reg::VT
pr_diag::VT
du_diag::VT
l_diag::VT
u_diag::VT
l_lower::VT
u_lower::VT
diag_hess::VT
# KKT system
aug_com::MT
# Info
ind_ineq::VI
ind_lb::VI
ind_ub::VI
# Linear Solver
linear_solver::LS
# Buffers
etc::Dict{Symbol, Any}
end

function create_kkt_system(
::Type{DenseKKTSystem},
cb::AbstractCallback{T,VT},
ind_cons,
linear_solver::Type;
opt_linear_solver=default_options(linear_solver),
hessian_approximation=ExactHessian,
) 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))
u_lower = fill!(VT(undef, nub), zero(T))

# Init!
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(hessian_approximation, cb, n)
_linear_solver = 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}(),
)
end

"""
DenseCondensedKKTSystem{T, VT, MT, QN} <: AbstractCondensedKKTSystem{T, VT, MT, QN}

Expand Down Expand Up @@ -204,110 +110,6 @@ function create_kkt_system(
)
end

# For templating
const AbstractDenseKKTSystem{T, VT, MT, QN} = Union{
DenseKKTSystem{T, VT, MT, QN},
DenseCondensedKKTSystem{T, VT, MT, QN},
}

#=
Generic functions
=#

function jtprod!(y::AbstractVector, kkt::AbstractDenseKKTSystem, x::AbstractVector)
nx = size(kkt.hess, 1)
ind_ineq = kkt.ind_ineq
ns = length(ind_ineq)
yx = view(y, 1:nx)
ys = view(y, 1+nx:nx+ns)
# / x
mul!(yx, kkt.jac', x)
# / s
ys .= -@view(x[ind_ineq])
return
end

function compress_jacobian!(kkt::AbstractDenseKKTSystem)
return
end

nnz_jacobian(kkt::AbstractDenseKKTSystem) = length(kkt.jac)

#=
DenseKKTSystem
=#

num_variables(kkt::DenseKKTSystem) = length(kkt.pr_diag)

function mul!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector)
symul!(y, kkt.aug_com, x)
end

# 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

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::VT, hess, jac, pr_diag, du_diag, diag_hess, ind_ineq, n, m, ns) where {T, VT <: AbstractMatrix{T}}
# 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 / variables
for i in 1:m, j in 1:n
dest[i + n + ns, j] = jac[i, j]
dest[j, i + n + ns] = jac[i, j]
end
# Transfer Jacobian / slacks
for j in 1:ns
is = ind_ineq[j]
dest[is + n + ns, j + n] = - one(T)
dest[j + n, is + n + ns] = - one(T)
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)

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

function compress_hessian!(kkt::DenseKKTSystem)
# Transfer diagonal term for future regularization
diag!(kkt.diag_hess, kkt.hess)
end

#=
DenseCondensedKKTSystem
=#

num_variables(kkt::DenseCondensedKKTSystem) = size(kkt.hess, 1)

function get_slack_regularization(kkt::DenseCondensedKKTSystem)
Expand Down
Loading
Loading