Initial infrastructure for Krylov-based linear solvers

Project.toml

@@ -4,6 +4,7 @@ authors = ["Mathieu Tanneau <mathieu.tanneau@gmail.com>"]
 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"

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,6 @@ 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ᵀ
+| [`LDLFact_SymQuasDef`](@ref) | `Real` | Augmented system | [LDLFactorizations.jl](https://github.com/JuliaSmoothOptimizers/LDLFactorizations.jl) | LDLᵀ
+| [`KrylovPDSolver`](@ref) | `Real` | Normal equations | [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) | Krylov
+| [`KrylovSQDSolver`](@ref) | `Real` | Augmented system | [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) | Krylov

docs/src/reference/kkt_solvers.md

@@ -65,4 +65,19 @@ Cholmod_SymPosDef
 
 ```@docs
 LDLFact_SymQuasDef
+```
+
+## Krylov
+
+!!! warning
+    Iterative methods are still an experimental feature.
+    Some numerical and performance issues should be expected.
+
+
+```@docs
+KrylovPDSolver
+```
+
+```@docs
+KrylovSQDSolver
 ```

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)

src/KKTSolver/krylov.jl

@@ -0,0 +1,239 @@
+import Krylov
+const LO = Krylov.LinearOperators
+
+abstract type KrylovSolver{T} <: AbstractKKTSolver{T} end
+
+# ==============================================================================
+#   KrylovPDSolver: 
+# ==============================================================================
+
+"""
+    KrylovPDSolver{T, F, Tv, Ta}
+
+Solves normal equations system with an iterative method `f::F`.
+
+```julia
+model.params.KKTOptions = Tulip.KKT.SolverOptions(
+    KrylovPDSolver, method=Krylov.cg
+)
+```
+"""
+mutable struct KrylovPDSolver{T, F, Tv, Ta} <: KrylovSolver{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 KrylovPDSolver(f::Function, A::Ta;
+        atol::T=sqrt(eps(T)),
+        rtol::T=sqrt(eps(T))
+    ) 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{KrylovPDSolver}, A; method, kwargs...) = KrylovPDSolver(method, A; kwargs...)
+
+backend(kkt::KrylovPDSolver) = "Krylov ($(kkt.f))"
+linear_system(::KrylovPDSolver) = "Normal equations"
+
+
+"""
+    update!(kkt::KrylovPDSolver, θ, regP, regD)
+
+Update diagonal scaling ``θ`` and primal-dual regularizations.
+"""
+function update!(
+    kkt::KrylovPDSolver{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::KrylovPDSolver, ξp, ξd)
+
+Solve the normal equations system using an iterative method.
+"""
+function solve!(
+    dx::Vector{Tv}, dy::Vector{Tv},
+    kkt::KrylovPDSolver{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
+
+
+# ==============================================================================
+#   KrylovSQDSolver: 
+# ==============================================================================
+
+"""
+    KrylovSQDSolver{T, F, Tv, Ta}
+
+Solves the augmented system with an iterative method `f::F`.
+
+```julia
+model.params.KKTOptions = Tulip.KKT.SolverOptions(
+    KrylovSQDSolver, method=Krylov.trimr
+)
+```
+"""
+mutable struct KrylovSQDSolver{T, F, Tv, Ta} <: KrylovSolver{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=sqrt(eps(T)),
+        rtol::T=sqrt(eps(T))
+    ) 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 normal equations system using an iterative 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)
+    D = Diagonal(d)
+
+    # Solve augmented system
+    # Currently assumes that Rd is a multiple of the identity matrix
+    _dy, _dx, stats = kkt.f(kkt.A, ξp, ξd,
+        N=D, τ = kkt.regD[1],
+        atol=kkt.atol, rtol=kkt.rtol
+    )
+
+    dx .= _dx
+    dy .= _dy
+
+    return nothing
+end

test/KKTSolver/KKTSolver.jl

@@ -2,4 +2,5 @@ const KKT = Tulip.KKT
 
 include("lapack.jl")
 include("cholmod.jl")
-include("ldlfact.jl")
+include("ldlfact.jl")
+include("krylov.jl")

test/KKTSolver/krylov.jl

@@ -0,0 +1,28 @@
+import Krylov
+
+@testset "Krylov" begin
+    for T in TvTYPES
+        @testset "$T" begin
+            for f in [Krylov.cg, Krylov.minres, Krylov.minres_qlp]
+                @testset "$f" begin
+                    A = SparseMatrixCSC{T, Int}([
+                        1 0 1 0;
+                        0 1 0 1
+                    ])
+                    kkt = KKT.KrylovPDSolver(f, A)
+                    KKT.run_ls_tests(A, kkt)
+                end
+            end
+            for f in [Krylov.tricg, Krylov.trimr]
+                @testset "$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