Decouple rand and eltype

docs/src/extends.md

@@ -4,11 +4,8 @@ Whereas this package already provides a large collection of common distributions
 
 Generally, you don't have to implement every API method listed in the documentation. This package provides a series of generic functions that turn a small number of internal methods into user-end API methods. What you need to do is to implement this small set of internal methods for your distributions.
 
-By default, `Discrete` sampleables have the support of type `Int` while `Continuous` sampleables have the support of type `Float64`. If this assumption does not hold for your new distribution or sampler, or its `ValueSupport` is neither `Discrete` nor `Continuous`, you should implement the `eltype` method in addition to the other methods listed below.
-
 **Note:** The methods that need to be implemented are different for distributions of different variate forms.
 
-
 ## Create a Sampler
 
 Unlike full-fledged distributions, a sampler, in general, only provides limited functionalities, mainly to support sampling.
@@ -18,60 +15,48 @@ Unlike full-fledged distributions, a sampler, in general, only provides limited
 To implement a univariate sampler, one can define a subtype (say `Spl`) of `Sampleable{Univariate,S}` (where `S` can be `Discrete` or `Continuous`), and provide a `rand` method, as
 
 ```julia
-function rand(rng::AbstractRNG, s::Spl)
+function Base.rand(rng::AbstractRNG, s::Spl)
     # ... generate a single sample from s
 end
 ```
 
-The package already implements a vectorized version of `rand!` and `rand` that repeatedly calls the scalar version to generate multiple samples; as wells as a one arg version that uses the default random number generator.
-
-### Multivariate Sampler
+The package already implements vectorized versions `rand!(rng::AbstractRNG, s::Spl, dims::Int...)` and `rand(rng::AbstractRNG, s::Spl, dims::Int...)` that repeatedly call the scalar version to generate multiple samples.
+Additionally, the package implements versions of these functions without the `rng::AbstractRNG` argument that use the default random number generator.
 
-To implement a multivariate sampler, one can define a subtype of `Sampleable{Multivariate,S}`, and provide both `length` and `_rand!` methods, as
+If there is a more efficient method to generate multiple samples, one should provide the following method
 
 ```julia
-Base.length(s::Spl) = ... # return the length of each sample
-
-function _rand!(rng::AbstractRNG, s::Spl, x::AbstractVector{T}) where T<:Real
-    # ... generate a single vector sample to x
+function Random.rand!(rng::AbstractRNG, s::Spl, x::AbstractArray{<:Real})
+    # ... generate multiple samples from s in x
 end
 ```
 
-This function can assume that the dimension of `x` is correct, and doesn't need to perform dimension checking.
+### Multivariate Sampler
 
-The package implements both `rand` and `rand!` as follows (which you don't need to implement in general):
+To implement a multivariate sampler, one can define a subtype of `Sampleable{Multivariate,S}`, and provide `length`, `rand`, and `rand!` methods, as
 
 ```julia
-function _rand!(rng::AbstractRNG, s::Sampleable{Multivariate}, A::DenseMatrix)
-    for i = 1:size(A,2)
-        _rand!(rng, s, view(A,:,i))
-    end
-    return A
-end
+Base.length(s::Spl) = ... # return the length of each sample
 
-function rand!(rng::AbstractRNG, s::Sampleable{Multivariate}, A::AbstractVector)
-    length(A) == length(s) ||
-        throw(DimensionMismatch("Output size inconsistent with sample length."))
-    _rand!(rng, s, A)
+function Base.rand(rng::AbstractRNG, s::Spl)
+    # ... generate a single vector sample from s
 end
 
-function rand!(rng::AbstractRNG, s::Sampleable{Multivariate}, A::DenseMatrix)
-    size(A,1) == length(s) ||
-        throw(DimensionMismatch("Output size inconsistent with sample length."))
-    _rand!(rng, s, A)
+@inline function Random.rand!(rng::AbstractRNG, s::Spl, x::AbstractVector{<:Real})
+    # `@inline` + `@boundscheck` allows users to skip bound checks by calling `@inbounds rand!(...)`
+    # Ref https://docs.julialang.org/en/v1/devdocs/boundscheck/#Eliding-bounds-checks
+    @boundscheck # ... check size (and possibly indices) of `x`
+    # ... generate a single vector sample from s in x
 end
-
-rand(rng::AbstractRNG, s::Sampleable{Multivariate,S}) where {S<:ValueSupport} =
-    _rand!(rng, s, Vector{eltype(S)}(length(s)))
-
-rand(rng::AbstractRNG, s::Sampleable{Multivariate,S}, n::Int) where {S<:ValueSupport} =
-    _rand!(rng, s, Matrix{eltype(S)}(length(s), n))
 ```
 
 If there is a more efficient method to generate multiple vector samples in a batch, one should provide the following method
 
 ```julia
-function _rand!(rng::AbstractRNG, s::Spl, A::DenseMatrix{T}) where T<:Real
+@inline function Random.rand!(rng::AbstractRNG, s::Spl, A::AbstractMatrix{<:Real})
+    # `@inline` + `@boundscheck` allows users to skip bound checks by calling `@inbounds rand!(...)`
+    # Ref https://docs.julialang.org/en/v1/devdocs/boundscheck/#Eliding-bounds-checks
+    @boundscheck # ... check size (and possibly indices) of `x`
     # ... generate multiple vector samples in batch
 end
 ```
@@ -80,17 +65,22 @@ Remember that each *column* of A is a sample.
 
 ### Matrix-variate Sampler
 
-To implement a multivariate sampler, one can define a subtype of `Sampleable{Multivariate,S}`, and provide both `size` and `_rand!` methods, as
+To implement a multivariate sampler, one can define a subtype of `Sampleable{Multivariate,S}`, and provide `size`, `rand`, and `rand!` methods, as
 
 ```julia
 Base.size(s::Spl) = ... # the size of each matrix sample
 
-function _rand!(rng::AbstractRNG, s::Spl, x::DenseMatrix{T}) where T<:Real
-    # ... generate a single matrix sample to x
+function Base.rand(rng::AbstractRNG, s::Spl)
+    # ... generate a single matrix sample from s
 end
-```
 
-Note that you can assume `x` has correct dimensions in `_rand!` and don't have to perform dimension checking, the generic `rand` and `rand!` will do dimension checking and array allocation for you.
+@inline function Random.rand!(rng::AbstractRNG, s::Spl, x::AbstractMatrix{<:Real})
+    # `@inline` + `@boundscheck` allows users to skip bound checks by calling `@inbounds rand!(...)`
+    # Ref https://docs.julialang.org/en/v1/devdocs/boundscheck/#Eliding-bounds-checks
+    @boundscheck # ... check size (and possibly indices) of `x`
+    # ... generate a single matrix sample from s in x
+end
+```
 
 ## Create a Distribution
 
@@ -106,7 +96,7 @@ A univariate distribution type should be defined as a subtype of `DiscreteUnivar
 
 The following methods need to be implemented for each univariate distribution type:
 
-- [`rand(::AbstractRNG, d::UnivariateDistribution)`](@ref)
+- [`Base.rand(::AbstractRNG, d::UnivariateDistribution)`](@ref)
 - [`sampler(d::Distribution)`](@ref)
 - [`logpdf(d::UnivariateDistribution, x::Real)`](@ref)
 - [`cdf(d::UnivariateDistribution, x::Real)`](@ref)
@@ -138,8 +128,8 @@ The following methods need to be implemented for each multivariate distribution
 
 - [`length(d::MultivariateDistribution)`](@ref)
 - [`sampler(d::Distribution)`](@ref)
-- [`eltype(d::Distribution)`](@ref)
-- [`Distributions._rand!(::AbstractRNG, d::MultivariateDistribution, x::AbstractArray)`](@ref)
+- [`Base.rand(::AbstractRNG, d::MultivariateDistribution)`](@ref)
+- [`Random.rand!(::AbstractRNG, d::MultivariateDistribution, x::AbstractVector{<:Real})`](@ref)
 - [`Distributions._logpdf(d::MultivariateDistribution, x::AbstractArray)`](@ref)
 
 Note that if there exist faster methods for batch evaluation, one should override `_logpdf!` and `_pdf!`.
@@ -161,6 +151,7 @@ A matrix-variate distribution type should be defined as a subtype of `DiscreteMa
 The following methods need to be implemented for each matrix-variate distribution type:
 
 - [`size(d::MatrixDistribution)`](@ref)
-- [`Distributions._rand!(rng::AbstractRNG, d::MatrixDistribution, A::AbstractMatrix)`](@ref)
+- [`Base.rand(rng::AbstractRNG, d::MatrixDistribution)`](@ref)
+- [`Random.rand!(rng::AbstractRNG, d::MatrixDistribution, A::AbstractMatrix{<:Real})`](@ref)
 - [`sampler(d::MatrixDistribution)`](@ref)
 - [`Distributions._logpdf(d::MatrixDistribution, x::AbstractArray)`](@ref)

docs/src/matrix.md

@@ -25,7 +25,8 @@ var(::MatrixDistribution)
 cov(::MatrixDistribution)
 pdf(d::MatrixDistribution, x::AbstractMatrix{<:Real})
 logpdf(d::MatrixDistribution, x::AbstractMatrix{<:Real})
-Distributions._rand!(::AbstractRNG, ::MatrixDistribution, A::AbstractMatrix)
+Base.rand(::AbstractRNG, ::MatrixDistribution)
+Random.rand!(::AbstractRNG, ::MatrixDistribution, A::AbstractMatrix{<:Real})
 ```
 
 ## Distributions

docs/src/multivariate.md

@@ -18,7 +18,6 @@ The methods listed below are implemented for each multivariate distribution, whi
 ```@docs
 length(::MultivariateDistribution)
 size(::MultivariateDistribution)
-eltype(::Type{MultivariateDistribution})
 mean(::MultivariateDistribution)
 var(::MultivariateDistribution)
 std(::MultivariateDistribution)
@@ -42,8 +41,8 @@ loglikelihood(::MultivariateDistribution, ::AbstractVector{<:Real})
 ### Sampling
 
 ```@docs
-rand(rng::AbstractRNG, ::MultivariateDistribution)
-rand!(rng::AbstractRNG, d::MultivariateDistribution, x::AbstractArray)
+Base.rand(rng::AbstractRNG, ::MultivariateDistribution)
+Random.rand!(rng::AbstractRNG, d::MultivariateDistribution, x::AbstractVector{<:Real})
 ```
 
 **Note:** In addition to these common methods, each multivariate distribution has its special methods, as introduced below.

docs/src/types.md

@@ -57,7 +57,6 @@ The basic functionalities that a sampleable object provides are to *retrieve inf
 length(::Sampleable)
 size(::Sampleable)
 nsamples(::Type{Sampleable}, ::Any)
-eltype(::Type{Sampleable})
 rand(::AbstractRNG, ::Sampleable)
 rand!(::AbstractRNG, ::Sampleable, ::AbstractArray)
 ```

src/censored.jl

@@ -112,8 +112,6 @@ function partype(d::Censored{<:UnivariateDistribution,<:ValueSupport,T}) where {
     return promote_type(partype(d.uncensored), T)
 end
 
-Base.eltype(::Type{<:Censored{D,S,T}}) where {D,S,T} = promote_type(T, eltype(D))
-
 #### Range and Support
 
 isupperbounded(d::LeftCensored) = isupperbounded(d.uncensored)

src/cholesky/lkjcholesky.jl

@@ -7,7 +7,7 @@ matrices (positive-definite matrices with ones on the diagonal).
 
 Variates or samples of the distribution are `LinearAlgebra.Cholesky` objects, as might
 be returned by `F = LinearAlgebra.cholesky(R)`, so that `Matrix(F) ≈ R` is a variate or
-sample of [`LKJ`](@ref). 
+sample of [`LKJ`](@ref).
 
 Sampling `LKJCholesky` is faster than sampling `LKJ`, and often having the correlation
 matrix in factorized form makes subsequent computations cheaper as well.
@@ -48,8 +48,8 @@ function LKJCholesky(d::Int, η::Real, _uplo::Union{Char,Symbol} = 'L'; check_ar
     )
     logc0 = lkj_logc0(d, η)
     uplo = _char_uplo(_uplo)
-    T = Base.promote_eltype(η, logc0)
-    return LKJCholesky(d, T(η), uplo, T(logc0))
+    _η, _logc0 = promote(η, logc0)
+    return LKJCholesky(d, _η, uplo, _logc0)
 end
 
 # adapted from LinearAlgebra.char_uplo
@@ -82,8 +82,6 @@ end
 #  Properties
 #  -----------------------------------------------------------------------------
 
-Base.eltype(::Type{LKJCholesky{T}}) where {T} = T
-
 function Base.size(d::LKJCholesky)
     p = d.d
     return (p, p)
@@ -134,7 +132,7 @@ function logkernel(d::LKJCholesky, R::LinearAlgebra.Cholesky)
     p == 1 && return c * log(first(factors))
     # assuming D = diag(factors) with length(D) = p,
     # logp = sum(i -> (c - i) * log(D[i]), 2:p)
-    logp = sum(Iterators.drop(enumerate(diagind(factors)), 1)) do (i, di) 
+    logp = sum(Iterators.drop(enumerate(diagind(factors)), 1)) do (i, di)
         return (c - i) * log(factors[di])
     end
     return logp
@@ -159,95 +157,104 @@ end
 #  -----------------------------------------------------------------------------
 
 function Base.rand(rng::AbstractRNG, d::LKJCholesky)
-    factors = Matrix{eltype(d)}(undef, size(d))
-    R = LinearAlgebra.Cholesky(factors, d.uplo, 0)
-    return _lkj_cholesky_onion_sampler!(rng, d, R)
-end
-function Base.rand(rng::AbstractRNG, d::LKJCholesky, dims::Dims)
     p = d.d
+    η = d.η
     uplo = d.uplo
-    T = eltype(d)
-    TM = Matrix{T}
-    Rs = Array{LinearAlgebra.Cholesky{T,TM}}(undef, dims)
-    for i in eachindex(Rs)
-        factors = TM(undef, p, p)
-        Rs[i] = R = LinearAlgebra.Cholesky(factors, uplo, 0)
-        _lkj_cholesky_onion_sampler!(rng, d, R)
+    factors = Matrix{float(partype(d))}(undef, p, p)
+    _lkj_cholesky_onion_tri!(rng, factors, p, η, uplo)
+    return LinearAlgebra.Cholesky(factors, uplo, 0)
+end
+function rand(rng::AbstractRNG, d::LKJCholesky, dims::Dims)
+    let rng=rng, d=d
+        map(_ -> rand(rng, d), CartesianIndices(dims))
     end
-    return Rs
 end
 
-Random.rand!(d::LKJCholesky, R::LinearAlgebra.Cholesky) = Random.rand!(default_rng(), d, R)
-function Random.rand!(rng::AbstractRNG, d::LKJCholesky, R::LinearAlgebra.Cholesky)
-    return _lkj_cholesky_onion_sampler!(rng, d, R)
+Base.@propagate_inbounds function rand!(d::LKJCholesky, R::LinearAlgebra.Cholesky{<:Real})
+    return rand!(default_rng(), d, R)
+end
+@inline function rand!(rng::AbstractRNG, d::LKJCholesky, R::LinearAlgebra.Cholesky{<:Real})
+    @boundscheck size(R.factors) == size(d)
+    _lkj_cholesky_onion_tri!(rng, R.factors, d.d, d.η, R.uplo)
+    return R
 end
 
-function Random.rand!(
+function rand!(
     rng::AbstractRNG,
     d::LKJCholesky,
     Rs::AbstractArray{<:LinearAlgebra.Cholesky{T,TM}},
     allocate::Bool,
-) where {T,TM}
+) where {T<:Real,TM}
     p = d.d
-    uplo = d.uplo
+    η = d.η
     if allocate
+        uplo = d.uplo
         for i in eachindex(Rs)
-            Rs[i] = _lkj_cholesky_onion_sampler!(
-                rng,
-                d,
-                LinearAlgebra.Cholesky(TM(undef, p, p), uplo, 0),
+            Rs[i] = LinearAlgebra.Cholesky(
+                _lkj_cholesky_onion_tri!(
+                    rng,
+                    TM(undef, p, p),
+                    p,
+                    η,
+                    uplo,
+                ),
+                uplo,
+                0,
             )
         end
     else
-        for i in eachindex(Rs)
-            _lkj_cholesky_onion_sampler!(rng, d, Rs[i])
+        for R in Rs
+            _lkj_cholesky_onion_tri!(rng, R.factors, p, η, R.uplo)
         end
     end
     return Rs
 end
-function Random.rand!(
+function rand!(
     rng::AbstractRNG,
     d::LKJCholesky,
     Rs::AbstractArray{<:LinearAlgebra.Cholesky{<:Real}},
 )
     allocate = any(!isassigned(Rs, i) for i in eachindex(Rs)) || any(R -> size(R, 1) != d.d, Rs)
-    return Random.rand!(rng, d, Rs, allocate)
+    return rand!(rng, d, Rs, allocate)
 end
 
 #
 # onion method
 #
 
-function _lkj_cholesky_onion_sampler!(
-    rng::AbstractRNG,
-    d::LKJCholesky,
-    R::LinearAlgebra.Cholesky,
-)
-    if R.uplo === 'U'
-        _lkj_cholesky_onion_tri!(rng, R.factors, d.d, d.η, Val(:U))
-    else
-        _lkj_cholesky_onion_tri!(rng, R.factors, d.d, d.η, Val(:L))
-    end
-    return R
-end
-
 function _lkj_cholesky_onion_tri!(
     rng::AbstractRNG,
-    A::AbstractMatrix,
+    A::AbstractMatrix{<:Real},
     d::Int,
     η::Real,
-    ::Val{uplo},
-) where {uplo}
+    uplo::Char,
+)
+    # Currently requires one-based indexing
+    # TODO: Generalize to more general indices
+    Base.require_one_based_indexing(A)
+    @assert size(A) == (d, d)
+
+    # Special case: Distribution collapses to a Dirac distribution at the identity matrix
+    # We handle this separately, to increase performance and since `rand(Beta(Inf, Inf))` is not well defined.
+    if η == Inf
+        if uplo === 'L'
+            copyto!(LowerTriangular(A), I)
+        else
+            copyto!(UpperTriangular(A), I)
+        end
+        return A
+    end
+
     # Section 3.2 in LKJ (2009 JMA)
     # reformulated to incrementally construct Cholesky factor as mentioned in Section 5
     # equivalent steps in algorithm in reference are marked.
-    @assert size(A) == (d, d)
+
     A[1, 1] = 1
     d > 1 || return A
     β = η + (d - 2)//2
     #  1. Initialization
     w0 = 2 * rand(rng, Beta(β, β)) - 1
-    if uplo === :L
+    if uplo === 'L'
         A[2, 1] = w0
     else
         A[1, 2] = w0
@@ -261,7 +268,7 @@ function _lkj_cholesky_onion_tri!(
         y = rand(rng, Beta(k//2, β))
         #  (c)-(e)
         # w is directionally uniform vector of length √y
-        w = @views uplo === :L ? A[k + 1, 1:k] : A[1:k, k + 1]
+        w = @views uplo === 'L' ? A[k + 1, 1:k] : A[1:k, k + 1]
         Random.randn!(rng, w)
         rmul!(w, sqrt(y) / norm(w))
         # normalize so new row/column has unit norm