Implementing cg with axpy! and axpby!

Author: fverdugo

PartitionedSolvers/src/krylov_solvers.jl

@@ -206,7 +206,8 @@ function pcg_step(x,ws,b,phase=:start;kwargs...)
         phase = :advance
         copyto!(r,b)
         mul!(s,A,x)
-        r .-= s
+        axpy!(-one(eltype(s)),s,r)
+        #r .-= s
         current = norm(r)
         target = max(reltol*current,abstol)
         dx .= zero(eltype(dx))
@@ -218,12 +219,15 @@ function pcg_step(x,ws,b,phase=:start;kwargs...)
     ρ_prev = ρ
     ρ = dot(r,u)
     β = ρ / ρ_prev
-    dx .= u .+ β .* dx
+    axpby!(one(eltype(u)),u,β,dx)
+    #dx .= u .+ β .* dx
     s = u
     mul!(s,A,dx)
     α = ρ / dot(s,dx)
-    x .+= α .* dx
-    r .-= α .* s
+    axpy!(α,dx,x)
+    #x .+= α .* dx
+    axpy!(-α,s,r)
+    #r .-= α .* s
     current = norm(r)
     iteration += 1
     current = norm(r)

PartitionedSolvers/test/krylov_solvers_tests.jl

@@ -25,6 +25,23 @@ s = PS.solve(s)
 
 z = IterativeSolvers.cg(A,b;verbose=true,Pl)
 
+parts_per_dir = (2,2,2)
+ranks = PA.DebugArray(LinearIndices((prod(parts_per_dir),)))
 
+args = PA.laplacian_fem(nodes_per_dir,parts_per_dir,ranks)
+A = PA.psparse(args...) |> fetch
+x = PA.pones(axes(A,2))
+b = A*x
+
+y = similar(x)
+y .= 0
+p = PS.linear_problem(y,A,b)
+Pl = PS.preconditioner(PS.amg,p)
+s = PS.pcg(p;verbose=true,Pl)
+s = PS.solve(s)
+@test x ≈ PS.solution(p)
+s = PS.update(s,matrix=2*A)
+s = PS.solve(s)
+@test x ≈ 2*PS.solution(p)
 
 end # module

src/p_vector.jl

@@ -1524,3 +1524,14 @@ function renumber(a::PVector,row_partition_2;renumber_local_indices=Val(true))
     PVector(values,row_partition_2)
 end
 
+function LinearAlgebra.axpy!(a::Number,x::PVector,y::PVector)
+    y .+= a .* x
+    y
+end
+
+function LinearAlgebra.axpby!(a::Number,x::PVector,b::Number,y::PVector)
+    y .= a .* x .+ b .* y
+    y
+end
+
+