Skip to content

Commit

Permalink
Unify and generalize rand!, logpdf and pdf
Browse files Browse the repository at this point in the history
  • Loading branch information
@devmotion
devmotion committed Nov 14, 2021
1 parent 5c7a82a commit ccf4f8e
Show file tree
Hide file tree
Showing 10 changed files with 286 additions and 366 deletions.
99 changes: 96 additions & 3 deletions src/genericrand.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
115 changes: 1 addition & 114 deletions src/matrixvariates.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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
113 changes: 3 additions & 110 deletions src/multivariates.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}) =
Expand All @@ -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

Expand Down Expand Up @@ -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",
Expand Down
Loading

0 comments on commit ccf4f8e

Please sign in to comment.