From c4aa64c343ad4c96985504e487540dc968d27a56 Mon Sep 17 00:00:00 2001 From: Martin Trapp Date: Sat, 13 Apr 2019 18:20:10 +0200 Subject: [PATCH] Work in progress. --- src/Turing.jl | 11 ++- src/inference/Inference.jl | 110 +++------------------------ src/inference/hmc.jl | 46 ++---------- src/inference/runners.jl | 147 +++++++++++++++++++------------------ 4 files changed, 100 insertions(+), 214 deletions(-) diff --git a/src/Turing.jl b/src/Turing.jl index 21f8c9e12..3b9ac6049 100644 --- a/src/Turing.jl +++ b/src/Turing.jl @@ -73,9 +73,16 @@ hash(s::Selector) = hash(s.gid) abstract type AbstractRunner end #abstract type AbstractSampler end +abstract type SampleFromDistribution end + +struct SampleFromPrior <: SampleFromDistribution end +_rand(dist::Distribution) = rand(dist) +_rand(dist::distribution, n::Int) = rand(dist, n) + +struct SampleFromUniform <: SampleFromDistribution end +_rand(dist::Distribution) = init(dist) +_rand(dist::Distributionm, n::Int) = init(dist, n) -struct SampleFromUniform <: AbstractRunner end -struct SampleFromPrior <: AbstractRunner end struct ComputeLogJointDensity <: AbstractRunner end struct ComputeLogDensity <: AbstractRunner end struct ParticleFiltering <: AbstractRunner end diff --git a/src/inference/Inference.jl b/src/inference/Inference.jl index 1d91cdead..95a704536 100644 --- a/src/inference/Inference.jl +++ b/src/inference/Inference.jl @@ -102,114 +102,28 @@ include("gibbs.jl") require_gradient(spl::Sampler) = false require_particles(spl::Sampler) = false -assume(spl::Sampler, dist::Distribution) = -error("Turing.assume: unmanaged inference algorithm: $(typeof(spl))") - -observe(spl::Sampler, weight::Float64) = -error("Turing.observe: unmanaged inference algorithm: $(typeof(spl))") - -## Default definitions for assume, observe, when sampler = nothing. -function assume(spl::A, - dist::Distribution, - vn::VarName, - vi::VarInfo) where {A<:Union{SampleFromPrior, SampleFromUniform}} - - if haskey(vi, vn) - r = vi[vn] - else - r = isa(spl, SampleFromUniform) ? init(dist) : rand(dist) - push!(vi, vn, r, dist) - end - # NOTE: The importance weight is not correctly computed here because - # r is genereated from some uniform distribution which is different from the prior - # acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn))) +# ############################# # +# Assume and observe functions. # +# ############################# # - r, logpdf_with_trans(dist, r, istrans(vi, vn)) +# Those functions have to be implemented for each new sampler. +function assume(spl::Sampler, dist::Distribution, vn::VarName, vi::VarInfo) + @error "[assume]: Unmanaged inference algorithm: $(typeof(spl))" end -function assume(spl::A, - dists::Vector{T}, - vn::VarName, - var::Any, - vi::VarInfo) where {T<:Distribution, A<:Union{SampleFromPrior, SampleFromUniform}} - - @assert length(dists) == 1 "Turing.assume only support vectorizing i.i.d distribution" - dist = dists[1] - n = size(var)[end] - - vns = map(i -> copybyindex(vn, "[$i]"), 1:n) - - if haskey(vi, vns[1]) - rs = vi[vns] - else - rs = isa(spl, SampleFromUniform) ? init(dist, n) : rand(dist, n) - - if isa(dist, UnivariateDistribution) || isa(dist, MatrixDistribution) - for i = 1:n - push!(vi, vns[i], rs[i], dist) - end - @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched" - var = rs - elseif isa(dist, MultivariateDistribution) - for i = 1:n - push!(vi, vns[i], rs[:,i], dist) - end - if isa(var, Vector) - @assert length(var) == size(rs)[2] "Turing.assume: variable and random number dimension unmatched" - for i = 1:n - var[i] = rs[:,i] - end - elseif isa(var, Matrix) - @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched" - var = rs - else - @error("Turing.assume: unsupported variable container"); error() - end - end - end - - # acclogp!(vi, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1])))) - - var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))) - +function assume(spl::Sampler, dists::Vector{<:Distribution}, vn::VarName, var, vi::VarInfo) + @error "[assume]: Unmanaged inference algorithm: $(typeof(spl))" end - -@inline observe(::Nothing, dist, value, vi) = observe(SampleFromPrior(), dist, value, vi) - -function observe(spl::A, - dist::Distribution, - value::Any, - vi::VarInfo) where {A<:Union{SampleFromPrior, SampleFromUniform}} - - vi.num_produce += one(vi.num_produce) - Turing.DEBUG && @debug "dist = $dist" - Turing.DEBUG && @debug "value = $value" - - # acclogp!(vi, logpdf(dist, value)) - return logpdf(dist, value) +function observe(spl::Sampler, dist::Distribution, value, vi::VarInfo) + @error "[observe]: Unmanaged inference algorithm: $(typeof(spl))" end -function observe(spl::A, - dists::Vector{T}, - value::Any, - vi::VarInfo) where {T<:Distribution, A<:Union{SampleFromPrior, SampleFromUniform}} - - @assert length(dists) == 1 "Turing.observe only support vectorizing i.i.d distribution" - dist = dists[1] - @assert isa(dist, UnivariateDistribution) || isa(dist, MultivariateDistribution) "Turing.observe: vectorizing matrix distribution is not supported" - if isa(dist, UnivariateDistribution) # only univariate distributions support broadcast operation (logpdf.) by Distributions.jl - # acclogp!(vi, sum(logpdf.(Ref(dist), value))) - sum(logpdf.(Ref(dist), value)) - else - # acclogp!(vi, sum(logpdf(dist, value))) - sum(logpdf(dist, value)) - end - +function observe(spl::Sampler, dists::Vector{<:Distribution}, values, vi::VarInfo) + @error "[observe]: Unmanaged inference algorithm: $(typeof(spl))" end - ############## # Utilities # ############## diff --git a/src/inference/hmc.jl b/src/inference/hmc.jl index f591b5452..5b4d7f54a 100644 --- a/src/inference/hmc.jl +++ b/src/inference/hmc.jl @@ -263,50 +263,14 @@ function step(model, spl::Sampler{<:Hamiltonian}, vi::VarInfo, is_first::Val{fal end function assume(spl::Sampler{<:Hamiltonian}, dist::Distribution, vn::VarName, vi::VarInfo) - Turing.DEBUG && @debug "assuming..." + # Why do we need this here and not if dist is a vector of dists? updategid!(vi, vn, spl) - r = vi[vn] - # acclogp!(vi, logpdf_with_trans(dist, r, istrans(vi, vn))) - # r - Turing.DEBUG && @debug "dist = $dist" - Turing.DEBUG && @debug "vn = $vn" - Turing.DEBUG && @debug "r = $r" "typeof(r)=$(typeof(r))" - r, logpdf_with_trans(dist, r, istrans(vi, vn)) + return assume(ComputeLogJointDensity(), dist, vn, vi) end function assume(spl::Sampler{<:Hamiltonian}, dists::Vector{<:Distribution}, vn::VarName, var::Any, vi::VarInfo) - @assert length(dists) == 1 "[observe] Turing only support vectorizing i.i.d distribution" - dist = dists[1] - n = size(var)[end] - - vns = map(i -> copybyindex(vn, "[$i]"), 1:n) - - rs = vi[vns] # NOTE: inside Turing the Julia conversion should be sticked to - - # acclogp!(vi, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1])))) - - if isa(dist, UnivariateDistribution) || isa(dist, MatrixDistribution) - @assert size(var) == size(rs) "Turing.assume variable and random number dimension unmatched" - var = rs - elseif isa(dist, MultivariateDistribution) - if isa(var, Vector) - @assert length(var) == size(rs)[2] "Turing.assume variable and random number dimension unmatched" - for i = 1:n - var[i] = rs[:,i] - end - elseif isa(var, Matrix) - @assert size(var) == size(rs) "Turing.assume variable and random number dimension unmatched" - var = rs - else - error("[Turing] unsupported variable container") - end - end - - var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))) + return assume(ComputeLogJointDensity(), dists, vn, var, vi) end -observe(spl::Sampler{<:Hamiltonian}, d::Distribution, value::Any, vi::VarInfo) = - observe(nothing, d, value, vi) - -observe(spl::Sampler{<:Hamiltonian}, ds::Vector{<:Distribution}, value::Any, vi::VarInfo) = - observe(nothing, ds, value, vi) +observe(spl::Sampler{<:Hamiltonian}, d, value, vi::VarInfo) = + observe(ComputeLogJointDensity(), d, value, vi) diff --git a/src/inference/runners.jl b/src/inference/runners.jl index 7324b33c7..083c5b4be 100644 --- a/src/inference/runners.jl +++ b/src/inference/runners.jl @@ -1,11 +1,11 @@ function _getdist(dists::Vector{<:Distribution}) - @assert length(dists) == 1 "[observe] Turing only support vectorizing iid distribution." + @assert length(dists) == 1 "[_getdist] Turing only support vectorizing iid distribution." return first(dists) end -## -# Default definition when runner = nothing -## +#################### +# runner = nothing # +#################### function observe(::Nothing, dist::Distribution, value, vi::VarInfo) vi.num_produce += one(vi.num_produce) @@ -13,99 +13,105 @@ function observe(::Nothing, dist::Distribution, value, vi::VarInfo) end function observe(::Nothing, dists::Vector{<:UnivariateDistribution}, values, vi::VarInfo) + # TODO: We should probably increase num_produce here. dist = _getdist(dists) return sum(logpdf.(dist, values)) end # NOTE: this is necessary as we cannot use broadcasting for MV dists. function observe(::Nothing, dists::Vector{<:MultivariateDistribution}, values, vi::VarInfo) + # TODO: We should probably increase num_produce here. dist = _getdist(dists) return sum(logpdf(dist, values)) end -## -# Sample from prior -## +############################### +# Sample from prior / uniform # +############################### -function assume(spl::SampleFromPrior, dist::Distribution, vn::VarName, vi::VarInfo) +function assume(spl::SampleFromDistribution, dist::Distribution, vn::VarName, vi::VarInfo) - if haskey(vi, vn) - r = vi[vn] - else - r = rand(dist) - push!(vi, vn, r, dist) + if !haskey(vi, vn) + push!(vi, vn, _rand(dist), dist) end + r = vi[vn] return r, logpdf_with_trans(dist, r, istrans(vi, vn)) end -function assume(spl::SampleFromUniform, dist::Distribution, vn::VarName, vi::VarInfo) +function assume(spl::SampleFromDistribution, + dists::Vector{<:Union{UnivariateDistribution, MatrixDistribution}}, + vn::VarName, + var, + vi::VarInfo) + + dist = _getdist(dists) + n = size(var)[end] + + vns = map(i -> copybyindex(vn, "[$i]"), 1:n) + + if haskey(vi, first(vns)) + rs = vi[vns] + else + rs = _rand(dist, n) + @assert size(var) == size(rs) "[assume]: Variable and random number dimension unmatched" + + for i = 1:n + @inbounds push!(vi, vns[i], rs[i], dist) + end + end - if haskey(vi, vn) - r = vi[vn] + if dist isa UnivariateDistribution + var = rs else - r = init(dist) - push!(vi, vn, r, dist) + @inbounds var[:] .= rs[:] end - # NOTE: The importance weight is not correctly computed here because - # r is genereated from some uniform distribution which is different from the prior - return r, logpdf_with_trans(dist, r, istrans(vi, vn)) + return var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))) end -function assume(spl::A, - dists::Vector{T}, - vn::VarName, - var::Any, - vi::VarInfo) where {T<:Distribution, A<:Union{SampleFromPrior, SampleFromUniform}} +function assume(spl::SampleFromDistribution, + dists::Vector{<:MultivariateDistribution}, + vn::VarName, + var::AbstractArray, + vi::VarInfo) - @assert length(dists) == 1 "Turing.assume only support vectorizing i.i.d distribution" - dist = dists[1] + dist = _getdist(dists) n = size(var)[end] vns = map(i -> copybyindex(vn, "[$i]"), 1:n) - if haskey(vi, vns[1]) + if haskey(vi, first(vns)) rs = vi[vns] else - rs = isa(spl, SampleFromUniform) ? init(dist, n) : rand(dist, n) + rs = _rand(dist, n) + @assert size(var) == size(rs) "[assume]: Variable and random number dimension unmatched" - if isa(dist, UnivariateDistribution) || isa(dist, MatrixDistribution) - for i = 1:n - push!(vi, vns[i], rs[i], dist) - end - @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched" - var = rs - elseif isa(dist, MultivariateDistribution) + @inbounds begin for i = 1:n push!(vi, vns[i], rs[:,i], dist) end - if isa(var, Vector) + + if var isa Vector @assert length(var) == size(rs)[2] "Turing.assume: variable and random number dimension unmatched" for i = 1:n var[i] = rs[:,i] end - elseif isa(var, Matrix) - @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched" - var = rs else - @error("Turing.assume: unsupported variable container"); error() + @assert size(var) == size(rs) "Turing.assume: variable and random number dimension unmatched" + var[:] .= rs[:] end end end - # acclogp!(vi, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1])))) - - var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))) - + return var, sum(logpdf_with_trans(dist, rs, istrans(vi, vns[1]))) end -@inline observe(spl::SampleFromPrior, dist, value, vi) = observe(nothing, dist, value, vi) -@inline observe(spl::SampleFromUniform, dist, value, vi) = observe(nothing, dist, value, vi) +@inline observe(spl::SampleFromDistribution, dist, value, vi) = observe(nothing, dist, value, vi) -### -# Functions for runner to compute the log joint. -### +################################# +# Compute the log joint Runner. # +################################# function assume(spl::ComputeLogJointDensity, dist::Distribution, vn::VarName, vi::VarInfo) @assert haskey(vi, vn) @@ -114,36 +120,31 @@ function assume(spl::ComputeLogJointDensity, dist::Distribution, vn::VarName, vi end function assume(spl::ComputeLogJointDensity, - dists::AbstractVector{<:Distribution}, - vn::VarName, - var, - vi::VarInfo) + dists::AbstractVector{<:Distribution}, + vn::VarName, + var, + vi::VarInfo) - @assert length(dists) == 1 "[assume] Turing only supports vectorizing iid distributions" + dist = _getdist(dists) + n = size(var)[end] - dist = first(dist) - N = size(var)[end] + vns = map(i -> copybyindex(vn, "[$i]"), 1:n) + @assert all(haskey.(vi, vn)) - vns = map(i -> copybyindex(vn, "[$i]"), 1:N) - rs = vi[vns] + @inbounds begin + rs = vi[vns] - if isa(dist, UnivariateDistribution) || isa(dist, MatrixDistribution) - @assert size(var) == size(rs) "[assume] Variable and random number dimension unmatched." - var = rs - elseif isa(dist, MultivariateDistribution) - if isa(var, Vector) - @assert length(var) == size(rs)[2] "[assume] Variable and random number dimension unmatched." - for i = 1:N - @inbounds var[i] = rs[:,i] + if (dist isa MultivariateDistribution) && (var isa AbstractVector) + @assert length(var) == last(size(rs)) + for i in 1:n + inbounds var[i][:] .= rs[:,i] end - elseif isa(var, Matrix) - @assert size(var) == size(rs) "[assume] Variable and random number dimension unmatched." + elseif dist isa UnivariateDistribution var = rs else - @error "[assume] Unsupported variable container." + @assert size(var) == size(rs) "[assume] Variable and random number dimension unmatched." + var[:] .= rs[:] end - else - @error "[assume] Unsupported distribution type." end return var, sum(logpdf_with_trans(dist, rs, istrans(vi, first(vns))))