diff --git a/src/LinearMaps.jl b/src/LinearMaps.jl index 373cf5b5..c67ed14b 100644 --- a/src/LinearMaps.jl +++ b/src/LinearMaps.jl @@ -5,6 +5,7 @@ export ⊗, kronsum, ⊕ using LinearAlgebra using SparseArrays +import Base.Broadcast: materialize! if VERSION < v"1.2-" import Base: has_offset_axes diff --git a/src/blockmap.jl b/src/blockmap.jl index f5caad7a..32ead314 100644 --- a/src/blockmap.jl +++ b/src/blockmap.jl @@ -1,19 +1,91 @@ -struct BlockMap{T,As<:Tuple{Vararg{LinearMap}},Rs<:Tuple{Vararg{Int}},Rranges<:Tuple{Vararg{UnitRange{Int}}},Cranges<:Tuple{Vararg{UnitRange{Int}}}} <: LinearMap{T} +struct BlockMap{T,As,Rs<:Tuple{Vararg{Int}},Rranges<:Tuple{Vararg{UnitRange{Int}}},Cranges<:Tuple{Vararg{UnitRange{Int}}}} <: LinearMap{T} maps::As rows::Rs rowranges::Rranges colranges::Cranges + # Store BlockMap size explicitly, to allow south-east corner to be + # empty. + M::Int + N::Int + function BlockMap{T,R,S}(maps::R, rows::S) where {T, R<:Tuple{Vararg{LinearMap}}, S<:Tuple{Vararg{Int}}} for A in maps promote_type(T, eltype(A)) == T || throw(InexactError()) end rowranges, colranges = rowcolranges(maps, rows) - return new{T,R,S,typeof(rowranges),typeof(colranges)}(maps, rows, rowranges, colranges) + M,N = last(last(rowranges)), last(last(colranges)) + return new{T,R,S,typeof(rowranges),typeof(colranges)}(maps, rows, rowranges, colranges, M, N) end + + BlockMap{T}(maps::As, rows::Rs, rowranges::Rranges, colranges::Cranges, + M::Integer, N::Integer) where {T,As,Rs,Rranges,Cranges} = + new{T,As,Rs,Rranges,Cranges}(maps, rows, rowranges, colranges, M, N) end BlockMap{T}(maps::As, rows::S) where {T,As<:Tuple{Vararg{LinearMap}},S} = BlockMap{T,As,S}(maps, rows) +""" + BlockMap{T}(M, N, (m, n), maps, is, js) + +Constructor for a `BlockMap` of `eltype(T)` for the case when all rows +have the same column dimensions; these are specified using the vectors +`m` and `n`. Only those blocks `maps` that are actually occupied need +to be specified, their coordinates are given by the vectors `is` and `js`. +""" +function BlockMap{T}(M::Int, N::Int, (m,n), maps, is, js) where T + all(m .> 0) && all(n .> 0) || + throw(ArgumentError("Block sizes must be positive integers")) + + sum(m) == M && sum(n) == N || + throw(DimensionMismatch("Cumulative block dimensions $(sum(m))×$(sum(n)) do not agree with overall size $(M)×$(N)")) + + length(maps) == length(is) == length(js) || + throw(ArgumentError("Must provide block coordinates for $(length(maps)) blocks")) + + allunique(zip(is, js)) || + throw(ArgumentError("Not all block indices unique")) + + for (A,i,j) in zip(maps,is,js) + promote_type(T, eltype(A)) == T || throw(InexactError()) + 0 < i ≤ length(m) || throw(ArgumentError("Invalid block row $(i) ∉ 1:$(length(m))")) + 0 < j ≤ length(n) || throw(ArgumentError("Invalid block column $(j) ∉ 1:$(length(n))")) + size(A) == (m[i],n[j]) || + throw(DimensionMismatch("Block of size $(size(A)) does not match size at $((i,j)): $((m[i],n[j]))")) + end + + rows = zeros(Int, length(m)) + colranges = () + colstarts = 1 .+ vcat(0,cumsum(n)) + for p = sortperm(collect(zip(is,js))) + A,i,j = maps[p],is[p],js[p] + rows[i] += 1 + colranges = (colranges..., colstarts[j]:colstarts[j+1]-1) + end + rows = (rows...,) + + rowranges = () + i = 1 + for (mm,r) in zip(m,rows) + iprev = i + i += mm + rowranges = (rowranges...,iprev:i-1) + end + + maps = map(convert_to_lmaps_, maps) + + BlockMap{T}(maps, rows, rowranges, colranges, M, N) +end + +function Base.show(io::IO, B::BlockMap{T}) where T + nrows = length(B.rows) + # One block in one row may actually span multiple block columns + # for other rows, but it's nice to get a sense of the structure + # anyway. + ncols = maximum(B.rows) + M,N = size(B) + write(io, "$(nrows)×$(ncols)-blocked $(M)×$(N) BlockMap{$T}") +end + MulStyle(A::BlockMap) = MulStyle(A.maps...) function check_dim(A::LinearMap, dim, n) @@ -49,7 +121,7 @@ function rowcolranges(maps, rows) return rowranges::NTuple{length(rows), UnitRange{Int}}, colranges::NTuple{length(maps), UnitRange{Int}} end -Base.size(A::BlockMap) = (last(last(A.rowranges)), last(last(A.colranges))) +Base.size(A::BlockMap) = (A.M, A.N) ############ # concatenation @@ -306,6 +378,7 @@ LinearAlgebra.adjoint(A::BlockMap) = AdjointMap(A) maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges mapind = 0 @views @inbounds for (row, yi) in zip(rows, yinds) + iszero(row) && continue yrow = selectdim(y, 1, yi) mapind += 1 mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, β) @@ -377,6 +450,29 @@ for Atype in (AbstractVector, AbstractMatrix) end end +############ +# materialization into matrices +############ + +function materialize!(M::MT, A::BlockMap) where {MT<:AbstractArray} + axes(M) == axes(A) || + throw(DimensionMismatch("Cannot materialize BlockMap of size $(size(A)) into matrix of size $(size(M))")) + + M .= false + + maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges + mapind = 0 + @views @inbounds for (row, yi) in zip(rows, yinds) + Mrow = selectdim(M, 1, yi) + for _ in 1:row + mapind +=1 + copyto!(selectdim(Mrow, 2, xinds[mapind]), convert(Matrix, maps[mapind])) + end + end + + M +end + ############ # show methods ############ diff --git a/src/conversion.jl b/src/conversion.jl index f3e9a45e..49402658 100644 --- a/src/conversion.jl +++ b/src/conversion.jl @@ -44,7 +44,7 @@ function Base.convert(::Type{Matrix}, Aλ::CompositeMap{<:Any,<:Tuple{UniformSca return A.lmap*J.λ end -Base.Matrix(A::BlockMap) = hvcat(A.rows, convert.(AbstractMatrix, A.maps)...) +Base.Matrix(A::BlockMap) = materialize!(Matrix{eltype(A)}(undef, size(A)), A) # sparse: create sparse matrix representation of LinearMap function SparseArrays.sparse(A::LinearMap{T}) where {T} @@ -73,4 +73,4 @@ end Base.convert(::Type{SparseMatrixCSC}, A::LinearMap) = sparse(A) Base.convert(::Type{SparseMatrixCSC}, A::WrappedMap) = convert(SparseMatrixCSC, A.lmap) -Base.convert(::Type{SparseMatrixCSC}, A::BlockMap) = hvcat(A.rows, convert.(SparseMatrixCSC, A.maps)...) +Base.convert(::Type{SparseMatrixCSC}, A::BlockMap) = materialize!(spzeros(eltype(A), size(A)...), A) diff --git a/src/linearcombination.jl b/src/linearcombination.jl index 7f556918..8718d917 100644 --- a/src/linearcombination.jl +++ b/src/linearcombination.jl @@ -1,4 +1,4 @@ -struct LinearCombination{T, As<:Tuple{Vararg{LinearMap}}} <: LinearMap{T} +struct LinearCombination{T, As} <: LinearMap{T} maps::As function LinearCombination{T, As}(maps::As) where {T, As} N = length(maps) @@ -13,6 +13,33 @@ end LinearCombination{T}(maps::As) where {T, As} = LinearCombination{T, As}(maps) +function Base.show(io::IO, A::LinearCombination) + write(io, "LinearMaps.LinearCombination(") + n = length(A.maps) + if get(io, :limit, true) && n > 10 + s = 1:5 + e = n-5:n + if e[1]-s[end] > 1 + write(io, "(") + for i in s + show(io, A.maps[i]) + write(io, ", ") + end + write(io, "…") + for i in e + write(io, ", ") + show(io, A.maps[i]) + end + write(io, ")") + else + show(io, A.maps) + end + else + show(io, A.maps) + end + write(io, ")") +end + MulStyle(A::LinearCombination) = MulStyle(A.maps...) # basic methods @@ -81,6 +108,9 @@ for Atype in (AbstractVector, AbstractMatrix) end end +# For tuple-like storage of the maps (default), we recurse on the tail +# of the tuple. + @inline _mul!(::FiveArg, y, A::LinearCombination, x, α::Number) = __mul!(y, Base.tail(A.maps), x, α, nothing) @inline function _mul!(::ThreeArg, y, A::LinearCombination, x, α::Number) @@ -93,6 +123,22 @@ end @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 vector-like storage of the maps, we simply loop over the maps. + +@inline function _mul!(::FiveArg, y, A::LinearCombination{<:Any,<:AbstractVector}, x, α::Number) + for i = 2:length(A.maps) + mul!(y, A.maps[i], x, α, true) + end + y +end +@inline function _mul!(mulstyle::ThreeArg, y, A::LinearCombination{<:Any,<:AbstractVector}, x, α::Number) + z = similar(y) + for i = 2:length(A.maps) + muladd!(mulstyle, y, A.maps[i], x, α, z) + end + y +end + @inline muladd!(::FiveArg, y, A, x, α, _) = mul!(y, A, x, α, true) @inline function muladd!(::ThreeArg, y, A, x, α, z) mul!(z, A, x)