enhance wrapped distributions

Author: alyst

src/distribution_wrappers.jl

@@ -2,11 +2,34 @@ using Distributions: Distributions
 using Bijectors: Bijectors
 using Distributions: Univariate, Multivariate, Matrixvariate
 
+"""
+Base type for distribution wrappers.
+"""
+abstract type WrappedDistribution{variate,support,Td<:Distribution{variate,support}} <:
+        Distribution{variate,support}
+end
+
+wrapped_dist_type(::Type{<:WrappedDistribution{<:Any,<:Any,Td}}) where Td = Td
+wrapped_dist_type(d::WrappedDistribution) = wrapped_dist_type(d)
+
+wrapped_dist(d::WrappedDistribution) = d.dist
+
+Base.length(d::WrappedDistribution{<:Multivariate}) = length(wrapped_dist(d))
+Base.size(d::WrappedDistribution{<:Multivariate}) = size(wrapped_dist(d))
+Base.eltype(::Type{T}) where T <: WrappedDistribution = eltype(wrapped_dist_type(T))
+Base.eltype(d::WrappedDistribution) = eltype(wrapped_dist_type(d))
+
+Distributions.rand(rng::Random.AbstractRNG, d::WrappedDistribution) = rand(rng, wrapped_dist(d))
+Distributions.minimum(d::WrappedDistribution) = minimum(wrapped_dist(d))
+Distributions.maximum(d::WrappedDistribution) = maximum(wrapped_dist(d))
+
+Bijectors.bijector(d::WrappedDistribution) = bijector(wrapped_dist(d))
+
 """
 A named distribution that carries the name of the random variable with it.
 """
 struct NamedDist{variate,support,Td<:Distribution{variate,support},Tv<:VarName} <:
-       Distribution{variate,support}
+       WrappedDistribution{variate,support,Td}
     dist::Td
     name::Tv
 end
@@ -24,21 +47,23 @@ function Distributions.loglikelihood(dist::NamedDist, x::AbstractArray{<:Real})
     return Distributions.loglikelihood(dist.dist, x)
 end
 
+"""
+Wrapper around distribution `Td` that suppresses `logpdf()` calculation.
+
+Note that *SampleFromPrior* would still sample from `Td`.
+"""
 struct NoDist{variate,support,Td<:Distribution{variate,support}} <:
-       Distribution{variate,support}
+       WrappedDistribution{variate,support,Td}
     dist::Td
 end
 NoDist(dist::NamedDist) = NamedDist(NoDist(dist.dist), dist.name)
 
-Distributions.rand(rng::Random.AbstractRNG, d::NoDist) = rand(rng, d.dist)
 Distributions.logpdf(d::NoDist{<:Univariate}, ::Real) = 0
 Distributions.logpdf(d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}) = 0
 function Distributions.logpdf(d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real})
     return zeros(Int, size(x, 2))
 end
 Distributions.logpdf(d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}) = 0
-Distributions.minimum(d::NoDist) = minimum(d.dist)
-Distributions.maximum(d::NoDist) = maximum(d.dist)
 
 Bijectors.logpdf_with_trans(d::NoDist{<:Univariate}, ::Real) = 0
 Bijectors.logpdf_with_trans(d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}) = 0