diff --git a/Project.toml b/Project.toml index 18480bd8..e7ce26f0 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ authors = ["Mathieu Tanneau "] version = "0.5.1" [deps] +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" @@ -15,6 +16,7 @@ SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] +Krylov = "0.5.2" LDLFactorizations = "^0.5" MathOptInterface = "~0.9.5" QPSReader = "0.2" diff --git a/docs/src/manual/linear_systems.md b/docs/src/manual/linear_systems.md index d3779704..1f493291 100644 --- a/docs/src/manual/linear_systems.md +++ b/docs/src/manual/linear_systems.md @@ -38,7 +38,7 @@ The augmented system above can be reduced to the positive-definite _normal equat \Delta x &= (Θ^{-1} + R_{p})^{-1} (A^{T} \Delta y - \xi_{d}) \end{array} ``` -When selected, this reduction is transparent to the interior-point algorithm. +If available and when selected, this reduction is transparent to the interior-point algorithm. To enable the use of fast external libraries and/or specialized routines, the resolution of linear systems is performed by an [`AbstractKKTSolver`] object. @@ -52,4 +52,11 @@ Here is a list of currently supported linear solvers: | [`Dense_SymPosDef`](@ref) | `Real` | Normal equations | Dense / LAPACK | Cholesky | [`Cholmod_SymQuasDef`](@ref) | `Float64` | Augmented system | CHOLMOD | LDLᵀ | [`Cholmod_SymPosDef`](@ref) | `Float64` | Normal equations | CHOLMOD | Cholesky -| [`LDLFact_SymQuasDef`](@ref) | `Real` | Augmented system | [LDLFactorizations.jl](https://github.com/JuliaSmoothOptimizers/LDLFactorizations.jl) | LDLᵀ \ No newline at end of file +| [`LDLFact_SymQuasDef`](@ref) | `Real` | Augmented system | [LDLFactorizations.jl](https://github.com/JuliaSmoothOptimizers/LDLFactorizations.jl) | LDLᵀ +| [`KrylovSPDSolver`](@ref) | `Real` | Normal equations | [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) | Krylov +| [`KrylovSIDSolver`](@ref) | `Real` | Augmented system[^1] | [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) | Krylov +| [`KrylovSQDSolver`](@ref) | `Real` | Augmented system[^1] | [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) | Krylov + +[^1]: [`KrylovSIDSolver`](@ref)s view the augmented system as a symmetric indefinite system, + while [`KrylovSQDSolver`](@ref)s exploit its 2x2 structure and quasi-definite property. + See the reference documentation for more details. \ No newline at end of file diff --git a/docs/src/reference/kkt_solvers.md b/docs/src/reference/kkt_solvers.md index ee7c1ab6..d8377449 100644 --- a/docs/src/reference/kkt_solvers.md +++ b/docs/src/reference/kkt_solvers.md @@ -65,4 +65,23 @@ Cholmod_SymPosDef ```@docs LDLFact_SymQuasDef +``` + +## Krylov + +!!! warning + Iterative methods are still an experimental feature. + Some numerical and performance issues should be expected. + + +```@docs +KrylovSPDSolver +``` + +```@docs +KrylovSIDSolver +``` + +```@docs +KrylovSQDSolver ``` \ No newline at end of file diff --git a/src/KKTSolver/KKTSolver.jl b/src/KKTSolver/KKTSolver.jl index c292a4ca..c5bdef54 100644 --- a/src/KKTSolver/KKTSolver.jl +++ b/src/KKTSolver/KKTSolver.jl @@ -107,6 +107,7 @@ include("test.jl") include("lapack.jl") include("cholmod.jl") include("ldlfact.jl") +include("krylov.jl") """ default_options(Tv) diff --git a/src/KKTSolver/krylov.jl b/src/KKTSolver/krylov.jl new file mode 100644 index 00000000..952476d2 --- /dev/null +++ b/src/KKTSolver/krylov.jl @@ -0,0 +1,427 @@ +import Krylov +const LO = Krylov.LinearOperators + +""" + AbstractKrylovSolver{T} + +Abstract type for Kyrlov-based linear solvers. +""" +abstract type AbstractKrylovSolver{T} <: AbstractKKTSolver{T} end + +# ============================================================================== +# KrylovSPDSolver: +# ============================================================================== + +""" + KrylovSPDSolver{T, F, Tv, Ta} + +Krylov-based **S**ymmetric **P**ositive-**D**efinite (SPD) linear solver. + +Applies a Krylov method to the normal equations systems, then recovers a solution +to the augmented system. +The selected Krylov method must therefore handle symmetric positive-definite systems. +Suitable methods are CG or CR. + +A `KrylovSPDSolver` is selected as follows +```julia +model.params.KKTOptions = Tulip.KKT.SolverOptions( + KrylovSPDSolver, method=Krylov.cg, + atol=1e-12, rtol=1e-12 +) +``` + +The `method` argument is a function `f::F` whose signature must match +```julia +dy, _ = f(S, ξ; atol, rtol) +``` +where `S`, `ξ` and `dy` are the normal equations system's +left- and right-hand side, and solution vector, respectively. +`S` may take the form of a matrix, or of a suitable linear operator. +""" +mutable struct KrylovSPDSolver{T, F, Tv, Ta} <: AbstractKrylovSolver{T} + m::Int + n::Int + + f::F # Krylov function + atol::T # absolute tolerance + rtol::T # relative tolerance + + # Problem data + A::Ta # Constraint matrix + θ::Tv # scaling + regP::Tv # primal regularization + regD::Tv # dual regularization + + function KrylovSPDSolver(f::Function, A::Ta; + atol::T=eps(T)^(3 // 4), + rtol::T=eps(T)^(3 // 4) + ) where{T, Ta <: AbstractMatrix{T}} + F = typeof(f) + m, n = size(A) + + return new{T, F, Vector{T}, Ta}( + m, n, + f, atol, rtol, + A, ones(T, n), ones(T, n), ones(T, m) + ) + end +end + +setup(::Type{KrylovSPDSolver}, A; method, kwargs...) = KrylovSPDSolver(method, A; kwargs...) + +backend(kkt::KrylovSPDSolver) = "Krylov ($(kkt.f))" +linear_system(::KrylovSPDSolver) = "Normal equations" + + +""" + update!(kkt::KrylovSPDSolver, θ, regP, regD) + +Update diagonal scaling ``θ`` and primal-dual regularizations. +""" +function update!( + kkt::KrylovSPDSolver{Tv}, + θ::AbstractVector{Tv}, + regP::AbstractVector{Tv}, + regD::AbstractVector{Tv} +) where{Tv<:Real} + # Sanity checks + length(θ) == kkt.n || throw(DimensionMismatch( + "θ has length $(length(θ)) but linear solver is for n=$(kkt.n)." + )) + length(regP) == kkt.n || throw(DimensionMismatch( + "regP has length $(length(regP)) but linear solver has n=$(kkt.n)" + )) + length(regD) == kkt.m || throw(DimensionMismatch( + "regD has length $(length(regD)) but linear solver has m=$(kkt.m)" + )) + + # Update diagonal scaling + kkt.θ .= θ + # Update regularizers + kkt.regP .= regP + kkt.regD .= regD + + return nothing +end + +""" + solve!(dx, dy, kkt::KrylovSPDSolver, ξp, ξd) + +Solve the normal equations system using the selected Krylov method. +""" +function solve!( + dx::Vector{Tv}, dy::Vector{Tv}, + kkt::KrylovSPDSolver{Tv}, + ξp::Vector{Tv}, ξd::Vector{Tv} +) where{Tv<:Real} + m, n = kkt.m, kkt.n + + # Setup + d = one(Tv) ./ (kkt.θ .+ kkt.regP) + D = Diagonal(d) + + # Set-up right-hand side + ξ_ = ξp .+ kkt.A * (D * ξd) + + # Form linear operator S = A * D * A' + Rd + # Here we pre-allocate the intermediary and final vectors + v1 = zeros(Tv, n) + v2 = zeros(Tv, m) + opS = LO.LinearOperator(Tv, m, m, true, true, + w -> begin + mul!(v1, kkt.A', w) + v1 .*= d + mul!(v2, kkt.A, v1) + v2 .+= kkt.regD .* w + v2 + end + ) + + # Solve normal equations + _dy, stats = kkt.f(opS, ξ_, atol=kkt.atol, rtol=kkt.rtol) + dy .= _dy + + # Recover dx + dx .= D * (kkt.A' * dy - ξd) + + return nothing +end + + +# ============================================================================== +# KrylovSIDSolver: +# ============================================================================== + +""" + KrylovSIDSolver{T, F, Tv, Ta} + +Krylov-based **S**ymmetric **I**n**D**efinite (SID) linear solver. + +Applies a Krylov method to solve the augmented system, +without exploiting its quasi-definiteness properties. +The selected Krylov method must therefore handle symmetric indefinite systems. +Suitable methods are MINRES or MINRES-QLP. + +A `KrylovSIDSolver` is selected as follows +```julia +model.params.KKTOptions = Tulip.KKT.SolverOptions( + KrylovSIDSolver, method=Krylov.minres, + atol=1e-12, rtol=1e-12 +) +``` + +The `method` argument is a function `f::F` whose signature must match +```julia +Δ, _ = f(K, ξ; atol, rtol) +``` +where `K`, `ξ` and `Δ` are the augmented system's +left- and right-hand side, and solution vector, respectively. +`K` may take the form of a matrix, or of a suitable linear operator. +""" +mutable struct KrylovSIDSolver{T, F, Tv, Ta} <: AbstractKrylovSolver{T} + m::Int + n::Int + + f::F # Krylov function + atol::T # absolute tolerance + rtol::T # relative tolerance + + # Problem data + A::Ta # Constraint matrix + θ::Tv # scaling + regP::Tv # primal regularization + regD::Tv # dual regularization + + function KrylovSIDSolver(f::Function, A::Ta; + atol::T=eps(T)^(3 // 4), + rtol::T=eps(T)^(3 // 4) + ) where{T, Ta <: AbstractMatrix{T}} + F = typeof(f) + m, n = size(A) + + return new{T, F, Vector{T}, Ta}( + m, n, + f, atol, rtol, + A, ones(T, n), ones(T, n), ones(T, m) + ) + end +end + +setup(::Type{KrylovSIDSolver}, A; method, kwargs...) = KrylovSIDSolver(method, A; kwargs...) + +backend(kkt::KrylovSIDSolver) = "Krylov ($(kkt.f))" +linear_system(::KrylovSIDSolver) = "Augmented system" + + +""" + update!(kkt::KrylovSIDSolver, θ, regP, regD) + +Update diagonal scaling ``θ`` and primal-dual regularizations. +""" +function update!( + kkt::KrylovSIDSolver{Tv}, + θ::AbstractVector{Tv}, + regP::AbstractVector{Tv}, + regD::AbstractVector{Tv} +) where{Tv<:Real} + # Sanity checks + length(θ) == kkt.n || throw(DimensionMismatch( + "θ has length $(length(θ)) but linear solver is for n=$(kkt.n)." + )) + length(regP) == kkt.n || throw(DimensionMismatch( + "regP has length $(length(regP)) but linear solver has n=$(kkt.n)" + )) + length(regD) == kkt.m || throw(DimensionMismatch( + "regD has length $(length(regD)) but linear solver has m=$(kkt.m)" + )) + + # Update diagonal scaling + kkt.θ .= θ + # Update regularizers + kkt.regP .= regP + kkt.regD .= regD + + return nothing +end + +""" + solve!(dx, dy, kkt::KrylovSIDSolver, ξp, ξd) + +Solve the augmented system using the selected Krylov method. +""" +function solve!( + dx::Vector{Tv}, dy::Vector{Tv}, + kkt::KrylovSIDSolver{Tv}, + ξp::Vector{Tv}, ξd::Vector{Tv} +) where{Tv<:Real} + m, n = kkt.m, kkt.n + + # Setup + A = kkt.A + D = Diagonal(-(kkt.θ .+ kkt.regP)) # D = Θ⁻¹ + Rp + + # Set-up right-hand side + ξ = [ξd; ξp] + + # Form linear operator K = [-(Θ⁻¹ + Rp) Aᵀ; A Rd] + # Here we pre-allocate the final vector + z = zeros(Tv, m+n) + opK = LO.LinearOperator(Tv, n+m, n+m, true, false, + w -> begin + @views z1 = z[1:n] + @views z2 = z[n+1:n+m] + + @views w1 = w[1:n] + @views w2 = w[n+1:n+m] + + # z1 = -(Θ⁻¹ + Rp) w1 + A'w2 + mul!(z1, D, w1, one(Tv), zero(Tv)) # z1 = -(Θ⁻¹ + Rp) w1 + mul!(z1, A', w2, one(Tv), one(Tv)) # z1 += A'w2 + + # z2 = A w1 + Rd w2 + mul!(z2, A, w1, one(Tv), zero(Tv)) # z2 = A w1 + mul!(z2, Diagonal(kkt.regD), w2, one(Tv), one(Tv)) # z2 += Rd * w2 + + z + end + ) + + # Solve augmented system + Δ, stats = kkt.f(opK, ξ, atol=kkt.atol, rtol=kkt.rtol) + + # Recover dx, dy + dx .= Δ[1:n] + dy .= Δ[n+1:n+m] + + return nothing +end + + +# ============================================================================== +# KrylovSQDSolver: +# ============================================================================== + +""" + KrylovSQDSolver{T, F, Tv, Ta} + +Krylov-based **S**ymmetric **Q**uasi-**D**efinite (SQD) linear solver. + +Applies a Krylov method to solve the augmented system, taking advantage of +its 2x2 block structure and quasi-definiteness. +The selected Krylov method must therefore handle 2x2 symmetric quasi-definite systems. +Suitable methods are TriCG and TriMR. + +A `KrylovSIDSolver` is selected as follows +```julia +model.params.KKTOptions = Tulip.KKT.SolverOptions( + KrylovSQDSolver, method=Krylov.trimr, + atol=1e-12, rtol=1e-12 +) +``` + +The `method` argument is a function `f::F` whose signature must match +``` +dy, dx, _ = f(A, ξp, ξd; M, N, atol, rtol) +``` +where the augmented system is of the form +``` +[ -N⁻¹ Aᵀ ] [dx] = [ξd] +[ A M⁻¹] [dy] = [ξp] +``` +i.e., ``N = (Θ^{-1} + R_{p})^{-1}`` and ``M = R_{d}^{-1}``. +""" +mutable struct KrylovSQDSolver{T, F, Tv, Ta} <: AbstractKrylovSolver{T} + m::Int + n::Int + + f::F # Krylov function + atol::T # absolute tolerance + rtol::T # relative tolerance + + # Problem data + A::Ta # Constraint matrix + θ::Tv # scaling + regP::Tv # primal regularization + regD::Tv # dual regularization + + function KrylovSQDSolver(f::Function, A::Ta; + atol::T=eps(T)^(3 // 4), + rtol::T=eps(T)^(3 // 4) + ) where{T, Ta <: AbstractMatrix{T}} + F = typeof(f) + m, n = size(A) + + return new{T, F, Vector{T}, Ta}( + m, n, + f, atol, rtol, + A, ones(T, n), ones(T, n), ones(T, m) + ) + end +end + +setup(::Type{KrylovSQDSolver}, A; method, kwargs...) = KrylovSQDSolver(method, A; kwargs...) + +backend(kkt::KrylovSQDSolver) = "Krylov ($(kkt.f))" +linear_system(::KrylovSQDSolver) = "Augmented system" + + +""" + update!(kkt::KrylovSQDSolver, θ, regP, regD) + +Update diagonal scaling ``θ`` and primal-dual regularizations. +""" +function update!( + kkt::KrylovSQDSolver{Tv}, + θ::AbstractVector{Tv}, + regP::AbstractVector{Tv}, + regD::AbstractVector{Tv} +) where{Tv<:Real} + # Sanity checks + length(θ) == kkt.n || throw(DimensionMismatch( + "θ has length $(length(θ)) but linear solver is for n=$(kkt.n)." + )) + length(regP) == kkt.n || throw(DimensionMismatch( + "regP has length $(length(regP)) but linear solver has n=$(kkt.n)" + )) + length(regD) == kkt.m || throw(DimensionMismatch( + "regD has length $(length(regD)) but linear solver has m=$(kkt.m)" + )) + + # Update diagonal scaling + kkt.θ .= θ + # Update regularizers + kkt.regP .= regP + kkt.regD .= regD + + return nothing +end + +""" + solve!(dx, dy, kkt::KrylovSQDSolver, ξp, ξd) + +Solve the augmented system using the selected Kyrlov method. +""" +function solve!( + dx::Vector{Tv}, dy::Vector{Tv}, + kkt::KrylovSQDSolver{Tv}, + ξp::Vector{Tv}, ξd::Vector{Tv} +) where{Tv<:Real} + m, n = kkt.m, kkt.n + + # Setup + d = one(Tv) ./ (kkt.θ .+ kkt.regP) + N = Diagonal(d) + + M = Diagonal(one(Tv) ./ kkt.regD) + + # Solve augmented system + _dy, _dx, stats = kkt.f(kkt.A, ξp, ξd, + M=M, N=N, + atol=kkt.atol, rtol=kkt.rtol + ) + + dx .= _dx + dy .= _dy + + return nothing +end \ No newline at end of file diff --git a/test/KKTSolver/KKTSolver.jl b/test/KKTSolver/KKTSolver.jl index 29d91627..8f3b8137 100644 --- a/test/KKTSolver/KKTSolver.jl +++ b/test/KKTSolver/KKTSolver.jl @@ -2,4 +2,5 @@ const KKT = Tulip.KKT include("lapack.jl") include("cholmod.jl") -include("ldlfact.jl") \ No newline at end of file +include("ldlfact.jl") +include("krylov.jl") \ No newline at end of file diff --git a/test/KKTSolver/krylov.jl b/test/KKTSolver/krylov.jl new file mode 100644 index 00000000..be643be0 --- /dev/null +++ b/test/KKTSolver/krylov.jl @@ -0,0 +1,41 @@ +import Krylov + +@testset "Krylov" begin + for T in TvTYPES + @testset "$T" begin + # Kyrlov solvers for symmetric positive-definite systems + for f in [Krylov.cg, Krylov.minres, Krylov.minres_qlp] + @testset "PD: $f" begin + A = SparseMatrixCSC{T, Int}([ + 1 0 1 0; + 0 1 0 1 + ]) + kkt = KKT.KrylovSPDSolver(f, A) + KKT.run_ls_tests(A, kkt) + end + end + # Kyrlov solvers for symmetric indefinite systems + for f in [Krylov.minres, Krylov.minres_qlp] + @testset "SID: $f" begin + A = SparseMatrixCSC{T, Int}([ + 1 0 1 0; + 0 1 0 1 + ]) + kkt = KKT.KrylovSIDSolver(f, A) + KKT.run_ls_tests(A, kkt) + end + end + # Kyrlov solvers for symmetric quasi-definite systems + for f in [Krylov.tricg, Krylov.trimr] + @testset "SQD: $f" begin + A = SparseMatrixCSC{T, Int}([ + 1 0 1 0; + 0 1 0 1 + ]) + kkt = KKT.KrylovSQDSolver(f, A) + KKT.run_ls_tests(A, kkt) + end + end + end + end +end \ No newline at end of file