Skip to content

Commit

Permalink
Relaxed type requirements on maps storage for LinearCombination and B…
Browse files Browse the repository at this point in the history
…lockMap
  • Loading branch information
jagot committed Feb 12, 2020
1 parent 8e3b35b commit ce8e7c5
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 5 deletions.
10 changes: 6 additions & 4 deletions src/blockmap.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
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
Expand Down Expand Up @@ -39,6 +39,9 @@ function BlockMap{T}(M::Int, N::Int, (m,n), maps, is, js) where T
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"))

Expand All @@ -50,9 +53,6 @@ function BlockMap{T}(M::Int, N::Int, (m,n), maps, is, js) where T
throw(DimensionMismatch("Block of size $(size(A)) does not match size at $((i,j)): $((m[i],n[j]))"))
end

# It would be nice to do just maps = map(LinearMap, maps)
maps = (map(A -> A isa LinearMap ? A : LinearMap(A), maps)...,)

rows = zeros(Int, length(m))
colranges = ()
colstarts = 1 .+ vcat(0,cumsum(n))
Expand All @@ -71,6 +71,8 @@ function BlockMap{T}(M::Int, N::Int, (m,n), maps, is, js) where T
rowranges = (rowranges...,iprev:i-1)
end

maps = map(convert_to_lmaps_, maps)

BlockMap{T}(maps, rows, rowranges, colranges, M, N)
end

Expand Down
48 changes: 47 additions & 1 deletion src/linearcombination.jl
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down

0 comments on commit ce8e7c5

Please sign in to comment.