diff --git a/docs/make.jl b/docs/make.jl index f95b3f60b..c5ab4409e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -17,6 +17,7 @@ makedocs( "reshape.md", "cholesky.md", "mixture.md", + "product.md", "order_statistics.md", "convolution.md", "fit.md", diff --git a/docs/src/multivariate.md b/docs/src/multivariate.md index c4e7c1764..35eff32ca 100644 --- a/docs/src/multivariate.md +++ b/docs/src/multivariate.md @@ -58,7 +58,6 @@ MvNormalCanon MvLogitNormal MvLogNormal Dirichlet -Product ``` ## Addition Methods @@ -105,15 +104,6 @@ params{D<:Distributions.AbstractMvLogNormal}(::Type{D},m::AbstractVector,S::Abst Distributions._logpdf(d::MultivariateDistribution, x::AbstractArray) ``` -## Product distributions - -```@docs -Distributions.product_distribution -``` - -Using `product_distribution` is advised to construct product distributions. -For some distributions, it constructs a special multivariate type. - ## Index ```@index diff --git a/docs/src/product.md b/docs/src/product.md new file mode 100644 index 000000000..9a01b6cd2 --- /dev/null +++ b/docs/src/product.md @@ -0,0 +1,27 @@ +# Product Distributions + +Product distributions are joint distributions of multiple independent distributions. +It is recommended to use `product_distribution` to construct product distributions. +Depending on the type of the argument, it may construct a different distribution type. + +## Multivariate products + +```@docs +Distributions.product_distribution(::AbstractArray{<:Distribution{<:ArrayLikeVariate}}) +Distributions.product_distribution(::AbstractVector{<:Normal}) +Distributions.ProductDistribution +Distributions.Product +``` + +## NamedTuple-variate products + +```@docs +Distributions.product_distribution(::NamedTuple{<:Any,<:Tuple{Distribution,Vararg{Distribution}}}) +Distributions.ProductNamedTupleDistribution +``` + +## Index + +```@index +Pages = ["product.md"] +``` diff --git a/src/Distributions.jl b/src/Distributions.jl index 8d344c415..3e72a8f35 100644 --- a/src/Distributions.jl +++ b/src/Distributions.jl @@ -41,6 +41,7 @@ export Multivariate, Matrixvariate, CholeskyVariate, + NamedTupleVariate, Discrete, Continuous, Sampleable, @@ -297,6 +298,7 @@ include("univariates.jl") include("edgeworth.jl") include("multivariates.jl") include("matrixvariates.jl") +include("namedtuple/productnamedtuple.jl") include("cholesky/lkjcholesky.jl") include("samplers.jl") diff --git a/src/common.jl b/src/common.jl index 31a218e19..4fdb36ecb 100644 --- a/src/common.jl +++ b/src/common.jl @@ -16,6 +16,12 @@ const Univariate = ArrayLikeVariate{0} const Multivariate = ArrayLikeVariate{1} const Matrixvariate = ArrayLikeVariate{2} +""" +`F <: NamedTupleVariate{K}` specifies that the variate or a sample is of type +`NamedTuple{K}`. +""" +struct NamedTupleVariate{K} <: VariateForm end + """ `F <: CholeskyVariate` specifies that the variate or a sample is of type `LinearAlgebra.Cholesky`. diff --git a/src/multivariate/product.jl b/src/multivariate/product.jl index 3bb0f0d9e..1227e5ba1 100644 --- a/src/multivariate/product.jl +++ b/src/multivariate/product.jl @@ -10,6 +10,10 @@ An N dimensional `MultivariateDistribution` constructed from a vector of N indep ```julia Product(Uniform.(rand(10), 1)) # A 10-dimensional Product from 10 independent `Uniform` distributions. ``` + +!!! note + `Product` is deprecated and will be removed in the next breaking release. + Use [`product_distribution`](@ref) instead. """ struct Product{ S<:ValueSupport, diff --git a/src/namedtuple/productnamedtuple.jl b/src/namedtuple/productnamedtuple.jl new file mode 100644 index 000000000..853159ee3 --- /dev/null +++ b/src/namedtuple/productnamedtuple.jl @@ -0,0 +1,155 @@ +""" + ProductNamedTupleDistribution{Tnames,Tdists,S<:ValueSupport,eltypes} <: + Distribution{NamedTupleVariate{Tnames},S} + +A distribution of `NamedTuple`s, constructed from a `NamedTuple` of independent named +distributions. + +Users should use [`product_distribution`](@ref) to construct a product distribution of +independent distributions instead of constructing a `ProductNamedTupleDistribution` +directly. + +# Examples + +```jldoctest ProductNamedTuple; setup = :(using Distributions, Random; Random.seed!(832)) +julia> d = product_distribution((x=Normal(), y=Dirichlet([2, 4]))) +ProductNamedTupleDistribution{(:x, :y)}( +x: Normal{Float64}(μ=0.0, σ=1.0) +y: Dirichlet{Int64, Vector{Int64}, Float64}(alpha=[2, 4]) +) + + +julia> nt = rand(d) +(x = 1.5155385995160346, y = [0.533531876438439, 0.466468123561561]) + +julia> pdf(d, nt) +0.13702825691074877 + +julia> mode(d) # mode of marginals +(x = 0.0, y = [0.25, 0.75]) + +julia> mean(d) # mean of marginals +(x = 0.0, y = [0.3333333333333333, 0.6666666666666666]) + +julia> var(d) # var of marginals +(x = 1.0, y = [0.031746031746031744, 0.031746031746031744]) +``` +""" +struct ProductNamedTupleDistribution{Tnames,Tdists,S<:ValueSupport,eltypes} <: + Distribution{NamedTupleVariate{Tnames},S} + dists::NamedTuple{Tnames,Tdists} +end +function ProductNamedTupleDistribution( + dists::NamedTuple{K,V} +) where {K,V<:Tuple{Distribution,Vararg{Distribution}}} + vs = _product_valuesupport(values(dists)) + eltypes = _product_namedtuple_eltype(values(dists)) + return ProductNamedTupleDistribution{K,V,vs,eltypes}(dists) +end + +_gentype(d::UnivariateDistribution) = eltype(d) +_gentype(d::Distribution{<:ArrayLikeVariate{S}}) where {S} = Array{eltype(d),S} +function _gentype(d::Distribution{CholeskyVariate}) + T = eltype(d) + return LinearAlgebra.Cholesky{T,Matrix{T}} +end +function _gentype(d::ProductNamedTupleDistribution{K}) where {K} + return NamedTuple{K,Tuple{map(_gentype, values(d.dists))...}} +end +_gentype(::Distribution) = Any + +_product_namedtuple_eltype(dists) = typejoin(map(_gentype, dists)...) + +function Base.show(io::IO, d::ProductNamedTupleDistribution) + return show_multline(io, d, collect(pairs(d.dists))) +end + +function distrname(::ProductNamedTupleDistribution{K}) where {K} + return "ProductNamedTupleDistribution{$K}" +end + +""" + product_distribution(dists::NamedTuple{K,Tuple{Vararg{Distribution}}}) where {K} + +Create a distribution of `NamedTuple`s as a product distribution of independent named +distributions. + +The function falls back to constructing a [`ProductNamedTupleDistribution`](@ref) +distribution but specialized methods can be defined. +""" +function product_distribution( + dists::NamedTuple{<:Any,<:Tuple{Distribution,Vararg{Distribution}}} +) + return ProductNamedTupleDistribution(dists) +end + +# Properties + +Base.eltype(::Type{<:ProductNamedTupleDistribution{<:Any,<:Any,<:Any,T}}) where {T} = T + +Base.minimum(d::ProductNamedTupleDistribution) = map(minimum, d.dists) + +Base.maximum(d::ProductNamedTupleDistribution) = map(maximum, d.dists) + +function insupport(dist::ProductNamedTupleDistribution{K}, x::NamedTuple{K}) where {K} + return all(map(insupport, dist.dists, x)) +end + +# Evaluation + +function pdf(dist::ProductNamedTupleDistribution{K}, x::NamedTuple{K}) where {K} + return exp(logpdf(dist, x)) +end + +function logpdf(dist::ProductNamedTupleDistribution{K}, x::NamedTuple{K}) where {K} + return sum(map(logpdf, dist.dists, x)) +end + +function loglikelihood(dist::ProductNamedTupleDistribution{K}, x::NamedTuple{K}) where {K} + return logpdf(dist, x) +end + +function loglikelihood( + dist::ProductNamedTupleDistribution{K}, xs::AbstractArray{<:NamedTuple{K}} +) where {K} + return sum(Base.Fix1(loglikelihood, dist), xs) +end + +# Statistics + +mode(d::ProductNamedTupleDistribution) = map(mode, d.dists) + +mean(d::ProductNamedTupleDistribution) = map(mean, d.dists) + +var(d::ProductNamedTupleDistribution) = map(var, d.dists) + +std(d::ProductNamedTupleDistribution) = map(std, d.dists) + +entropy(d::ProductNamedTupleDistribution) = sum(entropy, values(d.dists)) + +function kldivergence( + d1::ProductNamedTupleDistribution{K}, d2::ProductNamedTupleDistribution{K} +) where {K} + return sum(map(kldivergence, d1.dists, d2.dists)) +end + +# Sampling + +function sampler(d::ProductNamedTupleDistribution{K,<:Any,S}) where {K,S} + samplers = map(sampler, d.dists) + Tsamplers = typeof(values(samplers)) + return ProductNamedTupleSampler{K,Tsamplers,S}(samplers) +end + +function Base.rand(rng::AbstractRNG, d::ProductNamedTupleDistribution{K}) where {K} + return NamedTuple{K}(map(Base.Fix1(rand, rng), d.dists)) +end +function Base.rand( + rng::AbstractRNG, d::ProductNamedTupleDistribution{K}, dims::Dims +) where {K} + return convert(AbstractArray{<:NamedTuple{K}}, _rand(rng, sampler(d), dims)) +end + +function _rand!(rng::AbstractRNG, d::ProductNamedTupleDistribution, xs::AbstractArray) + return _rand!(rng, sampler(d), xs) +end diff --git a/src/samplers.jl b/src/samplers.jl index 794f2bff4..fa667fff2 100644 --- a/src/samplers.jl +++ b/src/samplers.jl @@ -24,7 +24,9 @@ for fname in ["aliastable.jl", "vonmises.jl", "vonmisesfisher.jl", "discretenonparametric.jl", - "categorical.jl"] + "categorical.jl", + "productnamedtuple.jl", + ] include(joinpath("samplers", fname)) end diff --git a/src/samplers/productnamedtuple.jl b/src/samplers/productnamedtuple.jl new file mode 100644 index 000000000..792ca8569 --- /dev/null +++ b/src/samplers/productnamedtuple.jl @@ -0,0 +1,21 @@ +struct ProductNamedTupleSampler{Tnames,Tsamplers,S<:ValueSupport} <: + Sampleable{NamedTupleVariate{Tnames},S} + samplers::NamedTuple{Tnames,Tsamplers} +end + +function Base.rand(rng::AbstractRNG, spl::ProductNamedTupleSampler{K}) where {K} + return NamedTuple{K}(map(Base.Fix1(rand, rng), spl.samplers)) +end + +function _rand(rng::AbstractRNG, spl::ProductNamedTupleSampler, dims::Dims) + return map(CartesianIndices(dims)) do _ + return rand(rng, spl) + end +end + +function _rand!(rng::AbstractRNG, spl::ProductNamedTupleSampler, xs::AbstractArray) + for i in eachindex(xs) + xs[i] = rand(rng, spl) + end + return xs +end diff --git a/test/namedtuple/productnamedtuple.jl b/test/namedtuple/productnamedtuple.jl new file mode 100644 index 000000000..2f121fd33 --- /dev/null +++ b/test/namedtuple/productnamedtuple.jl @@ -0,0 +1,195 @@ +using Distributions +using Distributions: ProductNamedTupleDistribution, ProductNamedTupleSampler +using LinearAlgebra +using Random +using Test + +@testset "ProductNamedTupleDistribution" begin + @testset "Constructor" begin + nt = (x=Normal(1.0, 2.0), y=Normal(3.0, 4.0)) + d = @inferred ProductNamedTupleDistribution(nt) + @test d isa ProductNamedTupleDistribution + @test d.dists === nt + @test Distributions.variate_form(typeof(d)) === NamedTupleVariate{(:x, :y)} + @test Distributions.value_support(typeof(d)) === Continuous + + nt = ( + x=Normal(), + y=Dirichlet(10, 1.0), + z=DiscreteUniform(1, 10), + w=LKJCholesky(3, 2.0), + ) + d = @inferred ProductNamedTupleDistribution(nt) + @test d isa ProductNamedTupleDistribution + @test d.dists === nt + @test Distributions.variate_form(typeof(d)) === NamedTupleVariate{(:x, :y, :z, :w)} + @test Distributions.value_support(typeof(d)) === Continuous + end + + @testset "product_distribution" begin + nt = (x=Normal(1.0, 2.0), y=Normal(3.0, 4.0)) + d = @inferred product_distribution(nt) + @test d === ProductNamedTupleDistribution(nt) + + nt = ( + x=Normal(), + y=Dirichlet(10, 1.0), + z=DiscreteUniform(1, 10), + w=LKJCholesky(3, 2.0), + ) + d = @inferred product_distribution(nt) + @test d === ProductNamedTupleDistribution(nt) + end + + @testset "show" begin + d = ProductNamedTupleDistribution((x=Gamma(1.0, 2.0), y=Normal())) + @test repr(d) == """ + ProductNamedTupleDistribution{(:x, :y)}( + x: Gamma{Float64}(α=1.0, θ=2.0) + y: Normal{Float64}(μ=0.0, σ=1.0) + ) + """ + end + + @testset "Properties" begin + @testset "eltype" begin + @testset for nt in [ + (x=Normal(1.0, 2.0), y=Normal(3.0, 4.0)), + (x=Normal(), y=Gamma()), + (x=Bernoulli(),), + (x=Normal(), y=Bernoulli()), + (w=LKJCholesky(3, 2.0),), + ( + x=Normal(), + y=Dirichlet(10, 1.0), + z=DiscreteUniform(1, 10), + w=LKJCholesky(3, 2.0), + ), + (x = product_distribution((x=Normal(), y=Gamma())),), + ] + d = ProductNamedTupleDistribution(nt) + @test eltype(d) === eltype(rand(d)) + end + end + + @testset "minimum" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=MvNormal(ones(5))) + d = ProductNamedTupleDistribution(nt) + @test @inferred(minimum(d)) == + (x=minimum(nt.x), y=minimum(nt.y), z=minimum(nt.z)) + end + + @testset "maximum" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=MvNormal(ones(5))) + d = ProductNamedTupleDistribution(nt) + @test @inferred(maximum(d)) == + (x=maximum(nt.x), y=maximum(nt.y), z=maximum(nt.z)) + end + + @testset "insupport" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=Dirichlet(5, 1.0)) + d = ProductNamedTupleDistribution(nt) + x = (x=rand(nt.x), y=rand(nt.y), z=rand(nt.z)) + @test @inferred(insupport(d, x)) + @test_throws MethodError insupport(d, NamedTuple{(:y, :z, :x)}(x)) + @test_throws MethodError insupport(d, NamedTuple{(:x, :y)}(x)) + @test !insupport(d, merge(x, (x=NaN,))) + @test !insupport(d, merge(x, (y=-1,))) + @test !insupport(d, merge(x, (z=fill(0.25, 4),))) + end + end + + @testset "Evaluation" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=Dirichlet(5, 1.0), w=Bernoulli()) + d = ProductNamedTupleDistribution(nt) + x = (x=rand(nt.x), y=rand(nt.y), z=rand(nt.z), w=rand(nt.w)) + @test @inferred(logpdf(d, x)) == + logpdf(nt.x, x.x) + logpdf(nt.y, x.y) + logpdf(nt.z, x.z) + logpdf(nt.w, x.w) + @test @inferred(pdf(d, x)) == exp(logpdf(d, x)) + @test @inferred(loglikelihood(d, x)) == logpdf(d, x) + xs = [(x=rand(nt.x), y=rand(nt.y), z=rand(nt.z), w=rand(nt.w)) for _ in 1:10] + @test @inferred(loglikelihood(d, xs)) == sum(logpdf.(Ref(d), xs)) + end + + @testset "Statistics" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=MvNormal(1.0:5.0), w=Poisson(100)) + d = ProductNamedTupleDistribution(nt) + @test @inferred(mode(d)) == (x=mode(nt.x), y=mode(nt.y), z=mode(nt.z), w=mode(nt.w)) + @test @inferred(mean(d)) == (x=mean(nt.x), y=mean(nt.y), z=mean(nt.z), w=mean(nt.w)) + @test @inferred(var(d)) == (x=var(nt.x), y=var(nt.y), z=var(nt.z), w=var(nt.w)) + @test @inferred(entropy(d)) == + entropy(nt.x) + entropy(nt.y) + entropy(nt.z) + entropy(nt.w) + + d1 = ProductNamedTupleDistribution((x=Normal(1.0, 2.0), y=Gamma())) + d2 = ProductNamedTupleDistribution((x=Normal(), y=Gamma(2.0, 3.0))) + @test kldivergence(d1, d2) == + kldivergence(d1.dists.x, d2.dists.x) + kldivergence(d1.dists.y, d2.dists.y) + + d3 = ProductNamedTupleDistribution((x=Normal(1.0, 2.0), y=Gamma(6.0, 7.0))) + @test std(d3) == (x=std(d3.dists.x), y=std(d3.dists.y)) + end + + @testset "Sampler" begin + nt1 = (x=Normal(1.0, 2.0), y=Gamma(), z=Dirichlet(5, 1.0), w=Bernoulli()) + d1 = ProductNamedTupleDistribution(nt1) + # sampler(::Gamma) is type-unstable + spl = @inferred ProductNamedTupleSampler{(:x, :y, :z, :w)} sampler(d1) + @test spl.samplers == (; (k => sampler(v) for (k, v) in pairs(nt1))...) + rng = MersenneTwister(973) + x1 = @inferred rand(rng, d1) + rng = MersenneTwister(973) + x2 = ( + x=rand(rng, nt1.x), y=rand(rng, nt1.y), z=rand(rng, nt1.z), w=rand(rng, nt1.w) + ) + @test x1 == x2 + x3 = rand(rng, d1) + @test x3 != x1 + + # sampler should now be type-stable + nt2 = (x=Normal(1.0, 2.0), z=Dirichlet(5, 1.0), w=Bernoulli()) + d2 = ProductNamedTupleDistribution(nt2) + @inferred sampler(d2) + end + + @testset "Sampling" begin + @testset "rand" begin + nt = (x=Normal(1.0, 2.0), y=Gamma(), z=Dirichlet(5, 1.0), w=Bernoulli()) + d = ProductNamedTupleDistribution(nt) + rng = MersenneTwister(973) + x1 = @inferred rand(rng, d) + @test eltype(x1) === eltype(d) + rng = MersenneTwister(973) + x2 = ( + x=rand(rng, nt.x), y=rand(rng, nt.y), z=rand(rng, nt.z), w=rand(rng, nt.w) + ) + @test x1 == x2 + x3 = rand(rng, d) + @test x3 != x1 + + # not completely type-inferrable due to sampler(::Gamma) being type-unstable + xs1 = @inferred Vector{<:NamedTuple{(:x, :y, :z, :w)}} rand(rng, d, 10) + @test length(xs1) == 10 + @test all(insupport.(Ref(d), xs1)) + + xs2 = @inferred Array{<:NamedTuple{(:x, :y, :z, :w)},3} rand(rng, d, (2, 3, 4)) + @test size(xs2) == (2, 3, 4) + @test all(insupport.(Ref(d), xs2)) + + nt2 = (x=Normal(1.0, 2.0), z=Dirichlet(5, 1.0), w=Bernoulli()) + d2 = ProductNamedTupleDistribution(nt2) + # now type-inferrable + @inferred rand(rng, d2, 10) + @inferred rand(rng, d2, 2, 3, 4) + end + + @testset "rand!" begin + d = ProductNamedTupleDistribution(( + x=Normal(1.0, 2.0), y=Gamma(), z=Dirichlet(5, 1.0), w=Bernoulli() + )) + x = rand(d) + xs = Array{typeof(x)}(undef, (2, 3, 4)) + rand!(d, xs) + @test all(insupport.(Ref(d), xs)) + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index c7d73f454..7eb909471 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -62,6 +62,7 @@ const tests = [ "qq", "univariate/continuous/pgeneralizedgaussian", "product", + "namedtuple/productnamedtuple", "univariate/discrete/discretenonparametric", "univariate/continuous/chernoff", "univariate_bounds", # extra file compared to /src