From 5f7528d987d143173834cae247aa2708c21202c4 Mon Sep 17 00:00:00 2001
From: Sebastian Callh <sebastian.callh@gmail..com>
Date: Sat, 24 Jun 2023 09:58:06 +0200
Subject: [PATCH] Invert preconditioner using matrix inversion lemma

---
 src/IterGP.jl              |  1 +
 src/lowrankplusdiagonal.jl | 16 ++++++++++++++++
 src/preconditioners.jl     |  4 ++--
 3 files changed, 19 insertions(+), 2 deletions(-)
 create mode 100644 src/lowrankplusdiagonal.jl

diff --git a/src/IterGP.jl b/src/IterGP.jl
index 68b3142..d558b92 100644
--- a/src/IterGP.jl
+++ b/src/IterGP.jl
@@ -12,6 +12,7 @@ using AbstractGPs: AbstractGP, FiniteGP, MeanFunction, ZeroMean
 export CholeskyPolicy, ConjugateGradientPolicy
 export NoPreconditioner, CholeskyPreconditioner
 
+include("lowrankplusdiagonal.jl")
 include("preconditioners.jl")
 include("policies.jl")
 include("actors.jl")
diff --git a/src/lowrankplusdiagonal.jl b/src/lowrankplusdiagonal.jl
new file mode 100644
index 0000000..8993dce
--- /dev/null
+++ b/src/lowrankplusdiagonal.jl
@@ -0,0 +1,16 @@
+"""
+Type for matrices on the form A = D + UU' that can be inverted using the matrix inversion lemma.
+"""
+struct LowRankPlusDiagonal{TD <: AbstractMatrix, TU <: AbstractMatrix}
+    D::TD # Diagonal component
+    U::TU # Low-rank component
+end
+
+function LinearAlgebra.inv(A::LowRankPlusDiagonal)
+    (;D, U) = A
+    V = U'
+    Dinv = inv(D)
+    B = I + V*Dinv*U
+    Binv = inv(B)
+    Dinv - Dinv*U*Binv*V*Dinv
+end
diff --git a/src/preconditioners.jl b/src/preconditioners.jl
index 51c804e..d850fc2 100644
--- a/src/preconditioners.jl
+++ b/src/preconditioners.jl
@@ -13,8 +13,8 @@ function (p::CholeskyPreconditioner)(K, Σy)
         K′ .= K′ - Iᵢ*Iᵢ'
         L[:,i] = Iᵢ
     end
-    L*L' + Σy
+    LowRankPlusDiagonal(Σy, L)
 end
 
 struct NoPreconditioner <: AbstractPreconditioner end
-function (p::NoPreconditioner)(_, _) I end
\ No newline at end of file
+function (p::NoPreconditioner)(_, _) I end