Added linalg mul! overloads for ArrayPartition with tests

Author: Lewih

src/array_partition.jl

@@ -370,3 +370,18 @@ function LinearAlgebra._swap_rows!(B::ArrayPartition, i::Integer, j::Integer)
     B[i], B[j] = B[j], B[i]
     return B
 end
+
+# linalg mul! overloads for ArrayPartition
+function LinearAlgebra.mul!(C::T, A::T, B::AbstractArray) where T<:ArrayPartition
+    @assert length(C.x) == length(A.x)
+    for index = 1:length(C.x)
+        mul!(C.x[index], A.x[index], B)
+    end
+end
+
+function LinearAlgebra.mul!(C::T, A::T, B::T) where T<:ArrayPartition
+    @assert length(C.x) == length(A.x) == length(B.x)
+    for index = 1:length(C.x)
+        mul!(C.x[index], A.x[index], B.x[index])
+    end
+end

test/linalg.jl

@@ -1,5 +1,5 @@
 using RecursiveArrayTools, Test, Random
-using LinearAlgebra
+using LinearAlgebra, SparseArrays
 
 n, m = 5, 6
 bb = rand(n), rand(m)
@@ -26,3 +26,24 @@ for ff in (lu, svd, qr)
     @test ldiv!(FF, bbb) === bbb
     @test A*bbb ≈ b
 end
+
+# linalg mul! overloads
+n, m, l = 5, 6, 7
+bb = rand(n, n), rand(m, n), rand(l, n)
+cc = rand(n), rand(n), rand(n)
+dd = rand(n), rand(m), rand(l)
+b = ArrayPartition(bb)
+c = ArrayPartition(cc)
+d = ArrayPartition(dd)
+A = rand(n)
+for T in (Array{Float64}, Array{ComplexF64}, sparse, )
+    B = T(A)
+    mul!(d, b, A)
+    for i = 1:length(c.x)
+        @test d.x[i] == b.x[i] * A
+    end
+    mul!(d, b, c)
+    for i = 1:length(d.x)
+        @test d.x[i] == b.x[i] * c.x[i]
+    end
+end