add multiplication with vectors method (#32)

* add multiplication method for BlockDiagonal*Vector

* allow FillArrays v0.8

* bump patch version

* reduce code duplication

* code style as per review

Project.toml

@@ -1,14 +1,14 @@
 name = "BlockDiagonals"
 uuid = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
-FillArrays = "0.6, 0.7"
+FillArrays = "0.6, 0.7, 0.8"
 julia = "1"
 
 [extras]

src/base_maths.jl

@@ -76,23 +76,31 @@ function Base.:*(B1::BlockDiagonal, B2::BlockDiagonal)
     end
 end
 
-function _check_matmul_dims(A::AbstractMatrix, B::AbstractMatrix)
+function _check_matmul_dims(A::AbstractMatrix, B::AbstractVecOrMat)
     # match error message from LinearAlgebra
     size(A, 2) == size(B, 1) || throw(DimensionMismatch(
         "A has dimensions $(size(A)) but B has dimensions $(size(B))"
     ))
 end
 
-function Base.:*(B::BlockDiagonal{T}, M::AbstractMatrix) where T
-    _check_matmul_dims(B, M)
+_mulblocksizes(bblocks, ::AbstractVector) = size.(bblocks, 1)
+function _mulblocksizes(bblocks, M::AbstractMatrix)
+    return zip(size.(bblocks, 1), Base.Iterators.repeated(size(M, 2), length(bblocks)))
+end
+
+Base.:*(B::BlockDiagonal, x::AbstractVector) = _mul(B, x)
+Base.:*(B::BlockDiagonal, X::AbstractMatrix) = _mul(B, X)
+
+function _mul(B::BlockDiagonal{T}, x::AbstractVecOrMat) where {T}
+    _check_matmul_dims(B, x)
     bblocks = blocks(B)
-    new_blocksizes = zip(size.(bblocks, 1), fill(size(M, 2), length(bblocks)))
+    new_blocksizes = _mulblocksizes(bblocks, x)
     d = similar.(bblocks, T, new_blocksizes)
     ed = 0
     @inbounds @views for (p, block) in enumerate(bblocks)
         st = ed + 1  # start
         ed += size(block, 2)  # end
-        mul!(d[p], block, M[st:ed, :])
+        mul!(d[p], block, selectdim(x, 1, st:ed))
     end
     return reduce(vcat, d)
 end