From 30b4a2234ccef7dcc43265a25c6f06d0f0761e71 Mon Sep 17 00:00:00 2001
From: Abel <abel.s.siqueira@gmail.com>
Date: Wed, 16 Oct 2019 17:56:36 -0300
Subject: [PATCH] Update variants to not convert unless necessary

Closes #126
---
 src/variants.jl       | 65 ++++++++++++++++++++++++++++++++++---------
 test/test_cg.jl       |  7 +++++
 test/test_variants.jl |  4 +--
 3 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/src/variants.jl b/src/variants.jl
index e3bed7506..92d456667 100644
--- a/src/variants.jl
+++ b/src/variants.jl
@@ -1,41 +1,80 @@
 using LinearAlgebra
 
 # Define a generic linear operator with preallocation
-function preallocated_LinearOperator(A)
+function preallocated_LinearOperator(A :: AbstractMatrix, T)
   (n, m) = size(A)
-  Ap = zeros(n)
-  Atq = zeros(m)
-  return LinearOperator(n, m, false, false, p -> mul!(Ap, A, p),
+  Ap = zeros(T, n)
+  Atq = zeros(T, m)
+  return LinearOperator(T, n, m, false, false, p -> mul!(Ap, A, p),
                         q -> mul!(Atq, transpose(A), q), q -> mul!(Atq, transpose(A), q))
 end
 
 # Variants where A is a matrix without specific properties
 for fn in (:cgls, :cgne, :craig, :craigmr, :crls, :crmr, :lslq, :lsmr, :lsqr, :dqgmres, :diom, :cgs, :bilq)
   @eval begin
-    $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number} =
-      $fn(preallocated_LinearOperator(A), convert(Vector{Float64}, b), args...; kwargs...)
+    function $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number}
+      S = promote_type(TA, Tb)
+      if S <: Integer || S <: Complex{<: Integer}
+        S = promote_type(S, Float64)
+      end
+      $fn(preallocated_LinearOperator(A, S), convert(Vector{S}, b), args...; kwargs...)
+    end
+
+    $fn(A :: AbstractMatrix{T}, b :: SparseVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), convert(Vector{T}, b), args...; kwargs...)
+
+    $fn(A :: AbstractMatrix{T}, b :: AbstractVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), b, args...; kwargs...)
   end
 end
 
 # Variants for USYMQR
 for fn in (:usymqr,)
   @eval begin
-    $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, c :: AbstractVector{Tc}, args...; kwargs...) where {TA, Tb, Tc <: Number} =
-      $fn(preallocated_LinearOperator(A), convert(Vector{Float64}, b), convert(Vector{Float64}, c), args...; kwargs...)
+    function $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, c :: AbstractVector{Tc}, args...; kwargs...) where {TA, Tb, Tc <: Number}
+      S = promote_type(TA, Tb, Tc)
+      if S <: Integer || S <: Complex{<: Integer}
+        S = promote_type(S, Float64)
+      end
+      $fn(preallocated_LinearOperator(A, S), convert(Vector{S}, b), convert(Vector{S}, c), args...; kwargs...)
+    end
+
+    $fn(A :: AbstractMatrix{T}, b :: SparseVector{T}, c :: SparseVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), convert(Vector{T}, b), convert(Vector{T}, c), args...; kwargs...)
+
+    $fn(A :: AbstractMatrix{T}, b :: AbstractVector{T}, c :: SparseVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), b, convert(Vector{T}, c), args...; kwargs...)
+
+    $fn(A :: AbstractMatrix{T}, b :: SparseVector{T}, c :: AbstractVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), convert(Vector{T}, b), c, args...; kwargs...)
+
+    $fn(A :: AbstractMatrix{T}, b :: AbstractVector{T}, c :: AbstractVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_LinearOperator(A, T), b, c, args...; kwargs...)
   end
 end
 
 # Define a symmetric linear operator with preallocation
-function preallocated_symmetric_LinearOperator(A)
+function preallocated_symmetric_LinearOperator(A :: AbstractMatrix, T)
   (n, m) = size(A)
-  Ap = zeros(n)
-  return LinearOperator(n, m, true, true, p -> mul!(Ap, A, p))
+  Ap = zeros(T, n)
+  return LinearOperator(T, n, m, true, true, p -> mul!(Ap, A, p))
 end
 
 # Variants where A must be symmetric
 for fn in (:cg_lanczos, :cg_lanczos_shift_seq, :cg, :cr, :minres, :minres_qlp, :symmlq)
   @eval begin
-    $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number} =
-      $fn(preallocated_symmetric_LinearOperator(A), convert(Vector{Float64}, b), args...; kwargs...)
+    function $fn(A :: AbstractMatrix{TA}, b :: AbstractVector{Tb}, args...; kwargs...) where {TA, Tb <: Number}
+      S = promote_type(TA, Tb)
+      if S <: Integer || S <: Complex{<: Integer}
+        S = promote_type(S, Float64)
+      end
+      $fn(preallocated_symmetric_LinearOperator(A, S), convert(Vector{S}, b), args...; kwargs...)
+    end
+
+    $fn(A :: AbstractMatrix{T}, b :: SparseVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_symmetric_LinearOperator(A, T), convert(Vector{T}, b), args...; kwargs...)
+
+    $fn(A :: AbstractMatrix{T}, b :: AbstractVector{T}, args...; kwargs...) where T <: AbstractFloat =
+      $fn(preallocated_symmetric_LinearOperator(A, T), b, args...; kwargs...)
   end
 end
diff --git a/test/test_cg.jl b/test/test_cg.jl
index 6b234f57f..677418420 100644
--- a/test/test_cg.jl
+++ b/test/test_cg.jl
@@ -70,6 +70,13 @@ function test_cg()
   @printf("CG: Relative residual: %8.1e\n", resid);
   @test(resid <= cg_tol);
   @test(stats.solved);
+
+  # Test that precision is not lost (#126)
+  A = rand(BigFloat, 3, 3)
+  A = A * A'
+  b = rand(BigFloat, 3)
+  x = cg(A, b)[1]
+  @test eltype(x) == BigFloat
 end
 
 test_cg()
diff --git a/test/test_variants.jl b/test/test_variants.jl
index 76a9792e1..9ca1806b0 100644
--- a/test/test_variants.jl
+++ b/test/test_variants.jl
@@ -6,9 +6,9 @@ function test_variants()
              :lsqr, :minres, :symmlq, :dqgmres, :diom, :cgs, :usymqr,
              :minres_qlp)
     @printf("%s ", string(fn))
-    for TA in (Int32, Int64, Float32, Float64)
+    for TA in (Int32, Int64, Float32, Float64, BigFloat)
       for IA in (Int32, Int64)
-        for Tb in (Int32, Int64, Float32, Float64)
+        for Tb in (Int32, Int64, Float32, Float64, BigFloat)
           for Ib in (Int32, Int64)
             A_dense = Matrix{TA}(I, 5, 5)
             A_sparse = convert(SparseMatrixCSC{TA,IA}, A_dense)