Initial infrasctructure for preconditioners

Author: mtanneau

src/KKT/KKT.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(T)

src/KKT/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

src/KKT/krylov.jl renamed to src/KKT/Krylov/krylov.jl

@@ -46,6 +46,9 @@ mutable struct KrylovSPD{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 KrylovSPD{T, F, Tv, Ta} <: AbstractKrylovSolver{T}
 
     function KrylovSPD(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
 
-
 # ==============================================================================
 #   KrylovSID:
 # ==============================================================================
@@ -424,4 +447,7 @@ function solve!(
     dy .= _dy
 
     return nothing
-end
+end
+
+# Code for pre-conditioners
+include("preconditioners.jl")

src/KKT/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")