Pointpriors

docs/src/api.md

@@ -130,6 +130,17 @@ For a chain of samples, one can compute the pointwise log-likelihoods of each ob
 pointwise_loglikelihoods
 ```
 
+Similarly, one can compute the pointwise log-priors of each sampled random variable
+with [`varwise_logpriors`](@ref).
+Differently from `pointwise_loglikelihoods` it reports only a
+single value for `.~` assignements. 
+If one needs to access the parts for single indices, one can
+reformulate the model to use an explicit loop instead.
+
+```@docs
+varwise_logpriors
+```
+
 For converting a chain into a format that can more easily be fed into a `Model` again, for example using `condition`, you can use [`value_iterator_from_chain`](@ref).
 
 ```@docs

src/DynamicPPL.jl

@@ -116,6 +116,7 @@ export AbstractVarInfo,
     logprior,
     logjoint,
     pointwise_loglikelihoods,
+    pointwise_logdensities,
     condition,
     decondition,
     fix,
@@ -181,14 +182,15 @@ include("varinfo.jl")
 include("simple_varinfo.jl")
 include("context_implementations.jl")
 include("compiler.jl")
-include("loglikelihoods.jl")
+include("pointwise_logdensities.jl")
 include("submodel_macro.jl")
 include("test_utils.jl")
 include("transforming.jl")
 include("logdensityfunction.jl")
 include("model_utils.jl")
 include("extract_priors.jl")
 include("values_as_in_model.jl")
+include("deprecated.jl")
 
 include("debug_utils.jl")
 using .DebugUtils

src/deprecated.jl

@@ -0,0 +1,9 @@
+# https://invenia.github.io/blog/2022/06/17/deprecating-in-julia/
+
+Base.@deprecate pointwise_loglikelihoods(model::Model, chain, keytype) pointwise_logdensities(
+    model::Model, LikelihoodContext(), chain, keytype)
+
+Base.@deprecate pointwise_loglikelihoods(
+    model::Model, varinfo::AbstractVarInfo) pointwise_logdensities(
+    model::Model, varinfo, LikelihoodContext())
+

src/loglikelihoods.jl → src/pointwise_logdensities.jl

@@ -1,83 +1,83 @@
 # Context version
-struct PointwiseLikelihoodContext{A,Ctx} <: AbstractContext
-    loglikelihoods::A
+struct PointwiseLogdensityContext{A,Ctx} <: AbstractContext
+    logdensities::A
     context::Ctx
 end
 
-function PointwiseLikelihoodContext(
+function PointwiseLogdensityContext(
     likelihoods=OrderedDict{VarName,Vector{Float64}}(),
-    context::AbstractContext=LikelihoodContext(),
+    context::AbstractContext=DefaultContext(),
 )
-    return PointwiseLikelihoodContext{typeof(likelihoods),typeof(context)}(
+    return PointwiseLogdensityContext{typeof(likelihoods),typeof(context)}(
         likelihoods, context
     )
 end
 
-NodeTrait(::PointwiseLikelihoodContext) = IsParent()
-childcontext(context::PointwiseLikelihoodContext) = context.context
-function setchildcontext(context::PointwiseLikelihoodContext, child)
-    return PointwiseLikelihoodContext(context.loglikelihoods, child)
+NodeTrait(::PointwiseLogdensityContext) = IsParent()
+childcontext(context::PointwiseLogdensityContext) = context.context
+function setchildcontext(context::PointwiseLogdensityContext, child)
+    return PointwiseLogdensityContext(context.logdensities, child)
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{VarName,Vector{Float64}}},
+    context::PointwiseLogdensityContext{<:AbstractDict{VarName,Vector{Float64}}},
     vn::VarName,
     logp::Real,
 )
-    lookup = context.loglikelihoods
+    lookup = context.logdensities
     ℓ = get!(lookup, vn, Float64[])
     return push!(ℓ, logp)
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{VarName,Float64}},
+    context::PointwiseLogdensityContext{<:AbstractDict{VarName,Float64}},
     vn::VarName,
     logp::Real,
 )
-    return context.loglikelihoods[vn] = logp
+    return context.logdensities[vn] = logp
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{String,Vector{Float64}}},
+    context::PointwiseLogdensityContext{<:AbstractDict{String,Vector{Float64}}},
     vn::VarName,
     logp::Real,
 )
-    lookup = context.loglikelihoods
+    lookup = context.logdensities
     ℓ = get!(lookup, string(vn), Float64[])
     return push!(ℓ, logp)
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{String,Float64}},
+    context::PointwiseLogdensityContext{<:AbstractDict{String,Float64}},
     vn::VarName,
     logp::Real,
 )
-    return context.loglikelihoods[string(vn)] = logp
+    return context.logdensities[string(vn)] = logp
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{String,Vector{Float64}}},
+    context::PointwiseLogdensityContext{<:AbstractDict{String,Vector{Float64}}},
     vn::String,
     logp::Real,
 )
-    lookup = context.loglikelihoods
+    lookup = context.logdensities
     ℓ = get!(lookup, vn, Float64[])
     return push!(ℓ, logp)
 end
 
 function Base.push!(
-    context::PointwiseLikelihoodContext{<:AbstractDict{String,Float64}},
+    context::PointwiseLogdensityContext{<:AbstractDict{String,Float64}},
     vn::String,
     logp::Real,
 )
-    return context.loglikelihoods[vn] = logp
+    return context.logdensities[vn] = logp
 end
 
-function tilde_observe!!(context::PointwiseLikelihoodContext, right, left, vi)
+function tilde_observe!!(context::PointwiseLogdensityContext, right, left, vi)
     # Defer literal `observe` to child-context.
     return tilde_observe!!(context.context, right, left, vi)
 end
-function tilde_observe!!(context::PointwiseLikelihoodContext, right, left, vn, vi)
+function tilde_observe!!(context::PointwiseLogdensityContext, right, left, vn, vi)
     # Need the `logp` value, so we cannot defer `acclogp!` to child-context, i.e.
     # we have to intercept the call to `tilde_observe!`.
     logp, vi = tilde_observe(context.context, right, left, vi)
@@ -88,17 +88,17 @@ function tilde_observe!!(context::PointwiseLikelihoodContext, right, left, vn, v
     return left, acclogp!!(vi, logp)
 end
 
-function dot_tilde_observe!!(context::PointwiseLikelihoodContext, right, left, vi)
+function dot_tilde_observe!!(context::PointwiseLogdensityContext, right, left, vi)
     # Defer literal `observe` to child-context.
     return dot_tilde_observe!!(context.context, right, left, vi)
 end
-function dot_tilde_observe!!(context::PointwiseLikelihoodContext, right, left, vn, vi)
+function dot_tilde_observe!!(context::PointwiseLogdensityContext, right, left, vn, vi)
     # Need the `logp` value, so we cannot defer `acclogp!` to child-context, i.e.
     # we have to intercept the call to `dot_tilde_observe!`.
 
     # We want to treat `.~` as a collection of independent observations,
     # hence we need the `logp` for each of them. Broadcasting the univariate
-    # `tilde_obseve` does exactly this.
+    # `tilde_observe` does exactly this.
     logps = _pointwise_tilde_observe(context.context, right, left, vi)
 
     # Need to unwrap the `vn`, i.e. get one `VarName` for each entry in `left`.
@@ -129,8 +129,47 @@ function _pointwise_tilde_observe(
     end
 end
 
+function tilde_assume(context::PointwiseLogdensityContext, right::Distribution, vn, vi)
+    #@info "PointwiseLogdensityContext tilde_assume called for $vn"
+    value, logp, vi = tilde_assume(context.context, right, vn, vi)
+    push!(context, vn, logp)
+    return value, logp, vi
+end
+
+function dot_tilde_assume(context::PointwiseLogdensityContext, right, left, vns, vi)
+    #@info "PointwiseLogdensityContext dot_tilde_assume called for $vns"
+    value, logp, vi_new = dot_tilde_assume(context.context, right, left, vns, vi)
+    # dispatch recording of log-densities based on type of right
+    logps = record_dot_tilde_assume(context, right, left, vns, vi, logp)
+    sum(logps) ≈ logp || error("Expected sum of individual logp equal origina, but differed sum($(join(logps, ","))) != $logp_orig")
+    return value, logp, vi
+end
+
+function record_dot_tilde_assume(context::PointwiseLogdensityContext, right::Distribution, left, vns, vi, logp)
+    # forward to tilde_assume for each variable
+    map(vns) do vn
+        value_i, logp_i, vi_i = tilde_assume(context, right, vn, vi)
+        logp_i
+    end
+end
+
+function record_dot_tilde_assume(context::PointwiseLogdensityContext, rights::AbstractVector{<:Distribution}, left, vns, vi, logp)
+    # forward to tilde_assume for each variable and distribution
+    map(vns, rights) do vn, right
+        # use current context to record vn
+        value_i, logp_i, vi_i = tilde_assume(context, right, vn, vi)
+        logp_i
+    end
+end
+
+() -> begin
+    # code that generates julia-repl in docstring below
+    # using DynamicPPL, Turing
+    # TODO when Turing version that is compatible with DynamicPPL 0.29 becomes available
+end
+
 """
-    pointwise_loglikelihoods(model::Model, chain::Chains, keytype = String)
+    pointwise_logdensities(model::Model, chain::Chains, keytype = String)
 
 Runs `model` on each sample in `chain` returning a `OrderedDict{String, Matrix{Float64}}`
 with keys corresponding to symbols of the observations, and values being matrices
@@ -184,21 +223,21 @@ julia> model = demo(randn(3), randn());
 
 julia> chain = sample(model, MH(), 10);
 
-julia> pointwise_loglikelihoods(model, chain)
+julia> pointwise_logdensities(model, chain)
 OrderedDict{String,Array{Float64,2}} with 4 entries:
   "xs[1]" => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
   "xs[2]" => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
   "xs[3]" => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
   "y"     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
 
-julia> pointwise_loglikelihoods(model, chain, String)
+julia> pointwise_logdensities(model, chain, String)
 OrderedDict{String,Array{Float64,2}} with 4 entries:
   "xs[1]" => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
   "xs[2]" => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
   "xs[3]" => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
   "y"     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
 
-julia> pointwise_loglikelihoods(model, chain, VarName)
+julia> pointwise_logdensities(model, chain, VarName)
 OrderedDict{VarName,Array{Float64,2}} with 4 entries:
   xs[1] => [-1.42932; -2.68123; … ; -1.66333; -1.66333]
   xs[2] => [-1.6724; -0.861339; … ; -1.62359; -1.62359]
@@ -217,41 +256,46 @@ julia> @model function demo(x)
 
 julia> m = demo([1.0, ]);
 
-julia> ℓ = pointwise_loglikelihoods(m, VarInfo(m)); first(ℓ[@varname(x[1])])
+julia> ℓ = pointwise_logdensities(m, VarInfo(m)); first(ℓ[@varname(x[1])])
 -1.4189385332046727
 
 julia> m = demo([1.0; 1.0]);
 
-julia> ℓ = pointwise_loglikelihoods(m, VarInfo(m)); first.((ℓ[@varname(x[1])], ℓ[@varname(x[2])]))
+julia> ℓ = pointwise_logdensities(m, VarInfo(m)); first.((ℓ[@varname(x[1])], ℓ[@varname(x[2])]))
 (-1.4189385332046727, -1.4189385332046727)
 ```
 
 """
-function pointwise_loglikelihoods(model::Model, chain, keytype::Type{T}=String) where {T}
+function pointwise_logdensities(model::Model, chain, 
+    context::AbstractContext=DefaultContext(), keytype::Type{T}=String) where {T}
     # Get the data by executing the model once
     vi = VarInfo(model)
-    context = PointwiseLikelihoodContext(OrderedDict{T,Vector{Float64}}())
+    point_context = PointwiseLogdensityContext(OrderedDict{T,Vector{Float64}}(), context)
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     for (sample_idx, chain_idx) in iters
         # Update the values
         setval!(vi, chain, sample_idx, chain_idx)
 
         # Execute model
-        model(vi, context)
+        model(vi, point_context)
     end
 
     niters = size(chain, 1)
     nchains = size(chain, 3)
-    loglikelihoods = OrderedDict(
+    logdensities = OrderedDict(
         varname => reshape(logliks, niters, nchains) for
-        (varname, logliks) in context.loglikelihoods
+        (varname, logliks) in point_context.logdensities
     )
-    return loglikelihoods
+    return logdensities
 end
 
-function pointwise_loglikelihoods(model::Model, varinfo::AbstractVarInfo)
-    context = PointwiseLikelihoodContext(OrderedDict{VarName,Vector{Float64}}())
-    model(varinfo, context)
-    return context.loglikelihoods
+function pointwise_logdensities(model::Model,   
+    varinfo::AbstractVarInfo, context::AbstractContext=DefaultContext())
+    point_context = PointwiseLogdensityContext(
+        OrderedDict{VarName,Vector{Float64}}(), context)
+    model(varinfo, point_context)
+    return point_context.logdensities
 end
+
+