fix vector-based CompositeMap mul kernel (#175)

Author: dkarrasch

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearMaps"
 uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
-version = "3.6.0"
+version = "3.6.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

docs/src/history.md

@@ -8,7 +8,9 @@
   cases where these higher-order `LinearMap`s are constructed from many maps where a tuple
   backend may get inefficient or impose hard work for the compiler at construction.
   The default behavior, however, does not change, and construction of vector-based
-  `LinearMap`s requires usage of the unexported constructors ("expert usage").
+  `LinearMap`s requires usage of the unexported constructors ("expert usage"), except for
+  constructions like `sum([A, B, C])` or `prod([A, B, C])` (`== C*B*A`), where `A`, `B` and
+  `C` are of some `LinearMap` type.
 
 ## What's new in v3.5
 

src/composition.jl

@@ -205,7 +205,7 @@ function _resize(dest::AbstractMatrix, sz::Tuple{<:Integer,<:Integer})
 end
 
 function _compositemul!(y::AbstractVecOrMat,
-                        A::CompositeMap,
+                        A::CompositeMap{<:Any,<:LinearMapTuple},
                         x::AbstractVecOrMat,
                         source = similar(y, (size(A.maps[1],1), size(x)[2:end]...)),
                         dest = similar(y, (size(A.maps[2],1), size(x)[2:end]...)))
@@ -219,3 +219,40 @@ function _compositemul!(y::AbstractVecOrMat,
     _unsafe_mul!(y, A.maps[N], source)
     return y
 end
+
+function _compositemul!(y::AbstractVecOrMat,
+                        A::CompositeMap{<:Any,<:LinearMapVector},
+                        x::AbstractVecOrMat)
+    N = length(A.maps)
+    if N == 1
+        return _unsafe_mul!(y, A.maps[1], x)
+    elseif N == 2
+        return _compositemul2!(y, A, x)
+    else
+        return _compositemulN!(y, A, x)
+    end
+end
+
+function _compositemul2!(y::AbstractVecOrMat,
+                        A::CompositeMap{<:Any,<:LinearMapVector},
+                        x::AbstractVecOrMat,
+                        source = similar(y, (size(A.maps[1],1), size(x)[2:end]...)))
+    _unsafe_mul!(source, A.maps[1], x)
+    _unsafe_mul!(y, A.maps[2], source)
+    return y
+end
+function _compositemulN!(y::AbstractVecOrMat,
+                            A::CompositeMap{<:Any,<:LinearMapVector},
+                            x::AbstractVecOrMat,
+                            source = similar(y, (size(A.maps[1],1), size(x)[2:end]...)),
+                            dest = similar(y, (size(A.maps[2],1), size(x)[2:end]...)))
+    N = length(A.maps)
+    _unsafe_mul!(source, A.maps[1], x)
+    for n in 2:N-1
+        dest = _resize(dest, (size(A.maps[n],1), size(x)[2:end]...))
+        _unsafe_mul!(dest, A.maps[n], source)
+        dest, source = source, dest # alternate dest and source
+    end
+    _unsafe_mul!(y, A.maps[N], source)
+    return y
+end

test/composition.jl

@@ -148,4 +148,11 @@ using LinearMaps: LinearMapVector, LinearMapTuple
     @test @inferred isposdef(LinearMaps.CompositeMap{Float64}([B, B, B, B, B']))
     @test @inferred isposdef(B * B) # even case for empty tuple test
     @test @inferred isposdef(LinearMaps.CompositeMap{Float64}([B, B]))
+
+    M = LinearMap(cumsum, 3)
+    for i in 1:4
+        P = prod(fill(M, i))
+        @test P isa LinearMaps.CompositeMap{<:Any,<:LinearMapVector}
+        @test P * ones(3) == (LowerTriangular(ones(3,3))^i) * ones(3)
+    end
 end