Implement std(::MultivariateDistribution) (#1352)

* Implement std(::MultivariateDistribution)

* Update src/multivariates.jl

* Add `std` to docs

* Add `std(::Dirichlet)`

* Safer default implementation

* More tests

---------

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

Author: bicycle1885

docs/src/multivariate.md

@@ -21,6 +21,7 @@ size(::MultivariateDistribution)
 eltype(::Type{MultivariateDistribution})
 mean(::MultivariateDistribution)
 var(::MultivariateDistribution)
+std(::MultivariateDistribution)
 cov(::MultivariateDistribution)
 cor(::MultivariateDistribution)
 entropy(::MultivariateDistribution)

ext/DistributionsTestExt.jl

@@ -33,6 +33,7 @@ function Distributions.TestUtils.test_mvnormal(
     @test length(μ) == d
     @test size(Σ) == (d, d)
     @test var(g) ≈ diag(Σ)
+    @test std(g) ≈ sqrt.(diag(Σ))
     @test entropy(g) ≈ 0.5 * logdet(2π * ℯ * Σ)
     ldcov = logdetcov(g)
     @test ldcov ≈ logdet(Σ)

src/multivariate/product.jl

@@ -51,6 +51,7 @@ function _logpdf(d::Product, x::AbstractVector{<:Real})
 end
 
 mean(d::Product) = mean.(d.v)
+std(d::Product) = std.(d.v)
 var(d::Product) = var.(d.v)
 cov(d::Product) = Diagonal(var(d))
 entropy(d::Product) = sum(entropy, d.v)

src/multivariates.jl

@@ -62,6 +62,13 @@ Compute the vector of element-wise variances for distribution `d`.
 """
 var(d::MultivariateDistribution)
 
+"""
+    std(d::MultivariateDistribution)
+
+Compute the vector of element-wise standard deviations for distribution `d`.
+"""
+std(d::MultivariateDistribution) = sqrt!!(var(d))
+
 """
     entropy(d::MultivariateDistribution)
 

src/product.jl

@@ -83,6 +83,8 @@ cov(d::ProductDistribution) = Diagonal(vec(var(d)))
 ## For product distributions of univariate distributions
 mean(d::ArrayOfUnivariateDistribution) = map(mean, d.dists)
 mean(d::VectorOfUnivariateDistribution{<:Tuple}) = collect(map(mean, d.dists))
+std(d::ArrayOfUnivariateDistribution) = map(std, d.dists)
+std(d::VectorOfUnivariateDistribution{<:Tuple}) = collect(map(std, d.dists))
 var(d::ArrayOfUnivariateDistribution) = map(var, d.dists)
 var(d::VectorOfUnivariateDistribution{<:Tuple}) = collect(map(var, d.dists))
 

src/utils.jl

@@ -97,6 +97,14 @@ isunitvec(v::AbstractVector) = (norm(v) - 1.0) < 1.0e-12
 isprobvec(p::AbstractVector{<:Real}) =
     all(x -> x ≥ zero(x), p) && isapprox(sum(p), one(eltype(p)))
 
+sqrt!!(x::AbstractVector{<:Real}) = map(sqrt, x)
+function sqrt!!(x::Vector{<:Real})
+    for i in eachindex(x)
+        x[i] = sqrt(x[i])
+    end
+    return x
+end
+
 # get a type wide enough to represent all a distributions's parameters
 # (if the distribution is parametric)
 # if the distribution is not parametric, we need this to be a float so that

test/multivariate/dirichlet.jl

@@ -26,6 +26,7 @@ rng = MersenneTwister(123)
         @test mode(d) ≈ fill(1/3, 3)
         @test cov(d)  ≈ [8 -4 -4; -4 8 -4; -4 -4 8] / (36 * 7)
         @test var(d)  ≈ diag(cov(d))
+        @test std(d)  ≈ sqrt.(var(d))
 
         r = Vector{Float64}(undef, 3)
         Distributions.dirichlet_mode!(r, d.alpha, d.alpha0)
@@ -65,6 +66,7 @@ rng = MersenneTwister(123)
         @test mean(d) ≈ v / sum(v)
         @test cov(d)  ≈ [8 -2 -6; -2 5 -3; -6 -3 9] / (36 * 7)
         @test var(d)  ≈ diag(cov(d))
+        @test std(d)  ≈ sqrt.(var(d))
 
         @test pdf(d, [0.2, 0.3, 0.5]) ≈ 3
         @test pdf(d, [0.4, 0.5, 0.1]) ≈ 0.24

test/multivariate/dirichletmultinomial.jl

@@ -33,11 +33,13 @@ d = DirichletMultinomial(10, α)
 @test mean(d) ≈ α * (d.n / d.α0)
 p = d.α / d.α0
 @test var(d)  ≈ d.n * (d.n + d.α0) / (1 + d.α0) .* p .* (1.0 .- p)
+@test std(d) ≈ sqrt.(var(d))
 x = func[2](d, 10_000)
 
 # test statistics with mle fit
 d = fit(DirichletMultinomial, x)
 @test isapprox(mean(d), vec(mean(x, dims=2)), atol=.5)
+@test isapprox(std(d) , vec(std(x, dims=2)) , atol=.5)
 @test isapprox(var(d) , vec(var(x, dims=2)) , atol=.5)
 @test isapprox(cov(d) , cov(x, dims=2)      , atol=.5)
 

test/multivariate/multinomial.jl

@@ -19,6 +19,7 @@ using Test
     @test length(d) == 3
     @test d.n == nt
     @test mean(d) ≈ T[2., 5., 3.]
+    @test std(d)  ≈ T[sqrt(1.6), sqrt(2.5), sqrt(2.1)]
     @test var(d)  ≈ T[1.6, 2.5, 2.1]
     @test cov(d)  ≈ T[1.6 -1.0 -0.6; -1.0 2.5 -1.5; -0.6 -1.5 2.1]
 

test/multivariate/mvlognormal.jl

@@ -32,6 +32,7 @@ function test_mvlognormal(g::MvLogNormal, n_tsamples::Int=10^6,
     @test length(s) == d
     @test size(S) == (d, d)
     @test s          ≈ diag(S)
+    @test std(g)     ≈ sqrt.(diag(S)) 
     @test md         ≈ exp.(mean(g.normal))
     @test mn         ≈ exp.(mean(g.normal) .+ var(g.normal)/2)
     @test mo         ≈ exp.(mean(g.normal) .- var(g.normal))

test/multivariate/product.jl

@@ -23,6 +23,7 @@ using Distributions: Product
     @test eltype(d_product) === eltype(ds[1])
     @test @inferred(logpdf(d_product, x)) ≈ sum(logpdf.(ds, x))
     @test mean(d_product) == mean.(ds)
+    @test std(d_product) == std.(ds)
     @test var(d_product) == var.(ds)
     @test cov(d_product) == Diagonal(var.(ds))
     @test entropy(d_product) ≈ sum(entropy.(ds))
@@ -46,6 +47,7 @@ end
     @test eltype(d_product) === eltype(ds[1])
     @test @inferred(logpdf(d_product, x)) ≈ sum(logpdf.(ds, x))
     @test mean(d_product) == mean.(ds)
+    @test std(d_product) == std.(ds)
     @test var(d_product) == var.(ds)
     @test cov(d_product) == Diagonal(var.(ds))
     @test entropy(d_product) == sum(entropy.(ds))
@@ -76,6 +78,7 @@ end
         @test eltype(d_product) === eltype(ds[1])
         @test @inferred(logpdf(d_product, x)) ≈ sum(logpdf.(ds, x))
         @test mean(d_product) == mean.(ds)
+        @test std(d_product) == std.(ds)
         @test var(d_product) == var.(ds)
         @test cov(d_product) == Diagonal(var.(ds))
         @test entropy(d_product) == sum(entropy.(ds))

test/product.jl

@@ -25,6 +25,7 @@ using LinearAlgebra
         @test length(d_product) == length(ds)
         @test eltype(d_product) === eltype(ds[1])
         @test mean(d_product) == mean.(ds)
+        @test std(d_product) == std.(ds)
         @test var(d_product) == var.(ds)
         @test cov(d_product) == Diagonal(var.(ds))
         @test entropy(d_product) ≈ sum(entropy.(ds))
@@ -65,6 +66,7 @@ end
         @test length(d_product) == length(ds)
         @test eltype(d_product) === eltype(ds[1])
         @test @inferred(mean(d_product)) == mean.(ds)
+        @test @inferred(std(d_product)) == std.(ds)
         @test @inferred(var(d_product)) == var.(ds)
         @test @inferred(cov(d_product)) == Diagonal(var.(ds))
         @test @inferred(entropy(d_product)) == sum(entropy.(ds))
@@ -115,6 +117,7 @@ end
             @test length(d_product) == length(ds)
             @test eltype(d_product) === eltype(ds[1])
             @test @inferred(mean(d_product)) == mean.(ds)
+            @test @inferred(std(d_product)) == std.(ds)
             @test @inferred(var(d_product)) == var.(ds)
             @test @inferred(cov(d_product)) == Diagonal(var.(ds))
             @test @inferred(entropy(d_product)) == sum(entropy.(ds))
@@ -150,6 +153,7 @@ end
     @test length(d_product) == 3
     @test eltype(d_product) === Float64
     @test @inferred(mean(d_product)) == mean.(ds_vec)
+    @test @inferred(std(d_product)) == std.(ds_vec)
     @test @inferred(var(d_product)) == var.(ds_vec)
     @test @inferred(cov(d_product)) == Diagonal(var.(ds_vec))
     @test @inferred(entropy(d_product)) == sum(entropy.(ds_vec))