in-place

Author: vpuri3

src/common.jl

@@ -1,12 +1,14 @@
-struct LinearCache{TA,Tb,Tp,Talg,Tc,Tr,Tl}
+struct LinearCache{TA,Tb,Tu,Tp,Talg,Tc,Tr,Tl}
     A::TA
     b::Tb
+    u::Tu
     p::Tp
     alg::Talg
     cacheval::Tc
     isfresh::Bool
     Pr::Tr
     Pl::Tl
+#   k::Tk # iteration count
 end
 
 function set_A(cache, A)
@@ -39,23 +41,28 @@ function SciMLBase.init(
     alias_b = false,
     kwargs...,
 )
-    @unpack A, b, p = prob
+    @unpack A, b, u0, p = prob
     if alg isa LUFactorization
         fact = lu_instance(A)
         Tfact = typeof(fact)
     else
         fact = nothing
         Tfact = Any
     end
-    Pr = nothing
-    Pl = nothing
+    Pr = LinearAlgebra.I
+    Pl = LinearAlgebra.I
 
     A = alias_A ? A : copy(A)
     b = alias_b ? b : copy(b)
 
+    if u0 == nothing
+        u0 = zero(b)
+    end
+
     cache = LinearCache{
         typeof(A),
         typeof(b),
+        typeof(u0),
         typeof(p),
         typeof(alg),
         Tfact,
@@ -64,6 +71,7 @@ function SciMLBase.init(
     }(
         A,
         b,
+        u0,
         p,
         alg,
         fact,

src/factorization.jl

@@ -16,7 +16,7 @@ function SciMLBase.solve(cache::LinearCache, alg::LUFactorization)
     cache.A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("LU is not defined for $(typeof(prob.A))")
     cache = set_cacheval(cache, lu!(cache.A, alg.pivot))
-    ldiv!(cache.cacheval, cache.b)
+    ldiv!(cache.u,cache.cacheval, cache.b)
 end
 
 struct QRFactorization{P} <: SciMLLinearSolveAlgorithm
@@ -40,7 +40,7 @@ function SciMLBase.solve(cache::LinearCache, alg::QRFactorization)
         cache,
         qr!(cache.A.A, alg.pivot; blocksize = alg.blocksize),
     )
-    ldiv!(cache.cacheval, cache.b)
+    ldiv!(cache.u,cache.cacheval, cache.b)
 end
 
 struct SVDFactorization{A} <: SciMLLinearSolveAlgorithm
@@ -54,5 +54,5 @@ function SciMLBase.solve(cache::LinearCache, alg::SVDFactorization)
     cache.A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("SVD is not defined for $(typeof(prob.A))")
     cache = set_cacheval(cache, svd!(cache.A; full = alg.full, alg = alg.alg))
-    ldiv!(cache.cacheval, cache.b)
+    ldiv!(cache.u,cache.cacheval, cache.b)
 end

src/krylov.jl

@@ -9,9 +9,9 @@ function KrylovJL(args...; solver = gmres, kwargs...)
 end
 
 function SciMLBase.solve(cache::LinearCache, alg::KrylovJL,args...;kwargs...)
-    @unpack A, b, Pl,Pr = cache
-    x, stats = alg.solver(A, b, args...; M=Pl, N=Pr, kwargs...)
-    resid = A * x - b
+    @unpack A, b, u, Pr, Pl = cache
+    u, stats = alg.solver(A, b, args...; M=Pl, N=Pr, kwargs...)
+    resid = A * u - b
     retcode = stats.solved ? :Success : :Failure
-    return x #SciMLBase.build_solution(prob, alg, x, resid; retcode = retcode)
+    return u #SciMLBase.build_solution(prob, alg, x, resid; retcode = retcode)
 end

test/runtests.jl

@@ -20,5 +20,8 @@ using Test
     # make algorithm callable - interoperable with DiffEq ecosystem
     @test A * LUFactorization()(x,A,b) ≈ b 
     @test A * KrylovJL()(x,A,b) ≈ b 
-    @test_broken A * x ≈ b # in place
+
+    # in place
+    KrylovJL()(x,A,b)
+    @test A * x ≈ b
 end