Sampling from MvNormal with BlockDiagonal covariance is slow

[mzgubic]

using Distributions
using LinearAlgebra
using PDMats
using BlockDiagonals

blocksize = 2
nblocks = 100
ϕ = rand(blocksize, blocksize)
means = rand(nblocks*blocksize)
cov_pd = PDMat(BlockDiagonal([ϕ'ϕ for _ in 1:nblocks]));
cov_m = Matrix(Matrix(cov_pd));

mvnm = MvNormal(means, cov_m)
mvnb = MvNormal(means, cov_pd)

julia> @time rand(mvnm, 100);
  0.000289 seconds (3 allocations: 156.312 KiB)

julia> @time rand(mvnb, 100);
  0.570064 seconds (1.97 M allocations: 630.919 MiB, 13.55% gc time)

Upon some inspection with @profview it turns out that time is mostly spent in getindex

and the stacktrace (by putting in an error manually) reveals that lmul!(LowerTriangular{BlockDiagonal}, Matrix) is where the indexing calls happen:

ERROR: 
Stacktrace:
  [1] error()
    @ Base ./error.jl:44
  [2] getindex
    @ ~/JuliaEnvs/BlockDiagonals.jl/src/blockdiagonal.jl:140 [inlined]
  [3] lmul!(A::LowerTriangular{Float64, BlockDiagonal{Float64, Adjoint{Float64, Matrix{Float64}}}}, B::Matrix{Float64})
    @ LinearAlgebra /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/triangular.jl:937
  [4] unwhiten!
    @ ~/.julia/packages/PDMats/ZW0lN/src/generics.jl:38 [inlined]
  [5] unwhiten!
    @ ~/.julia/packages/PDMats/ZW0lN/src/generics.jl:29 [inlined]
  [6] _rand!(rng::Random._GLOBAL_RNG, d::MvNormal{Float64, PDMat{Float64, BlockDiagonal{Float64, Matrix{Float64}}}, Vector{Float64}}, x::Matrix{Float64})
    @ Distributions ~/.julia/packages/Distributions/0Nl1l/src/multivariate/mvnormal.jl:277
  [7] rand!
    @ ~/.julia/packages/Distributions/0Nl1l/src/genericrand.jl:108 [inlined]
  [8] rand
    @ ~/.julia/packages/Distributions/0Nl1l/src/multivariates.jl:23 [inlined]
  [9] rand(s::MvNormal{Float64, PDMat{Float64, BlockDiagonal{Float64, Matrix{Float64}}}, Vector{Float64}}, dims::Int64)
    @ Distributions ~/.julia/packages/Distributions/0Nl1l/src/genericrand.jl:22
 [10] top-level scope
    @ ./timing.jl:263 [inlined]
 [11] top-level scope
    @ ./REPL[198]:0

[mjp98]

This appears to be resolvable by specialising lmul! with something like:

function LinearAlgebra.lmul!(B::LowerTriangular{<:Any,<:BlockDiagonal}, vm::StridedVecOrMat)
    row_i = 1
    # BlockDiagonals with non-square blocks
    if !all(BlockDiagonals.is_square, blocks(parent(B)))
        return lmul!(LowerTriangular(Matrix(B)), vm) # Fallback on the generic LinearAlgebra method
    end
    for block in blocks(parent(B))
        nrow = size(block, 1)
        @views lmul!(LowerTriangular(block), vm[row_i:(row_i + nrow - 1), :])
        row_i += nrow
    end
    vm
end

which results in

julia> @time rand(mvnm, 100);
  0.000152 seconds (3 allocations: 156.312 KiB)

julia> @time rand(mvnb, 100);
  0.000134 seconds (5 allocations: 158.062 KiB)

Will put together a PR. May be worth looking through other common LinearAlgebra methods that could do with specialisation.