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,
+    varwise_logpriors,
     condition,
     decondition,
     fix,
@@ -182,6 +183,7 @@ include("simple_varinfo.jl")
 include("context_implementations.jl")
 include("compiler.jl")
 include("loglikelihoods.jl")
+include("logpriors_var.jl")
 include("submodel_macro.jl")
 include("test_utils.jl")
 include("transforming.jl")

src/context_implementations.jl

@@ -425,6 +425,15 @@ function dot_assume(
     var::AbstractMatrix,
     vns::AbstractVector{<:VarName},
     vi::AbstractVarInfo,
+)
+    r, lp, vi = dot_assume_vec(dist, var, vns, vi)
+    return r, sum(lp), vi
+end
+function dot_assume_vec(
+    dist::MultivariateDistribution,
+    var::AbstractMatrix,
+    vns::AbstractVector{<:VarName},
+    vi::AbstractVarInfo,
 )
     @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
@@ -434,7 +443,7 @@ function dot_assume(
     #
     # in which case `var` will have `undef` elements, even if `m` is present in `vi`.
     r = vi[vns, dist]
-    lp = sum(zip(vns, eachcol(r))) do (vn, ri)
+    lp = map(zip(vns, eachcol(r))) do (vn, ri)
         return Bijectors.logpdf_with_trans(dist, ri, istrans(vi, vn))
     end
     return r, lp, vi
@@ -455,21 +464,29 @@ function dot_assume(
 end
 
 function dot_assume(
+    dist::Union{Distribution,AbstractArray{<:Distribution}}, var::AbstractArray, vns::AbstractArray{<:VarName}, vi
+)
+    # possibility to acesss the single logpriors
+    r, lp, vi = dot_assume_vec(dist, var, vns, vi)
+    return r, sum(lp), vi
+end
+
+function dot_assume_vec(
     dist::Distribution, var::AbstractArray, vns::AbstractArray{<:VarName}, vi
 )
     r = getindex.((vi,), vns, (dist,))
-    lp = sum(Bijectors.logpdf_with_trans.((dist,), r, istrans.((vi,), vns)))
+    lp = Bijectors.logpdf_with_trans.((dist,), r, istrans.((vi,), vns))
     return r, lp, vi
 end
 
-function dot_assume(
+function dot_assume_vec(
     dists::AbstractArray{<:Distribution},
     var::AbstractArray,
     vns::AbstractArray{<:VarName},
     vi,
 )
     r = getindex.((vi,), vns, dists)
-    lp = sum(Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns)))
+    lp = Bijectors.logpdf_with_trans.(dists, r, istrans.((vi,), vns))
     return r, lp, vi
 end
 

src/logpriors_var.jl

@@ -0,0 +1,201 @@
+"""
+Context that records logp after tilde_assume!!
+for each VarName used by [`varwise_logpriors`](@ref).
+"""
+struct VarwisePriorContext{A,Ctx} <: AbstractContext
+    logpriors::A
+    context::Ctx
+end
+
+function VarwisePriorContext(
+    logpriors=OrderedDict{Symbol,Float64}(),
+    context::AbstractContext=DynamicPPL.PriorContext(),
+    #OrderedDict{Symbol,Vector{Float64}}(),PriorContext()),
+)
+    return VarwisePriorContext{typeof(logpriors),typeof(context)}(
+        logpriors, context
+    )
+end
+
+NodeTrait(::VarwisePriorContext) = IsParent()
+childcontext(context::VarwisePriorContext) = context.context
+function setchildcontext(context::VarwisePriorContext, child)
+    return VarwisePriorContext(context.logpriors, child)
+end
+
+function tilde_assume(context::VarwisePriorContext, right, vn, vi)
+    #@info "VarwisePriorContext tilde_assume!! called for $vn"
+    value, logp, vi = tilde_assume(context.context, right, vn, vi)
+    #sym = DynamicPPL.getsym(vn)
+    new_context = acc_logp!(context, vn, logp)
+    return value, logp, vi
+end
+
+function dot_tilde_assume(context::VarwisePriorContext, right, left, vn, vi)
+    #@info "VarwisePriorContext dot_tilde_assume!! called for $vn"
+    # @show vn, left, right, typeof(context).name
+    value, logp, vi = dot_tilde_assume(context.context, right, left, vn, vi)
+    new_context = acc_logp!(context, vn, logp)
+    return value, logp, vi
+end
+
+
+function tilde_observe(context::VarwisePriorContext, right, left, vi)
+    # Since we are evaluating the prior, the log probability of all the observations
+    # is set to 0. This has the effect of ignoring the likelihood.
+    return 0.0, vi
+    #tmp = tilde_observe(context.context, SampleFromPrior(), right, left, vi)
+    #return tmp
+end
+
+function acc_logp!(context::VarwisePriorContext, vn::Union{VarName,AbstractVector{<:VarName}}, logp)
+    #sym = DynamicPPL.getsym(vn)  # leads to duplicates
+    # if vn is a Vector leads to Symbol("VarName{:s, IndexLens{Tuple{Int64}}}[s[1], s[2]]")    
+    sym = Symbol(vn) 
+    context.logpriors[sym] = logp
+    return (context)
+end
+
+
+# """
+#     pointwise_logpriors(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
+# of shape `(num_chains, num_samples)`.
+
+# `keytype` specifies what the type of the keys used in the returned `OrderedDict` are.
+# Currently, only `String` and `VarName` are supported.
+
+# # Notes
+# Say `y` is a `Vector` of `n` i.i.d. `Normal(μ, σ)` variables, with `μ` and `σ`
+# both being `<:Real`. Then the *observe* (i.e. when the left-hand side is an
+# *observation*) statements can be implemented in three ways:
+# 1. using a `for` loop:
+# ```julia
+# for i in eachindex(y)
+#     y[i] ~ Normal(μ, σ)
+# end
+# ```
+# 2. using `.~`:
+# ```julia
+# y .~ Normal(μ, σ)
+# ```
+# 3. using `MvNormal`:
+# ```julia
+# y ~ MvNormal(fill(μ, n), σ^2 * I)
+# ```
+
+# In (1) and (2), `y` will be treated as a collection of `n` i.i.d. 1-dimensional variables,
+# while in (3) `y` will be treated as a _single_ n-dimensional observation.
+
+# This is important to keep in mind, in particular if the computation is used
+# for downstream computations.
+
+# # Examples
+# ## From chain
+# ```julia-repl
+# julia> using DynamicPPL, Turing
+
+# julia> @model function demo(xs, y)
+#            s ~ InverseGamma(2, 3)
+#            m ~ Normal(0, √s)
+#            for i in eachindex(xs)
+#                xs[i] ~ Normal(m, √s)
+#            end
+
+#            y ~ Normal(m, √s)
+#        end
+# demo (generic function with 1 method)
+
+# julia> model = demo(randn(3), randn());
+
+# julia> chain = sample(model, MH(), 10);
+
+# julia> pointwise_logpriors(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_logpriors(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_logpriors(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]
+#   xs[3] => [-1.42862; -2.67573; … ; -1.66251; -1.66251]
+#   y     => [-1.51265; -0.914129; … ; -1.5499; -1.5499]
+# ```
+
+# ## Broadcasting
+# Note that `x .~ Dist()` will treat `x` as a collection of
+# _independent_ observations rather than as a single observation.
+
+# ```jldoctest; setup = :(using Distributions)
+# julia> @model function demo(x)
+#            x .~ Normal()
+#        end;
+
+# julia> m = demo([1.0, ]);
+
+# julia> ℓ = pointwise_logpriors(m, VarInfo(m)); first(ℓ[@varname(x[1])])
+# -1.4189385332046727
+
+# julia> m = demo([1.0; 1.0]);
+
+# julia> ℓ = pointwise_logpriors(m, VarInfo(m)); first.((ℓ[@varname(x[1])], ℓ[@varname(x[2])]))
+# (-1.4189385332046727, -1.4189385332046727)
+# ```
+
+# """
+function varwise_logpriors(
+    model::Model, varinfo::AbstractVarInfo,
+    context::AbstractContext=PriorContext()
+)
+#    top_context = VarwisePriorContext(OrderedDict{Symbol,Float64}(), context)
+    top_context = VarwisePriorContext(OrderedDict{Symbol,Float64}(), context)
+    model(varinfo, top_context)
+    return top_context.logpriors
+end
+
+function varwise_logpriors(model::Model, chain::AbstractChains,
+    context::AbstractContext=PriorContext();
+    top_context::VarwisePriorContext{T} = VarwisePriorContext(OrderedDict{Symbol,Float64}(), context)
+    ) where T
+    # pass top-context as keyword to allow adapt Number type of log-prior
+    get_values = (vi) -> begin
+        model(vi, top_context)
+        values(top_context.logpriors)
+    end
+    arr =  map_model(get_values, model, chain)
+    par_names = collect(keys(top_context.logpriors))
+    return(arr, par_names)
+end
+
+function map_model(get_values, model::Model, chain::AbstractChains)
+    niters = size(chain, 1)
+    nchains = size(chain, 3)
+    vi = VarInfo(model)
+    iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+    # initialize the array by the first result
+    (sample_idx, chain_idx), iters2 = Iterators.peel(iters)
+    setval!(vi, chain, sample_idx, chain_idx)
+    values1 = get_values(vi)
+    arr = Array{eltype(values1)}(undef, niters, length(values1), nchains)
+    arr[sample_idx, :, chain_idx] .= values1
+    #(sample_idx, chain_idx), iters3 = Iterators.peel(iters2)
+    for (sample_idx, chain_idx) in iters2
+        # Update the values
+        setval!(vi, chain, sample_idx, chain_idx)
+        values_i = get_values(vi)
+        arr[sample_idx, :, chain_idx] .= values_i
+    end
+    return(arr)
+end

src/test_utils.jl

@@ -1042,4 +1042,37 @@ function test_context_interface(context)
     end
 end
 
+"""
+Context that multiplies each log-prior by mod
+used to test whether varwise_logpriors respects child-context.
+"""
+struct TestLogModifyingChildContext{T,Ctx} <: DynamicPPL.AbstractContext
+    mod::T
+    context::Ctx
+end
+function TestLogModifyingChildContext(
+    mod=1.2,
+    context::DynamicPPL.AbstractContext=DynamicPPL.DefaultContext(),
+        #OrderedDict{VarName,Vector{Float64}}(),PriorContext()),
+)
+    return TestLogModifyingChildContext{typeof(mod),typeof(context)}(
+        mod, context
+    )
+end
+DynamicPPL.NodeTrait(::TestLogModifyingChildContext) = DynamicPPL.IsParent()
+DynamicPPL.childcontext(context::TestLogModifyingChildContext) = context.context
+function DynamicPPL.setchildcontext(context::TestLogModifyingChildContext, child)
+    return TestLogModifyingChildContext(context.mod, child)
+end
+function DynamicPPL.tilde_assume(context::TestLogModifyingChildContext, right, vn, vi)
+    #@info "TestLogModifyingChildContext tilde_assume!! called for $vn"
+    value, logp, vi = DynamicPPL.tilde_assume(context.context, right, vn, vi)
+    return value, logp*context.mod, vi
+end
+function DynamicPPL.dot_tilde_assume(context::TestLogModifyingChildContext, right, left, vn, vi)
+    #@info "TestLogModifyingChildContext dot_tilde_assume!! called for $vn"
+    value, logp, vi = DynamicPPL.dot_tilde_assume(context.context, right, left, vn, vi)
+    return value, logp*context.mod, vi
+end
+
 end

test/Project.toml

@@ -14,6 +14,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -26,10 +27,10 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-Accessors = "0.1"
 ADTypes = "0.2, 1"
 AbstractMCMC = "5"
 AbstractPPL = "0.8.2"
+Accessors = "0.1"
 Bijectors = "0.13"
 Compat = "4.3.0"
 Distributions = "0.25"