Relaxed type requirements on maps storage for LinearCombination and BlockMap

Author: jagot

src/blockmap.jl

@@ -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
@@ -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"))
 
@@ -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))
@@ -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
 

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)