Skip to content

Commit

Permalink
kind of merge block map enhancements into this branch
Browse files Browse the repository at this point in the history
  • Loading branch information
@dkarrasch
dkarrasch committed Nov 18, 2019
1 parent 0c21402 commit 450bfec
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 38 deletions.
166 changes: 129 additions & 37 deletions src/blockmap.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,75 +297,85 @@ LinearAlgebra.transpose(A::BlockMap) = TransposeMap(A)
LinearAlgebra.adjoint(A::BlockMap) = AdjointMap(A)

############
# multiplication with vectors
# multiplication helper functions
############

function A_mul_B!(y::AbstractVector, A::BlockMap, x::AbstractVector)
require_one_based_indexing(y, x)
m, n = size(A)
@boundscheck (m == length(y) && n == length(x)) || throw(DimensionMismatch("A_mul_B!"))
@inline function _blockmul!(y, A::BlockMap, x, α, β)
maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges
mapind = 0
@views @inbounds for rowind in 1:length(rows)
yrow = y[yinds[rowind]]
yrow = selectdim(y, 1, yinds[rowind])
mapind += 1
A_mul_B!(yrow, maps[mapind], x[xinds[mapind]])
mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, β)
for colind in 2:rows[rowind]
mapind +=1
mul!(yrow, maps[mapind], x[xinds[mapind]], true, true)
mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, true)
end
end
return y
end

function At_mul_B!(y::AbstractVector, A::BlockMap, x::AbstractVector)
require_one_based_indexing(y, x)
m, n = size(A)
@boundscheck (n == length(y) && m == length(x)) || throw(DimensionMismatch("At_mul_B!"))
@inline function _transblockmul!(y, A::BlockMap, x, α, β, transform)
maps, rows, xinds, yinds = A.maps, A.rows, A.rowranges, A.colranges
mapind = 0
# first block row (rowind = 1) of A, meaning first block column of A', fill all of y
@views @inbounds begin
xcol = x[xinds[1]]
xcol = selectdim(x, 1, xinds[1])
for colind in 1:rows[1]
mapind +=1
A_mul_B!(y[yinds[mapind]], transpose(maps[mapind]), xcol)
mul!(selectdim(y, 1, yinds[mapind]), transform(maps[mapind]), xcol, α, β)
end
# subsequent block rows of A, add results to corresponding parts of y
for rowind in 2:length(rows)
xcol = x[xinds[rowind]]
xcol = selectdim(x, 1, xinds[rowind])
for colind in 1:rows[rowind]
mapind +=1
mul!(y[yinds[mapind]], transpose(maps[mapind]), xcol, true, true)
mul!(selectdim(y, 1, yinds[mapind]), transform(maps[mapind]), xcol, α, true)
end
end
end
return y
end

function Ac_mul_B!(y::AbstractVector, A::BlockMap, x::AbstractVector)
############
# multiplication with vectors
############

Base.@propagate_inbounds function LinearAlgebra.mul!(y::AbstractVector, A::BlockMap, x::AbstractVector,
α::Number=true, β::Number=false)
require_one_based_indexing(y, x)
m, n = size(A)
@boundscheck (n == length(y) && m == length(x)) || throw(DimensionMismatch("At_mul_B!"))
maps, rows, xinds, yinds = A.maps, A.rows, A.rowranges, A.colranges
mapind = 0
# first block row (rowind = 1) of A, fill all of y
@views @inbounds begin
xcol = x[xinds[1]]
for colind in 1:rows[1]
mapind +=1
A_mul_B!(y[yinds[mapind]], adjoint(maps[mapind]), xcol)
end
# subsequent block rows of A, add results to corresponding parts of y
for rowind in 2:length(rows)
xcol = x[xinds[rowind]]
for colind in 1:rows[rowind]
mapind +=1
mul!(y[yinds[mapind]], adjoint(maps[mapind]), xcol, true, true)
end
end
@boundscheck check_dim_mul(y, A, x)
return _blockmul!(y, A, x, α, β)
end

for (maptype, transform) in ((:(TransposeMap{<:Any,<:BlockMap}), :transpose), (:(AdjointMap{<:Any,<:BlockMap}), :adjoint))
@eval Base.@propagate_inbounds function LinearAlgebra.mul!(y::AbstractVector, wrapA::$maptype, x::AbstractVector,
α::Number=true, β::Number=false)
require_one_based_indexing(y, x)
@boundscheck check_dim_mul(y, wrapA, x)
return _transblockmul!(y, wrapA.lmap, x, α, β, $transform)
end
end

############
# multiplication with matrices
############

Base.@propagate_inbounds function LinearAlgebra.mul!(Y::AbstractMatrix, A::BlockMap, X::AbstractMatrix,
α::Number=true, β::Number=false)
require_one_based_indexing(Y, X)
@boundscheck check_dim_mul(Y, A, X)
maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges
return _blockmul!(Y, A, X, α, β)
end

for (maptype, transform) in ((:(TransposeMap{<:Any,<:BlockMap}), :transpose), (:(AdjointMap{<:Any,<:BlockMap}), :adjoint))
@eval Base.@propagate_inbounds function LinearAlgebra.mul!(Y::AbstractMatrix, wrapA::$maptype, X::AbstractMatrix,
α::Number=true, β::Number=false)
require_one_based_indexing(Y, X)
@boundscheck check_dim_mul(Y, wrapA, X)
return _transblockmul!(Y, wrapA.lmap, X, α, β, $transform)
end
return y
end

############
Expand All @@ -386,3 +396,85 @@ end
# show(io, T)
# print(io, '}')
# end

############
# BlockDiagonalMap
############

struct BlockDiagonalMap{T,As<:Tuple{Vararg{LinearMap}}} <: LinearMap{T}
maps::As
rowranges::Vector{UnitRange{Int}}
colranges::Vector{UnitRange{Int}}
function BlockDiagonalMap{T,As}(maps::As) where {T, As<:Tuple{Vararg{LinearMap}}}
for A in maps
promote_type(T, eltype(A)) == T || throw(InexactError())
end
# row ranges
inds = vcat(1, size.(maps, 1)...)
cumsum!(inds, inds)
rowranges = map(i -> inds[i]:inds[i+1]-1, 1:length(maps))
# column ranges
inds[2:end] .= size.(maps, 2)
cumsum!(inds, inds)
colranges = map(i -> inds[i]:inds[i+1]-1, 1:length(maps))
return new{T,As}(maps, rowranges, colranges)
end
end

BlockDiagonalMap{T}(maps::As) where {T,As<:Tuple{Vararg{LinearMap}}} =
BlockDiagonalMap{T,As}(maps)
BlockDiagonalMap(maps::LinearMap...) =
BlockDiagonalMap{promote_type(map(eltype, maps)...)}(maps)

for k in 1:8 # is 8 sufficient?
Is = ntuple(n->:($(Symbol(:A,n))::AbstractMatrix), Val(k-1))
# yields (:A1, :A2, :A3, ..., :A(k-1))
L = :($(Symbol(:A,k))::LinearMap)
# yields :Ak
mapargs = ntuple(n -> :(LinearMap($(Symbol(:A,n)))), Val(k-1))
# yields (:LinearMap(A1), :LinearMap(A2), ..., :LinearMap(A(k-1)))

@eval begin
SparseArrays.blockdiag($(Is...), $L, As::Union{LinearMap,AbstractMatrix}...) =
BlockDiagonalMap($(mapargs...), $(Symbol(:A,k)), convert_to_lmaps(As...)...)
function Base.cat($(Is...), $L, As::Union{LinearMap,AbstractMatrix}...; dims::Dims{2})
if dims == (1,2)
return BlockDiagonalMap($(mapargs...), $(Symbol(:A,k)), convert_to_lmaps(As...)...)
else
throw(ArgumentError("dims keyword in cat of LinearMaps must be (1,2)"))
end
end
end
end

Base.size(A::BlockDiagonalMap) = (last(A.rowranges[end]), last(A.colranges[end]))

LinearAlgebra.issymmetric(A::BlockDiagonalMap) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::BlockDiagonalMap{<:Real}) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::BlockDiagonalMap) = all(ishermitian, A.maps)

LinearAlgebra.adjoint(A::BlockDiagonalMap{T}) where {T} = BlockDiagonalMap{T}(map(adjoint, A.maps))
LinearAlgebra.transpose(A::BlockDiagonalMap{T}) where {T} = BlockDiagonalMap{T}(map(transpose, A.maps))

Base.:(==)(A::BlockDiagonalMap, B::BlockDiagonalMap) = (eltype(A) == eltype(B) && A.maps == B.maps)

Base.@propagate_inbounds function LinearAlgebra.mul!(y::AbstractVector, A::BlockDiagonalMap, x::AbstractVector, α::Number=true, β::Number=false)
require_one_based_indexing(y, x)
@boundscheck check_dim_mul(y, A, x)
return _blockscaling!(y, A, x, α, β)
end

Base.@propagate_inbounds function LinearAlgebra.mul!(Y::AbstractMatrix, A::BlockDiagonalMap, X::AbstractMatrix,
α::Number=true, β::Number=false)
require_one_based_indexing(Y, X)
@boundscheck check_dim_mul(Y, A, X)
return _blockscaling!(Y, A, X, α, β)
end

@inline function _blockscaling!(y, A::BlockDiagonalMap, x, α, β)
maps, yinds, xinds = A.maps, A.rowranges, A.colranges
@views @inbounds for i in 1:length(maps)
mul!(selectdim(y, 1, yinds[i]), maps[i], selectdim(x, 1, xinds[i]), α, β)
end
return y
end
2 changes: 1 addition & 1 deletion test/runtests.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ include("uniformscalingmap.jl")

include("numbertypes.jl")

# include("blockmap.jl")
include("blockmap.jl")

include("kronecker.jl")

Expand Down

0 comments on commit 450bfec

Please sign in to comment.