Add COCG method for complex symmetric linear systems

src/cg.jl

@@ -1,21 +1,22 @@
 import Base: iterate
 using Printf
-export cg, cg!, CGIterable, PCGIterable, cg_iterator!, CGStateVariables
+export cg, cg!, cocg, cocg!, CGIterable, PCGIterable, cg_iterator!, CGStateVariables
 
-mutable struct CGIterable{matT, solT, vecT, numT <: Real}
+mutable struct CGIterable{matT, solT, vecT, numT <: Real, paramT <: Number, dotT <: AbstractDot}
     A::matT
     x::solT
     r::vecT
     c::vecT
     u::vecT
     tol::numT
     residual::numT
-    prev_residual::numT
+    ρ_prev::paramT
     maxiter::Int
     mv_products::Int
+    dotproduct::dotT
 end
 
-mutable struct PCGIterable{precT, matT, solT, vecT, numT <: Real, paramT <: Number}
+mutable struct PCGIterable{precT, matT, solT, vecT, numT <: Real, paramT <: Number, dotT <: AbstractDot}
     Pl::precT
     A::matT
     x::solT
@@ -24,9 +25,10 @@ mutable struct PCGIterable{precT, matT, solT, vecT, numT <: Real, paramT <: Numb
     u::vecT
     tol::numT
     residual::numT
-    ρ::paramT
+    ρ_prev::paramT
     maxiter::Int
     mv_products::Int
+    dotproduct::dotT
 end
 
 @inline converged(it::Union{CGIterable, PCGIterable}) = it.residual ≤ it.tol
@@ -47,18 +49,20 @@ function iterate(it::CGIterable, iteration::Int=start(it))
     end
 
     # u := r + βu (almost an axpy)
-    β = it.residual^2 / it.prev_residual^2
+    ρ = isa(it.dotproduct, ConjugatedDot) ?  it.residual^2 : _norm(it.r, it.dotproduct)^2
+    β = ρ / it.ρ_prev
+
     it.u .= it.r .+ β .* it.u
 
     # c = A * u
     mul!(it.c, it.A, it.u)
-    α = it.residual^2 / dot(it.u, it.c)
+    α = ρ / _dot(it.u, it.c, it.dotproduct)
 
     # Improve solution and residual
+    it.ρ_prev = ρ
     it.x .+= α .* it.u
     it.r .-= α .* it.c
 
-    it.prev_residual = it.residual
     it.residual = norm(it.r)
 
     # Return the residual at item and iteration number as state
@@ -78,18 +82,17 @@ function iterate(it::PCGIterable, iteration::Int=start(it))
     # Apply left preconditioner
     ldiv!(it.c, it.Pl, it.r)
 
-    ρ_prev = it.ρ
-    it.ρ = dot(it.c, it.r)
-
     # u := c + βu (almost an axpy)
-    β = it.ρ / ρ_prev
+    ρ = _dot(it.r, it.c, it.dotproduct)
+    β = ρ / it.ρ_prev
     it.u .= it.c .+ β .* it.u
 
     # c = A * u
     mul!(it.c, it.A, it.u)
-    α = it.ρ / dot(it.u, it.c)
+    α = ρ / _dot(it.u, it.c, it.dotproduct)
 
     # Improve solution and residual
+    it.ρ_prev = ρ
     it.x .+= α .* it.u
     it.r .-= α .* it.c
 
@@ -122,7 +125,8 @@ function cg_iterator!(x, A, b, Pl = Identity();
                       reltol::Real = sqrt(eps(real(eltype(b)))),
                       maxiter::Int = size(A, 2),
                       statevars::CGStateVariables = CGStateVariables(zero(x), similar(x), similar(x)),
-                      initially_zero::Bool = false)
+                      initially_zero::Bool = false,
+                      dotproduct::AbstractDot = ConjugatedDot())
     u = statevars.u
     r = statevars.r
     c = statevars.c
@@ -143,14 +147,12 @@ function cg_iterator!(x, A, b, Pl = Identity();
     # Return the iterable
     if isa(Pl, Identity)
         return CGIterable(A, x, r, c, u,
-            tolerance, residual, one(residual),
-            maxiter, mv_products
-        )
+            tolerance, residual, one(eltype(r)),
+            maxiter, mv_products, dotproduct)
     else
         return PCGIterable(Pl, A, x, r, c, u,
-            tolerance, residual, one(eltype(x)),
-            maxiter, mv_products
-        )
+            tolerance, residual, one(eltype(r)),
+            maxiter, mv_products, dotproduct)
     end
 end
 
@@ -211,6 +213,7 @@ function cg!(x, A, b;
              statevars::CGStateVariables = CGStateVariables(zero(x), similar(x), similar(x)),
              verbose::Bool = false,
              Pl = Identity(),
+             dotproduct::AbstractDot = ConjugatedDot(),
              kwargs...)
     history = ConvergenceHistory(partial = !log)
     history[:abstol] = abstol
@@ -219,7 +222,7 @@ function cg!(x, A, b;
 
     # Actually perform CG
     iterable = cg_iterator!(x, A, b, Pl; abstol = abstol, reltol = reltol, maxiter = maxiter,
-                            statevars = statevars, kwargs...)
+                            statevars = statevars, dotproduct = dotproduct, kwargs...)
     if log
         history.mvps = iterable.mv_products
     end
@@ -237,3 +240,18 @@ function cg!(x, A, b;
 
     log ? (iterable.x, history) : iterable.x
 end
+
+"""
+    cocg(A, b; kwargs...) -> x, [history]
+
+Same as [`cocg!`](@ref), but allocates a solution vector `x` initialized with zeros.
+"""
+cocg(A, b; kwargs...) = cocg!(zerox(A, b), A, b; initially_zero = true, kwargs...)
+
+"""
+    cocg!(x, A, b; kwargs...) -> x, [history]
+
+Same as [`cg!`](@ref), but uses the unconjugated dot product instead of the usual,
+conjugated dot product.
+"""
+cocg!(x, A, b; kwargs...) = cg!(x, A, b; dotproduct = UnconjugatedDot(), kwargs...)

src/common.jl

@@ -1,6 +1,6 @@
 import LinearAlgebra: ldiv!, \
 
-export Identity
+export Identity, ConjugatedDot, UnconjugatedDot
 
 #### Type-handling
 """
@@ -30,3 +30,16 @@ struct Identity end
 \(::Identity, x) = copy(x)
 ldiv!(::Identity, x) = x
 ldiv!(y, ::Identity, x) = copyto!(y, x)
+
+"""
+Conjugated and unconjugated dot products
+"""
+abstract type AbstractDot end
+struct ConjugatedDot <: AbstractDot end
+struct UnconjugatedDot <: AbstractDot end
+
+_norm(x, ::ConjugatedDot) = norm(x)
+_dot(x, y, ::ConjugatedDot) = dot(x, y)
+
+_norm(x, ::UnconjugatedDot) = sqrt(sum(xₖ->xₖ^2, x))
+_dot(x, y, ::UnconjugatedDot) = transpose(@view(x[:])) * @view(y[:])  # allocating, but faster than sum(prod, zip(x,y))

test/cg.jl

@@ -24,15 +24,15 @@ Random.seed!(1234321)
 @testset "Small full system" begin
     n = 10
 
-    @testset "Matrix{$T}" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    @testset "Matrix{$T}, conjugated dot product" for T in (Float32, Float64, ComplexF32, ComplexF64)
         A = rand(T, n, n)
         A = A' * A + I
         b = rand(T, n)
         reltol = √eps(real(T))
 
         x,ch = cg(A, b; reltol=reltol, maxiter=2n, log=true)
         @test isa(ch, ConvergenceHistory)
-        @test norm(A*x - b) / norm(b) ≤ reltol
+        @test A*x ≈ b rtol=reltol
         @test ch.isconverged
 
         # If you start from the exact solution, you should converge immediately
@@ -50,6 +50,32 @@ Random.seed!(1234321)
         x0 = cg(A, zeros(T, n))
         @test x0 == zeros(T, n)
     end
+
+    @testset "Matrix{$T}, unconjugated dot product" for T in (Float32, Float64, ComplexF32, ComplexF64)
+        A = rand(T, n, n)
+        A = A + transpose(A) + 15I
+        x = ones(T, n)
+        b = A * x
+
+        reltol = √eps(real(T))
+
+        # Solve without preconditioner
+        x1, his1 = cocg(A, b, reltol = reltol, maxiter = 100, log = true)
+        @test isa(his1, ConvergenceHistory)
+        @test A*x1 ≈ b rtol=reltol
+
+        # With an initial guess
+        x_guess = rand(T, n)
+        x2, his2 = cocg!(x_guess, A, b, reltol = reltol, maxiter = 100, log = true)
+        @test isa(his2, ConvergenceHistory)
+        @test x2 == x_guess
+        @test A*x2 ≈ b rtol=reltol
+
+        # Do an exact LU decomp of a nearby matrix
+        F = lu(A + rand(T, n, n))
+        x3, his3 = cocg(A, b, Pl = F, maxiter = 100, reltol = reltol, log = true)
+        @test A*x3 ≈ b rtol=reltol
+    end
 end
 
 @testset "Sparse Laplacian" begin
@@ -64,24 +90,24 @@ end
     @testset "SparseMatrixCSC{$T, $Ti}" for T in (Float64, Float32), Ti in (Int64, Int32)
         xCG = cg(A, rhs; reltol=reltol, maxiter=100)
         xJAC = cg(A, rhs; Pl=P, reltol=reltol, maxiter=100)
-        @test norm(A * xCG - rhs) ≤ reltol
-        @test norm(A * xJAC - rhs) ≤ reltol
+        @test A*xCG ≈ rhs rtol=reltol
+        @test A*xJAC ≈ rhs rtol=reltol
     end
 
     Af = LinearMap(A)
     @testset "Function" begin
         xCG = cg(Af, rhs; reltol=reltol, maxiter=100)
         xJAC = cg(Af, rhs; Pl=P, reltol=reltol, maxiter=100)
-        @test norm(A * xCG - rhs) ≤ reltol
-        @test norm(A * xJAC - rhs) ≤ reltol
+        @test A*xCG ≈ rhs rtol=reltol
+        @test A*xJAC ≈ rhs rtol=reltol
     end
 
     @testset "Function with specified starting guess" begin
         x0 = randn(size(rhs))
         xCG, hCG = cg!(copy(x0), Af, rhs; abstol=abstol, reltol=0.0, maxiter=100, log=true)
         xJAC, hJAC = cg!(copy(x0), Af, rhs; Pl=P, abstol=abstol, reltol=0.0, maxiter=100, log=true)
-        @test norm(A * xCG - rhs) ≤ reltol
-        @test norm(A * xJAC - rhs) ≤ reltol
+        @test A*xCG ≈ rhs rtol=reltol
+        @test A*xJAC ≈ rhs rtol=reltol
         @test niters(hJAC) == niters(hCG)
     end
 end

test/common.jl

@@ -27,6 +27,20 @@ end
     @test ldiv!(y, P, copy(x)) == x
 end
 
+@testset "Vector{$T}, conjugated and unconjugated dot products" for T in (ComplexF32, ComplexF64)
+    n = 100
+    x = rand(T, n)
+    y = rand(T, n)
+
+    # Conjugated dot product
+    @test IterativeSolvers._norm(x, ConjugatedDot()) ≈ sqrt(x'x)
+    @test IterativeSolvers._dot(x, y, ConjugatedDot()) ≈ x'y
+
+    # Unonjugated dot product
+    @test IterativeSolvers._norm(x, UnconjugatedDot()) ≈ sqrt(transpose(x) * x)
+    @test IterativeSolvers._dot(x, y, UnconjugatedDot()) ≈ transpose(x) * y
+end
+
 end
 
 DocMeta.setdocmeta!(IterativeSolvers, :DocTestSetup, :(using IterativeSolvers); recursive=true)