Square Kronecker products and sums

Author: dkarrasch

docs/src/types.md

@@ -72,6 +72,16 @@ kronsum
 LinearMaps.:⊕
 ```
 
+There exist alternative constructors of Kronecker products and sums for square factors and
+summands, respectively. These are designed for cases of 3 or more arguments, and
+benchmarking intended use cases for comparison with `KroneckerMap` and `KroneckerSumMap`
+is recommended.
+
+```@docs
+squarekron
+sumkronsum
+```
+
 ### `BlockMap` and `BlockDiagonalMap`
 
 Types for representing block (diagonal) maps lazily.
@@ -116,6 +126,7 @@ Base.:*(::AbstractMatrix,::LinearMap)
 LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::AbstractVector)
 LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::AbstractVector,::Number,::Number)
 LinearAlgebra.mul!(::AbstractMatrix,::AbstractMatrix,::LinearMap)
+LinearAlgebra.mul!(::AbstractMatrix,::AbstractMatrix,::LinearMap,::Number,::Number)
 LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::Number)
 LinearAlgebra.mul!(::AbstractMatrix,::LinearMap,::Number,::Number,::Number)
 *(::LinearAlgebra.AdjointAbsVec,::LinearMap)

src/LinearMaps.jl

@@ -1,7 +1,7 @@
 module LinearMaps
 
 export LinearMap
-export ⊗, kronsum, ⊕
+export ⊗, squarekron, kronsum, ⊕, sumkronsum
 export FillMap
 export InverseMap
 
@@ -17,6 +17,7 @@ abstract type LinearMap{T} end
 
 const MapOrVecOrMat{T} = Union{LinearMap{T}, AbstractVecOrMat{T}}
 const MapOrMatrix{T} = Union{LinearMap{T}, AbstractMatrix{T}}
+const TransposeAbsVecOrMat{T} = Transpose{T,<:AbstractVecOrMat}
 const RealOrComplex = Union{Real, Complex}
 
 const LinearMapTuple = Tuple{Vararg{LinearMap}}
@@ -77,10 +78,12 @@ function check_dim_mul(C, A, B)
     return nothing
 end
 
+_issquare(A) = size(A, 1) == size(A, 2)
+
 _front(As::Tuple) = Base.front(As)
-_front(As::AbstractVector) = @inbounds @views As[1:end-1]
+_front(As::AbstractVector) = @inbounds @views As[begin:end-1]
 _tail(As::Tuple) = Base.tail(As)
-_tail(As::AbstractVector) = @inbounds @views As[2:end]
+_tail(As::AbstractVector) = @inbounds @views As[begin+1:end]
 
 _combine(A::LinearMap, B::LinearMap) = tuple(A, B)
 _combine(A::LinearMap, Bs::LinearMapTuple) = tuple(A, Bs...)
@@ -330,9 +333,9 @@ function _generic_map_mul!(Y, A, s::Number, α, β)
     return Y
 end
 
-include("left.jl") # left multiplication by a transpose or adjoint vector
 include("transpose.jl") # transposing linear maps
 include("wrappedmap.jl") # wrap a matrix of linear map in a new type, thereby allowing to alter its properties
+include("left.jl") # left multiplication by a matrix/transpose or adjoint vector
 include("uniformscalingmap.jl") # the uniform scaling map, to be able to make linear combinations of LinearMap objects and multiples of I
 include("linearcombination.jl") # defining linear combinations of linear maps
 include("scaledmap.jl") # multiply by a (real or complex) scalar

src/fillmap.jl

@@ -21,7 +21,7 @@ Base.size(A::FillMap) = A.size
 MulStyle(A::FillMap) = FiveArg()
 LinearAlgebra.issymmetric(A::FillMap) = A.size[1] == A.size[2]
 LinearAlgebra.ishermitian(A::FillMap) = isreal(A.λ) && A.size[1] == A.size[2]
-LinearAlgebra.isposdef(A::FillMap) = (size(A, 1) == size(A, 2) == 1 && isposdef(A.λ))
+LinearAlgebra.isposdef(A::FillMap) = (LinearAlgebra.checksquare(A) == 1 && isposdef(A.λ))
 Base.:(==)(A::FillMap, B::FillMap) = A.λ == B.λ && A.size == B.size
 
 LinearAlgebra.adjoint(A::FillMap) = FillMap(adjoint(A.λ), reverse(A.size))

src/kronecker.jl

@@ -79,6 +79,43 @@ for k in 3:8 # is 8 sufficient?
         kron($(mapargs...), $(Symbol(:A, k)), convert_to_lmaps(As...)...)
 end
 
+"""
+    squarekron(A₁::MapOrMatrix, A₂::MapOrMatrix, A₃::MapOrMatrix, Aᵢ::MapOrMatrix...)::CompositeMap
+
+Construct a (lazy) representation of the Kronecker product `⨂ᵢ₌₁ⁿ Aᵢ` of at least 3 _square_
+Kronecker factors. In contrast to [`kron`](@ref), this function assumes that all Kronecker
+factors are square, and makes use of the following identity:
+
+```
+\\bigotimes_{i=1}^n A_i = \\prod_{i=1}^n I_1 \\otimes \\ldots \\otimes I_{i-1} \\otimes A_i \\otimes I_{i+1} \\otimes \\ldots \\otimes I_n
+```
+
+where ```I_k``` is an identity matrix of the size of ```A_k```. By associativity, the
+Kronecker product of the identity operators may be combined to larger identity operators
+```I_{1:i-1}``` and ```I_{i+1:n}```, which yields
+
+```
+\\bigotimes_{i=1}^n A_i = \\prod_{i=1}^n I_{1:i-1} \\otimes A_i \\otimes I_{i+1:n}
+```
+
+i.e., a `CompositeMap` where each factor is a Kronecker product consisting of three maps:
+outer `UniformScalingMap`s and the respective Kronecker factor. This representation is
+expected to yield significantly faster multiplication (and reduce memory allocation)
+compared to [`kron`](@ref), but benchmarking intended use cases is highly recommended.
+"""
+function squarekron(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...)
+    maps = (A, B, C, Ds...)
+    T = promote_type(map(eltype, maps)...)
+    all(_issquare, maps) || throw(ArgumentError("operators need to be square in Kronecker sums"))
+    ns = map(a -> size(a, 1), maps)
+    firstmap = first(maps) ⊗ UniformScalingMap(true, prod(ns[2:end]))
+    lastmap  = UniformScalingMap(true, prod(ns[1:end-1])) ⊗ last(maps)
+    middlemaps = prod(enumerate(maps[2:end-1])) do (i, map)
+        UniformScalingMap(true, prod(ns[1:i])) ⊗ map ⊗ UniformScalingMap(true, prod(ns[i+2:end]))
+    end
+    return firstmap * middlemaps * lastmap
+end
+
 struct KronPower{p}
     function KronPower(p::Integer)
         p > 1 || throw(ArgumentError("the Kronecker power is only defined for exponents larger than 1, got $k"))
@@ -114,74 +151,90 @@ Base.:(==)(A::KroneckerMap, B::KroneckerMap) =
 # multiplication helper functions
 #################
 
-@inline function _kronmul!(y, B, x, A, T)
-    ma, na = size(A)
-    mb, nb = size(B)
-    X = reshape(x, (nb, na))
-    Y = reshape(y, (mb, ma))
-    if B isa UniformScalingMap
-        _unsafe_mul!(Y, X, transpose(A))
-        lmul!(B.λ, y)
+@inline function _kronmul!(Y, B, X, A)
+    # minimize intermediate memory allocation
+    if size(B, 2) * size(A, 1) <= size(B, 1) * size(A, 2)
+        temp = similar(Y, (size(B, 2), size(A, 1) ))
+        _unsafe_mul!(temp, X, transpose(A))
+        _unsafe_mul!(Y, B, temp)
     else
-        temp = similar(Y, (ma, nb))
-        _unsafe_mul!(temp, A, copy(transpose(X)))
-        _unsafe_mul!(Y, B, transpose(temp))
+        temp = similar(Y, (size(B, 1), size(A, 2)))
+        _unsafe_mul!(temp, B, X)
+        _unsafe_mul!(Y, temp, transpose(A))
     end
-    return y
+    return Y
 end
-@inline function _kronmul!(y, B, x, A::UniformScalingMap, _)
-    ma, na = size(A)
-    mb, nb = size(B)
-    iszero(A.λ) && return fill!(y, zero(eltype(y)))
-    X = reshape(x, (nb, na))
-    Y = reshape(y, (mb, ma))
+@inline function _kronmul!(Y, B::UniformScalingMap, X, A)
+    _unsafe_mul!(Y, X, transpose(A))
+    !isone(B.λ) && lmul!(B.λ, Y)
+    return Y
+end
+@inline function _kronmul!(Y, B, X, A::UniformScalingMap)
     _unsafe_mul!(Y, B, X)
-    !isone(A.λ) && rmul!(y, A.λ)
-    return y
+    !isone(A.λ) && rmul!(Y, A.λ)
+    return Y
 end
-@inline function _kronmul!(y, B, x, A::VecOrMatMap, _)
-    ma, na = size(A)
-    mb, nb = size(B)
-    X = reshape(x, (nb, na))
-    Y = reshape(y, (mb, ma))
+# disambiguation (cannot occur)
+@inline function _kronmul!(Y, B::UniformScalingMap, X, A::UniformScalingMap)
+    mul!(parent(Y), A.λ * B.λ, parent(X))
+    return Y
+end
+@inline function _kronmul!(Y, B, X, A::VecOrMatMap)
     At = transpose(A.lmap)
-    if B isa UniformScalingMap
-        # the following is (perhaps due to the reshape?) faster than
-        # _unsafe_mul!(Y, B * X, At)
-        _unsafe_mul!(Y, X, At)
-        lmul!(B.λ, y)
-    elseif nb*ma <= mb*na
+    if size(B, 2) * size(A, 1) <= size(B, 1) * size(A, 2)
         _unsafe_mul!(Y, B, X * At)
     else
         _unsafe_mul!(Y, Matrix(B * X), At)
     end
-    return y
+    return Y
+end
+@inline function _kronmul!(Y, B::UniformScalingMap, X, A::VecOrMatMap)
+    _unsafe_mul!(Y, X, transpose(A.lmap))
+    !isone(B.λ) && lmul!(B.λ, Y)
+    return Y
 end
+
 const VectorMap{T} = WrappedMap{T,<:AbstractVector}
 const AdjOrTransVectorMap{T} = WrappedMap{T,<:LinearAlgebra.AdjOrTransAbsVec}
-@inline _kronmul!(y, B::AdjOrTransVectorMap, x, a::VectorMap, _) = mul!(y, a.lmap, B.lmap * x)
 
 #################
 # multiplication with vectors
 #################
 
 const KroneckerMap2{T} = KroneckerMap{T, <:Tuple{LinearMap, LinearMap}}
-
+const OuterProductMap{T} = KroneckerMap{T, <:Tuple{VectorMap, AdjOrTransVectorMap}}
+function _unsafe_mul!(y::AbstractVecOrMat, L::OuterProductMap, x::AbstractVector)
+    a, bt = L.maps
+    mul!(y, a.lmap, bt.lmap * x)
+end
 function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerMap2, x::AbstractVector)
     require_one_based_indexing(y)
     A, B = L.maps
-    _kronmul!(y, B, x, A, eltype(L))
+    ma, na = size(A)
+    mb, nb = size(B)
+    X = reshape(x, (nb, na))
+    Y = reshape(y, (mb, ma))
+    _kronmul!(Y, B, X, A)
     return y
 end
 function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerMap, x::AbstractVector)
     require_one_based_indexing(y)
     maps = L.maps
     if length(maps) == 2 # reachable only for L.maps::Vector
-        @inbounds _kronmul!(y, maps[2], x, maps[1], eltype(L))
+        A, B = maps
+        ma, na = size(A)
+        mb, nb = size(B)
+        X = reshape(x, (nb, na))
+        Y = reshape(y, (mb, ma))
+        _kronmul!(Y, B, X, A)
     else
         A = first(maps)
         B = KroneckerMap{eltype(L)}(_tail(maps))
-        _kronmul!(y, B, x, A, eltype(L))
+        ma, na = size(A)
+        mb, nb = size(B)
+        X = reshape(x, (nb, na))
+        Y = reshape(y, (mb, ma))
+        _kronmul!(Y, B, X, A)
     end
     return y
 end
@@ -225,7 +278,7 @@ struct KroneckerSumMap{T, As<:Tuple{LinearMap, LinearMap}} <: LinearMap{T}
     maps::As
     function KroneckerSumMap{T}(maps::Tuple{LinearMap,LinearMap}) where {T}
         A1, A2 = maps
-        (size(A1, 1) == size(A1, 2) && size(A2, 1) == size(A2, 2)) ||
+        (_issquare(A1) && _issquare(A2)) ||
             throw(ArgumentError("operators need to be square in Kronecker sums"))
         for TA in Base.Iterators.map(eltype, maps)
             promote_type(T, TA) == T ||
@@ -269,6 +322,66 @@ kronsum(A::MapOrMatrix, B::MapOrMatrix) =
 kronsum(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...) =
     kronsum(A, kronsum(B, C, Ds...))
 
+"""
+    sumkronsum(A, B)::LinearCombination
+    sumkronsum(A, B, Cs...)::LinearCombination
+
+Construct a (lazy) representation of the Kronecker sum `A⊕B` of two or more square
+objects of type `LinearMap` or `AbstractMatrix`. This function makes use of the following
+representation of Kronecker sums:
+
+```
+\\bigoplus_{i=1}^n A_i = \\sum_{i=1}^n I_1 \\otimes \\ldots \\otimes I_{i-1} \\otimes A_i \\otimes I_{i+1} \\otimes \\ldots \\otimes I_n
+```
+
+where ```I_k``` is the identity operator of the size of ```A_k```. By associativity, the
+Kronecker product of the identity operators may be combined to larger identity operators
+```I_{1:i-1}``` and ```I_{i+1:n}```, which yields
+
+```
+\\bigoplus_{i=1}^n A_i = \\sum_{i=1}^n I_{1:i-1} \\otimes A_i \\otimes I_{i+1:n}
+```
+
+i.e., a `LinearCombination` where each summand is a Kronecker product consisting of three
+maps: outer `UniformScalingMap`s and the respective Kronecker factor. This representation is
+expected to yield significantly faster multiplication (and reduce memory allocation)
+compared to [`kronsum`](@ref), especially for 3 or more Kronecker summands, but
+benchmarking intended use cases is highly recommended.
+
+# Examples
+```jldoctest; setup=(using LinearAlgebra, SparseArrays, LinearMaps)
+julia> J = LinearMap(I, 2) # 2×2 identity map
+2×2 LinearMaps.UniformScalingMap{Bool} with scaling factor: true
+
+julia> E = spdiagm(-1 => trues(1)); D = LinearMap(E + E' - 2I);
+
+julia> Δ₁ = kron(D, J) + kron(J, D); # discrete 2D-Laplace operator, Kronecker sum
+
+julia> Δ₂ = sumkronsum(D, D);
+
+julia> Δ₃ = D^⊕(2);
+
+julia> Matrix(Δ₁) == Matrix(Δ₂) == Matrix(Δ₃)
+true
+```
+"""
+function sumkronsum(A::MapOrMatrix, B::MapOrMatrix)
+    LinearAlgebra.checksquare(A, B)
+    A ⊗ UniformScalingMap(true, size(B,1)) + UniformScalingMap(true, size(A,1)) ⊗ B
+end
+function sumkronsum(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...)
+    maps = (A, B, C, Ds...)
+    all(_issquare, maps) || throw(ArgumentError("operators need to be square in Kronecker sums"))
+    ns = map(a -> size(a, 1), maps)
+    n = length(maps)
+    firstmap = first(maps) ⊗ UniformScalingMap(true, prod(ns[2:end]))
+    lastmap  = UniformScalingMap(true, prod(ns[1:end-1])) ⊗ last(maps)
+    middlemaps = sum(enumerate(Base.front(Base.tail(maps)))) do (i, map)
+        UniformScalingMap(true, prod(ns[1:i])) ⊗ map ⊗ UniformScalingMap(true, prod(ns[i+2:end]))
+    end
+    return firstmap + middlemaps + lastmap
+end
+
 struct KronSumPower{p}
     function KronSumPower(p::Integer)
         p > 1 || throw(ArgumentError("the Kronecker sum power is only defined for exponents larger than 1, got $k"))
@@ -280,14 +393,24 @@ end
     ⊕(k::Integer)
 
 Construct a lazy representation of the `k`-th Kronecker sum power `A^⊕(k) = A ⊕ A ⊕ ... ⊕ A`,
-where `A` can be a square `AbstractMatrix` or a `LinearMap`.
+where `A` can be a square `AbstractMatrix` or a `LinearMap`. This calls [`sumkronsum`](@ref)
+on the `k`-tuple `(A, ..., A)`.
+
+# Example
+```jldoctest
+julia> Matrix([1 0; 0 1]^⊕(2))
+4×4 Matrix{Int64}:
+ 2  0  0  0
+ 0  2  0  0
+ 0  0  2  0
+ 0  0  0  2
 """
 ⊕(k::Integer) = KronSumPower(k)
 
 ⊕(a, b, c...) = kronsum(a, b, c...)
 
 Base.:(^)(A::MapOrMatrix, ::KronSumPower{p}) where {p} =
-    kronsum(ntuple(n -> convert(LinearMap, A), Val(p))...)
+    sumkronsum(ntuple(n -> convert(LinearMap, A), Val(p))...)
 
 Base.size(A::KroneckerSumMap, i) = prod(size.(A.maps, i))
 Base.size(A::KroneckerSumMap) = (size(A, 1), size(A, 2))
@@ -305,10 +428,10 @@ Base.:(==)(A::KroneckerSumMap, B::KroneckerSumMap) =
 
 function _unsafe_mul!(y::AbstractVecOrMat, L::KroneckerSumMap, x::AbstractVector)
     A, B = L.maps
-    ma, na = size(A)
-    mb, nb = size(B)
-    X = reshape(x, (nb, na))
-    Y = reshape(y, (nb, na))
+    a = size(A, 1)
+    b = size(B, 1)
+    X = reshape(x, (b, a))
+    Y = reshape(y, (b, a))
     _unsafe_mul!(Y, X, transpose(A))
     _unsafe_mul!(Y, B, X, true, true)
     return y