Reducing allocations for repeatedly generating MvNormal distributions

General Usage Performance
1

I have to repeatedly generate few random numbers from many different bivariate normal distributions. I want to do this with as few allocations as possible. A simplified example of my code:

using Distributions
using Random, BenchmarkTools
function generate_measurement_model(x,θ)
    Σ = [x[1]*θ[1] x[1]*θ[2]; x[1]*θ[2] x[1]*θ[1]]
    μ = [x[1]^2 + x[2]^2; x[1]*θ[1]]
    MvNormal(μ,Σ)
end
function measure!(y,x,θ)
    measurement_model = generate_measurement_model(x,θ)
    rand!(measurement_model,y)
    return nothing
end
function main!(y_particles,x_particles,θ)
    length_particle = size(y_particles,1)
    num_particles = size(y_particles,2)
    for k = 1:num_particles
        measure!(@view(y_particles[:,k]),x_particles[:,k],θ)
    end
    return nothing
end
Random.seed!(123);
n = 100_000
y_particles = rand(2,n)
x_particles = rand(2,n)
θ = [1.0; 0.1]
main!(y_particles,x_particles,θ)
@btime main!(y_particles,x_particles,θ) #49ms 69MiB

I have tried to get rid of as many allocations as possible, but did not get rid of everything.
Particularly cholesky seems to still allocate quite a bit, in my naive opinion. I would appreciate any advice.

using Distributions
using LinearAlgebra
using PDMats
using Random, BenchmarkTools
function generate_measurement_model(x,θ,μ,Σ,Σ_copy)
    Σ[1,1] = x[1]*θ[1]
    Σ[1,2] = x[1]*θ[2]
    Σ[2,1] = x[1]*θ[2]
    Σ[2,2] = x[1]*θ[1]
    copy!(Σ_copy,Σ)
    Σ_chol = cholesky!(Σ_copy)
    Σ_PDM = PDMat(Σ,Σ_chol)
    μ[1] = x[1]^2 + x[2]^2
    μ[2] = x[1]*θ[1]
    MvNormal(μ,Σ_PDM)
end
function measure!(y,x,θ)
    measurement_model = generate_measurement_model(x,θ)
    rand!(measurement_model,y)
    return nothing
end
function main!(y_particles,x_particles,θ,μ,Σ,Σ_copy)
    n = size(y_particles,2)
    for k = 1:n
        measure!(@view(y_particles[:,k]),@view(x_particles[:,k]),θ)
    end
end
Random.seed!(123);
n = 100_000
y_particles = rand(2,n)
x_particles = rand(2,n)
θ = [1.0; 0.1]
μ = Vector{Float64}(undef,2)
Σ = Matrix{Float64}(undef,2,2)
Σ_copy = similar(Σ)
generate_measurement_model(x,θ) = generate_measurement_model(x,θ,μ::Vector{Float64},Σ::Matrix{Float64},Σ_copy::Matrix{Float64})
main!(y_particles,x_particles,θ,μ,Σ,Σ_copy) 
@btime main!(y_particles,x_particles,θ,μ,Σ,Σ_copy)  #31ms 27MiB
@code_warntype main!(y_particles,x_particles,θ,μ,Σ,Σ_copy)
2

I would normally suggest StaticArrays, but unfortunately PDMat does not accept an SMatrix. You could make a PR to PDMats.jl fixing this, and use static vectors throughout, that should be very fast.

Incidentally, don’t use [a; b] for vectors. You really want [a, b].

1 Like
3

Thank you for the tips, I will give StaticArrays a shot.