Unify and generalize rand!, logpdf and pdf

Author: devmotion

src/genericrand.jl

@@ -26,6 +26,32 @@ rand(s::Sampleable, dim1::Int, moredims::Int...) =
 rand(rng::AbstractRNG, s::Sampleable, dim1::Int, moredims::Int...) =
     rand(rng, s, (dim1, moredims...))
 
+# default fallback (redefined for univariate distributions)
+function rand(rng::AbstractRNG, s::Sampleable{<:ArrayLikeVariate,Continuous})
+    return @inbounds rand!(rng, sampler(s), Array{float(eltype(s))}(undef, size(s)))
+end
+function rand(rng::AbstractRNG, s::Sampleable{<:ArrayLikeVariate,Discrete})
+    return @inbounds rand!(rng, sampler(s), Array{eltype(s)}(undef, size(s)))
+end
+
+# multiple samples (redefined for univariate distributions)
+function rand(
+    rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{N},Discrete}, dims::Dims,
+) where {N}
+    sz = size(s)
+    ax = map(Base.OneTo, dims)
+    out = [Array{eltype(s),N}(undef, sz) for _ in Iterators.product(ax)]
+    return @inbounds rand!(rng, sampler(s), out)
+end
+function rand(
+    rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{N},Continuous}, dims::Dims,
+) where {N}
+    sz = size(s)
+    ax = map(Base.OneTo, dims)
+    out = [Array{float(eltype(s)),N}(undef, sz) for _ in Iterators.product(ax)]
+    return @inbounds rand!(rng, sampler(s), out)
+end
+
 """
     rand!([rng::AbstractRNG,] s::Sampleable, A::AbstractArray)
 
@@ -40,11 +66,78 @@ form as specified above. The rules are summarized as below:
   matrices with each element for a sample matrix.
 """
 function rand! end
-rand!(s::Sampleable, X::AbstractArray{<:AbstractArray}, allocate::Bool) =
-    rand!(GLOBAL_RNG, s, X, allocate)
-rand!(s::Sampleable, X::AbstractArray) = rand!(GLOBAL_RNG, s, X)
+Base.@propagate_inbounds rand!(s::Sampleable, X::AbstractArray) = rand!(GLOBAL_RNG, s, X)
 rand!(rng::AbstractRNG, s::Sampleable, X::AbstractArray) = _rand!(rng, s, X)
 
+# default definitions for arraylike variates
+@inline function rand!(
+    rng::AbstractRNG,
+    s::Sampleable{ArrayLikeVariate{N}},
+    x::AbstractArray{<:Real,N},
+) where {N}
+    @boundscheck begin
+        size(x) == size(s) || throw(DimensionMismatch("inconsistent array dimensions"))
+    end
+    # the function barrier fixes performance issues if `sampler(s)` is type unstable
+    return _rand!(rng, sampler(s), x)
+end
+
+@inline function rand!(
+    rng::AbstractRNG,
+    s::Sampleable{ArrayLikeVariate{N}},
+    x::AbstractArray{<:Real,M},
+) where {N,M}
+    @boundscheck begin
+        M > N ||
+            throw(DimensionMismatch(
+                "number of dimensions of `x` ($M) must be greater than number of dimensions of `s` ($N)"
+            ))
+        ntuple(i -> size(x, i), Val(N)) == size(s) ||
+            throw(DimensionMismatch("inconsistent array dimensions"))
+    end
+    return _rand!(rng, sampler(s), x)
+end
+
+function _rand!(
+    rng::AbstractRNG,
+    s::Sampleable{<:ArrayLikeVariate},
+    x::AbstractArray{<:Real},
+)
+    @inbounds for xi in eachvariate(x, variate_form(typeof(s)))
+        rand!(rng, s, xi)
+    end
+    return x
+end
+
+Base.@propagate_inbounds function rand!(
+    rng::AbstractRNG,
+    s::Sampleable{ArrayLikeVariate{N}},
+    x::AbstractArray{<:AbstractArray{<:Real,N}},
+) where {N}
+    # the function barrier fixes performance issues if `sampler(s)` is type unstable
+    return _rand!(rng, sampler(s), x)
+end
+
+Base.@propagate_inbounds function rand!(
+    rng::AbstractRNG,
+    s::Sampleable{ArrayLikeVariate{N}},
+    x::AbstractArray{<:AbstractArray{<:Real,N}},
+) where {N}
+    # the function barrier fixes performance issues if `sampler(s)` is type unstable
+    return _rand!(rng, sampler(s), x)
+end
+
+Base.@propagate_inbounds function _rand!(
+    rng::AbstractRNG,
+    s::Sampleable{ArrayLikeVariate{N}},
+    x::AbstractArray{<:AbstractArray{<:Real,N}},
+) where {N}
+    for i in eachindex(x)
+        rand!(rng, s, @inbounds(x[i]))
+    end
+    return x
+end
+
 """
     sampler(d::Distribution) -> Sampleable
     sampler(s::Sampleable) -> s

src/matrixvariates.jl

@@ -93,17 +93,6 @@ _rand!(::AbstractRNG, ::MatrixDistribution, A::AbstractMatrix)
 
 ## sampling
 
-# multivariate with pre-allocated 3D array
-function _rand!(rng::AbstractRNG, s::Sampleable{Matrixvariate},
-                m::AbstractArray{<:Real, 3})
-    @boundscheck (size(m, 1), size(m, 2)) == (size(s, 1), size(s, 2)) ||
-        throw(DimensionMismatch("Output size inconsistent with matrix size."))
-    smp = sampler(s)
-    for i in Base.OneTo(size(m,3))
-        _rand!(rng, smp, view(m,:,:,i))
-    end
-    return m
-end
 
 # multiple matrix-variates with pre-allocated array of maybe pre-allocated matrices
 rand!(rng::AbstractRNG, s::Sampleable{Matrixvariate},
@@ -127,120 +116,18 @@ function rand!(rng::AbstractRNG, s::Sampleable{Matrixvariate},
     return X
 end
 
-# multiple matrix-variates, must allocate array of arrays
-rand(rng::AbstractRNG, s::Sampleable{Matrixvariate}, dims::Dims) =
-    rand!(rng, s, Array{Matrix{eltype(s)}}(undef, dims), true)
-rand(rng::AbstractRNG, s::Sampleable{Matrixvariate,Continuous}, dims::Dims) =
-    rand!(rng, s, Array{Matrix{float(eltype(s))}}(undef, dims), true)
-
-# single matrix-variate, must allocate one matrix
-rand(rng::AbstractRNG, s::Sampleable{Matrixvariate}) =
-    _rand!(rng, s, Matrix{eltype(s)}(undef, size(s)))
-rand(rng::AbstractRNG, s::Sampleable{Matrixvariate,Continuous}) =
-    _rand!(rng, s, Matrix{float(eltype(s))}(undef, size(s)))
-
-# single matrix-variate with pre-allocated matrix
-function rand!(rng::AbstractRNG, s::Sampleable{Matrixvariate},
-               A::AbstractMatrix{<:Real})
-    @boundscheck size(A) == size(s) ||
-        throw(DimensionMismatch("Output size inconsistent with matrix size."))
-    return _rand!(rng, s, A)
-end
-
 # pdf & logpdf
 
 _logpdf(d::MatrixDistribution, X::AbstractMatrix{<:Real}) = logkernel(d, X) + d.logc0
 
-_pdf(d::MatrixDistribution, x::AbstractMatrix{<:Real}) = exp(_logpdf(d, x))
-
-"""
-    logpdf(d::MatrixDistribution, AbstractMatrix)
-
-Compute the logarithm of the probability density at the input matrix `x`.
-"""
-function logpdf(d::MatrixDistribution, x::AbstractMatrix{<:Real})
-    size(x) == size(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _logpdf(d, x)
-end
-
-"""
-    pdf(d::MatrixDistribution, x::AbstractArray)
-
-Compute the probability density at the input matrix `x`.
-"""
-function pdf(d::MatrixDistribution, x::AbstractMatrix{<:Real})
-    size(x) == size(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _pdf(d, x)
-end
-
-function _logpdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{<:AbstractMatrix{<:Real}})
-    for i = 1:length(X)
-        r[i] = logpdf(d, X[i])
-    end
-    return r
-end
-
-function _pdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
-    for i = 1:length(X)
-        r[i] = pdf(d, X[i])
-    end
-    return r
-end
-
-function logpdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
-    length(X) == length(r) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _logpdf!(r, d, X)
-end
-
-function pdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
-    length(X) == length(r) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _pdf!(r, d, X)
-end
-
-function logpdf(d::MatrixDistribution, X::AbstractArray{<:AbstractMatrix{<:Real}})
-    map(Base.Fix1(logpdf, d), X)
-end
-
-function pdf(d::MatrixDistribution, X::AbstractArray{<:AbstractMatrix{<:Real}})
-    map(Base.Fix1(pdf, d), X)
-end
-
-"""
-    _logpdf(d::MatrixDistribution, x::AbstractArray)
-
-Evaluate logarithm of pdf value for a given sample `x`. This function need not perform dimension checking.
-"""
-_logpdf(d::MatrixDistribution, x::AbstractArray)
-
-"""
-    loglikelihood(d::MatrixDistribution, x::AbstractArray)
-
-The log-likelihood of distribution `d` with respect to all samples contained in array `x`.
-
-Here, `x` can be a matrix of size `size(d)`, a three-dimensional array with `size(d, 1)`
-rows and `size(d, 2)` columns, or an array of matrices of size `size(d)`.
-"""
-loglikelihood(d::MatrixDistribution, X::AbstractMatrix{<:Real}) = logpdf(d, X)
-function loglikelihood(d::MatrixDistribution, X::AbstractArray{<:Real,3})
-    (size(X, 1), size(X, 2)) == size(d) || throw(DimensionMismatch("Inconsistent array dimensions."))
-    return sum(i -> _logpdf(d, view(X, :, :, i)), axes(X, 3))
-end
-function loglikelihood(d::MatrixDistribution, X::AbstractArray{<:AbstractMatrix{<:Real}})
-    return sum(x -> logpdf(d, x), X)
-end
-
 #  for testing
 is_univariate(d::MatrixDistribution) = size(d) == (1, 1)
 check_univariate(d::MatrixDistribution) = is_univariate(d) || throw(ArgumentError("not 1 x 1"))
 
 ##### Specific distributions #####
 
 for fname in ["wishart.jl", "inversewishart.jl", "matrixnormal.jl",
-              "matrixreshaped.jl", "matrixtdist.jl", "matrixbeta.jl", 
+              "matrixreshaped.jl", "matrixtdist.jl", "matrixbeta.jl",
               "matrixfdist.jl", "lkj.jl"]
     include(joinpath("matrix", fname))
 end

src/multivariates.jl

@@ -24,25 +24,6 @@ a vector of length `dim(d)` or a matrix with `dim(d)` rows.
 """
 rand!(rng::AbstractRNG, d::MultivariateDistribution, x::AbstractArray)
 
-# multivariate with pre-allocated array
-function _rand!(rng::AbstractRNG, s::Sampleable{Multivariate}, m::AbstractMatrix)
-    @boundscheck size(m, 1) == length(s) ||
-        throw(DimensionMismatch("Output size inconsistent with sample length."))
-    smp = sampler(s)
-    for i in Base.OneTo(size(m,2))
-        _rand!(rng, smp, view(m,:,i))
-    end
-    return m
-end
-
-# single multivariate with pre-allocated vector
-function rand!(rng::AbstractRNG, s::Sampleable{Multivariate},
-               v::AbstractVector{<:Real})
-    @boundscheck length(v) == length(s) ||
-        throw(DimensionMismatch("Output size inconsistent with sample length."))
-    _rand!(rng, s, v)
-end
-
 # multiple multivariates with pre-allocated array of maybe pre-allocated vectors
 rand!(rng::AbstractRNG, s::Sampleable{Multivariate},
       X::AbstractArray{<:AbstractVector}) =
@@ -65,22 +46,11 @@ function rand!(rng::AbstractRNG, s::Sampleable{Multivariate},
     return X
 end
 
-# multiple multivariate, must allocate matrix or array of vectors
-rand(s::Sampleable{Multivariate}, n::Int) = rand(GLOBAL_RNG, s, n)
+# multiple multivariate, must allocate matrix
 rand(rng::AbstractRNG, s::Sampleable{Multivariate}, n::Int) =
-    _rand!(rng, s, Matrix{eltype(s)}(undef, length(s), n))
+    @inbounds rand!(rng, sampler(s), Matrix{eltype(s)}(undef, length(s), n))
 rand(rng::AbstractRNG, s::Sampleable{Multivariate,Continuous}, n::Int) =
-    _rand!(rng, s, Matrix{float(eltype(s))}(undef, length(s), n))
-rand(rng::AbstractRNG, s::Sampleable{Multivariate}, dims::Dims) =
-    rand!(rng, s, Array{Vector{eltype(s)}}(undef, dims), true)
-rand(rng::AbstractRNG, s::Sampleable{Multivariate,Continuous}, dims::Dims) =
-    rand!(rng, s, Array{Vector{float(eltype(s))}}(undef, dims), true)
-
-# single multivariate, must allocate vector
-rand(rng::AbstractRNG, s::Sampleable{Multivariate}) =
-    _rand!(rng, s, Vector{eltype(s)}(undef, length(s)))
-rand(rng::AbstractRNG, s::Sampleable{Multivariate,Continuous}) =
-    _rand!(rng, s, Vector{float(eltype(s))}(undef, length(s)))
+    @inbounds rand!(rng, sampler(s), Matrix{float(eltype(s))}(undef, length(s), n))
 
 ## domain
 
@@ -193,83 +163,6 @@ Return the logarithm of probability density evaluated at `x`.
 """
 logpdf(d::MultivariateDistribution, x::AbstractArray)
 
-_pdf(d::MultivariateDistribution, X::AbstractVector) = exp(_logpdf(d, X))
-
-function logpdf(d::MultivariateDistribution, X::AbstractVector)
-    length(X) == length(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _logpdf(d, X)
-end
-
-function pdf(d::MultivariateDistribution, X::AbstractVector)
-    length(X) == length(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _pdf(d, X)
-end
-
-function _logpdf!(r::AbstractArray, d::MultivariateDistribution, X::AbstractMatrix)
-    for i in 1 : size(X,2)
-        @inbounds r[i] = logpdf(d, view(X,:,i))
-    end
-    return r
-end
-
-function _pdf!(r::AbstractArray, d::MultivariateDistribution, X::AbstractMatrix)
-    for i in 1 : size(X,2)
-        @inbounds r[i] = pdf(d, view(X,:,i))
-    end
-    return r
-end
-
-function logpdf!(r::AbstractArray, d::MultivariateDistribution, X::AbstractMatrix)
-    size(X) == (length(d), length(r)) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _logpdf!(r, d, X)
-end
-
-function pdf!(r::AbstractArray, d::MultivariateDistribution, X::AbstractMatrix)
-    size(X) == (length(d), length(r)) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    _pdf!(r, d, X)
-end
-
-function logpdf(d::MultivariateDistribution, X::AbstractMatrix)
-    size(X, 1) == length(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    map(i -> _logpdf(d, view(X, :, i)), axes(X, 2))
-end
-
-function pdf(d::MultivariateDistribution, X::AbstractMatrix)
-    size(X, 1) == length(d) ||
-        throw(DimensionMismatch("Inconsistent array dimensions."))
-    map(i -> _pdf(d, view(X, :, i)), axes(X, 2))
-end
-
-"""
-    _logpdf{T<:Real}(d::MultivariateDistribution, x::AbstractArray)
-
-Evaluate logarithm of pdf value for a given vector `x`. This function need not perform dimension checking.
-Generally, one does not need to implement `pdf` (or `_pdf`) as fallback methods are provided in `src/multivariates.jl`.
-"""
-_logpdf(d::MultivariateDistribution, x::AbstractArray)
-
-"""
-    loglikelihood(d::MultivariateDistribution, x::AbstractArray)
-
-The log-likelihood of distribution `d` with respect to all samples contained in array `x`.
-
-Here, `x` can be a vector of length `dim(d)`, a matrix with `dim(d)` rows, or an array of
-vectors of length `dim(d)`.
-"""
-loglikelihood(d::MultivariateDistribution, X::AbstractVector{<:Real}) = logpdf(d, X)
-function loglikelihood(d::MultivariateDistribution, X::AbstractMatrix{<:Real})
-    size(X, 1) == length(d) || throw(DimensionMismatch("Inconsistent array dimensions."))
-    return sum(i -> _logpdf(d, view(X, :, i)), 1:size(X, 2))
-end
-function loglikelihood(d::MultivariateDistribution, X::AbstractArray{<:AbstractVector})
-    return sum(x -> logpdf(d, x), X)
-end
-
 ##### Specific distributions #####
 
 for fname in ["dirichlet.jl",