Merge pull request #105 from vpuri3/vp-krylovkit

KrylovKit wrapper

Author: ChrisRackauckas

src/LinearSolve.jl

@@ -50,6 +50,7 @@ export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
        UMFPACKFactorization, KLUFactorization
 export KrylovJL, KrylovJL_CG, KrylovJL_GMRES, KrylovJL_BICGSTAB, KrylovJL_MINRES,
        IterativeSolversJL, IterativeSolversJL_CG, IterativeSolversJL_GMRES,
-       IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES
+       IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES,
+       KrylovKitJL, KrylovKitJL_CG, KrylovKitJL_GMRES
 
 end

src/iterative_wrappers.jl

@@ -231,3 +231,43 @@ function purge_history!(iter::IterativeSolvers.GMRESIterable, x, b)
   iter.β = iter.residual.current
   nothing
 end
+
+## KrylovKit.jl
+
+struct KrylovKitJL{F,A,I,K} <: AbstractKrylovSubspaceMethod
+    KrylovAlg::F
+    gmres_restart::I
+    args::A
+    kwargs::K
+end
+
+function KrylovKitJL(args...;
+                     KrylovAlg = KrylovKit.GMRES, gmres_restart = 0,
+                     kwargs...)
+    return KrylovJL(KrylovAlg, gmres_restart, args, kwargs)
+end
+
+KrylovKitJL_CG(args...;kwargs...) =
+    KrylovKitJL(args...; KrylovAlg=KrylovKit.CG, kwargs...)
+KrylovKitJL_GMRES(args...;kwargs...) =
+    KrylovKitJL(args...; KrylovAlg=KrylovKit.GMRES, kwargs...)
+
+function SciMLBase.solve(cache::LinearCache, alg::KrylovKitJL, kwargs...)
+
+    atol      = float(cache.abstol)
+    rtol      = float(cache.reltol)
+    maxiter   = cache.maxiters
+    verbosity = cache.verbose ? 1 : 0
+    krylovdim = (alg.gmres_restart == 0) ? min(20, size(A,1)) : alg.gmres_restart
+
+    kwargs = (atol=atol, rtol=rtol, maxiter=maxiter, verbosity=verbosity,
+              krylovdim = krylovdim, alg.kwargs...)
+
+    x, info = KrylovKit.linsolve(cache.A, cache.b, cache.u, alg.KrylovAlg)
+
+    copy!(cache.u, x)
+    resid = info.normres
+    retcode = info.converged == 1 ? :Default : :DidNotConverge
+    iters = info.numiter
+    return SciMLBase.build_linear_solution(alg, cache.u, resid, cache; retcode = retcode, iters = iters)
+end

test/basictests.jl

@@ -161,6 +161,19 @@ end
     end
 end
 
+@testset "KrylovKit" begin
+    kwargs = (;gmres_restart=5,)
+    for alg in (
+                ("Default", KrylovKitJL(kwargs...)),
+                ("CG", KrylovKitJL_CG(kwargs...)),
+                ("GMRES",KrylovKitJL_GMRES(kwargs...)),
+               )
+        @testset "$(alg[1])" begin
+            test_interface(alg[2], prob1, prob2)
+        end
+    end
+end
+
 @testset "PardisoJL" begin
     @test_throws UndefVarError alg = PardisoJL()