Extension needed for StaticArrays

[astrozot]

I believe we should consider writing an extension for StaticArrays, so that multivariate distributions with "static" parameters works as expected. For example, as noted in #992, rand(MvNormal(μ::SVector, Σ)) should return an SVector. Similarly, when building a MvNormal with just the covariance as a static matrix, the mean should be set to a vanishing static vector.

I think the work should be not too complicated, but the gain is obvious as we would avoid allocations when using these distributions with small static vectors. For example, these few lines could solve the two issues I just mentioned:

Distributions.MvNormal(Σ::SMatrix{N, N, T}) where {N,T<:Real} = 
    Distributions.MvNormal(zero(SVector{N,T}), Σ)

Distributions.rand(rng::Random.AbstractRNG, d::Distributions.MvNormal{T,P,<:SVector{N}}) where {T, P, N} = 
    Distributions.unwhiten(d.Σ, randn(rng, SVector{N, T})) .+ d.μ

Moreover, I propose that when generating multiple random vectors from static-array multivariate distributions, the result is not a matrix, but a vector of static vectors:

Distributions.rand(rng::Random.AbstractRNG, d::Distributions.MvNormal{T,P,<:SVector{N}}, n::Int) where {T,P,N} =
    [Distributions.unwhiten(d.Σ, randn(rng, SVector{N,T})) .+ d.μ for _ in 1:n]

I believe this is much more natural; also, conversion to a matrix (and vice-versa) is trivial (see this section in the documentation of StaticArrays).

If there is agreement on this I would be happy helping to implement the extension.

[devmotion]

Similarly, when building a MvNormal with just the covariance as a static matrix, the mean should be set to a vanishing static vector.

There's nothing to be done in this case. If you only specify a covariance matrix, the mean is set to a non-allocating FillArrays.Zeros (for which eg broadcasting operations are also optimized).

[devmotion]

Moreover, I propose that when generating multiple random vectors from static-array multivariate distributions, the result is not a matrix, but a vector of static vectors:

This cannot be done as it would break any downstream package relying on the API of rand for multivariate distributions. The behaviour of rand must be consistent regardless of the parameter type.

That is, if this seems generally desirable, it must be changed for all parameter types in a breaking release. But with the currently documented API it's not possible.

[astrozot]

Yes, I agree, this is a breaking change which would require a new (breaking) release.

In my opinion it is worth doing it, as it saves allocations, and thus makes programs faster, in the common situation of multivariate distributions with a small dimensionality. Also, I think that the proposed behavior is what a user using StaticArrays as parameters expects.

[GaneshPatil7517]

Hi @astrozot
I’d like to work on this issue and contribute a fix/extension. Could you please assign it to me?

I’ll start by:

• Reproducing the behavior with StaticArrays parameters
• Studying the current API design to ensure compatibility
• Prototyping the extension and sharing updates via PR

Looking forward to contributing. Thanks!

[astrozot]

Hi @GaneshPatil7517 , thank you for your interest in this issue and for your offer to contribute. I would be happy to assign this issue to you, but I do not have the rights to. Hopefully someone with the sufficient rights can read your request and answer positively to you.

[GaneshPatil7517]

Hi @astrozot ,
Thanks for the clarification and for responding. I understand the permission limitations.
I’ll wait for a maintainer/org admin with sufficient rights to assign the issue.
Meanwhile, I’ll start studying the codebase so I can begin work quickly once it’s assigned.
Appreciate your time and support!