Skip to content

Commit

Permalink
Square Kronecker products and sums (#183)
Browse files Browse the repository at this point in the history
  • Loading branch information
dkarrasch authored Jun 27, 2022
1 parent 10b6901 commit e39cbc1
Show file tree
Hide file tree
Showing 6 changed files with 233 additions and 74 deletions.
5 changes: 5 additions & 0 deletions docs/src/history.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,11 @@
argument to the constructor (see the docstring for details). Note that `A` must be
compatible with the solver: `A` can, for example, be a factorization, or another
`LinearMap` in combination with an iterative solver.
* New constructors for lazy representations of Kronecker products ([`squarekron`](@ref))
and sums ([`sumkronsum`](@ref)) for _square_ factors and summands, respectively, are
introduced. They target cases with 3 or more factors/summands, and benchmarking intended
use cases for comparison with `KroneckerMap` (constructed via [`Base.kron`](@ref)) and
`KroneckerSumMap` (constructed via [`kronsum`](@ref)) is recommended.

## What's new in v3.7

Expand Down
10 changes: 10 additions & 0 deletions docs/src/types.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
6 changes: 4 additions & 2 deletions src/LinearMaps.jl
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
module LinearMaps

export LinearMap
export , kronsum,
export , squarekron, kronsum, , sumkronsum
export FillMap
export InverseMap

Expand Down Expand Up @@ -78,6 +78,8 @@ 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[begin:end-1]
_tail(As::Tuple) = Base.tail(As)
Expand Down Expand Up @@ -331,7 +333,7 @@ end

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 transpose or adjoint vector
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
Expand Down
217 changes: 173 additions & 44 deletions src/kronecker.jl
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,45 @@ for k in 3:8 # is 8 sufficient?
kron($(mapargs...), $(Symbol(:A, k)), convert_to_lmaps(As...)...)
end

@doc raw"""
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[^1]:
```math
\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
```math
\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.
[^1]: Fernandes, P. and Plateau, B. and Stewart, W. J. ["Efficient Descriptor-Vector Multiplications in Stochastic Automata Networks"](https://doi.org/10.1145/278298.278303), _Journal of the ACM_, 45(3), 381–414, 1998.
"""
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"))
Expand Down Expand Up @@ -114,74 +153,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, L::OuterProductMap, x::AbstractVector)
a, bt = L.maps
mul!(y, a.lmap, bt.lmap * x)
end
function _unsafe_mul!(y, 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, 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
Expand Down Expand Up @@ -225,7 +280,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 ||
Expand Down Expand Up @@ -269,6 +324,68 @@ kronsum(A::MapOrMatrix, B::MapOrMatrix) =
kronsum(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...) =
kronsum(A, kronsum(B, C, Ds...))

@doc raw"""
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[^1]:
```math
\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
```math
\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
```
[^1]: Fernandes, P. and Plateau, B. and Stewart, W. J. ["Efficient Descriptor-Vector Multiplications in Stochastic Automata Networks"](https://doi.org/10.1145/278298.278303), _Journal of the ACM_, 45(3), 381–414, 1998.
"""
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"))
Expand All @@ -280,14 +397,26 @@ 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)` for `k ≥ 3`.
# 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{2}) =
kronsum(convert(LinearMap, A), convert(LinearMap, A))
Base.:(^)(A::MapOrMatrix, ::KronSumPower{p}) where {p} =
kronsum(ntuple(n -> convert(LinearMap, A), Val(p))...)
sumkronsum(ntuple(_ -> 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))
Expand All @@ -305,10 +434,10 @@ Base.:(==)(A::KroneckerSumMap, B::KroneckerSumMap) =

function _unsafe_mul!(y, 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
Expand Down
2 changes: 1 addition & 1 deletion src/left.jl
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ function _unsafe_mul!(X, Y::TransposeAbsVecOrMat, A::LinearMap)
return X
end
# unwrap WrappedMaps
_unsafe_mul!(X, Y::AbstractMatrix, A::WrappedMap) = mul!(X, Y, A.lmap)
_unsafe_mul!(X, Y::AbstractMatrix, A::WrappedMap) = _unsafe_mul!(X, Y, A.lmap)
# disambiguation
_unsafe_mul!(X, Y::TransposeAbsVecOrMat, A::WrappedMap) = _unsafe_mul!(X, Y, A.lmap)

Expand Down
Loading

2 comments on commit e39cbc1

@dkarrasch
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/63202

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v3.8.0 -m "<description of version>" e39cbc17d5b3c0d556e9fe9f251374f62c637826
git push origin v3.8.0

Please sign in to comment.