diff --git a/Project.toml b/Project.toml index 726a74343..bfa13d956 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DynamicPPL" uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8" -version = "0.19.3" +version = "0.19.4" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" @@ -8,7 +8,9 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -22,7 +24,9 @@ AbstractPPL = "0.5.1" BangBang = "0.3" Bijectors = "0.5.2, 0.6, 0.7, 0.8, 0.9, 0.10" ChainRulesCore = "0.9.7, 0.10, 1" +ConstructionBase = "1" Distributions = "0.23.8, 0.24, 0.25" +DocStringExtensions = "0.8" MacroTools = "0.5.6" Setfield = "0.7.1, 0.8" ZygoteRules = "0.2" diff --git a/bors.toml b/bors.toml index 302d9e1d5..273e80ce6 100644 --- a/bors.toml +++ b/bors.toml @@ -17,4 +17,5 @@ required_approvals = 1 use_squash_merge = true # Uncomment this to use a four hour timeout. # The default is one hour. -timeout_sec = 14400 +# timeout_sec = 14400 +timeout_sec = 43200 diff --git a/docs/src/api.md b/docs/src/api.md index 133b86e9b..809e6c49e 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -103,8 +103,14 @@ NamedDist DynamicPPL provides several demo models and helpers for testing samplers in the `DynamicPPL.TestUtils` submodule. ```@docs -DynamicPPL.TestUtils.test_sampler_demo_models +DynamicPPL.TestUtils.test_sampler +DynamicPPL.TestUtils.test_sampler_on_demo_models DynamicPPL.TestUtils.test_sampler_continuous +DynamicPPL.TestUtils.marginal_mean_of_samples +``` + +```@docs +DynamicPPL.TestUtils.DEMO_MODELS ``` For every demo model, one can define the true log prior, log likelihood, and log joint probabilities. @@ -115,6 +121,20 @@ DynamicPPL.TestUtils.loglikelihood_true DynamicPPL.TestUtils.logjoint_true ``` +And in the case where the model includes constrained variables, it can also be useful to define + +```@docs +DynamicPPL.TestUtils.logprior_true_with_logabsdet_jacobian +DynamicPPL.TestUtils.logjoint_true_with_logabsdet_jacobian +``` + +Finally, the following methods can also be of use: + +```@docs +DynamicPPL.TestUtils.varnames +DynamicPPL.TestUtils.posterior_mean +``` + ## Advanced ### Variable names diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index 4a44fa869..86d4e0def 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.jl @@ -12,6 +12,8 @@ using MacroTools: MacroTools using Setfield: Setfield using ZygoteRules: ZygoteRules +using DocStringExtensions + using Random: Random import Base: diff --git a/src/context_implementations.jl b/src/context_implementations.jl index 20c4af446..6271b5d8c 100644 --- a/src/context_implementations.jl +++ b/src/context_implementations.jl @@ -55,7 +55,7 @@ end function tilde_assume(context::PriorContext{<:NamedTuple}, right, vn, vi) if haskey(context.vars, getsym(vn)) vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end return tilde_assume(PriorContext(), right, vn, vi) end @@ -64,7 +64,7 @@ function tilde_assume( ) if haskey(context.vars, getsym(vn)) vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end return tilde_assume(rng, PriorContext(), sampler, right, vn, vi) end @@ -72,7 +72,7 @@ end function tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, vn, vi) if haskey(context.vars, getsym(vn)) vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end return tilde_assume(LikelihoodContext(), right, vn, vi) end @@ -86,7 +86,7 @@ function tilde_assume( ) if haskey(context.vars, getsym(vn)) vi = setindex!!(vi, vectorize(right, get(context.vars, vn)), vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end return tilde_assume(rng, LikelihoodContext(), sampler, right, vn, vi) end @@ -194,7 +194,7 @@ end # fallback without sampler function assume(dist::Distribution, vn::VarName, vi) - r = vi[vn] + r = vi[vn, dist] return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn)), vi end @@ -211,16 +211,21 @@ function assume( if sampler isa SampleFromUniform || is_flagged(vi, vn, "del") unset_flag!(vi, vn, "del") r = init(rng, dist, sampler) - vi[vn] = vectorize(dist, r) - settrans!(vi, false, vn) + vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r)) setorder!(vi, vn, get_num_produce(vi)) else - r = vi[vn] + # Otherwise we just extract it. + r = vi[vn, dist] end else r = init(rng, dist, sampler) - push!!(vi, vn, r, dist, sampler) - settrans!(vi, false, vn) + if istrans(vi) + push!!(vi, vn, link(dist, r), dist, sampler) + # By default `push!!` sets the transformed flag to `false`. + settrans!!(vi, true, vn) + else + push!!(vi, vn, r, dist, sampler) + end end return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn)), vi @@ -286,7 +291,7 @@ function dot_tilde_assume(context::LikelihoodContext{<:NamedTuple}, right, left, var = get(context.vars, vn) _right, _left, _vns = unwrap_right_left_vns(right, var, vn) set_val!(vi, _vns, _right, _left) - settrans!.(Ref(vi), false, _vns) + settrans!!.((vi,), false, _vns) dot_tilde_assume(LikelihoodContext(), _right, _left, _vns, vi) else dot_tilde_assume(LikelihoodContext(), right, left, vn, vi) @@ -305,19 +310,20 @@ function dot_tilde_assume( var = get(context.vars, vn) _right, _left, _vns = unwrap_right_left_vns(right, var, vn) set_val!(vi, _vns, _right, _left) - settrans!.(Ref(vi), false, _vns) + settrans!!.((vi,), false, _vns) dot_tilde_assume(rng, LikelihoodContext(), sampler, _right, _left, _vns, vi) else dot_tilde_assume(rng, LikelihoodContext(), sampler, right, left, vn, vi) end end + function dot_tilde_assume(context::LikelihoodContext, right, left, vn, vi) - return dot_assume(NoDist.(right), left, vn, vi) + return dot_assume(nodist(right), left, vn, vi) end function dot_tilde_assume( rng::Random.AbstractRNG, context::LikelihoodContext, sampler, right, left, vn, vi ) - return dot_assume(rng, sampler, NoDist.(right), vn, left, vi) + return dot_assume(rng, sampler, nodist(right), vn, left, vi) end # `PriorContext` @@ -326,7 +332,7 @@ function dot_tilde_assume(context::PriorContext{<:NamedTuple}, right, left, vn, var = get(context.vars, vn) _right, _left, _vns = unwrap_right_left_vns(right, var, vn) set_val!(vi, _vns, _right, _left) - settrans!.(Ref(vi), false, _vns) + settrans!!.((vi,), false, _vns) dot_tilde_assume(PriorContext(), _right, _left, _vns, vi) else dot_tilde_assume(PriorContext(), right, left, vn, vi) @@ -345,7 +351,7 @@ function dot_tilde_assume( var = get(context.vars, vn) _right, _left, _vns = unwrap_right_left_vns(right, var, vn) set_val!(vi, _vns, _right, _left) - settrans!.(Ref(vi), false, _vns) + settrans!!.((vi,), false, _vns) dot_tilde_assume(rng, PriorContext(), sampler, _right, _left, _vns, vi) else dot_tilde_assume(rng, PriorContext(), sampler, right, left, vn, vi) @@ -383,14 +389,14 @@ function dot_assume( vns::AbstractVector{<:VarName}, vi::AbstractVarInfo, ) - @assert length(dist) == size(var, 1) + @assert length(dist) == size(var, 1) "dimensionality of `var` ($(size(var, 1))) is incompatible with dimensionality of `dist` $(length(dist))" # NOTE: We cannot work with `var` here because we might have a model of the form # # m = Vector{Float64}(undef, n) # m .~ Normal() # # in which case `var` will have `undef` elements, even if `m` is present in `vi`. - r = vi[vns] + r = vi[vns, dist] lp = sum(zip(vns, eachcol(r))) do (vn, ri) return Bijectors.logpdf_with_trans(dist, ri, istrans(vi, vn)) end @@ -412,19 +418,21 @@ function dot_assume( end function dot_assume( - dists::Union{Distribution,AbstractArray{<:Distribution}}, + dist::Distribution, var::AbstractArray, vns::AbstractArray{<:VarName}, vi +) + r = getindex.((vi,), vns, (dist,)) + lp = sum(Bijectors.logpdf_with_trans.((dist,), r, istrans.((vi,), vns))) + return r, lp, vi +end + +function dot_assume( + dists::AbstractArray{<:Distribution}, var::AbstractArray, vns::AbstractArray{<:VarName}, vi, ) - # NOTE: We cannot work with `var` here because we might have a model of the form - # - # m = Vector{Float64}(undef, n) - # m .~ Normal() - # - # in which case `var` will have `undef` elements, even if `m` is present in `vi`. - r = reshape(vi[vec(vns)], size(vns)) - lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1]))) + r = getindex.((vi,), vns, dists) + lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns))) return r, lp, vi end @@ -438,7 +446,7 @@ function dot_assume( ) r = get_and_set_val!(rng, vi, vns, dists, spl) # Make sure `r` is not a matrix for multivariate distributions - lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans(vi, vns[1]))) + lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns))) return r, lp, vi end function dot_assume(rng, spl::Sampler, ::Any, ::AbstractArray{<:VarName}, ::Any, ::Any) @@ -462,19 +470,23 @@ function get_and_set_val!( r = init(rng, dist, spl, n) for i in 1:n vn = vns[i] - vi[vn] = vectorize(dist, r[:, i]) - settrans!(vi, false, vn) + vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r[:, i])) setorder!(vi, vn, get_num_produce(vi)) end else - r = vi[vns] + r = vi[vns, dist] end else r = init(rng, dist, spl, n) for i in 1:n vn = vns[i] - push!!(vi, vn, r[:, i], dist, spl) - settrans!(vi, false, vn) + if istrans(vi) + push!!(vi, vn, Bijectors.link(dist, r[:, i]), dist, spl) + # `push!!` sets the trans-flag to `false` by default. + settrans!!(vi, true, vn) + else + push!!(vi, vn, r[:, i], dist, spl) + end end end return r @@ -496,12 +508,13 @@ function get_and_set_val!( for i in eachindex(vns) vn = vns[i] dist = dists isa AbstractArray ? dists[i] : dists - vi[vn] = vectorize(dist, r[i]) - settrans!(vi, false, vn) + vi[vn] = vectorize(dist, maybe_link(vi, vn, dist, r[i])) setorder!(vi, vn, get_num_produce(vi)) end else - r = reshape(vi[vec(vns)], size(vns)) + # r = reshape(vi[vec(vns)], size(vns)) + r_raw = getindex_raw(vi, vec(vns)) + r = maybe_invlink.((vi,), vns, dists, reshape(r_raw, size(vns))) end else f = (vn, dist) -> init(rng, dist, spl) @@ -511,8 +524,13 @@ function get_and_set_val!( # 1. Figure out the broadcast size and use a `foreach`. # 2. Define an anonymous function which returns `nothing`, which # we then broadcast. This will allocate a vector of `nothing` though. - push!!.(Ref(vi), vns, r, dists, Ref(spl)) - settrans!.(Ref(vi), false, vns) + if istrans(vi) + push!!.((vi,), vns, link.((vi,), vns, dists, r), dists, (spl,)) + # `push!!` sets the trans-flag to `false` by default. + settrans!!.((vi,), true, vns) + else + push!!.((vi,), vns, r, dists, (spl,)) + end end return r end diff --git a/src/distribution_wrappers.jl b/src/distribution_wrappers.jl index 4045cc089..d8968a68e 100644 --- a/src/distribution_wrappers.jl +++ b/src/distribution_wrappers.jl @@ -13,6 +13,9 @@ end NamedDist(dist::Distribution, name::Symbol) = NamedDist(dist, VarName{name}()) +Base.length(dist::NamedDist) = Base.length(dist.dist) +Base.size(dist::NamedDist) = Base.size(dist.dist) + Distributions.logpdf(dist::NamedDist, x::Real) = Distributions.logpdf(dist.dist, x) function Distributions.logpdf(dist::NamedDist, x::AbstractArray{<:Real}) return Distributions.logpdf(dist.dist, x) @@ -24,12 +27,20 @@ function Distributions.loglikelihood(dist::NamedDist, x::AbstractArray{<:Real}) return Distributions.loglikelihood(dist.dist, x) end +Bijectors.bijector(d::NamedDist) = Bijectors.bijector(d.dist) + struct NoDist{variate,support,Td<:Distribution{variate,support}} <: Distribution{variate,support} dist::Td end NoDist(dist::NamedDist) = NamedDist(NoDist(dist.dist), dist.name) +nodist(dist::Distribution) = NoDist(dist) +nodist(dists::AbstractArray) = nodist.(dists) + +Base.length(dist::NoDist) = Base.length(dist.dist) +Base.size(dist::NoDist) = Base.size(dist.dist) + 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 @@ -40,9 +51,21 @@ 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 -function Bijectors.logpdf_with_trans(d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real}) +Bijectors.logpdf_with_trans(d::NoDist{<:Univariate}, ::Real, ::Bool) = 0 +function Bijectors.logpdf_with_trans( + d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}, ::Bool +) + return 0 +end +function Bijectors.logpdf_with_trans( + d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real}, ::Bool +) return zeros(Int, size(x, 2)) end -Bijectors.logpdf_with_trans(d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}) = 0 +function Bijectors.logpdf_with_trans( + d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}, ::Bool +) + return 0 +end + +Bijectors.bijector(d::NoDist) = Bijectors.bijector(d.dist) diff --git a/src/simple_varinfo.jl b/src/simple_varinfo.jl index 5cecda4b2..5b9edefdf 100644 --- a/src/simple_varinfo.jl +++ b/src/simple_varinfo.jl @@ -1,11 +1,19 @@ +abstract type AbstractTransformation end + +struct NoTransformation <: AbstractTransformation end +struct DefaultTransformation <: AbstractTransformation end + """ - SimpleVarInfo{NT,T} <: AbstractVarInfo + $(TYPEDEF) A simple wrapper of the parameters with a `logp` field for accumulation of the logdensity. Currently only implemented for `NT<:NamedTuple` and `NT<:Dict`. +# Fields +$(FIELDS) + # Notes The major differences between this and `TypedVarInfo` are: 1. `SimpleVarInfo` does not require linearization. @@ -16,7 +24,7 @@ The major differences between this and `TypedVarInfo` are: # Examples ## General usage -```jldoctest; setup=:(using Distributions) +```jldoctest simplevarinfo-general; setup=:(using Distributions) julia> using StableRNGs julia> @model function demo() @@ -78,6 +86,60 @@ ERROR: KeyError: key x[1:2] not found [...] ``` +You can also sample in _transformed_ space: + +```jldoctest simplevarinfo-general +julia> @model demo_constrained() = x ~ Exponential() +demo_constrained (generic function with 2 methods) + +julia> m = demo_constrained(); + +julia> _, vi = DynamicPPL.evaluate!!(m, SimpleVarInfo(), ctx); + +julia> vi[@varname(x)] # (✓) 0 ≤ x < ∞ +1.8632965762164932 + +julia> _, vi = DynamicPPL.evaluate!!(m, DynamicPPL.settrans!!(SimpleVarInfo(), true), ctx); + +julia> vi[@varname(x)] # (✓) -∞ < x < ∞ +-0.21080155351918753 + +julia> xs = [last(DynamicPPL.evaluate!!(m, DynamicPPL.settrans!!(SimpleVarInfo(), true), ctx))[@varname(x)] for i = 1:10]; + +julia> any(xs .< 0) # (✓) Positive probability mass on negative numbers! +true + +julia> # And with `Dict` of course! + _, vi = DynamicPPL.evaluate!!(m, DynamicPPL.settrans!!(SimpleVarInfo(Dict()), true), ctx); + +julia> vi[@varname(x)] # (✓) -∞ < x < ∞ +0.6225185067787314 + +julia> xs = [last(DynamicPPL.evaluate!!(m, DynamicPPL.settrans!!(SimpleVarInfo(), true), ctx))[@varname(x)] for i = 1:10]; + +julia> any(xs .< 0) # (✓) Positive probability mass on negative numbers! +true +``` + +Evaluation in transformed space of course also works: + +```jldoctest simplevarinfo-general +julia> vi = DynamicPPL.settrans!!(SimpleVarInfo((x = -1.0,)), true) +Transformed SimpleVarInfo((x = -1.0,), 0.0) + +julia> # (✓) Positive probability mass on negative numbers! + getlogp(last(DynamicPPL.evaluate!!(m, vi, DynamicPPL.DefaultContext()))) +-1.3678794411714423 + +julia> # While if we forget to make indicate that it's transformed: + vi = DynamicPPL.settrans!!(SimpleVarInfo((x = -1.0,)), false) +SimpleVarInfo((x = -1.0,), 0.0) + +julia> # (✓) No probability mass on negative numbers! + getlogp(last(DynamicPPL.evaluate!!(m, vi, DynamicPPL.DefaultContext()))) +-Inf +``` + ## Indexing Using `NamedTuple` as underlying storage. @@ -126,14 +188,26 @@ ERROR: type NamedTuple has no field b [...] ``` """ -struct SimpleVarInfo{NT,T} <: AbstractVarInfo +struct SimpleVarInfo{NT,T,C<:AbstractTransformation} <: AbstractVarInfo + "underlying representation of the realization represented" values::NT + "holds the accumulated log-probability" logp::T + "represents whether it assumes variables to be transformed" + transformation::C end -SimpleVarInfo{T}(θ) where {T<:Real} = SimpleVarInfo{typeof(θ),T}(θ, zero(T)) -SimpleVarInfo{T}(; kwargs...) where {T<:Real} = SimpleVarInfo{T}(NamedTuple(kwargs)) -SimpleVarInfo(; kwargs...) = SimpleVarInfo{Float64}(NamedTuple(kwargs)) +SimpleVarInfo(values, logp) = SimpleVarInfo(values, logp, NoTransformation()) + +function SimpleVarInfo{T}(θ) where {T<:Real} + return SimpleVarInfo(θ, zero(T)) +end +function SimpleVarInfo{T}(; kwargs...) where {T<:Real} + return SimpleVarInfo{T}(NamedTuple(kwargs)) +end +function SimpleVarInfo(; kwargs...) + return SimpleVarInfo{Float64}(NamedTuple(kwargs)) +end SimpleVarInfo(θ) = SimpleVarInfo{Float64}(θ) # Constructor from `Model`. @@ -158,8 +232,8 @@ function BangBang.empty!!(vi::SimpleVarInfo) end getlogp(vi::SimpleVarInfo) = vi.logp -setlogp!!(vi::SimpleVarInfo, logp) = SimpleVarInfo(vi.values, logp) -acclogp!!(vi::SimpleVarInfo, logp) = SimpleVarInfo(vi.values, getlogp(vi) + logp) +setlogp!!(vi::SimpleVarInfo, logp) = Setfield.@set vi.logp = logp +acclogp!!(vi::SimpleVarInfo, logp) = Setfield.@set vi.logp = getlogp(vi) + logp """ keys(vi::SimpleVarInfo) @@ -167,6 +241,7 @@ acclogp!!(vi::SimpleVarInfo, logp) = SimpleVarInfo(vi.values, getlogp(vi) + logp Return an iterator of keys present in `vi`. """ Base.keys(vi::SimpleVarInfo) = keys(vi.values) +Base.keys(vi::SimpleVarInfo{<:NamedTuple}) = map(k -> VarName{k}(), keys(vi.values)) function setlogp!!(vi::SimpleVarInfo{<:Any,<:Ref}, logp) vi.logp[] = logp @@ -179,10 +254,24 @@ function acclogp!!(vi::SimpleVarInfo{<:Any,<:Ref}, logp) end function Base.show(io::IO, ::MIME"text/plain", svi::SimpleVarInfo) + if !(svi.transformation isa NoTransformation) + print(io, "Transformed ") + end + return print(io, "SimpleVarInfo(", svi.values, ", ", svi.logp, ")") end # `NamedTuple` +function Base.getindex(vi::SimpleVarInfo, vn::VarName, dist::Distribution) + return maybe_invlink(vi, vn, dist, getindex(vi, vn)) +end +function Base.getindex(vi::SimpleVarInfo, vns::Vector{<:VarName}, dist::Distribution) + vals_linked = mapreduce(vcat, vns) do vn + getindex(vi, vn, dist) + end + return reconstruct(dist, vals_linked, length(vns)) +end + Base.getindex(vi::SimpleVarInfo, vn::VarName) = get(vi.values, vn) # `Dict` @@ -221,9 +310,23 @@ Base.getindex(vi::SimpleVarInfo, vns::Vector{<:VarName}) = map(Base.Fix1(getinde Base.getindex(vi::SimpleVarInfo, spl::SampleFromPrior) = vi.values Base.getindex(vi::SimpleVarInfo, spl::SampleFromUniform) = vi.values + # TODO: Should we do better? Base.getindex(vi::SimpleVarInfo, spl::Sampler) = vi.values +# Since we don't perform any transformations in `getindex` for `SimpleVarInfo` +# we simply call `getindex` in `getindex_raw`. +getindex_raw(vi::SimpleVarInfo, vn::VarName) = vi[vn] +function getindex_raw(vi::SimpleVarInfo, vn::VarName, dist::Distribution) + return reconstruct(dist, getindex_raw(vi, vn)) +end +getindex_raw(vi::SimpleVarInfo, vns::Vector{<:VarName}) = vi[vns] +function getindex_raw(vi::SimpleVarInfo, vns::Vector{<:VarName}, dist::Distribution) + # `reconstruct` expects a flattened `Vector` regardless of the type of `dist`, so we `vcat` everything. + vals = mapreduce(Base.Fix1(getindex_raw, vi), vcat, vns) + return reconstruct(dist, vals, length(vns)) +end + Base.haskey(vi::SimpleVarInfo, vn::VarName) = _haskey(vi.values, vn) function _haskey(nt::NamedTuple, vn::VarName) # LHS: Ensure that `nt` indeed has the property we want. @@ -259,7 +362,7 @@ end function BangBang.setindex!!(vi::SimpleVarInfo, val, vn::VarName) # For `NamedTuple` we treat the symbol in `vn` as the _property_ to set. - return SimpleVarInfo(set!!(vi.values, vn, val), vi.logp) + return Setfield.@set vi.values = set!!(vi.values, vn, val) end # TODO: Specialize to handle certain cases, e.g. a collection of `VarName` with @@ -290,7 +393,7 @@ function BangBang.setindex!!(vi::SimpleVarInfo{<:AbstractDict}, val, vn::VarName vn_key = VarName(vn, keylens) BangBang.setindex!!(dict, set!!(dict[vn_key], child, val), vn_key) end - return SimpleVarInfo(dict_new, vi.logp) + return Setfield.@set vi.values = dict_new end # `NamedTuple` @@ -325,8 +428,8 @@ function BangBang.push!!( return vi end -const SimpleOrThreadSafeSimple{T,V} = Union{ - SimpleVarInfo{T,V},ThreadSafeVarInfo{<:SimpleVarInfo{T,V}} +const SimpleOrThreadSafeSimple{T,V,C} = Union{ + SimpleVarInfo{T,V,C},ThreadSafeVarInfo{<:SimpleVarInfo{T,V,C}} } # Necessary for `matchingvalue` to work properly. @@ -337,79 +440,90 @@ function Base.eltype( end # Context implementations -function assume(dist::Distribution, vn::VarName, vi::SimpleOrThreadSafeSimple) - left = vi[vn] - return left, Distributions.loglikelihood(dist, left), vi -end - +# NOTE: Evaluations, i.e. those without `rng` are shared with other +# implementations of `AbstractVarInfo`. function assume( rng::Random.AbstractRNG, - sampler::SampleFromPrior, + sampler::Union{SampleFromPrior,SampleFromUniform}, dist::Distribution, vn::VarName, vi::SimpleOrThreadSafeSimple, ) value = init(rng, dist, sampler) - vi = BangBang.push!!(vi, vn, value, dist, sampler) - return value, Distributions.loglikelihood(dist, value), vi -end - -function dot_assume( - dist::MultivariateDistribution, - var::AbstractMatrix, - vns::AbstractVector{<:VarName}, - vi::SimpleOrThreadSafeSimple, -) - @assert length(dist) == size(var, 1) - # NOTE: We cannot work with `var` here because we might have a model of the form - # - # m = Vector{Float64}(undef, n) - # m .~ Normal() - # - # in which case `var` will have `undef` elements, even if `m` is present in `vi`. - value = vi[vns] - lp = sum(zip(vns, eachcol(value))) do (vn, val) - return Distributions.logpdf(dist, val) - end - return value, lp, vi + # Transform if we're working in unconstrained space. + value_raw = maybe_link(vi, vn, dist, value) + vi = BangBang.push!!(vi, vn, value_raw, dist, sampler) + return value, Bijectors.logpdf_with_trans(dist, value, istrans(vi, vn)), vi end function dot_assume( + rng, + spl::Union{SampleFromPrior,SampleFromUniform}, dists::Union{Distribution,AbstractArray{<:Distribution}}, - var::AbstractArray, vns::AbstractArray{<:VarName}, + var::AbstractArray, vi::SimpleOrThreadSafeSimple, ) - # NOTE: We cannot work with `var` here because we might have a model of the form - # - # m = Vector{Float64}(undef, n) - # m .~ Normal() - # - # in which case `var` will have `undef` elements, even if `m` is present in `vi`. - value = vi[vns] - lp = sum(Distributions.logpdf.(dists, value)) + f = (vn, dist) -> init(rng, dist, spl) + value = f.(vns, dists) + + # Transform if we're working in transformed space. + value_raw = if dists isa Distribution + maybe_link.((vi,), vns, (dists,), value) + else + maybe_link.((vi,), vns, dists, value) + end + + # Update `vi` + vi = BangBang.setindex!!(vi, value_raw, vns) + + # Compute logp. + lp = sum(Bijectors.logpdf_with_trans.(dists, value, istrans.((vi,), vns))) return value, lp, vi end function dot_assume( rng, spl::Union{SampleFromPrior,SampleFromUniform}, - dists::Union{Distribution,AbstractArray{<:Distribution}}, - vns::AbstractArray{<:VarName}, - var::AbstractArray, + dist::MultivariateDistribution, + vns::AbstractVector{<:VarName}, + var::AbstractMatrix, vi::SimpleOrThreadSafeSimple, ) - f = (vn, dist) -> init(rng, dist, spl) - value = f.(vns, dists) - vi = BangBang.setindex!!(vi, value, vns) - lp = sum(Distributions.logpdf.(dists, value)) + @assert length(dist) == size(var, 1) "dimensionality of `var` ($(size(var, 1))) is incompatible with dimensionality of `dist` $(length(dist))" + + # r = get_and_set_val!(rng, vi, vns, dist, spl) + n = length(vns) + value = init(rng, dist, spl, n) + + # Update `vi`. + for (vn, val) in zip(vns, eachcol(value)) + val_linked = maybe_link(vi, vn, dist, val) + vi = BangBang.setindex!!(vi, val_linked, vn) + end + + # Compute logp. + lp = sum(Bijectors.logpdf_with_trans(dist, value, istrans(vi))) return value, lp, vi end # HACK: Allows us to re-use the implementation of `dot_tilde`, etc. for literals. increment_num_produce!(::SimpleOrThreadSafeSimple) = nothing -settrans!(vi::SimpleOrThreadSafeSimple, trans::Bool, vn::VarName) = nothing -istrans(::SimpleVarInfo, vn::VarName) = false + +# NOTE: We don't implement `settrans!!(vi, trans, vn)`. +function settrans!!(vi::SimpleVarInfo, trans) + return settrans!!(vi, trans ? DefaultTransformation() : NoTransformation()) +end +function settrans!!(vi::SimpleVarInfo, transformation::AbstractTransformation) + return Setfield.@set vi.transformation = transformation +end +function settrans!!(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, trans) + return Setfield.@set vi.varinfo = settrans!!(vi.varinfo, trans) +end + +istrans(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation) +istrans(vi::SimpleVarInfo, vn::VarName) = istrans(vi) +istrans(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, vn::VarName) = istrans(vi.varinfo, vn) """ values_as(varinfo[, Type]) diff --git a/src/test_utils.jl b/src/test_utils.jl index ca0fabc9a..ef314fa91 100644 --- a/src/test_utils.jl +++ b/src/test_utils.jl @@ -6,10 +6,35 @@ using LinearAlgebra using Distributions using Test +using Random: Random +using Bijectors: Bijectors +using Setfield: Setfield + """ - logprior_true(model, θ) + varname_leaves(vn::VarName, val) -Return the `logprior` of `model` for `θ`. +Return iterator over all varnames that are represented by `vn` on `val`, +e.g. `varname_leaves(@varname(x), rand(2))` results in an iterator over `[@varname(x[1]), @varname(x[2])]`. +""" +varname_leaves(vn::VarName, val::Real) = [vn] +function varname_leaves(vn::VarName, val::AbstractArray{<:Union{Real,Missing}}) + return ( + VarName(vn, DynamicPPL.getlens(vn) ∘ Setfield.IndexLens(Tuple(I))) for + I in CartesianIndices(val) + ) +end +function varname_leaves(vn::VarName, val::AbstractArray) + return Iterators.flatten( + varname_leaves( + VarName(vn, DynamicPPL.getlens(vn) ∘ Setfield.IndexLens(Tuple(I))), val[I] + ) for I in CartesianIndices(val) + ) +end + +""" + logprior_true(model, args...) + +Return the `logprior` of `model` for `args`. This should generally be implemented by hand for every specific `model`. @@ -18,9 +43,9 @@ See also: [`logjoint_true`](@ref), [`loglikelihood_true`](@ref). function logprior_true end """ - loglikelihood_true(model, θ) + loglikelihood_true(model, args...) -Return the `loglikelihood` of `model` for `θ`. +Return the `loglikelihood` of `model` for `args`. This should generally be implemented by hand for every specific `model`. @@ -29,11 +54,11 @@ See also: [`logjoint_true`](@ref), [`logprior_true`](@ref). function loglikelihood_true end """ - logjoint_true(model, θ) + logjoint_true(model, args...) -Return the `logjoint` of `model` for `θ`. +Return the `logjoint` of `model` for `args`. -Defaults to `logprior_true(model, θ) + loglikelihood_true(model, θ)`. +Defaults to `logprior_true(model, args...) + loglikelihood_true(model, args..)`. This should generally be implemented by hand for every specific `model` so that the returned value can be used as a ground-truth for testing things like: @@ -49,202 +74,561 @@ function logjoint_true(model::Model, args...) return logprior_true(model, args...) + loglikelihood_true(model, args...) end -# A collection of models for which the mean-of-means for the posterior should -# be same. +""" + logjoint_true_with_logabsdet_jacobian(model::Model, args...) + +Return a tuple `(args_unconstrained, logjoint)` of `model` for `args`. + +Unlike [`logjoint_true`](@ref), the returned logjoint computation includes the +log-absdet-jacobian adjustment, thus computing logjoint for the unconstrained variables. + +Note that `args` are assumed be in the support of `model`, while `args_unconstrained` +will be unconstrained. + +This should generally not be implemented directly, instead one should implement +[`logprior_true_with_logabsdet_jacobian`](@ref) for a given `model`. + +See also: [`logjoint_true`](@ref), [`logprior_true_with_logabsdet_jacobian`](@ref). +""" +function logjoint_true_with_logabsdet_jacobian(model::Model, args...) + args_unconstrained, lp = logprior_true_with_logabsdet_jacobian(model, args...) + return args_unconstrained, lp + loglikelihood_true(model, args...) +end + +""" + logprior_true_with_logabsdet_jacobian(model::Model, args...) + +Return a tuple `(args_unconstrained, logprior_unconstrained)` of `model` for `args...`. + +Unlike [`logprior_true`](@ref), the returned logprior computation includes the +log-absdet-jacobian adjustment, thus computing logprior for the unconstrained variables. + +Note that `args` are assumed be in the support of `model`, while `args_unconstrained` +will be unconstrained. + +See also: [`logprior_true`](@ref). +""" +function logprior_true_with_logabsdet_jacobian end + +""" + varnames(model::Model) + +Return a collection of `VarName` as they are expected to appear in the model. + +Even though it is recommended to implement this by hand for a particular `Model`, +a default implementation using [`SimpleVarInfo{<:Dict}`](@ref) is provided. +""" +function varnames(model::Model) + return collect( + keys(last(DynamicPPL.evaluate!!(model, SimpleVarInfo(Dict()), SamplingContext()))) + ) +end + +""" + posterior_mean(model::Model) + +Return a `NamedTuple` compatible with `varnames(model)` where the values represent +the posterior mean under `model`. + +"Compatible" means that a `varname` from `varnames(model)` can be used to extract the +corresponding value using `get`, e.g. `get(posterior_mean(model), varname)`. +""" +function posterior_mean end + +""" + demo_dynamic_constraint() + +A model with variables `m` and `x` with `x` having support depending on `m`. +""" +@model function demo_dynamic_constraint() + m ~ Normal() + x ~ truncated(Normal(), m, Inf) + + return (m=m, x=x) +end + +function logprior_true(model::Model{typeof(demo_dynamic_constraint)}, m, x) + return logpdf(Normal(), m) + logpdf(truncated(Normal(), m, Inf), x) +end +function loglikelihood_true(model::Model{typeof(demo_dynamic_constraint)}, m, x) + return zero(float(eltype(m))) +end +function varnames(model::Model{typeof(demo_dynamic_constraint)}) + return [@varname(m), @varname(x)] +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dynamic_constraint)}, m, x +) + b_x = Bijectors.bijector(truncated(Normal(), m, Inf)) + x_unconstrained, Δlogp = Bijectors.with_logabsdet_jacobian(b_x, x) + return (m=m, x=x_unconstrained), logprior_true(model, m, x) - Δlogp +end + +# A collection of models for which the posterior should be "similar". +# Some utility methods for these. +function _demo_logprior_true_with_logabsdet_jacobian(model, s, m) + b = Bijectors.bijector(InverseGamma(2, 3)) + s_unconstrained = b.(s) + Δlogp = sum(Base.Fix1(Bijectors.logabsdetjac, b), s) + return (s=s_unconstrained, m=m), logprior_true(model, s, m) - Δlogp +end + @model function demo_dot_assume_dot_observe( - x=[10.0, 10.0], ::Type{TV}=Vector{Float64} + x=[1.5, 2.0], ::Type{TV}=Vector{Float64} ) where {TV} # `dot_assume` and `observe` + s = TV(undef, length(x)) m = TV(undef, length(x)) - m .~ Normal() - x ~ MvNormal(m, 0.25 * I) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + s .~ InverseGamma(2, 3) + m .~ Normal.(0, sqrt.(s)) + + x ~ MvNormal(m, Diagonal(s)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) +end +function logprior_true(model::Model{typeof(demo_dot_assume_dot_observe)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function logprior_true(model::Model{typeof(demo_dot_assume_dot_observe)}, m) - return loglikelihood(Normal(), m) +function loglikelihood_true(model::Model{typeof(demo_dot_assume_dot_observe)}, s, m) + return loglikelihood(MvNormal(m, Diagonal(s)), model.args.x) end -function loglikelihood_true(model::Model{typeof(demo_dot_assume_dot_observe)}, m) - return loglikelihood(MvNormal(m, 0.25 * I), model.args.x) +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_dot_observe)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_dot_assume_dot_observe)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end @model function demo_assume_index_observe( - x=[10.0, 10.0], ::Type{TV}=Vector{Float64} + x=[1.5, 2.0], ::Type{TV}=Vector{Float64} ) where {TV} # `assume` with indexing and `observe` + s = TV(undef, length(x)) + for i in eachindex(s) + s[i] ~ InverseGamma(2, 3) + end m = TV(undef, length(x)) for i in eachindex(m) - m[i] ~ Normal() + m[i] ~ Normal(0, sqrt(s[i])) end - x ~ MvNormal(m, 0.25 * I) + x ~ MvNormal(m, Diagonal(s)) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_assume_index_observe)}, m) - return loglikelihood(Normal(), m) +function logprior_true(model::Model{typeof(demo_assume_index_observe)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function loglikelihood_true(model::Model{typeof(demo_assume_index_observe)}, m) - return logpdf(MvNormal(m, 0.25 * I), model.args.x) +function loglikelihood_true(model::Model{typeof(demo_assume_index_observe)}, s, m) + return logpdf(MvNormal(m, Diagonal(s)), model.args.x) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_index_observe)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_assume_index_observe)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end -@model function demo_assume_multivariate_observe(x=[10.0, 10.0]) +@model function demo_assume_multivariate_observe(x=[1.5, 2.0]) # Multivariate `assume` and `observe` - m ~ MvNormal(zero(x), I) - x ~ MvNormal(m, 0.25 * I) + s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)]) + m ~ MvNormal(zero(x), Diagonal(s)) + x ~ MvNormal(m, Diagonal(s)) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) +end +function logprior_true(model::Model{typeof(demo_assume_multivariate_observe)}, s, m) + s_dist = product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)]) + m_dist = MvNormal(zero(model.args.x), Diagonal(s)) + return logpdf(s_dist, s) + logpdf(m_dist, m) end -function logprior_true(model::Model{typeof(demo_assume_multivariate_observe)}, m) - return logpdf(MvNormal(zero(model.args.x), I), m) +function loglikelihood_true(model::Model{typeof(demo_assume_multivariate_observe)}, s, m) + return logpdf(MvNormal(m, Diagonal(s)), model.args.x) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_multivariate_observe)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) end -function loglikelihood_true(model::Model{typeof(demo_assume_multivariate_observe)}, m) - return logpdf(MvNormal(m, 0.25 * I), model.args.x) +function varnames(model::Model{typeof(demo_assume_multivariate_observe)}) + return [@varname(s), @varname(m)] end @model function demo_dot_assume_observe_index( - x=[10.0, 10.0], ::Type{TV}=Vector{Float64} + x=[1.5, 2.0], ::Type{TV}=Vector{Float64} ) where {TV} # `dot_assume` and `observe` with indexing + s = TV(undef, length(x)) + s .~ InverseGamma(2, 3) m = TV(undef, length(x)) - m .~ Normal() + m .~ Normal.(0, sqrt.(s)) for i in eachindex(x) - x[i] ~ Normal(m[i], 0.5) + x[i] ~ Normal(m[i], sqrt(s[i])) end - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_dot_assume_observe_index)}, m) - return loglikelihood(Normal(), m) +function logprior_true(model::Model{typeof(demo_dot_assume_observe_index)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function loglikelihood_true(model::Model{typeof(demo_dot_assume_observe_index)}, m) - return sum(logpdf.(Normal.(m, 0.5), model.args.x)) +function loglikelihood_true(model::Model{typeof(demo_dot_assume_observe_index)}, s, m) + return sum(logpdf.(Normal.(m, sqrt.(s)), model.args.x)) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_observe_index)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_dot_assume_observe_index)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end # Using vector of `length` 1 here so the posterior of `m` is the same # as the others. -@model function demo_assume_dot_observe(x=[10.0]) +@model function demo_assume_dot_observe(x=[1.5, 2.0]) # `assume` and `dot_observe` - m ~ Normal() - x .~ Normal(m, 0.5) + s ~ InverseGamma(2, 3) + m ~ Normal(0, sqrt(s)) + x .~ Normal(m, sqrt(s)) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_assume_dot_observe)}, m) - return logpdf(Normal(), m) +function logprior_true(model::Model{typeof(demo_assume_dot_observe)}, s, m) + return logpdf(InverseGamma(2, 3), s) + logpdf(Normal(0, sqrt(s)), m) end -function loglikelihood_true(model::Model{typeof(demo_assume_dot_observe)}, m) - return sum(logpdf.(Normal.(m, 0.5), model.args.x)) +function loglikelihood_true(model::Model{typeof(demo_assume_dot_observe)}, s, m) + return sum(logpdf.(Normal.(m, sqrt.(s)), model.args.x)) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_dot_observe)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_assume_dot_observe)}) + return [@varname(s), @varname(m)] end @model function demo_assume_observe_literal() # `assume` and literal `observe` - m ~ MvNormal(zeros(2), I) - [10.0, 10.0] ~ MvNormal(m, 0.25 * I) + s ~ product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)]) + m ~ MvNormal(zeros(2), Diagonal(s)) + [1.5, 2.0] ~ MvNormal(m, Diagonal(s)) - return (; m=m, x=[10.0, 10.0], logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=[1.5, 2.0], logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_assume_observe_literal)}, m) - return logpdf(MvNormal(zeros(2), I), m) +function logprior_true(model::Model{typeof(demo_assume_observe_literal)}, s, m) + s_dist = product_distribution([InverseGamma(2, 3), InverseGamma(2, 3)]) + m_dist = MvNormal(zeros(2), Diagonal(s)) + return logpdf(s_dist, s) + logpdf(m_dist, m) end -function loglikelihood_true(model::Model{typeof(demo_assume_observe_literal)}, m) - return logpdf(MvNormal(m, 0.25 * I), [10.0, 10.0]) +function loglikelihood_true(model::Model{typeof(demo_assume_observe_literal)}, s, m) + return logpdf(MvNormal(m, Diagonal(s)), [1.5, 2.0]) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_observe_literal)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_assume_observe_literal)}) + return [@varname(s), @varname(m)] end @model function demo_dot_assume_observe_index_literal(::Type{TV}=Vector{Float64}) where {TV} # `dot_assume` and literal `observe` with indexing + s = TV(undef, 2) m = TV(undef, 2) - m .~ Normal() - for i in eachindex(m) - 10.0 ~ Normal(m[i], 0.5) - end + s .~ InverseGamma(2, 3) + m .~ Normal.(0, sqrt.(s)) + + 1.5 ~ Normal(m[1], sqrt(s[1])) + 2.0 ~ Normal(m[2], sqrt(s[2])) - return (; m=m, x=fill(10.0, length(m)), logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=[1.5, 2.0], logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_dot_assume_observe_index_literal)}, m) - return loglikelihood(Normal(), m) +function logprior_true(model::Model{typeof(demo_dot_assume_observe_index_literal)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function loglikelihood_true(model::Model{typeof(demo_dot_assume_observe_index_literal)}, m) - return sum(logpdf.(Normal.(m, 0.5), fill(10.0, length(m)))) +function loglikelihood_true( + model::Model{typeof(demo_dot_assume_observe_index_literal)}, s, m +) + return sum(logpdf.(Normal.(m, sqrt.(s)), [1.5, 2.0])) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_observe_index_literal)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_dot_assume_observe_index_literal)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end @model function demo_assume_literal_dot_observe() # `assume` and literal `dot_observe` - m ~ Normal() - [10.0] .~ Normal(m, 0.5) + s ~ InverseGamma(2, 3) + m ~ Normal(0, sqrt(s)) + [1.5, 2.0] .~ Normal(m, sqrt(s)) - return (; m=m, x=[10.0], logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=[1.5, 2.0], logp=getlogp(__varinfo__)) +end +function logprior_true(model::Model{typeof(demo_assume_literal_dot_observe)}, s, m) + return logpdf(InverseGamma(2, 3), s) + logpdf(Normal(0, sqrt(s)), m) +end +function loglikelihood_true(model::Model{typeof(demo_assume_literal_dot_observe)}, s, m) + return loglikelihood(Normal(m, sqrt(s)), [1.5, 2.0]) end -function logprior_true(model::Model{typeof(demo_assume_literal_dot_observe)}, m) - return logpdf(Normal(), m) +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_literal_dot_observe)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) end -function loglikelihood_true(model::Model{typeof(demo_assume_literal_dot_observe)}, m) - return logpdf(Normal(m, 0.5), 10.0) +function varnames(model::Model{typeof(demo_assume_literal_dot_observe)}) + return [@varname(s), @varname(m)] end @model function _prior_dot_assume(::Type{TV}=Vector{Float64}) where {TV} + s = TV(undef, 2) + s .~ InverseGamma(2, 3) m = TV(undef, 2) - m .~ Normal() + m .~ Normal.(0, sqrt.(s)) - return m + return s, m end @model function demo_assume_submodel_observe_index_literal() # Submodel prior - @submodel m = _prior_dot_assume() - for i in eachindex(m) - 10.0 ~ Normal(m[i], 0.5) - end + @submodel s, m = _prior_dot_assume() + 1.5 ~ Normal(m[1], sqrt(s[1])) + 2.0 ~ Normal(m[2], sqrt(s[2])) - return (; m=m, x=[10.0], logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=[1.5, 2.0], logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_assume_submodel_observe_index_literal)}, m) - return loglikelihood(Normal(), m) +function logprior_true( + model::Model{typeof(demo_assume_submodel_observe_index_literal)}, s, m +) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end function loglikelihood_true( - model::Model{typeof(demo_assume_submodel_observe_index_literal)}, m + model::Model{typeof(demo_assume_submodel_observe_index_literal)}, s, m ) - return sum(logpdf.(Normal.(m, 0.5), 10.0)) + return sum(logpdf.(Normal.(m, sqrt.(s)), [1.5, 2.0])) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_assume_submodel_observe_index_literal)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_assume_submodel_observe_index_literal)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end -@model function _likelihood_dot_observe(m, x) - return x ~ MvNormal(m, 0.25 * I) +@model function _likelihood_mltivariate_observe(s, m, x) + return x ~ MvNormal(m, Diagonal(s)) end @model function demo_dot_assume_observe_submodel( - x=[10.0, 10.0], ::Type{TV}=Vector{Float64} + x=[1.5, 2.0], ::Type{TV}=Vector{Float64} ) where {TV} + s = TV(undef, length(x)) + s .~ InverseGamma(2, 3) m = TV(undef, length(x)) - m .~ Normal() + m .~ Normal.(0, sqrt.(s)) # Submodel likelihood - @submodel _likelihood_dot_observe(m, x) + @submodel _likelihood_mltivariate_observe(s, m, x) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) end -function logprior_true(model::Model{typeof(demo_dot_assume_observe_submodel)}, m) - return loglikelihood(Normal(), m) +function logprior_true(model::Model{typeof(demo_dot_assume_observe_submodel)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function loglikelihood_true(model::Model{typeof(demo_dot_assume_observe_submodel)}, m) - return logpdf(MvNormal(m, 0.25 * I), model.args.x) +function loglikelihood_true(model::Model{typeof(demo_dot_assume_observe_submodel)}, s, m) + return logpdf(MvNormal(m, Diagonal(s)), model.args.x) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_observe_submodel)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_dot_assume_observe_submodel)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end @model function demo_dot_assume_dot_observe_matrix( - x=fill(10.0, 2, 1), ::Type{TV}=Vector{Float64} + x=transpose([1.5 2.0;]), ::Type{TV}=Vector{Float64} ) where {TV} + s = TV(undef, length(x)) + s .~ InverseGamma(2, 3) m = TV(undef, length(x)) - m .~ Normal() + m .~ Normal.(0, sqrt.(s)) # Dotted observe for `Matrix`. - x .~ MvNormal(m, 0.25 * I) + x .~ MvNormal(m, Diagonal(s)) - return (; m=m, x=x, logp=getlogp(__varinfo__)) + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) +end +function logprior_true(model::Model{typeof(demo_dot_assume_dot_observe_matrix)}, s, m) + return loglikelihood(InverseGamma(2, 3), s) + sum(logpdf.(Normal.(0, sqrt.(s)), m)) end -function logprior_true(model::Model{typeof(demo_dot_assume_dot_observe_matrix)}, m) - return loglikelihood(Normal(), m) +function loglikelihood_true(model::Model{typeof(demo_dot_assume_dot_observe_matrix)}, s, m) + return sum(logpdf.(Normal.(m, sqrt.(s)), model.args.x)) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_dot_observe_matrix)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) end -function loglikelihood_true(model::Model{typeof(demo_dot_assume_dot_observe_matrix)}, m) - return loglikelihood(MvNormal(m, 0.25 * I), model.args.x) +function varnames(model::Model{typeof(demo_dot_assume_dot_observe_matrix)}) + return [@varname(s[1]), @varname(s[2]), @varname(m[1]), @varname(m[2])] end +@model function demo_dot_assume_matrix_dot_observe_matrix( + x=transpose([1.5 2.0;]), ::Type{TV}=Array{Float64} +) where {TV} + n = length(x) + d = length(x) ÷ 2 + s = TV(undef, d, 2) + s .~ product_distribution([InverseGamma(2, 3) for _ in 1:d]) + s_vec = vec(s) + m ~ MvNormal(zeros(n), Diagonal(s_vec)) + + # Dotted observe for `Matrix`. + x .~ MvNormal(m, Diagonal(s_vec)) + + return (; s=s, m=m, x=x, logp=getlogp(__varinfo__)) +end +function logprior_true( + model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}, s, m +) + n = length(model.args.x) + s_vec = vec(s) + return loglikelihood(InverseGamma(2, 3), s_vec) + + logpdf(MvNormal(zeros(n), Diagonal(s_vec)), m) +end +function loglikelihood_true( + model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}, s, m +) + return loglikelihood(MvNormal(m, Diagonal(vec(s))), model.args.x) +end +function logprior_true_with_logabsdet_jacobian( + model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}, s, m +) + return _demo_logprior_true_with_logabsdet_jacobian(model, s, m) +end +function varnames(model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}) + return [@varname(s[:, 1]), @varname(s[:, 2]), @varname(m)] +end + +const DemoModels = Union{ + Model{typeof(demo_dot_assume_dot_observe)}, + Model{typeof(demo_assume_index_observe)}, + Model{typeof(demo_assume_multivariate_observe)}, + Model{typeof(demo_dot_assume_observe_index)}, + Model{typeof(demo_assume_dot_observe)}, + Model{typeof(demo_assume_literal_dot_observe)}, + Model{typeof(demo_assume_observe_literal)}, + Model{typeof(demo_dot_assume_observe_index_literal)}, + Model{typeof(demo_assume_submodel_observe_index_literal)}, + Model{typeof(demo_dot_assume_observe_submodel)}, + Model{typeof(demo_dot_assume_dot_observe_matrix)}, + Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}, +} + +# We require demo models to have explict impleentations of `rand` since we want +# these to be considered as ground truth. +function Random.rand(rng::Random.AbstractRNG, ::Type{NamedTuple}, model::DemoModels) + return error("demo models requires explicit implementation of rand") +end + +const UnivariateAssumeDemoModels = Union{ + Model{typeof(demo_assume_dot_observe)},Model{typeof(demo_assume_literal_dot_observe)} +} +function posterior_mean(model::UnivariateAssumeDemoModels) + return (s=49 / 24, m=7 / 6) +end +function Random.rand( + rng::Random.AbstractRNG, ::Type{NamedTuple}, model::UnivariateAssumeDemoModels +) + s = rand(rng, InverseGamma(2, 3)) + m = rand(rng, Normal(0, sqrt(s))) + + return (s=s, m=m) +end + +const MultivariateAssumeDemoModels = Union{ + Model{typeof(demo_dot_assume_dot_observe)}, + Model{typeof(demo_assume_index_observe)}, + Model{typeof(demo_assume_multivariate_observe)}, + Model{typeof(demo_dot_assume_observe_index)}, + Model{typeof(demo_assume_observe_literal)}, + Model{typeof(demo_dot_assume_observe_index_literal)}, + Model{typeof(demo_assume_submodel_observe_index_literal)}, + Model{typeof(demo_dot_assume_observe_submodel)}, + Model{typeof(demo_dot_assume_dot_observe_matrix)}, + Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)}, +} +function posterior_mean(model::MultivariateAssumeDemoModels) + # Get some containers to fill. + vals = Random.rand(model) + + vals.s[1] = 19 / 8 + vals.m[1] = 3 / 4 + + vals.s[2] = 8 / 3 + vals.m[2] = 1 + + return vals +end +function Random.rand( + rng::Random.AbstractRNG, ::Type{NamedTuple}, model::MultivariateAssumeDemoModels +) + # Get template values from `model`. + retval = model(rng) + vals = (s=retval.s, m=retval.m) + # Fill containers with realizations from prior. + for i in LinearIndices(vals.s) + vals.s[i] = rand(rng, InverseGamma(2, 3)) + vals.m[i] = rand(rng, Normal(0, sqrt(vals.s[i]))) + end + + return vals +end + +""" +A collection of models corresponding to the posterior distribution defined by +the generative process + + s ~ InverseGamma(2, 3) + m ~ Normal(0, √s) + 1.5 ~ Normal(m, √s) + 2.0 ~ Normal(m, √s) + +or by + + s[1] ~ InverseGamma(2, 3) + s[2] ~ InverseGamma(2, 3) + m[1] ~ Normal(0, √s) + m[2] ~ Normal(0, √s) + 1.5 ~ Normal(m[1], √s[1]) + 2.0 ~ Normal(m[2], √s[2]) + +These are examples of a Normal-InverseGamma conjugate prior with Normal likelihood, +for which the posterior is known in closed form. + +In particular, for the univariate model (the former one): + + mean(s) == 49 / 24 + mean(m) == 7 / 6 + +And for the multivariate one (the latter one): + + mean(s[1]) == 19 / 8 + mean(m[1]) == 3 / 4 + mean(s[2]) == 8 / 3 + mean(m[2]) == 1 + +""" const DEMO_MODELS = ( demo_dot_assume_dot_observe(), demo_assume_index_observe(), @@ -257,64 +641,79 @@ const DEMO_MODELS = ( demo_assume_submodel_observe_index_literal(), demo_dot_assume_observe_submodel(), demo_dot_assume_dot_observe_matrix(), + demo_dot_assume_matrix_dot_observe_matrix(), ) -# TODO: Is this really the best/most convenient "default" test method? """ - test_sampler_demo_models(meanfunction, sampler, args...; kwargs...) + marginal_mean_of_samples(chain, varname) -Test that `sampler` produces the correct marginal posterior means on all models in `demo_models`. +Return the mean of variable represented by `varname` in `chain`. +""" +marginal_mean_of_samples(chain, varname) = mean(Array(chain[Symbol(varname)])) + +""" + test_sampler(models, sampler, args...; kwargs...) + +Test that `sampler` produces correct marginal posterior means on each model in `models`. -In short, this method iterators through `demo_models`, calls `AbstractMCMC.sample` on the -`model` and `sampler` to produce a `chain`, and then checks `meanfunction(chain)` against `target` -provided in `kwargs...`. +In short, this method iterates through `models`, calls `AbstractMCMC.sample` on the +`model` and `sampler` to produce a `chain`, and then checks `marginal_mean_of_samples(chain, vn)` +for every (leaf) varname `vn` against the corresponding value returned by +[`posterior_mean`](@ref) for each model. + +To change how comparison is done for a particular `chain` type, one can overload +[`marginal_mean_of_samples`](@ref) for the corresponding type. # Arguments -- `meanfunction`: A callable which computes the mean of the marginal means from the - chain resulting from the `sample` call. +- `models`: A collection of instaces of [`DynamicPPL.Model`](@ref) to test on. - `sampler`: The `AbstractMCMC.AbstractSampler` to test. - `args...`: Arguments forwarded to `sample`. # Keyword arguments -- `target`: Value to compare result of `meanfunction(chain)` to. - `atol=1e-1`: Absolute tolerance used in `@test`. - `rtol=1e-3`: Relative tolerance used in `@test`. - `kwargs...`: Keyword arguments forwarded to `sample`. """ -function test_sampler_demo_models( - meanfunction, - sampler::AbstractMCMC.AbstractSampler, - args...; - target=8.0, - atol=1e-1, - rtol=1e-3, - kwargs..., +function test_sampler( + models, sampler::AbstractMCMC.AbstractSampler, args...; atol=1e-1, rtol=1e-3, kwargs... ) - @testset "$(nameof(typeof(sampler))) on $(nameof(m))" for model in DEMO_MODELS + @testset "$(typeof(sampler)) on $(nameof(model))" for model in models chain = AbstractMCMC.sample(model, sampler, args...; kwargs...) - μ = meanfunction(chain) - @test μ ≈ target atol = atol rtol = rtol + target_values = posterior_mean(model) + for vn in varnames(model) + # We want to compare elementwise which can be achieved by + # extracting the leaves of the `VarName` and the corresponding value. + for vn_leaf in varname_leaves(vn, get(target_values, vn)) + target_value = get(target_values, vn_leaf) + chain_mean_value = marginal_mean_of_samples(chain, vn_leaf) + @test chain_mean_value ≈ target_value atol = atol rtol = rtol + end + end end end """ - test_sampler_continuous([meanfunction, ]sampler, args...; kwargs...) + test_sampler_on_demo_models(meanfunction, sampler, args...; kwargs...) -Test that `sampler` produces the correct marginal posterior means on all models in `demo_models`. +Test `sampler` on every model in [`DEMO_MODELS`](@ref). -As of right now, this is just an alias for [`test_sampler_demo_models`](@ref). +This is just a proxy for `test_sampler(meanfunction, DEMO_MODELS, sampler, args...; kwargs...)`. """ -function test_sampler_continuous( - meanfunction, sampler::AbstractMCMC.AbstractSampler, args...; kwargs... +function test_sampler_on_demo_models( + sampler::AbstractMCMC.AbstractSampler, args...; kwargs... ) - return test_sampler_demo_models(meanfunction, sampler, args...; kwargs...) + return test_sampler(DEMO_MODELS, sampler, args...; kwargs...) end +""" + test_sampler_continuous(sampler, args...; kwargs...) + +Test that `sampler` produces the correct marginal posterior means on all models in `demo_models`. + +As of right now, this is just an alias for [`test_sampler_on_demo_models`](@ref). +""" function test_sampler_continuous(sampler::AbstractMCMC.AbstractSampler, args...; kwargs...) - # Default for `MCMCChains.Chains`. - return test_sampler_continuous(sampler, args...; kwargs...) do chain - mean(Array(chain)) - end + return test_sampler_on_demo_models(sampler, args...; kwargs...) end end diff --git a/src/threadsafe.jl b/src/threadsafe.jl index 6f020a352..7c7dd13ac 100644 --- a/src/threadsafe.jl +++ b/src/threadsafe.jl @@ -62,8 +62,24 @@ islinked(vi::ThreadSafeVarInfo, spl::AbstractSampler) = islinked(vi.varinfo, spl getindex(vi::ThreadSafeVarInfo, spl::AbstractSampler) = getindex(vi.varinfo, spl) getindex(vi::ThreadSafeVarInfo, spl::SampleFromPrior) = getindex(vi.varinfo, spl) getindex(vi::ThreadSafeVarInfo, spl::SampleFromUniform) = getindex(vi.varinfo, spl) + getindex(vi::ThreadSafeVarInfo, vn::VarName) = getindex(vi.varinfo, vn) +function getindex(vi::ThreadSafeVarInfo, vn::VarName, dist::Distribution) + return getindex(vi.varinfo, vn, dist) +end getindex(vi::ThreadSafeVarInfo, vns::Vector{<:VarName}) = getindex(vi.varinfo, vns) +function getindex(vi::ThreadSafeVarInfo, vns::Vector{<:VarName}, dist::Distribution) + return getindex(vi.varinfo, vns, dist) +end + +getindex_raw(vi::ThreadSafeVarInfo, vn::VarName) = getindex_raw(vi.varinfo, vn) +function getindex_raw(vi::ThreadSafeVarInfo, vn::VarName, dist::Distribution) + return getindex_raw(vi.varinfo, vn, dist) +end +getindex_raw(vi::ThreadSafeVarInfo, vns::Vector{<:VarName}) = getindex_raw(vi.varinfo, vns) +function getindex_raw(vi::ThreadSafeVarInfo, vns::Vector{<:VarName}, dist::Distribution) + return getindex_raw(vi.varinfo, vns, dist) +end function BangBang.setindex!!(vi::ThreadSafeVarInfo, val, spl::AbstractSampler) return Setfield.@set vi.varinfo = BangBang.setindex!!(vi.varinfo, val, spl) diff --git a/src/utils.jl b/src/utils.jl index 821eba38e..ac9222818 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -183,6 +183,7 @@ vectorize(d::MatrixDistribution, r::AbstractMatrix{<:Real}) = copy(vec(r)) # otherwise we will have error for MatrixDistribution. # Note this is not the case for MultivariateDistribution so I guess this might be lack of # support for some types related to matrices (like PDMat). +reconstruct(d::UnivariateDistribution, val::Real) = val reconstruct(d::Distribution, val::AbstractVector) = reconstruct(size(d), val) reconstruct(::Tuple{}, val::AbstractVector) = val[1] reconstruct(s::NTuple{1}, val::AbstractVector) = copy(val) @@ -417,3 +418,11 @@ function splitlens(condition, lens) return current_parent, current_child, condition(current_parent) end + +# HACK(torfjelde): Avoids type-instability in `dot_assume` for `SimpleVarInfo`. +function BangBang.possible( + ::typeof(BangBang._setindex!), ::C, ::T, ::Colon, ::Integer +) where {C<:AbstractMatrix,T<:AbstractVector} + return BangBang.implements(setindex!, C) && + promote_type(eltype(C), eltype(T)) <: eltype(C) +end diff --git a/src/varinfo.jl b/src/varinfo.jl index 9ce0414d6..22728ba9a 100644 --- a/src/varinfo.jl +++ b/src/varinfo.jl @@ -358,12 +358,31 @@ Return the set of sampler selectors associated with `vn` in `vi`. getgid(vi::VarInfo, vn::VarName) = getmetadata(vi, vn).gids[getidx(vi, vn)] """ - settrans!(vi::VarInfo, trans::Bool, vn::VarName) + settrans!!(vi::VarInfo, trans::Bool, vn::VarName) Set the `trans` flag value of `vn` in `vi`. """ -function settrans!(vi::AbstractVarInfo, trans::Bool, vn::VarName) - return trans ? set_flag!(vi, vn, "trans") : unset_flag!(vi, vn, "trans") +function settrans!!(vi::AbstractVarInfo, trans::Bool, vn::VarName) + if trans + set_flag!(vi, vn, "trans") + else + unset_flag!(vi, vn, "trans") + end + + return vi +end + +""" + settrans!!(vi::AbstractVarInfo, trans) + +Return new instance of `vi` but with `istrans(vi, trans)` now evaluating to `true`. +""" +function settrans!!(vi::VarInfo, trans::Bool) + for vn in keys(vi) + settrans!!(vi, trans, vn) + end + + return vi end """ @@ -638,6 +657,14 @@ function setgid!(vi::VarInfo, gid::Selector, vn::VarName) return push!(getmetadata(vi, vn).gids[getidx(vi, vn)], gid) end +""" + istrans(vi::AbstractVarInfo) + +Return `true` if `vi` is working in unconstrained space, and `false` +if `vi` is assuming realizations to be in support of the corresponding distributions. +""" +istrans(vi::AbstractVarInfo) = false # `VarInfo` works in constrained space by default. + """ istrans(vi::VarInfo, vn::VarName) @@ -645,6 +672,9 @@ Return true if `vn`'s values in `vi` are transformed to Euclidean space, and fal they are in the support of `vn`'s distribution. """ istrans(vi::AbstractVarInfo, vn::VarName) = is_flagged(vi, vn, "trans") +function istrans(vi::AbstractVarInfo, vns::AbstractVector{<:VarName}) + return all(Base.Fix1(istrans, vi), vns) +end """ getlogp(vi::VarInfo) @@ -749,7 +779,7 @@ function link!(vi::UntypedVarInfo, spl::Sampler) vectorize(dist, Bijectors.link(dist, reconstruct(dist, getval(vi, vn)))), vn, ) - settrans!(vi, true, vn) + settrans!!(vi, true, vn) end else @warn("[DynamicPPL] attempt to link a linked vi") @@ -785,7 +815,7 @@ end ), vn, ) - settrans!(vi, true, vn) + settrans!!(vi, true, vn) end else @warn("[DynamicPPL] attempt to link a linked vi") @@ -816,7 +846,7 @@ function invlink!(vi::UntypedVarInfo, spl::AbstractSampler) vectorize(dist, Bijectors.invlink(dist, reconstruct(dist, getval(vi, vn)))), vn, ) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end else @warn("[DynamicPPL] attempt to invlink an invlinked vi") @@ -854,7 +884,7 @@ end ), vn, ) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end else @warn("[DynamicPPL] attempt to invlink an invlinked vi") @@ -866,6 +896,9 @@ end return expr end +maybe_link(vi, vn, dist, val) = istrans(vi, vn) ? Bijectors.link(dist, val) : val +maybe_invlink(vi, vn, dist, val) = istrans(vi, vn) ? Bijectors.invlink(dist, val) : val + """ islinked(vi::VarInfo, spl::Union{Sampler, SampleFromPrior}) @@ -904,23 +937,37 @@ distribution(s). If the value(s) is (are) transformed to the Euclidean space, it is (they are) transformed back. """ -function getindex(vi::AbstractVarInfo, vn::VarName) +getindex(vi::AbstractVarInfo, vn::VarName) = getindex(vi, vn, getdist(vi, vn)) +function getindex(vi::AbstractVarInfo, vn::VarName, dist::Distribution) @assert haskey(vi, vn) "[DynamicPPL] attempted to replay unexisting variables in VarInfo" - dist = getdist(vi, vn) - return if istrans(vi, vn) - Bijectors.invlink(dist, reconstruct(dist, getval(vi, vn))) - else - reconstruct(dist, getval(vi, vn)) - end + val = getindex_raw(vi, vn, dist) + return maybe_invlink(vi, vn, dist, val) end function getindex(vi::AbstractVarInfo, vns::Vector{<:VarName}) + # FIXME(torfjelde): Using `getdist(vi, first(vns))` won't be correct in cases + # such as `x .~ [Normal(), Exponential()]`. + # BUT we also can't fix this here because this will lead to "incorrect" + # behavior if `vns` arose from something like `x .~ MvNormal(zeros(2), I)`, + # where by "incorrect" we mean there exists pieces of code expecting this behavior. + return getindex(vi, vns, getdist(vi, first(vns))) +end +function getindex(vi::AbstractVarInfo, vns::Vector{<:VarName}, dist::Distribution) @assert haskey(vi, vns[1]) "[DynamicPPL] attempted to replay unexisting variables in VarInfo" - dist = getdist(vi, vns[1]) - return if istrans(vi, vns[1]) - Bijectors.invlink(dist, reconstruct(dist, getval(vi, vns), length(vns))) - else - reconstruct(dist, getval(vi, vns), length(vns)) + vals_linked = mapreduce(vcat, vns) do vn + getindex(vi, vn, dist) end + return reconstruct(dist, vals_linked, length(vns)) +end + +getindex_raw(vi::AbstractVarInfo, vn::VarName) = getindex_raw(vi, vn, getdist(vi, vn)) +function getindex_raw(vi::AbstractVarInfo, vn::VarName, dist::Distribution) + return reconstruct(dist, getval(vi, vn)) +end +function getindex_raw(vi::AbstractVarInfo, vns::Vector{<:VarName}) + return getindex_raw(vi, vns, getdist(vi, first(vns))) +end +function getindex_raw(vi::AbstractVarInfo, vns::Vector{<:VarName}, dist::Distribution) + return reconstruct(dist, getval(vi, vns), length(vns)) end """ @@ -1411,7 +1458,7 @@ function _setval_kernel!(vi::VarInfo, vn::VarName, values, keys) if !isempty(indices) val = reduce(vcat, values[indices]) setval!(vi, val, vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) end return indices @@ -1492,7 +1539,7 @@ function _setval_and_resample_kernel!(vi::VarInfo, vn::VarName, values, keys) if !isempty(indices) val = reduce(vcat, values[indices]) setval!(vi, val, vn) - settrans!(vi, false, vn) + settrans!!(vi, false, vn) else # Ensures that we'll resample the variable corresponding to `vn` if we run # the model on `vi` again. diff --git a/test/Project.toml b/test/Project.toml index 6a3726474..edba29e5a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -23,7 +23,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" AbstractMCMC = "2.1, 3.0, 4" AbstractPPL = "0.5.1" Bijectors = "0.9.5, 0.10" -Distributions = "0.25" +Distributions = "<0.25.65" DistributionsAD = "0.6.3" Documenter = "0.26.1, 0.27" ForwardDiff = "0.10.12" diff --git a/test/contexts.jl b/test/contexts.jl index 65629afec..edcf5d0f3 100644 --- a/test/contexts.jl +++ b/test/contexts.jl @@ -57,26 +57,6 @@ function remove_prefix(vn::VarName) ) end -""" - varnames(vn::VarName, val) - -Return iterator over all varnames that are represented by `vn` on `val`, -e.g. `varnames(@varname(x), rand(2))` results in an iterator over `[@varname(x[1]), @varname(x[2])]`. -""" -varnames(vn::VarName, val::Real) = [vn] -function varnames(vn::VarName, val::AbstractArray{<:Union{Real,Missing}}) - return ( - VarName(vn, getlens(vn) ∘ Setfield.IndexLens(Tuple(I))) for - I in CartesianIndices(val) - ) -end -function varnames(vn::VarName, val::AbstractArray) - return Iterators.flatten( - varnames(VarName(vn, getlens(vn) ∘ Setfield.IndexLens(Tuple(I))), val[I]) for - I in CartesianIndices(val) - ) -end - @testset "contexts.jl" begin child_contexts = [DefaultContext(), PriorContext(), LikelihoodContext()] @@ -185,7 +165,8 @@ end vn_without_prefix = remove_prefix(vn) # Let's check elementwise. - for vn_child in varnames(vn_without_prefix, val) + for vn_child in + DynamicPPL.TestUtils.varname_leaves(vn_without_prefix, val) if get(val, getlens(vn_child)) === missing @test contextual_isassumption(context, vn_child) else @@ -217,7 +198,8 @@ end # `ConditionContext` with the conditioned variable. vn_without_prefix = remove_prefix(vn) - for vn_child in varnames(vn_without_prefix, val) + for vn_child in + DynamicPPL.TestUtils.varname_leaves(vn_without_prefix, val) # `vn_child` should be in `context`. @test hasvalue_nested(context, vn_child) # Value should be the same as extracted above. diff --git a/test/distribution_wrappers.jl b/test/distribution_wrappers.jl index 350ce6014..8bb692783 100644 --- a/test/distribution_wrappers.jl +++ b/test/distribution_wrappers.jl @@ -9,5 +9,5 @@ @test minimum(nd) == -Inf @test maximum(nd) == Inf @test logpdf(nd, 15.0) == 0 - @test Bijectors.logpdf_with_trans(nd, 0) == 0 + @test Bijectors.logpdf_with_trans(nd, 30, true) == 0 end diff --git a/test/loglikelihoods.jl b/test/loglikelihoods.jl index 4d5003f03..b390997af 100644 --- a/test/loglikelihoods.jl +++ b/test/loglikelihoods.jl @@ -1,15 +1,14 @@ @testset "loglikelihoods.jl" begin - for m in DynamicPPL.TestUtils.DEMO_MODELS - vi = VarInfo(m) + @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS + example_values = rand(NamedTuple, m) - vns = vi.metadata.m.vns - if length(vns) == 1 && length(vi[vns[1]]) == 1 - # Only have one latent variable. - DynamicPPL.setval!(vi, [1.0], ["m"]) - else - DynamicPPL.setval!(vi, [1.0, 1.0], ["m[1]", "m[2]"]) + # Instantiate a `VarInfo` with the example values. + vi = VarInfo(m) + for vn in DynamicPPL.TestUtils.varnames(m) + vi = DynamicPPL.setindex!!(vi, get(example_values, vn), vn) end + # Compute the pointwise loglikelihoods. lls = pointwise_loglikelihoods(m, vi) if isempty(lls) @@ -17,14 +16,9 @@ continue end - loglikelihood = if length(keys(lls)) == 1 && length(m.args.x) == 1 - # Only have one observation, so we need to double it - # for comparison with other models. - 2 * sum(lls[first(keys(lls))]) - else - sum(sum, values(lls)) - end + loglikelihood = sum(sum, values(lls)) + loglikelihood_true = DynamicPPL.TestUtils.loglikelihood_true(m, example_values...) - @test loglikelihood ≈ -324.45158270528947 + @test loglikelihood ≈ loglikelihood_true end end diff --git a/test/model.jl b/test/model.jl index 78ef27a6e..8eabeeba9 100644 --- a/test/model.jl +++ b/test/model.jl @@ -112,6 +112,22 @@ end @test !any(map(x -> x isa DynamicPPL.AbstractVarInfo, call_retval)) end + @testset "Dynamic constraints" begin + model = DynamicPPL.TestUtils.demo_dynamic_constraint() + vi = VarInfo(model) + spl = SampleFromPrior() + link!(vi, spl) + + for i in 1:10 + # Sample with large variations. + r_raw = randn(length(vi[spl])) * 10 + vi[spl] = r_raw + @test vi[@varname(m)] == r_raw[1] + @test vi[@varname(x)] != r_raw[2] + model(vi) + end + end + @testset "rand" begin model = gdemo_default diff --git a/test/simple_varinfo.jl b/test/simple_varinfo.jl index 7e9346450..6a8c545ca 100644 --- a/test/simple_varinfo.jl +++ b/test/simple_varinfo.jl @@ -62,68 +62,112 @@ DynamicPPL.TestUtils.DEMO_MODELS # We might need to pre-allocate for the variable `m`, so we need # to see whether this is the case. - m = model().m - svi_nt = if m isa AbstractArray - SimpleVarInfo((m=similar(m),)) - else - SimpleVarInfo() - end + svi_nt = SimpleVarInfo(rand(NamedTuple, model)) svi_dict = SimpleVarInfo(VarInfo(model), Dict) - @testset "$(nameof(typeof(svi.values)))" for svi in (svi_nt, svi_dict) + @testset "$(nameof(typeof(DynamicPPL.values_as(svi))))" for svi in ( + svi_nt, + svi_dict, + DynamicPPL.settrans!!(svi_nt, true), + DynamicPPL.settrans!!(svi_dict, true), + ) # Random seed is set in each `@testset`, so we need to sample # a new realization for `m` here. - m = model().m + retval = model() ### Sampling ### # Sample a new varinfo! _, svi_new = DynamicPPL.evaluate!!(model, svi, SamplingContext()) - # If the `m[1]` varname doesn't exist, this is a univariate model. - # TODO: Find a better way of dealing with this that is not dependent - # on knowledge of internals of `model`. - isunivariate = !haskey(svi_new, @varname(m[1])) - # Realization for `m` should be different wp. 1. - if isunivariate - @test svi_new[@varname(m)] != m - else - @test svi_new[@varname(m[1])] != m[1] - @test svi_new[@varname(m[2])] != m[2] + for vn in DynamicPPL.TestUtils.varnames(model) + @test svi_new[vn] != get(retval, vn) end + # Logjoint should be non-zero wp. 1. @test getlogp(svi_new) != 0 ### Evaluation ### - # Sample some random testing values. - m_eval = if m isa AbstractArray - randn!(similar(m)) + values_eval_constrained = rand(NamedTuple, model) + if DynamicPPL.istrans(svi) + _values_prior, logpri_true = DynamicPPL.TestUtils.logprior_true_with_logabsdet_jacobian( + model, values_eval_constrained... + ) + values_eval, logπ_true = DynamicPPL.TestUtils.logjoint_true_with_logabsdet_jacobian( + model, values_eval_constrained... + ) + # Make sure that these two computation paths provide the same + # transformed values. + @test values_eval == _values_prior else - randn(eltype(m)) + logpri_true = DynamicPPL.TestUtils.logprior_true( + model, values_eval_constrained... + ) + logπ_true = DynamicPPL.TestUtils.logjoint_true( + model, values_eval_constrained... + ) + values_eval = values_eval_constrained end + # No logabsdet-jacobian correction needed for the likelihood. + loglik_true = DynamicPPL.TestUtils.loglikelihood_true( + model, values_eval_constrained... + ) + # Update the realizations in `svi_new`. - svi_eval = if isunivariate - DynamicPPL.setindex!!(svi_new, m_eval, @varname(m)) - else - DynamicPPL.setindex!!(svi_new, m_eval, [@varname(m[1]), @varname(m[2])]) + svi_eval = svi_new + for vn in DynamicPPL.TestUtils.varnames(model) + svi_eval = DynamicPPL.setindex!!(svi_eval, get(values_eval, vn), vn) end + # Reset the logp field. svi_eval = DynamicPPL.resetlogp!!(svi_eval) # Compute `logjoint` using the varinfo. logπ = logjoint(model, svi_eval) - # Extract the parameters from `svi_eval`. - m_vi = if isunivariate - svi_eval[@varname(m)] - else - svi_eval[[@varname(m[1]), @varname(m[2])]] + logpri = logprior(model, svi_eval) + loglik = loglikelihood(model, svi_eval) + + # Values should not have changed. + for vn in DynamicPPL.TestUtils.varnames(model) + @test svi_eval[vn] == get(values_eval, vn) end - # These should not have changed. - @test m_vi == m_eval - # Compute the true `logjoint` and compare. - logπ_true = DynamicPPL.TestUtils.logjoint_true(model, m_vi) + + # Compare log-probability computations. + @test logpri ≈ logpri_true + @test loglik ≈ loglik_true @test logπ ≈ logπ_true end end + + @testset "Dynamic constraints" begin + model = DynamicPPL.TestUtils.demo_dynamic_constraint() + + # Initialize. + svi = DynamicPPL.settrans!!(SimpleVarInfo(), true) + svi = last(DynamicPPL.evaluate!!(model, svi, SamplingContext())) + + # Sample with large variations in unconstrained space. + for i in 1:10 + for vn in keys(svi) + svi = DynamicPPL.setindex!!(svi, 10 * randn(), vn) + end + retval, svi = DynamicPPL.evaluate!!(model, svi, DefaultContext()) + @test retval.m == svi[@varname(m)] # `m` is unconstrained + @test retval.x ≠ svi[@varname(x)] # `x` is constrained depending on `m` + + retval_unconstrained, lp_true = DynamicPPL.TestUtils.logjoint_true_with_logabsdet_jacobian( + model, retval.m, retval.x + ) + + # Realizations from model should all be equal to the unconstrained realization. + for vn in DynamicPPL.TestUtils.varnames(model) + @test get(retval_unconstrained, vn) ≈ svi[vn] rtol = 1e-6 + end + + # `getlogp` should be equal to the logjoint with log-absdet-jac correction. + lp = getlogp(svi) + @test lp ≈ lp_true + end + end end diff --git a/test/turing/Project.toml b/test/turing/Project.toml index 41a873172..e8713dae2 100644 --- a/test/turing/Project.toml +++ b/test/turing/Project.toml @@ -1,10 +1,12 @@ [deps] +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" [compat] +Distributions = "<0.25.65" DynamicPPL = "0.19" Turing = "0.21" julia = "1.3" diff --git a/test/varinfo.jl b/test/varinfo.jl index 2f7816024..32c90bf47 100644 --- a/test/varinfo.jl +++ b/test/varinfo.jl @@ -271,4 +271,47 @@ DynamicPPL.setval_and_resample!(vi, vi.metadata.x.vals, ks) @test vals_prev == vi.metadata.x.vals end + + @testset "istrans" begin + @model demo_constrained() = x ~ truncated(Normal(), 0, Inf) + model = demo_constrained() + vn = @varname(x) + dist = truncated(Normal(), 0, Inf) + + ### `VarInfo` + # Need to run once since we can't specify that we want to _sample_ + # in the unconstrained space for `VarInfo` without having `vn` + # present in the `varinfo`. + ## `UntypedVarInfo` + vi = VarInfo() + vi = last(DynamicPPL.evaluate!!(model, vi, SamplingContext())) + vi = DynamicPPL.settrans!!(vi, true, vn) + # Sample in unconstrained space. + vi = last(DynamicPPL.evaluate!!(model, vi, SamplingContext())) + x = Bijectors.invlink(dist, DynamicPPL.getindex_raw(vi, vn)) + @test getlogp(vi) ≈ Bijectors.logpdf_with_trans(dist, x, true) + + ## `TypedVarInfo` + vi = VarInfo(model) + vi = DynamicPPL.settrans!!(vi, true, vn) + # Sample in unconstrained space. + vi = last(DynamicPPL.evaluate!!(model, vi, SamplingContext())) + x = Bijectors.invlink(dist, DynamicPPL.getindex_raw(vi, vn)) + @test getlogp(vi) ≈ Bijectors.logpdf_with_trans(dist, x, true) + + ### `SimpleVarInfo` + ## `SimpleVarInfo{<:NamedTuple}` + vi = DynamicPPL.settrans!!(SimpleVarInfo(), true) + # Sample in unconstrained space. + vi = last(DynamicPPL.evaluate!!(model, vi, SamplingContext())) + x = Bijectors.invlink(dist, DynamicPPL.getindex_raw(vi, vn)) + @test getlogp(vi) ≈ Bijectors.logpdf_with_trans(dist, x, true) + + ## `SimpleVarInfo{<:Dict}` + vi = DynamicPPL.settrans!!(SimpleVarInfo(Dict()), true) + # Sample in unconstrained space. + vi = last(DynamicPPL.evaluate!!(model, vi, SamplingContext())) + x = Bijectors.invlink(dist, DynamicPPL.getindex_raw(vi, vn)) + @test getlogp(vi) ≈ Bijectors.logpdf_with_trans(dist, x, true) + end end