diff --git a/src/KKTSolver/KKTSolver.jl b/src/KKTSolver/KKTSolver.jl index c5bdef54..faec3da5 100644 --- a/src/KKTSolver/KKTSolver.jl +++ b/src/KKTSolver/KKTSolver.jl @@ -107,7 +107,7 @@ include("test.jl") include("lapack.jl") include("cholmod.jl") include("ldlfact.jl") -include("krylov.jl") +include("Krylov/krylov.jl") """ default_options(Tv) diff --git a/src/KKTSolver/Krylov/jacobi.jl b/src/KKTSolver/Krylov/jacobi.jl new file mode 100644 index 00000000..97ced1a5 --- /dev/null +++ b/src/KKTSolver/Krylov/jacobi.jl @@ -0,0 +1,45 @@ +struct JacobiPreconditioner{T, Tv} <: AbstractPreconditioner{T} + d::Tv + + JacobiPreconditioner(d::Tv) where{T, Tv<:AbstractVector{T}} = new{T, Tv}(d) +end + +op(P::JacobiPreconditioner) = Diagonal(P.d) + +# Code for Jacobi preconditioners +function update_preconditioner(kkt::KrylovSPDSolver{T, F, Tv, Ta, Tp}) where{T, F, Tv, Ta<:SparseMatrixCSC, Tp<:JacobiPreconditioner{T, Tv}} + + A = kkt.A + d = kkt.P.d + + # S = A (Θ⁻¹ + Rp)⁻¹ A' + Rd, so its diagonal writes + # S[i, i] = Σ_j=1..n D[j, j] * A[i, j] ^2, + # where D = (Θ⁻¹ + Rp)⁻¹ + + _d = one(T) ./ (kkt.θ .+ kkt.regP) + + rows = rowvals(A) + vals = nonzeros(A) + m, n = size(A) + + d .= zero(T) + + for j = 1:n + _a = zero(T) + dj = _d[j] + for i in nzrange(A, j) + row = rows[i] + a_ij = vals[i] + + d[row] += dj * a_ij ^2 + end + end + + # I think putting this afterwards is more numerics-friendly + d .+= kkt.regD + + # Invert + d .\= one(T) + + return nothing +end \ No newline at end of file diff --git a/src/KKTSolver/krylov.jl b/src/KKTSolver/Krylov/krylov.jl similarity index 91% rename from src/KKTSolver/krylov.jl rename to src/KKTSolver/Krylov/krylov.jl index 952476d2..e9f760d7 100644 --- a/src/KKTSolver/krylov.jl +++ b/src/KKTSolver/Krylov/krylov.jl @@ -38,7 +38,7 @@ 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} +mutable struct KrylovSPDSolver{T, F, Tv, Ta, Tp} <: AbstractKrylovSolver{T} m::Int n::Int @@ -46,6 +46,9 @@ mutable struct KrylovSPDSolver{T, F, Tv, Ta} <: AbstractKrylovSolver{T} atol::T # absolute tolerance rtol::T # relative tolerance + P::Tp # Preconditioner + nitertot::Int # Total number of Krylov iterations + # Problem data A::Ta # Constraint matrix θ::Tv # scaling @@ -54,14 +57,26 @@ mutable struct KrylovSPDSolver{T, F, Tv, Ta} <: AbstractKrylovSolver{T} function KrylovSPDSolver(f::Function, A::Ta; atol::T=eps(T)^(3 // 4), - rtol::T=eps(T)^(3 // 4) + rtol::T=eps(T)^(3 // 4), + preconditioner::Symbol=:Identity ) where{T, Ta <: AbstractMatrix{T}} F = typeof(f) m, n = size(A) - return new{T, F, Vector{T}, Ta}( + if preconditioner == :Jacobi + P = JacobiPreconditioner(ones(T, m)) + elseif preconditioner == :Identity + P = IdentityPreconditioner() + else + @error "Unknown preconditioner: $preconditioner, using identity instead" + P = IdentityPreconditioner() + end + + return new{T, F, Vector{T}, Ta, typeof(P)}( m, n, f, atol, rtol, + P, + 0, A, ones(T, n), ones(T, n), ones(T, m) ) end @@ -101,6 +116,9 @@ function update!( kkt.regP .= regP kkt.regD .= regD + # Update Jacobi preconditioner + update_preconditioner(kkt) + return nothing end @@ -137,17 +155,22 @@ function solve!( end ) + opM = op(kkt.P) + # @info typeof(opM) extrema(kkt.P.d) + # Solve normal equations - _dy, stats = kkt.f(opS, ξ_, atol=kkt.atol, rtol=kkt.rtol) - dy .= _dy + _dy, stats = kkt.f(opS, ξ_, M = opM, atol=kkt.atol, rtol=kkt.rtol) - # Recover dx + # Book-keeping + kkt.nitertot += length(stats.residuals) - 1 + + # Recover dy, dx + dy .= _dy dx .= D * (kkt.A' * dy - ξd) return nothing end - # ============================================================================== # KrylovSIDSolver: # ============================================================================== @@ -424,4 +447,7 @@ function solve!( dy .= _dy return nothing -end \ No newline at end of file +end + +# Code for pre-conditioners +include("preconditioners.jl") \ No newline at end of file diff --git a/src/KKTSolver/Krylov/preconditioners.jl b/src/KKTSolver/Krylov/preconditioners.jl new file mode 100644 index 00000000..78bb4ff3 --- /dev/null +++ b/src/KKTSolver/Krylov/preconditioners.jl @@ -0,0 +1,9 @@ +abstract type AbstractPreconditioner{T} end + +update_preconditioner(::AbstractKKTSolver) = nothing + +struct IdentityPreconditioner end + +op(::IdentityPreconditioner) = LO.opEye() + +include("jacobi.jl") \ No newline at end of file