Skip to content

Commit

Permalink
Initial infrasctructure for preconditioners
Browse files Browse the repository at this point in the history
  • Loading branch information
mtanneau committed Sep 20, 2020
1 parent 05cb811 commit 2e2f49f
Show file tree
Hide file tree
Showing 4 changed files with 89 additions and 9 deletions.
2 changes: 1 addition & 1 deletion src/KKTSolver/KKTSolver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions src/KKTSolver/Krylov/jacobi.jl
Original file line number Diff line number Diff line change
@@ -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
42 changes: 34 additions & 8 deletions src/KKTSolver/krylov.jl → src/KKTSolver/Krylov/krylov.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,14 +38,17 @@ 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

f::F # Krylov function
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
Expand All @@ -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
Expand Down Expand Up @@ -101,6 +116,9 @@ function update!(
kkt.regP .= regP
kkt.regD .= regD

# Update Jacobi preconditioner
update_preconditioner(kkt)

return nothing
end

Expand Down Expand Up @@ -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:
# ==============================================================================
Expand Down Expand Up @@ -424,4 +447,7 @@ function solve!(
dy .= _dy

return nothing
end
end

# Code for pre-conditioners
include("preconditioners.jl")
9 changes: 9 additions & 0 deletions src/KKTSolver/Krylov/preconditioners.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
abstract type AbstractPreconditioner{T} end

update_preconditioner(::AbstractKKTSolver) = nothing

struct IdentityPreconditioner end

op(::IdentityPreconditioner) = LO.opEye()

include("jacobi.jl")

0 comments on commit 2e2f49f

Please sign in to comment.