Invert preconditioner using matrix inversion lemma

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")

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

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
+function (p::NoPreconditioner)(_, _) I end