Accept any AbstractVecOrMat when sampling from an MvNormal

Project.toml

@@ -1,7 +1,7 @@
 name = "Distributions"
 uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
 authors = ["JuliaStats"]
-version = "0.25.130"
+version = "0.25.131"
 
 [deps]
 AliasTables = "66dad0bd-aa9a-41b7-9441-69ab47430ed8"

src/multivariate/mvnormal.jl

@@ -256,22 +256,12 @@ gradlogpdf(d::MvNormal, x::AbstractVector{<:Real}) = -(d.Σ \ (x .- d.μ))
 
 # Sampling (for GenericMvNormal)
 
-function _rand!(rng::AbstractRNG, d::MvNormal, x::VecOrMat)
+function _rand!(rng::AbstractRNG, d::MvNormal, x::AbstractVecOrMat)
     unwhiten!(d.Σ, randn!(rng, x))
     x .+= d.μ
     return x
 end
 
-# Workaround: randn! only works for Array, but not generally for AbstractArray
-function _rand!(rng::AbstractRNG, d::MvNormal, x::AbstractVector)
-    for i in eachindex(x)
-        x[i] = randn(rng, eltype(x))
-    end
-    unwhiten!(d.Σ, x)
-    x .+= d.μ
-    return x
-end
-
 ### Affine transformations
 
 Base.:+(d::MvNormal, c::AbstractVector) = MvNormal(d.μ + c, d.Σ)

test/multivariate/mvnormal.jl

@@ -9,6 +9,7 @@ using Distributions
 using LinearAlgebra, Random, Test
 using SparseArrays
 using FillArrays
+using StableRNGs
 
 ###### General Testing
 
@@ -302,3 +303,46 @@ end
     x = rand(d)
     @test logpdf(d, x) ≈ logpdf(Normal(), x[1]) + logpdf(Normal(), x[2])
 end
+
+@testset "MvNormal: Sampling into an AbstractVecOrMat (#2086)" begin
+    μ = [1.0, 2.0, 3.0]
+    C = [4.0 -2.0 -1.0; -2.0 5.0 -1.0; -1.0 -1.0 6.0]
+
+    Σs = (PDMat(C), PDiagMat([1.2, 3.4, 2.6]), ScalMat(3, 2.0))
+    @testset "$(nameof(typeof(Σ)))" for Σ in Σs
+        d = MvNormal(μ, Σ)
+        x = rand!(StableRNG(1234), d, Matrix{Float64}(undef, 3, 6))
+
+        # A view of a whole matrix takes the same single `randn!` and `unwhiten!` call as
+        # the matrix, so the samples are exactly equal. Only RNGs without a bulk `randn!`
+        # method give exact equality: the ones in `Random` are defined for
+        # `MersenneTwister`/`Xoshiro` and `Array{Float64}`, which a view is not.
+        y = Matrix{Float64}(undef, 3, 6)
+        rand!(StableRNG(1234), d, view(y, :, :))
+        @test y == x
+
+        # A view of some columns is filled like a matrix of the same size, and no other
+        # column is touched.
+        y = zeros(3, 6)
+        rand!(StableRNG(1234), d, view(y, :, 3:4))
+        @test y[:, 3:4] ≈ rand!(StableRNG(1234), d, Matrix{Float64}(undef, 3, 2))
+        @test all(iszero, y[:, [1, 2, 5, 6]])
+
+        # The same for a view that is not strided, for which `unwhiten!` cannot use BLAS.
+        y = zeros(3, 6)
+        rand!(StableRNG(1234), d, view(y, :, [1, 3, 5]))
+        @test y[:, [1, 3, 5]] ≈ rand!(StableRNG(1234), d, Matrix{Float64}(undef, 3, 3))
+        @test all(iszero, y[:, [2, 4, 6]])
+
+        # Vectors do not need a separate method: `randn!` is generic.
+        y = zeros(3, 2)
+        rand!(StableRNG(1234), d, view(y, :, 1))
+        @test y[:, 1] ≈ rand!(StableRNG(1234), d, Vector{Float64}(undef, 3))
+        @test all(iszero, y[:, 2])
+
+        # A view without columns draws nothing.
+        y = zeros(3, 6)
+        rand!(StableRNG(1234), d, view(y, :, 1:0))
+        @test all(iszero, y)
+    end
+end