How to efficiently generate a distribution for every element in matrix

for example, I have two matrix which separately represent mu and std:

mu = [1.0 2.0; 3.0 4.0]
std = [0.1 0.2; 0.3 0.4]

I want to generate a matrix Normal Distribution like this:

dist = [Normal(1.0, 0.1) Normal(2.0, 0.2); Normal(3.0, 0.3) Normal(4.0, 0.4)]

Unfortunately, Distributions.jl does not support directly using matrix to generate this distribution, so the only method I can figure out to solve this is like bellow:

mu = [1.0 2.0; 3.0 4.0]
std = [0.1 0.2; 0.3 0.4]
dist = [Normal(mu[i, j], std[i, j]) for i in 1:size(mu, 1), j in 1:size(mu, 2)]

but I think this can be very low efficiency. because the matrix I need to deal with may be very big. so how can I improve the efficiency? :sob:

If you want to shorten the expression, you can use broadcasting, but I think the performance is the almost the same between your list comprehension and the broadcast:

using Distributions
using BenchmarkTools

mu = [1.0 2.0; 3.0 4.0]
sigma = [0.1 0.2; 0.3 0.4]

make_dist(mu, sigma) = dist = [Normal(mu[i, j], sigma[i, j]) for i in 1:size(mu, 1), j in 1:size(mu, 2)]

@btime make_dist($mu, $sigma)
#   38.013 ns (1 allocation: 128 bytes)

@btime Normal.($mu, $sigma)
#  32.173 ns (1 allocation: 128 bytes)

Note: If you want to sample a matrix with the corresponding random entries later, you can also broadcast rand.(dist), otherwise, you will pick the element of the matrix of distributions randomly.

PS: For larger matrices, the broadcast version seems to be consistently a bit faster than the array comprehension.

Oh, thank you very much. I don’t know β€œ.” can even work here, thanks!

1 Like