Fix _unsafe_mul! for custom maps (#110)

Author: dkarrasch

src/LinearMaps.jl

@@ -93,9 +93,10 @@ julia> A*x
 ```
 """
 function Base.:(*)(A::LinearMap, x::AbstractVector)
-    size(A, 2) == length(x) || throw(DimensionMismatch("linear map has dimensions ($mA,$nA), " *
-        "vector has length $mB"))
-    return _unsafe_mul!(similar(x, promote_type(eltype(A), eltype(x)), size(A, 1)), A, x)
+    m, n = size(A)
+    n == length(x) || throw(DimensionMismatch("linear map has dimensions ($m,$n), " *
+        "vector has length $(length(x))"))
+    return mul!(similar(x, promote_type(eltype(A), eltype(x)), m), A, x)
 end
 
 """
@@ -163,21 +164,23 @@ function mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α::Number,
 end
 
 function _generic_mapvec_mul!(y, A, x, α, β)
+    # this function needs to call mul! for, e.g.,  AdjointMap{...,<:CustomMap}
     if isone(α)
-        iszero(β) && (_unsafe_mul!(y, A, x); return y)
-        isone(β) && (y .+= A * x; return y)
-        # β != 0, 1
+        iszero(β) && return mul!(y, A, x)
         z = A * x
-        y .= y.*β .+ z
+        if isone(β)
+            y .+= z
+        else
+            y .= y.*β .+ z
+        end
         return y
     elseif iszero(α)
-        iszero(β) && (fill!(y, zero(eltype(y))); return y)
+        iszero(β) && return fill!(y, zero(eltype(y)))
         isone(β) && return y
         # β != 0, 1
-        rmul!(y, β)
-        return y
+        return rmul!(y, β)
     else # α != 0, 1
-        iszero(β) && (_unsafe_mul!(y, A, x); rmul!(y, α); return y)
+        iszero(β) && return rmul!(mul!(y, A, x), α)
         z = A * x
         if isone(β)
             y .+= z .* α

src/functionmap.jl

@@ -71,7 +71,7 @@ function Base.:(*)(A::AdjointMap{<:Any,<:FunctionMap}, x::AbstractVector)
         conj!(y)
         return y
     else
-        error("adjoint not implemented for $A")
+        error("adjoint not implemented for $(A.lmap)")
     end
 end
 function Base.:(*)(A::TransposeMap{<:Any,<:FunctionMap}, x::AbstractVector)
@@ -103,7 +103,7 @@ function Base.:(*)(A::TransposeMap{<:Any,<:FunctionMap}, x::AbstractVector)
         conj!(y)
         return y
     else
-        error("transpose not implemented for $A")
+        error("transpose not implemented for $(A.lmap)")
     end
 end
 
@@ -129,7 +129,7 @@ function _unsafe_mul!(y::AbstractVecOrMat, At::TransposeMap{<:Any,<:FunctionMap}
         conj!(y)
         return y
     else
-        error("transpose not implemented for $A")
+        error("transpose not implemented for $(A.lmap)")
     end
 end
 
@@ -144,6 +144,6 @@ function _unsafe_mul!(y::AbstractVecOrMat, Ac::AdjointMap{<:Any,<:FunctionMap},
         conj!(y)
         return y
     else
-        error("adjoint not implemented for $A")
+        error("adjoint not implemented for $(A.lmap)")
     end
 end

src/transpose.jl

@@ -51,7 +51,7 @@ Base.:(==)(A::LinearMap, B::AdjointMap)      = ishermitian(A) && B.lmap == A
 # multiplication with vector/matrices
 # # TransposeMap
 function _unsafe_mul!(y::AbstractVecOrMat, A::TransposeMap, x::AbstractVector)
-    issymmetric(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : error("transpose not implemented for $A")
+    issymmetric(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : error("transpose not implemented for $(A.lmap)")
 end
 function _unsafe_mul!(y::AbstractMatrix, A::TransposeMap, x::AbstractMatrix)
     issymmetric(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : _generic_mapmat_mul!(y, A, x)
@@ -64,7 +64,7 @@ function _unsafe_mul!(y::AbstractMatrix, A::TransposeMap, x::AbstractMatrix, α:
 end
 # # AdjointMap
 function _unsafe_mul!(y::AbstractVecOrMat, A::AdjointMap, x::AbstractVector)
-    ishermitian(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : error("adjoint not implemented for $A")
+    ishermitian(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : error("adjoint not implemented for $(A.lmap)")
 end
 function _unsafe_mul!(y::AbstractMatrix, A::AdjointMap, x::AbstractMatrix)
     ishermitian(A.lmap) ? _unsafe_mul!(y, A.lmap, x) : _generic_mapmat_mul!(y, A, x)

src/wrappedmap.jl

@@ -37,8 +37,6 @@ LinearAlgebra.ishermitian(A::WrappedMap) = A._ishermitian
 LinearAlgebra.isposdef(A::WrappedMap) = A._isposdef
 
 # multiplication with vectors & matrices
-Base.:(*)(A::WrappedMap, x::AbstractVector) = *(A.lmap, x)
-
 for (intype, outtype) in ((AbstractVector, AbstractVecOrMat), (AbstractMatrix, AbstractMatrix))
     @eval begin
         _unsafe_mul!(y::$outtype, A::WrappedMap, x::$intype) = _unsafe_mul!(y, A.lmap, x)

test/kronecker.jl

@@ -1,6 +1,6 @@
 using Test, LinearMaps, LinearAlgebra, SparseArrays
 
-@testset "kronecker products" begin
+@testset "Kronecker products and sums" begin
     @testset "Kronecker product" begin
         A = rand(ComplexF64, 3, 3)
         B = rand(ComplexF64, 2, 2)
@@ -70,6 +70,7 @@ using Test, LinearMaps, LinearAlgebra, SparseArrays
             @test_throws ArgumentError kronsum(LA, [B B]) # non-square map
             KSmat = kron(A, Matrix(I, 2, 2)) + kron(Matrix(I, 3, 3), B)
             @test Matrix(KS) ≈ Matrix(kron(A, LinearMap(I, 2)) + kron(LinearMap(I, 3), B))
+            @test KS * ones(size(KS, 2)) ≈ KSmat * ones(size(KS, 2))
             @test size(KS) == size(kron(A, Matrix(I, 2, 2)))
             for transform in (identity, transpose, adjoint)
                 @test Matrix(transform(KS)) ≈ transform(Matrix(KS)) ≈ transform(KSmat)

test/linearmaps.jl

@@ -99,3 +99,37 @@ LinearAlgebra.mul!(y::AbstractVector, A::Union{SimpleFunctionMap,SimpleComplexFu
     @test SparseMatrixCSC(F) == Fs == L
     @test Fs isa SparseMatrixCSC
 end
+
+struct MyFillMap{T} <: LinearMaps.LinearMap{T}
+    λ::T
+    size::Dims{2}
+    function MyFillMap(λ::T, dims::Dims{2}) where {T}
+        all(d -> d >= 0, dims) || throw(ArgumentError("dims of MyFillMap must be non-negative"))
+        promote_type(T, typeof(λ)) == T || throw(InexactError())
+        return new{T}(λ, dims)
+    end
+end
+Base.size(A::MyFillMap) = A.size
+function LinearAlgebra.mul!(y::AbstractVecOrMat, A::MyFillMap, x::AbstractVector)
+    LinearMaps.check_dim_mul(y, A, x)
+    return fill!(y, iszero(A.λ) ? zero(eltype(y)) : A.λ*sum(x))
+end
+function LinearAlgebra.mul!(y::AbstractVecOrMat, transA::LinearMaps.TransposeMap{<:Any,<:MyFillMap}, x::AbstractVector)
+    LinearMaps.check_dim_mul(y, transA, x)
+    λ = transA.lmap.λ
+    return fill!(y, iszero(λ) ? zero(eltype(y)) : transpose(λ)*sum(x))
+end
+
+@testset "transpose of new LinearMap type" begin
+    A = MyFillMap(5.0, (3, 3))
+    x = ones(3)
+    @test A * x == fill(15.0, 3)
+    @test mul!(zeros(3), A, x) == mul!(zeros(3), A, x, 1, 0) == fill(15.0, 3)
+    @test mul!(ones(3), A, x, 2, 2) == fill(32.0, 3)
+    @test mul!(ones(3,3), A, reshape(collect(1:9), 3, 3), 2, 2)[:,1] == fill(62.0, 3)
+    @test A'x == mul!(zeros(3), A', x)
+    for α in (true, false, 0, 1, randn()), β in (true, false, 0, 1, randn())
+        @test mul!(ones(3), A', x, α, β) == fill(β, 3) + fill(15α, 3)
+        @test mul!(ones(3, 2), A', [x x], α, β) == fill(β, 3, 2) + fill(15α, 3, 2)
+    end
+end