Rework conversions, speed-up matrix Kronecker products (#108)

Author: dkarrasch

src/LinearMaps.jl

@@ -68,6 +68,15 @@ convert_to_lmaps(A) = (convert_to_lmaps_(A),)
 @inline convert_to_lmaps(A, B, Cs...) =
     (convert_to_lmaps_(A), convert_to_lmaps_(B), convert_to_lmaps(Cs...)...)
 
+# The (internal) multiplication logic is as follows:
+#  - `*(A, x)` calls `mul!(y, A, x)` for appropriately-sized y
+#  - `mul!` checks consistency of the sizes, and calls `_unsafe_mul!`,
+#    which does not check sizes, but potentially one-based indexing if necessary
+#  - by default, `_unsafe_mul!` is redirected back to `mul!`
+#  - custom map types only need to implement 3-arg (vector) `mul!`, and
+#    everything else (5-arg multiplication, application to matrices,
+#    conversion to matrices) will just work
+
 """
     *(A::LinearMap, x::AbstractVector)::AbstractVector
 
@@ -200,7 +209,7 @@ function _generic_mapmat_mul!(Y, A, X, α=true, β=false)
     return Y
 end
 
-_unsafe_mul!(y, A::MapOrMatrix, x)       = mul!(y, A, x)
+_unsafe_mul!(y, A::MapOrMatrix, x) = mul!(y, A, x)
 _unsafe_mul!(y, A::AbstractMatrix, x, α, β) = mul!(y, A, x, α, β)
 function _unsafe_mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α, β)
     return _generic_mapvec_mul!(y, A, x, α, β)

src/conversion.jl

@@ -1,7 +1,6 @@
 # Matrix: create matrix representation of LinearMap
-function Base.Matrix(A::LinearMap)
+function Base.Matrix{T}(A::LinearMap) where {T}
     M, N = size(A)
-    T = eltype(A)
     mat = Matrix{T}(undef, (M, N))
     v = fill(zero(T), N)
     @inbounds for i in 1:N
@@ -11,8 +10,9 @@ function Base.Matrix(A::LinearMap)
     end
     return mat
 end
+Base.Matrix(A::LinearMap{T}) where {T} = Matrix{T}(A)
 Base.Array(A::LinearMap) = Matrix(A)
-Base.convert(::Type{Matrix}, A::LinearMap) = Matrix(A)
+Base.convert(::Type{T}, A::LinearMap) where {T<:Matrix} = T(A)
 Base.convert(::Type{Array}, A::LinearMap) = convert(Matrix, A)
 Base.convert(::Type{AbstractMatrix}, A::LinearMap) = convert(Matrix, A)
 Base.convert(::Type{AbstractArray}, A::LinearMap) = convert(AbstractMatrix, A)
@@ -42,21 +42,22 @@ function SparseArrays.sparse(A::LinearMap{T}) where {T}
     return SparseMatrixCSC(M, N, colptr, rowind, nzval)
 end
 Base.convert(::Type{SparseMatrixCSC}, A::LinearMap) = sparse(A)
+SparseArrays.SparseMatrixCSC(A::LinearMap) = sparse(A)
 
 # special cases
 
 # ScaledMap
-Base.Matrix(A::ScaledMap{<:Any,<:Any,<:MatrixMap}) = convert(Matrix, A.λ*A.lmap.lmap)
+Base.Matrix{T}(A::ScaledMap{<:Any,<:Any,<:MatrixMap}) where {T} = convert(Matrix{T}, A.λ*A.lmap.lmap)
 SparseArrays.sparse(A::ScaledMap{<:Any,<:Any,<:MatrixMap}) = convert(SparseMatrixCSC, A.λ*A.lmap.lmap)
 
 # UniformScalingMap
-Base.Matrix(J::UniformScalingMap) = Matrix(J.λ*I, size(J))
-Base.convert(::Type{AbstractMatrix}, J::UniformScalingMap) = Diagonal(fill(J.λ, size(J, 1)))
+Base.Matrix{T}(J::UniformScalingMap) where {T} = Matrix{T}(J.λ*I, size(J))
+Base.convert(::Type{AbstractMatrix}, J::UniformScalingMap) = Diagonal(fill(J.λ, J.M))
 
 # WrappedMap
-Base.Matrix(A::WrappedMap) = Matrix(A.lmap)
+Base.Matrix{T}(A::WrappedMap) where {T} = Matrix{T}(A.lmap)
+Base.convert(::Type{T}, A::WrappedMap) where {T<:Matrix} = convert(T, A.lmap)
 Base.convert(::Type{AbstractMatrix}, A::WrappedMap) = convert(AbstractMatrix, A.lmap)
-Base.convert(::Type{Matrix}, A::WrappedMap) = convert(Matrix, A.lmap)
 SparseArrays.sparse(A::WrappedMap) = sparse(A.lmap)
 Base.convert(::Type{SparseMatrixCSC}, A::WrappedMap) = convert(SparseMatrixCSC, A.lmap)
 
@@ -67,68 +68,90 @@ for (TT, T) in ((AdjointMap, adjoint), (TransposeMap, transpose))
 end
 
 # LinearCombination
-for (TT, T) in ((Type{Matrix}, Matrix), (Type{SparseMatrixCSC}, SparseMatrixCSC))
-    @eval function Base.convert(::$TT, ΣA::LinearCombination{<:Any,<:Tuple{Vararg{MatrixMap}}})
-        maps = ΣA.maps
-        mats = map(A->getfield(A, :lmap), maps)
-        return convert($T, sum(mats))
-    end
+function Base.Matrix{T}(ΣA::LinearCombination{<:Any,<:Tuple{Vararg{MatrixMap}}}) where {T}
+    maps = ΣA.maps
+    mats = map(A->getfield(A, :lmap), maps)
+    return Matrix{T}(sum(mats))
+end
+function SparseArrays.sparse(ΣA::LinearCombination{<:Any,<:Tuple{Vararg{MatrixMap}}})
+    maps = ΣA.maps
+    mats = map(A->getfield(A, :lmap), maps)
+    return convert(SparseMatrixCSC, sum(mats))
 end
 
 # CompositeMap
-for (TT, T) in ((Type{Matrix}, Matrix), (Type{SparseMatrixCSC}, SparseMatrixCSC))
-    @eval function Base.convert(::$TT, AB::CompositeMap{<:Any,<:Tuple{MatrixMap,MatrixMap}})
+function Base.Matrix{T}(AB::CompositeMap{<:Any,<:Tuple{MatrixMap,LinearMap}}) where {T}
+    B, A = AB.maps
+    require_one_based_indexing(B)
+    Y = Matrix{eltype(AB)}(undef, size(AB))
+    @views for i in 1:size(Y, 2)
+        _unsafe_mul!(Y[:, i], A, B.lmap[:, i])
+    end
+    return Y
+end
+for ((TA, fieldA), (TB, fieldB)) in (((MatrixMap, :lmap), (MatrixMap, :lmap)),
+                                     ((MatrixMap, :lmap), (UniformScalingMap, :λ)),
+                                     ((UniformScalingMap, :λ), (MatrixMap, :lmap)))
+    @eval function Base.convert(::Type{AbstractMatrix}, AB::CompositeMap{<:Any,<:Tuple{$TB,$TA}})
         B, A = AB.maps
-        return convert($T, A.lmap*B.lmap)
+        return A.$fieldA*B.$fieldB
     end
 end
-function Base.Matrix(λA::CompositeMap{<:Any,<:Tuple{MatrixMap,UniformScalingMap}})
+function Base.Matrix{T}(AB::CompositeMap{<:Any,<:Tuple{MatrixMap,MatrixMap}}) where {T}
+    B, A = AB.maps
+    return convert(Matrix{T}, A.lmap*B.lmap)
+end
+function SparseArrays.sparse(AB::CompositeMap{<:Any,<:Tuple{MatrixMap,MatrixMap}})
+    B, A = AB.maps
+    return convert(SparseMatrixCSC, A.lmap*B.lmap)
+end
+function Base.Matrix{T}(λA::CompositeMap{<:Any,<:Tuple{MatrixMap,UniformScalingMap}}) where {T}
     A, J = λA.maps
-    return convert(Matrix, J.λ*A.lmap)
+    return convert(Matrix{T}, J.λ*A.lmap)
 end
 function SparseArrays.sparse(λA::CompositeMap{<:Any,<:Tuple{MatrixMap,UniformScalingMap}})
     A, J = λA.maps
     return convert(SparseMatrixCSC, J.λ*A.lmap)
 end
-function Base.Matrix(Aλ::CompositeMap{<:Any,<:Tuple{UniformScalingMap,MatrixMap}})
+function Base.Matrix{T}(Aλ::CompositeMap{<:Any,<:Tuple{UniformScalingMap,MatrixMap}}) where {T}
     J, A = Aλ.maps
-    return convert(Matrix, A.lmap*J.λ)
+    return convert(Matrix{T}, A.lmap*J.λ)
 end
 function SparseArrays.sparse(Aλ::CompositeMap{<:Any,<:Tuple{UniformScalingMap,MatrixMap}})
     J, A = Aλ.maps
     return convert(SparseMatrixCSC, A.lmap*J.λ)
 end
 
 # BlockMap & BlockDiagonalMap
-Base.Matrix(A::BlockMap) = hvcat(A.rows, convert.(Matrix, A.maps)...)
+Base.Matrix{T}(A::BlockMap) where {T} = hvcat(A.rows, convert.(Matrix{T}, A.maps)...)
 Base.convert(::Type{AbstractMatrix}, A::BlockMap) = hvcat(A.rows, convert.(AbstractMatrix, A.maps)...)
-function Base.convert(::Type{SparseMatrixCSC}, A::BlockMap)
+function SparseArrays.sparse(A::BlockMap)
     return hvcat(
         A.rows,
         convert(SparseMatrixCSC, first(A.maps)),
         convert.(AbstractMatrix, Base.tail(A.maps))...
     )
 end
-Base.Matrix(A::BlockDiagonalMap) = cat(convert.(Matrix, A.maps)...; dims=(1,2))
+Base.Matrix{T}(A::BlockDiagonalMap) where {T} = cat(convert.(Matrix{T}, A.maps)...; dims=(1,2))
 Base.convert(::Type{AbstractMatrix}, A::BlockDiagonalMap) = sparse(A)
 function SparseArrays.sparse(A::BlockDiagonalMap)
     return blockdiag(convert.(SparseMatrixCSC, A.maps)...)
 end
 
 # KroneckerMap & KroneckerSumMap
-Base.Matrix(A::KroneckerMap) = kron(convert.(Matrix, A.maps)...)
+Base.Matrix{T}(A::KroneckerMap) where {T} = kron(convert.(Matrix{T}, A.maps)...)
 Base.convert(::Type{AbstractMatrix}, A::KroneckerMap) = kron(convert.(AbstractMatrix, A.maps)...)
 function SparseArrays.sparse(A::KroneckerMap)
     return kron(
         convert(SparseMatrixCSC, first(A.maps)),
         convert.(AbstractMatrix, Base.tail(A.maps))...
     )
 end
-function Base.Matrix(L::KroneckerSumMap)
+function Base.Matrix{T}(L::KroneckerSumMap) where {T}
     A, B = L.maps
     IA = Diagonal(ones(Bool, size(A, 1)))
     IB = Diagonal(ones(Bool, size(B, 1)))
-    return kron(Matrix(A), IB) + kron(IA, Matrix(B))
+    return kron(Matrix{T}(A), IB) + kron(IA, Matrix{T}(B))
 end
 function Base.convert(::Type{AbstractMatrix}, L::KroneckerSumMap)
     A, B = L.maps

src/kronecker.jl

@@ -115,13 +115,14 @@ Base.:(==)(A::KroneckerMap, B::KroneckerMap) = (eltype(A) == eltype(B) && A.maps
     end
     return y
 end
-@inline function _kronmul!(y, B::Union{MatrixMap,UniformScalingMap}, X, At::Union{MatrixMap,UniformScalingMap}, T)
+@inline function _kronmul!(y, B, X, At::Union{MatrixMap,UniformScalingMap}, T)
     na, ma = size(At)
     mb, nb = size(B)
-    if nb*ma < mb*na
-        _unsafe_mul!(reshape(y, (mb, ma)), B, convert(Matrix, X*At))
+    Y = reshape(y, (mb, ma))
+    if nb*ma < mb*na 
+        _unsafe_mul!(Y, B, Matrix(X*At))
     else
-        _unsafe_mul!(reshape(y, (mb, ma)), convert(Matrix, B*X), At isa MatrixMap ? At.lmap : At.λ)
+        _unsafe_mul!(Y, Matrix(B*X), At isa MatrixMap ? At.lmap : At.λ)
     end
     return y
 end

test/composition.jl

@@ -1,4 +1,4 @@
-using Test, LinearMaps, LinearAlgebra
+using Test, LinearMaps, LinearAlgebra, SparseArrays
 
 @testset "composition" begin
     F = @inferred LinearMap(cumsum, reverse ∘ cumsum ∘ reverse, 10; ismutating=false)
@@ -50,8 +50,9 @@ using Test, LinearMaps, LinearAlgebra
     R1 = rand(ComplexF64, 10, 10); L1 = LinearMap(R1)
     R2 = rand(ComplexF64, 10, 10); L2 = LinearMap(R2)
     R3 = rand(ComplexF64, 10, 10); L3 = LinearMap(R3)
-    CompositeR = prod(R -> LinearMap(R), [R1, R2, R3])
+    CompositeR = prod(LinearMap, [R1, R2, R3])
     @test @inferred L1 * L2 * L3 == CompositeR
+    @test Matrix(L1 * L2) ≈ sparse(L1 * L2) ≈ R1 * R2
     @test @inferred transpose(CompositeR) == transpose(L3) * transpose(L2) * transpose(L1)
     @test @inferred adjoint(CompositeR) == L3' * L2' * L1'
     @test @inferred adjoint(adjoint((CompositeR))) == CompositeR
@@ -61,6 +62,18 @@ using Test, LinearMaps, LinearAlgebra
     Lc = @inferred adjoint(LinearMap(CompositeR))
     @test Lc * v ≈ R3' * R2' * R1' * v
 
+    # convert to AbstractMatrix
+    for A in (LinearMap(sprandn(10, 10, 0.3)), LinearMap(rand()*I, 10))
+        for B in (LinearMap(sprandn(10, 10, 0.3)), LinearMap(rand()*I, 10))
+            AA = convert(AbstractMatrix, A*B)
+            if A isa LinearMaps.UniformScalingMap && B isa LinearMaps.UniformScalingMap
+                @test isdiag(AA)
+            else
+                @test issparse(AA)
+            end
+        end
+    end
+
     # test inplace operations
     w = similar(v)
     mul!(w, L, v)

test/linearcombination.jl

@@ -37,6 +37,7 @@ using Test, LinearMaps, LinearAlgebra, BenchmarkTools
     @test @inferred(LinearMaps.MulStyle(LC)) === matrixstyle
     @test @inferred(LinearMaps.MulStyle(LC + I)) === matrixstyle
     @test @inferred(LinearMaps.MulStyle(LC + 2.0*I)) === matrixstyle
+    @test sparse(LC) == Matrix(LC) == A+B
     v = rand(ComplexF64, 10)
     w = similar(v)
     b = @benchmarkable mul!($w, $M, $v)

test/linearmaps.jl

@@ -96,6 +96,6 @@ LinearAlgebra.mul!(y::AbstractVector, A::Union{SimpleFunctionMap,SimpleComplexFu
     @test FM == L
     @test F * v ≈ L * v
     Fs = sparse(F)
-    @test Fs == L
+    @test SparseMatrixCSC(F) == Fs == L
     @test Fs isa SparseMatrixCSC
 end

test/wrappedmap.jl

@@ -8,6 +8,7 @@ using Test, LinearMaps, LinearAlgebra
     L = @inferred LinearMap{Float64}(A)
     MA = @inferred LinearMap(SA)
     MB = @inferred LinearMap(SB)
+    @test eltype(Matrix{Complex{Float32}}(LinearMap(A))) <: Complex
     @test size(L) == size(A)
     @test @inferred !issymmetric(L)
     @test @inferred issymmetric(MA)