Simplify the tests of lsmr and remove the 5-argument mul! stuff since that is not in LinearAlgebra for dense matrices.

Author: haampie

src/lsmr.jl

@@ -107,10 +107,14 @@ function lsmr_method!(log::ConvergenceHistory, x, A, b, v, h, hbar;
     normArs = Tr[]
     conlim > 0 ? ctol = convert(Tr, inv(conlim)) : ctol = zero(Tr)
     # form the first vectors u and v (satisfy  β*u = b,  α*v = A'u)
-    u = mul!(b, A, x, -1, 1)
+    tmp_u = similar(b)
+    tmp_v = similar(v)
+    mul!(tmp_u, A, x)
+    b .-= tmp_u
+    u = b
     β = norm(u)
     u .*= inv(β)
-    mul!(v, adjoint(A), u, 1, 0)
+    mul!(v, adjoint(A), u)
     α = norm(v)
     v .*= inv(α)
 
@@ -158,12 +162,14 @@ function lsmr_method!(log::ConvergenceHistory, x, A, b, v, h, hbar;
         while iter < maxiter
             nextiter!(log,mvps=1)
             iter += 1
-            mul!(u, A, v, 1, -α)
+            mul!(tmp_u, A, v)
+            u .= tmp_u .+ u .* -α
             β = norm(u)
             if β > 0
                 log.mtvps+=1
                 u .*= inv(β)
-                mul!(v, adjoint(A), u, 1, -β)
+                mul!(tmp_v, adjoint(A), u)
+                v .= tmp_v .+ v .* -β
                 α = norm(v)
                 v .*= inv(α)
             end
@@ -278,11 +284,3 @@ function lsmr_method!(log::ConvergenceHistory, x, A, b, v, h, hbar;
     setconv(log, istop ∉ (3, 6, 7))
     x
 end
-
-function LinearAlgebra.mul!(y::AbstractVector, A::StridedVecOrMat, x::AbstractVector, α::Number, β::Number)
-    BLAS.gemm!('N', 'N', convert(eltype(y), α), A, x, convert(eltype(y), β), y)
-end
-
-function LinearAlgebra.mul!(y::AbstractVector, A::Adjoint{F, T}, x::AbstractVector, α::Number, β::Number) where {F, T <: StridedVecOrMat{F}}
-    BLAS.gemm!('T', 'N', convert(eltype(y), α), adjoint(A), x, convert(eltype(y), β), y)
-end

test/lsmr.jl

@@ -4,56 +4,17 @@ using LinearAlgebra
 using Random
 using SparseArrays
 
-import Base: size, eltype, similar, copyto!, fill!, length
+import Base: size, eltype, similar, copyto!, fill!, length, axes, getindex, setindex!
 import LinearAlgebra: norm, mul!, rmul!, lmul!
 
 # Type used in Dampenedtest
-# solve (A'A + diag(v).^2 ) x = b
-# using LSMR in the augmented space A' = [A ; diag(v)] b' = [b; zeros(size(A, 2)]
-mutable struct DampenedVector{Ty, Tx}
-    y::Ty
-    x::Tx
-end
-
-eltype(a::DampenedVector) = promote_type(eltype(a.y), eltype(a.x))
-norm(a::DampenedVector) = sqrt(norm(a.y)^2 + norm(a.x)^2)
-
-function Base.Broadcast.broadcast!(f::Tf, to::DampenedVector, from::DampenedVector, args...) where {Tf}
-    to.x .= f.(from.x, args...)
-    to.y .= f.(from.y, args...)
-    to
-end
-
-function copyto!(a::DampenedVector{Ty, Tx}, b::DampenedVector{Ty, Tx}) where {Ty, Tx}
-    copyto!(a.y, b.y)
-    copyto!(a.x, b.x)
-    a
-end
-
-function fill!(a::DampenedVector, α::Number)
-    fill!(a.y, α)
-    fill!(a.x, α)
-    a
-end
-
-lmul!(α::Number, a::DampenedVector) = rmul!(a, α)
-
-function rmul!(a::DampenedVector, α::Number)
-    rmul!(a.y, α)
-    rmul!(a.x, α)
-    a
-end
-
-similar(a::DampenedVector, T) = DampenedVector(similar(a.y, T), similar(a.x, T))
-length(a::DampenedVector) = length(a.y) + length(a.x)
-
-mutable struct DampenedMatrix{TA, Tx}
+# solve (A'A + diag(v).^2 ) x = A'b
+# using LSMR in the augmented space à = [A ; diag(v)] b̃ = [b; zeros(size(A, 2)]
+struct DampenedMatrix{Tv,TA<:AbstractMatrix{Tv},TD<:AbstractVector{Tv}} <: AbstractMatrix{Tv}
     A::TA
-    diagonal::Tx
+    diagonal::TD
 end
 
-eltype(A::DampenedMatrix) = promote_type(eltype(A.A), eltype(A.diagonal))
-
 function size(A::DampenedMatrix)
     m, n = size(A.A)
     l = length(A.diagonal)
@@ -63,36 +24,23 @@ end
 function size(A::DampenedMatrix, dim::Integer)
     m, n = size(A.A)
     l = length(A.diagonal)
-    dim == 1 ? (m + l) :
-    dim == 2 ? n : 1
+    dim == 1 ? (m + l) : (dim == 2 ? n : 1)
 end
 
-function mul!(b::DampenedVector{Ty, Tx}, mw::DampenedMatrix{TA, Tx}, a::Tx,
-    α::Number, β::Number) where {TA, Tx, Ty}
-    if β != 1.
-        if β == 0.
-            fill!(b, 0.)
-        else
-            rmul!(b, β)
-        end
-    end
-    mul!(b.y, mw.A, a, α, 1.0)
-    map!((z, x, y)-> z + α * x * y, b.x, b.x, a, mw.diagonal)
-    return b
+function mul!(y::AbstractVector{Tv}, mw::DampenedMatrix, x::AbstractVector{Tv}) where {Tv}
+    m₁ = size(mw.A, 1)
+    m₂ = size(mw, 1)
+    mul!(view(y, 1:m₁), mw.A, x)
+    y[m₁+1:m₂] .= mw.diagonal .* x
+    return y
 end
 
-function mul!(b::Tx, mw::Adjoint{DampenedMatrix{TA, Tx}}, a::DampenedVector{Ty, Tx},
-    α::Number, β::Number) where {TA, Tx, Ty}
-    if β != 1.
-        if β == 0.
-            fill!(b, 0.)
-        else
-            rmul!(b, β)
-        end
-    end
-    mul!(b, adjoint(mw.A), a.y, α, 1.0)
-    map!((z, x, y)-> z + α * x * y, b, b, a.x, mw.diagonal)
-    return b
+function mul!(y::AbstractVector, mw::Adjoint{Tv,<:DampenedMatrix}, x::AbstractVector) where {Tv}
+    m₁ = size(mw.parent.A, 1)
+    m₂ = size(mw.parent, 1)
+    mul!(y, adjoint(mw.parent.A), view(x, 1:m₁))
+    y .+= mw.parent.diagonal .* view(x, m₁+1:m₂)
+    return y
 end
 
 """
@@ -138,14 +86,14 @@ end
         @test norm(b - A * x) ≤ 1e-4
     end
 
-    # @testset "Dampened test" for (m, n) = ((10, 10), (20, 10))
-    #     # Test used to make sure A, b can be generic matrix / vector
-    #     b = rand(m)
-    #     A = rand(m, n)
-    #     v = rand(n)
-    #     Adampened = DampenedMatrix(A, v)
-    #     bdampened = DampenedVector(b, zeros(n))
-    #     x, ch = lsmr(Adampened, bdampened, log=true)
-    #     @test norm((A'A + Matrix(Diagonal(v)) .^ 2)x - A'b) ≤ 1e-3
-    # end
+    @testset "Dampened test" for (m, n) = ((10, 10), (20, 10))
+        # Test used to make sure A, b can be generic matrix / vector
+        b = rand(m)
+        A = rand(m, n)
+        v = rand(n)
+        A′ = DampenedMatrix(A, v)
+        b′ = [b; zeros(n)]
+        x, ch = lsmr(A′, b′, log=true)
+        @test norm((A'A + Matrix(Diagonal(v)) .^ 2)x - A'b) ≤ 1e-3
+    end
 end