make preconditioners part of the solver rather than a random extra

Author: oscardssmith

Project.toml

@@ -133,6 +133,7 @@ IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
+KrylovPreconditioners = "45d422c2-293f-44ce-8315-2cb988662dec"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 MultiFloats = "bdf0d083-296b-4888-a5b6-7498122e68a5"
@@ -146,4 +147,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays", "AllocCheck", "StableRNGs", "Zygote"]
+test = ["Aqua", "Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "KrylovPreconditioners", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices", "StaticArrays", "AllocCheck", "StableRNGs", "Zygote"]

src/common.jl

@@ -121,6 +121,8 @@ default_alias_b(::Any, ::Any, ::Any) = false
 default_alias_A(::AbstractKrylovSubspaceMethod, ::Any, ::Any) = true
 default_alias_b(::AbstractKrylovSubspaceMethod, ::Any, ::Any) = true
 
+DEFAULT_PRECS(A, p) = IdentityOperator(size(A)[1]), IdentityOperator(size(A)[2])
+
 function __init_u0_from_Ab(A, b)
     u0 = similar(b, size(A, 2))
     fill!(u0, false)
@@ -136,12 +138,12 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
         reltol = default_tol(real(eltype(prob.b))),
         maxiters::Int = length(prob.b),
         verbose::Bool = false,
-        Pl = IdentityOperator(size(prob.A)[1]),
-        Pr = IdentityOperator(size(prob.A)[2]),
+        Pl = nothing,
+        Pr = nothing,
         assumptions = OperatorAssumptions(issquare(prob.A)),
         sensealg = LinearSolveAdjoint(),
         kwargs...)
-    @unpack A, b, u0, p = prob
+    (;A, b, u0, p) = prob
 
     A = if alias_A || A isa SMatrix
         A
@@ -167,6 +169,18 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
     reltol = real(eltype(prob.b))(reltol)
     abstol = real(eltype(prob.b))(abstol)
 
+    precs = hasproperty(alg, :precs) ? alg.precs : DEFAULT_PRECS
+    _Pl, _Pr = precs(A, p)
+    if isnothing(Pl)
+        Pl = _Pl
+    else
+        @warn "passing Preconditioners at `init`/`solve` time is deprecated. Instead add a `precs` function to your algorithm."
+    end
+    if isnothing(Pr)
+        Pr = _Pr
+    else
+        @warn "passing Preconditioners at `init`/`solve` time is deprecated. Instead add a `precs` function to your algorithm."
+    end
     cacheval = init_cacheval(alg, A, b, u0_, Pl, Pr, maxiters, abstol, reltol, verbose,
         assumptions)
     isfresh = true
@@ -179,6 +193,33 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
     return cache
 end
 
+
+function SciMLBase.reinit!(cache::LinearCache;
+                           A = nothing,
+                           b = cache.b,
+                           u = cache.u,
+                           p = nothing,)
+    (; alg, cacheval, isfresh, abstol, reltol, maxiters, verbose, assumptions, sensealg) = cache
+
+    precs = hasproperty(alg, :precs) ? alg.precs : DEFAULT_PRECS
+    Pl, Pr = if isnothing(A) || isnothing(p)
+        (cache.Pl, cache.Pr)
+    else
+        if isnothing(A)
+            A = cache.A
+        end
+        if isnothing(p)
+            p = cache.p
+        end
+        precs(A, p)
+    end
+
+    return LinearCache{typeof(A), typeof(b), typeof(u), typeof(p), typeof(alg), typeof(cacheval),
+        typeof(Pl), typeof(Pr), typeof(reltol), typeof(assumptions.issq),
+        typeof(sensealg)}(A, b, u, p, alg, cacheval, isfresh, Pl, Pr, abstol, reltol,
+        maxiters, verbose, assumptions, sensealg)
+end
+
 function SciMLBase.solve(prob::LinearProblem, args...; kwargs...)
     return solve(prob, nothing, args...; kwargs...)
 end

src/iterative_wrappers.jl

@@ -10,19 +10,21 @@ KrylovJL(args...; KrylovAlg = Krylov.gmres!,
 
 A generic wrapper over the Krylov.jl krylov-subspace iterative solvers.
 """
-struct KrylovJL{F, I, A, K} <: AbstractKrylovSubspaceMethod
+struct KrylovJL{F, I, A, P, K} <: AbstractKrylovSubspaceMethod
     KrylovAlg::F
     gmres_restart::I
     window::I
+    precs::P
     args::A
     kwargs::K
 end
 
 function KrylovJL(args...; KrylovAlg = Krylov.gmres!,
         gmres_restart = 0, window = 0,
+        precs = DEFAULT_PRECS,
         kwargs...)
     return KrylovJL(KrylovAlg, gmres_restart, window,
-        args, kwargs)
+        precs, args, kwargs)
 end
 
 default_alias_A(::KrylovJL, ::Any, ::Any) = true
@@ -231,8 +233,7 @@ function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
         cache.isfresh = false
     end
 
-    M = cache.Pl
-    N = cache.Pr
+    M, N = cache.Pl, cache.Pr
 
     # use no-op preconditioner for Krylov.jl (LinearAlgebra.I) when M/N is identity
     M = _isidentity_struct(M) ? I : M
@@ -258,25 +259,21 @@ function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
     end
 
     args = (cacheval, cache.A, cache.b)
-    kwargs = (atol = atol, rtol = rtol, itmax = itmax, verbose = verbose,
+    kwargs = (atol = atol, rtol, itmax, verbose,
         ldiv = true, history = true, alg.kwargs...)
 
     if cache.cacheval isa Krylov.CgSolver
         N !== I &&
             @warn "$(alg.KrylovAlg) doesn't support right preconditioning."
-        Krylov.solve!(args...; M = M,
-            kwargs...)
+        Krylov.solve!(args...; M, kwargs...)
     elseif cache.cacheval isa Krylov.GmresSolver
-        Krylov.solve!(args...; M = M, N = N, restart = alg.gmres_restart > 0,
-            kwargs...)
+        Krylov.solve!(args...; M, N, restart = alg.gmres_restart > 0, kwargs...)
     elseif cache.cacheval isa Krylov.BicgstabSolver
-        Krylov.solve!(args...; M = M, N = N,
-            kwargs...)
+        Krylov.solve!(args...; M, N, kwargs...)
     elseif cache.cacheval isa Krylov.MinresSolver
         N !== I &&
             @warn "$(alg.KrylovAlg) doesn't support right preconditioning."
-        Krylov.solve!(args...; M = M,
-            kwargs...)
+        Krylov.solve!(args...; M, kwargs...)
     else
         Krylov.solve!(args...; kwargs...)
     end