add mulstyle feature (#76)

* refactor linear combination multiplication

* make LinearCombination multiplication recursive

Author: dkarrasch

src/LinearMaps.jl

@@ -19,6 +19,23 @@ const MapOrMatrix{T} = Union{LinearMap{T},AbstractMatrix{T}}
 
 Base.eltype(::LinearMap{T}) where {T} = T
 
+abstract type MulStyle end
+
+struct FiveArg <: MulStyle end
+struct ThreeArg <: MulStyle end
+
+MulStyle(::FiveArg, ::FiveArg) = FiveArg()
+MulStyle(::ThreeArg, ::FiveArg) = ThreeArg()
+MulStyle(::FiveArg, ::ThreeArg) = ThreeArg()
+MulStyle(::ThreeArg, ::ThreeArg) = ThreeArg()
+MulStyle(::LinearMap) = ThreeArg() # default
+@static if VERSION ≥ v"1.3.0-alpha.115"
+    MulStyle(::AbstractMatrix) = FiveArg()
+else
+    MulStyle(::AbstractMatrix) = ThreeArg()
+end
+MulStyle(A::LinearMap, As::LinearMap...) = MulStyle(MulStyle(A), MulStyle(As...))
+
 Base.isreal(A::LinearMap) = eltype(A) <: Real
 LinearAlgebra.issymmetric(::LinearMap) = false # default assumptions
 LinearAlgebra.ishermitian(A::LinearMap{<:Real}) = issymmetric(A)

src/blockmap.jl

@@ -14,6 +14,8 @@ end
 
 BlockMap{T}(maps::As, rows::S) where {T,As<:Tuple{Vararg{LinearMap}},S} = BlockMap{T,As,S}(maps, rows)
 
+MulStyle(A::BlockMap) = MulStyle(A.maps...)
+
 function check_dim(A::LinearMap, dim, n)
     n == size(A, dim) || throw(DimensionMismatch("Expected $n, got $(size(A, dim))"))
     return nothing

src/linearcombination.jl

@@ -13,6 +13,8 @@ end
 
 LinearCombination{T}(maps::As) where {T, As} = LinearCombination{T, As}(maps)
 
+MulStyle(A::LinearCombination) = MulStyle(A.maps...)
+
 # basic methods
 Base.size(A::LinearCombination) = size(A.maps[1])
 LinearAlgebra.issymmetric(A::LinearCombination) = all(issymmetric, A.maps) # sufficient but not necessary
@@ -61,84 +63,50 @@ Base.:(==)(A::LinearCombination, B::LinearCombination) = (eltype(A) == eltype(B)
 LinearAlgebra.transpose(A::LinearCombination) = LinearCombination{eltype(A)}(map(transpose, A.maps))
 LinearAlgebra.adjoint(A::LinearCombination)   = LinearCombination{eltype(A)}(map(adjoint, A.maps))
 
-# multiplication with vectors
-if VERSION < v"1.3.0-alpha.115"
-
-function A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector)
-    # no size checking, will be done by individual maps
-    A_mul_B!(y, A.maps[1], x)
-    l = length(A.maps)
-    if l>1
-        z = similar(y)
-        for n in 2:l
-            A_mul_B!(z, A.maps[n], x)
-            y .+= z
+# multiplication with vectors & matrices
+for Atype in (AbstractVector, AbstractMatrix)
+    @eval Base.@propagate_inbounds function LinearAlgebra.mul!(y::$Atype, A::LinearCombination, x::$Atype,
+                             α::Number=true, β::Number=false)
+        @boundscheck check_dim_mul(y, A, x)
+        if iszero(α) # trivial cases
+            iszero(β) && (fill!(y, zero(eltype(y))); return y)
+            isone(β) && return y
+            # β != 0, 1
+            rmul!(y, β)
+            return y
+        else
+            mul!(y, first(A.maps), x, α, β)
+            return _mul!(MulStyle(A), y, A, x, α, β)
         end
     end
-    return y
 end
 
-else # 5-arg mul! is available for matrices
-
-# map types that have an allocation-free 5-arg mul! implementation
-const FreeMap = Union{MatrixMap,UniformScalingMap}
-
-function A_mul_B!(y::AbstractVector, A::LinearCombination{T,As}, x::AbstractVector) where {T, As<:Tuple{Vararg{FreeMap}}}
-    # no size checking, will be done by individual maps
-    A_mul_B!(y, A.maps[1], x)
-    for n in 2:length(A.maps)
-        mul!(y, A.maps[n], x, true, true)
-    end
-    return y
-end
-function A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector)
-    # no size checking, will be done by individual maps
-    A_mul_B!(y, A.maps[1], x)
-    l = length(A.maps)
-    if l>1
-        z = similar(y)
-        for n in 2:l
-            An = A.maps[n]
-            if An isa FreeMap
-                mul!(y, An, x, true, true)
-            else
-                A_mul_B!(z, A.maps[n], x)
-                y .+= z
-            end
-        end
-    end
-    return y
+@inline _mul!(::FiveArg, y, A::LinearCombination, x, α::Number, β::Number) =
+    __mul!(y, Base.tail(A.maps), x, α, nothing)
+@inline function _mul!(::ThreeArg, y, A::LinearCombination, x, α::Number, β::Number)
+    z = similar(y)
+    __mul!(y, Base.tail(A.maps), x, α, z)
 end
 
-function LinearAlgebra.mul!(y::AbstractVector, A::LinearCombination{T,As}, x::AbstractVector, α::Number=true, β::Number=false) where {T, As<:Tuple{Vararg{FreeMap}}}
-    length(y) == size(A, 1) || throw(DimensionMismatch("mul!"))
-    if isone(α)
-        iszero(β) && (A_mul_B!(y, A, x); return y)
-        !isone(β) && rmul!(y, β)
-    elseif iszero(α)
-        iszero(β) && (fill!(y, zero(eltype(y))); return y)
-        isone(β) && return y
-        # β != 0, 1
-        rmul!(y, β)
-        return y
-    else # α != 0, 1
-        if iszero(β)
-            A_mul_B!(y, A, x)
-            rmul!(y, α)
-            return y
-        elseif !isone(β)
-            rmul!(y, β)
-        end # β-cases
-    end # α-cases
+@inline __mul!(y, As::Tuple{Vararg{LinearMap}}, x, α, z) =
+    __mul!(__mul!(y, first(As), x, α, z), Base.tail(As), x, α, z)
+@inline __mul!(y, A::Tuple{LinearMap}, x, α, z) = __mul!(y, first(A), x, α, z)
+@inline __mul!(y, A::LinearMap, x, α, z) = muladd!(MulStyle(A), y, A, x, α, z)
 
-    for An in A.maps
-        mul!(y, An, x, α, true)
+@inline muladd!(::FiveArg, y, A, x, α, _) = mul!(y, A, x, α, true)
+@inline function muladd!(::ThreeArg, y, A, x, α, z)
+    # TODO: replace by mul!(z, A, x)
+    A_mul_B!(z, A, x)
+    if isone(α)
+        y .+= z
+    else
+        y .+= z .* α
     end
     return y
 end
 
-end # VERSION
+A_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, A, x)
 
-At_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = A_mul_B!(y, transpose(A), x)
+At_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, transpose(A), x)
 
-Ac_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = A_mul_B!(y, adjoint(A), x)
+Ac_mul_B!(y::AbstractVector, A::LinearCombination, x::AbstractVector) = mul!(y, adjoint(A), x)

src/transpose.jl

@@ -5,6 +5,8 @@ struct AdjointMap{T, A<:LinearMap{T}} <: LinearMap{T}
     lmap::A
 end
 
+MulStyle(A::Union{TransposeMap,AdjointMap}) = MulStyle(A.lmap)
+
 # transposition behavior of LinearMap objects
 LinearAlgebra.transpose(A::TransposeMap) = A.lmap
 LinearAlgebra.adjoint(A::AdjointMap) = A.lmap

src/uniformscalingmap.jl

@@ -11,6 +11,8 @@ UniformScalingMap(λ::Number, M::Int, N::Int) =
 UniformScalingMap(λ::T, sz::Dims{2}) where {T} =
     (sz[1] == sz[2] ? UniformScalingMap(λ, sz[1]) : error("UniformScalingMap needs to be square"))
 
+MulStyle(::UniformScalingMap) = FiveArg()
+
 # properties
 Base.size(A::UniformScalingMap) = (A.M, A.M)
 Base.isreal(A::UniformScalingMap) = isreal(A.λ)
@@ -91,11 +93,22 @@ function _scaling!(y, J::UniformScalingMap, x, α::Number=true, β::Number=false
         rmul!(y, β)
         return y
     else # α != 0, 1
-        iszero(β) && (y .= λ .* x .* α; return y)
-        isone(β) && (y .+= λ .* x .* α; return y)
-        # β != 0, 1
-        y .= y .* β .+ λ .* x .* α
-        return y
+        if iszero(β)
+            iszero(λ) && return fill!(y, zero(eltype(y)))
+            isone(λ) && return y .= x .* α
+            y .= λ .* x .* α
+            return y
+        elseif isone(β)
+            iszero(λ) && return y
+            isone(λ) && return y .+= x .* α
+            y .+= λ .* x .* α
+            return y
+        else # β != 0, 1
+            iszero(λ) && (rmul!(y, β); return y)
+            isone(λ) && (y .= y .* β .+ x .* α; return y)
+            y .= y .* β .+ λ .* x .* α
+            return y
+        end
     end
 end
 

src/wrappedmap.jl

@@ -19,6 +19,8 @@ end
 
 const MatrixMap{T} = WrappedMap{T,<:AbstractMatrix}
 
+MulStyle(A::WrappedMap) = MulStyle(A.lmap)
+
 LinearAlgebra.transpose(A::MatrixMap{T}) where {T} =
     WrappedMap{T}(transpose(A.lmap); issymmetric=A._issymmetric, ishermitian=A._ishermitian, isposdef=A._isposdef)
 LinearAlgebra.adjoint(A::MatrixMap{T}) where {T} =

test/blockmap.jl

@@ -6,6 +6,7 @@ using Test, LinearMaps, LinearAlgebra
             A11 = rand(elty, 10, 10)
             A12 = rand(elty, 10, n2)
             L = @inferred hcat(LinearMap(A11), LinearMap(A12))
+            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
             @test L isa LinearMaps.BlockMap{elty}
             A = [A11 A12]
             x = rand(10+n2)
@@ -36,6 +37,7 @@ using Test, LinearMaps, LinearAlgebra
             A21 = rand(elty, 20, 10)
             L = @inferred vcat(LinearMap(A11), LinearMap(A21))
             @test L isa LinearMaps.BlockMap{elty}
+            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
             A = [A11; A21]
             x = rand(10)
             @test size(L) == size(A)
@@ -62,6 +64,7 @@ using Test, LinearMaps, LinearAlgebra
             A = [A11 A12; A21 A22]
             @inferred hvcat((2,2), LinearMap(A11), LinearMap(A12), LinearMap(A21), LinearMap(A22))
             L = [LinearMap(A11) LinearMap(A12); LinearMap(A21) LinearMap(A22)]
+            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
             @test @inferred !issymmetric(L)
             @test @inferred !ishermitian(L)
             x = rand(30)
@@ -102,12 +105,13 @@ using Test, LinearMaps, LinearAlgebra
             @test Matrix(adjoint(B)) == C'
         end
     end
-    
+
     @testset "adjoint/transpose" begin
         for elty in (Float32, Float64, ComplexF64), transform in (transpose, adjoint)
             A12 = rand(elty, 10, 10)
             A = [I A12; transform(A12) I]
             L = [I LinearMap(A12); transform(LinearMap(A12)) I]
+            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
             if elty <: Complex
                 if transform == transpose
                     @test @inferred issymmetric(L)

test/kronecker.jl

@@ -9,6 +9,7 @@ using Test, LinearMaps, LinearAlgebra
         LB = LinearMap(B)
         LK = @inferred kron(LA, LB)
         @test @inferred size(LK) == size(K)
+        @test LinearMaps.MulStyle(LK) === LinearMaps.ThreeArg()
         for i in (1, 2)
             @test @inferred size(LK, i) == size(K, i)
         end
@@ -31,6 +32,11 @@ using Test, LinearMaps, LinearAlgebra
         @test @inferred kron(LA, LB)' == @inferred kron(LA', LB')
         @test (@inferred kron(LA, B)) == (@inferred kron(LA, LB)) == (@inferred kron(A, LB))
         @test @inferred ishermitian(kron(LA'LA, LB'LB))
+        A = rand(2, 5); B = rand(4, 2)
+        K = @inferred kron(A, LinearMap(B))
+        @test Matrix(K) ≈ kron(A, B)
+        K = @inferred kron(LinearMap(B), A)
+        @test Matrix(K) ≈ kron(B, A)
         A = rand(3, 3); B = rand(2, 2); LA = LinearMap(A); LB = LinearMap(B)
         @test @inferred issymmetric(kron(LA'LA, LB'LB))
         @test @inferred ishermitian(kron(LA'LA, LB'LB))
@@ -59,7 +65,7 @@ using Test, LinearMaps, LinearAlgebra
                 @test Matrix(kronsum(transform(LA), transform(LB))) ≈ transform(KSmat)
                 @test Matrix(transform(LinearMap(kronsum(LA, LB)))) ≈ Matrix(transform(KS)) ≈ transform(KSmat)
             end
-            @inferred kronsum(A, A, LB)
+            @test @inferred(kronsum(A, A, LB)) == @inferred(⊕(A, A, B))
             @test Matrix(@inferred LA^⊕(3)) == Matrix(@inferred A^⊕(3)) ≈ Matrix(kronsum(LA, A, A))
             @test @inferred(kronsum(LA, LA, LB)) == @inferred(kronsum(LA, kronsum(LA, LB))) == @inferred(kronsum(A, A, B))
             @test Matrix(@inferred kronsum(A, B, A, B, A, B)) ≈ Matrix(@inferred kronsum(LA, LB, LA, LB, LA, LB))

test/linearcombination.jl

@@ -4,7 +4,7 @@ using Test, LinearMaps, LinearAlgebra, BenchmarkTools
     CS! = LinearMap{ComplexF64}(cumsum!,
                                 (y, x) -> (copyto!(y, x); reverse!(y); cumsum!(y, y)), 10;
                                 ismutating=true)
-    v = rand(10)
+    v = rand(ComplexF64, 10)
     u = similar(v)
     b = @benchmarkable mul!($u, $CS!, $v)
     @test run(b, samples=3).allocs == 0
@@ -13,12 +13,20 @@ using Test, LinearMaps, LinearAlgebra, BenchmarkTools
     @test mul!(u, L, v) ≈ n * cumsum(v)
     b = @benchmarkable mul!($u, $L, $v)
     @test run(b, samples=5).allocs <= 1
+    for α in (false, true, rand(ComplexF64)), β in (false, true, rand(ComplexF64))
+        @test mul!(copy(u), L, v, α, β) ≈ Matrix(L)*v*α + u*β
+    end
 
     A = 2 * rand(ComplexF64, (10, 10)) .- 1
     B = rand(ComplexF64, size(A)...)
     M = @inferred LinearMap(A)
     N = @inferred LinearMap(B)
+    @test @inferred(LinearMaps.MulStyle(M)) === matrixstyle
+    @test @inferred(LinearMaps.MulStyle(N)) === matrixstyle
     LC = @inferred M + N
+    @test @inferred(LinearMaps.MulStyle(LC)) === matrixstyle
+    @test @inferred(LinearMaps.MulStyle(LC + I)) === matrixstyle
+    @test @inferred(LinearMaps.MulStyle(LC + 2.0*I)) === matrixstyle
     v = rand(ComplexF64, 10)
     w = similar(v)
     b = @benchmarkable mul!($w, $M, $v)
@@ -27,10 +35,14 @@ using Test, LinearMaps, LinearAlgebra, BenchmarkTools
         b = @benchmarkable mul!($w, $LC, $v)
         @test run(b, samples=3).allocs == 0
         for α in (false, true, rand(ComplexF64)), β in (false, true, rand(ComplexF64))
-            b = @benchmarkable mul!($w, $LC, $v, $α, $β)
-            @test run(b, samples=3).allocs == 0
-            b = @benchmarkable mul!($w, $(LC + I), $v, $α, $β)
-            @test run(b, samples=3).allocs == 0
+            if testallocs
+                b = @benchmarkable mul!($w, $LC, $v, $α, $β)
+                @test run(b, samples=3).allocs == 0
+                b = @benchmarkable mul!($w, $(I + LC), $v, $α, $β)
+                @test run(b, samples=3).allocs == 0
+                b = @benchmarkable mul!($w, $(LC + I), $v, $α, $β)
+                @test run(b, samples=3).allocs == 0
+            end
             y = rand(ComplexF64, size(v))
             @test mul!(copy(y), LC, v, α, β) ≈ Matrix(LC)*v*α + y*β
             @test mul!(copy(y), LC+I, v, α, β) ≈ Matrix(LC + I)*v*α + y*β

test/runtests.jl

@@ -1,4 +1,9 @@
 using Test, LinearMaps
+import LinearMaps: FiveArg, ThreeArg
+
+const matrixstyle = VERSION ≥ v"1.3.0-alpha.115" ? FiveArg() : ThreeArg()
+
+const testallocs = true
 
 include("linearmaps.jl")