initial commit

Author: vpuri3

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"

src/LinearSolve.jl

@@ -8,6 +8,7 @@ using SparseArrays
 using SciMLBase: AbstractDiffEqOperator, AbstractLinearAlgorithm
 using Setfield
 using UnPack
+using FastBroadcast
 
 # wrap
 import Krylov

src/common.jl

@@ -44,6 +44,12 @@ function set_cacheval(cache::LinearCache, alg_cache)
     return cache
 end
 
+function set_prec(cache, Pl, Pr)
+    @set! cache.Pl = Pl
+    @set! cache.Pr = Pr
+    return cache
+end
+
 init_cacheval(alg::Union{SciMLLinearSolveAlgorithm,Nothing}, A, b, u) = nothing
 
 SciMLBase.init(prob::LinearProblem, args...; kwargs...) = SciMLBase.init(prob,nothing,args...;kwargs...)

src/wrappers.jl

@@ -1,18 +1,29 @@
 
-#TODO: composed preconditioners, preconditioner setter for cache,
-#   detailed tests for wrappers
-
 ## Preconditioners
 
 struct ScaleVector{T}
     s::T
     isleft::Bool
 end
 
-function LinearAlgebra.ldiv!(v::ScaleVector, x)
+function LinearAlgebra.ldiv!(P::ScaleVector, x)
+    P.s == one(eltype(P.s)) && return x
+
+    if P.isleft
+        @. x = x * P.s # @.. doesnt speed up scalar computation
+    else
+        @. x = x / P.s
+    end
 end
 
-function LinearAlgebra.ldiv!(y, v::ScaleVector, x)
+function LinearAlgebra.ldiv!(y, P::ScaleVector, x)
+    P.s == one(eltype(P.s)) && return y = x
+
+    if P.isleft
+        @. y = x / P.s
+    else
+        @. y = x * P.s
+    end
 end
 
 struct ComposePreconditioner{Ti,To}
@@ -21,12 +32,19 @@ struct ComposePreconditioner{Ti,To}
     isleft::Bool
 end
 
-function LinearAlgebra.ldiv!(v::ComposePreconditioner, x)
-    @unpack inner, outer, isleft = v
+function LinearAlgebra.ldiv!(P::ComposePreconditioner, x)
+    @unpack inner, outer, isleft = P
+    if isleft
+        ldiv!(outer, x)
+        ldiv!(inner, x)
+    else
+        ldiv!(inner, x)
+        ldiv!(outer, x)
+    end
 end
 
-function LinearAlgebra.ldiv!(y, v::ComposePreconditioner, x)
-    @unpack inner, outer, isleft = v
+function LinearAlgebra.ldiv!(y, P::ComposePreconditioner, x)
+    @unpack inner, outer, isleft = P
 end
 
 ## Krylov.jl
@@ -132,6 +150,9 @@ function SciMLBase.solve(cache::LinearCache, alg::KrylovJL; kwargs...)
         cache = set_cacheval(cache, solver)
     end
 
+#   Pl = ComposePreconditioner(alg.Pl, cache.Pl, true)
+#   Pr = ComposePreconditioner(alg.Pr, cache.Pr, false)
+
     atol    = cache.abstol
     rtol    = cache.reltol
     itmax   = cache.maxiters
@@ -204,8 +225,11 @@ IterativeSolversJL_MINRES(args...;kwargs...) =
 function init_cacheval(alg::IterativeSolversJL, cache::LinearCache)
     @unpack A, b, u = cache
 
-    Pl = (alg.Pl == LinearAlgebra.I) ? IterativeSolvers.Identity() : alg.Pl
-    Pr = (alg.Pr == LinearAlgebra.I) ? IterativeSolvers.Identity() : alg.Pr
+    Pl = alg.Pl
+    Pr = alg.Pr
+
+#   Pl = ComposePreconditioner(alg.Pl, cache.Pl, true)
+#   Pr = ComposePreconditioner(alg.Pr, cache.Pr, false)
 
     abstol  = cache.abstol
     reltol  = cache.reltol