From 2e2f49f66a2c899abfae98836ab2504b96c5ca46 Mon Sep 17 00:00:00 2001
From: mtanneau <mathieu.tanneau@gmail.com>
Date: Thu, 17 Sep 2020 21:35:57 -0400
Subject: [PATCH] Initial infrasctructure for preconditioners

---
 src/KKTSolver/KKTSolver.jl              |  2 +-
 src/KKTSolver/Krylov/jacobi.jl          | 45 +++++++++++++++++++++++++
 src/KKTSolver/{ => Krylov}/krylov.jl    | 42 ++++++++++++++++++-----
 src/KKTSolver/Krylov/preconditioners.jl |  9 +++++
 4 files changed, 89 insertions(+), 9 deletions(-)
 create mode 100644 src/KKTSolver/Krylov/jacobi.jl
 rename src/KKTSolver/{ => Krylov}/krylov.jl (91%)
 create mode 100644 src/KKTSolver/Krylov/preconditioners.jl

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