forward dot_tilde_assume to tilde_assume for Multivariate

Author: bgctw

src/pointwise_logdensities.jl

@@ -98,7 +98,7 @@ function dot_tilde_observe!!(context::PointwiseLogdensityContext, right, left, v
 
     # 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,8 @@ function _pointwise_tilde_observe(
     end
 end
 
-function tilde_assume(context::PointwiseLogdensityContext, right, vn, vi)
-    #@info "PointwiseLogdensityContext tilde_assume!! called for $vn"
+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
@@ -145,7 +145,7 @@ function dot_tilde_assume(context::PointwiseLogdensityContext, right, left, vns,
     return value, logp, vi
 end
 
-function record_dot_tilde_assume(context::PointwiseLogdensityContext, right::UnivariateDistribution, left, vns, vi, logp)
+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)
@@ -155,27 +155,13 @@ end
 
 function record_dot_tilde_assume(context::PointwiseLogdensityContext, rights::AbstractVector{<:Distribution}, left, vns, vi, logp)
     # forward to tilde_assume for each variable and distribution
-    logps = map(vns, rights) do vn, right
+    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
 
-function record_dot_tilde_assume(context::PointwiseLogdensityContext, right::MultivariateDistribution, left, vns, vi, logp)
-    #@info "PointwiseLogdensityContext record_dot_tilde_assume multivariate called for $vns"
-    # For multivariate distribution on the right there is only a single density.
-    # Need to construct a combined VarName.
-    # Assume that all vns have an IndexLens with a Colon at the first position
-    # and a single number at the second position.
-    indices = map(vn -> getoptic(vn).indices[2], vns)
-    indices_combined = (:,indices)
-    #indices = tuplejoin(map(vn -> getoptic(vn).indices[2], vns)...)
-    vn = VarName(first(vns), Accessors.IndexLens(indices_combined))
-    push!(context, vn, logp)
-    return logp
-end
-
 () -> begin
     # code that generates julia-repl in docstring below
     # using DynamicPPL, Turing

src/test_utils.jl

@@ -667,6 +667,62 @@ function varnames(model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)
     return [@varname(s[:, 1], true), @varname(s[:, 2], true), @varname(m)]
 end
 
+# model with truly two 3 columns in of d=2 (rather than ProductMatrix of single factor) to explore dot_tilde_assume with MultivariateDistribution
+@model function demo_dot_assume_matrix_dot_observe_matrix2(
+    x=transpose([1.5 2.0;1.6 2.1;1.45 2.05]), ::Type{TV}=Array{Float64}
+) where {TV}
+    d = size(x,1)
+    n = size(x,2)
+    s = TV(undef, d, n)
+    # for i in 1:n
+    #     s[:,i] ~ product_distribution([InverseGamma(2, 3) for _ in 1:d])
+    # end
+    s .~ product_distribution([InverseGamma(2, 3) for _ in 1:d])
+    m = TV(undef, d, n)
+    Sigma_x = Diagonal(s[:,1])
+    for i in 1:n
+        diag_s = Diagonal(s[:,i])
+        m[:,i] ~ MvNormal(zeros(d), diag_s)
+        x[:,i] ~ MvNormal(m[:,i], Sigma_x) 
+    end
+
+    return (; s=s, m=m, x=x, logp=getlogp(__varinfo__))
+end
+function logprior_true(
+    model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix2)}, s, m
+)
+    n = size(model.args.x,1)
+    d = size(model.args.x,2)
+    logd = map(1:d) do i_d
+        s_vec = Diagonal(s[:,i_d])
+        loglikelihood(InverseGamma(2, 3), s_vec) +
+        logpdf(MvNormal(zeros(n), Diagonal(s_vec)), m[:,i_d])
+    end
+    return sum(logd)
+end
+function loglikelihood_true(
+    model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix2)}, s, m
+)
+    n = size(model.args.x,2)
+    d = size(model.args.x,1)
+    s_vec = Diagonal(s[:,1])
+    logd = map(1:n) do i
+        loglikelihood(MvNormal(m[:,i], Diagonal(s_vec)), model.args.x[:,i])
+    end
+    return sum(logd)
+end
+function logprior_true_with_logabsdet_jacobian(
+    model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix2)}, s, m
+)
+    return _demo_logprior_true_with_logabsdet_jacobian(model, s, m)
+end
+function varnames(model::Model{typeof(demo_dot_assume_matrix_dot_observe_matrix2)})
+    s = m = zeros(2, 3) # used for varname concretization only
+    return [@varname(s[:, 1], true), @varname(s[:, 2], true), @varname(s[:, 3], true), 
+    @varname(m[:,1], true), @varname(m[:,2], true), @varname(m[:,3], true)]
+end
+
+
 @model function demo_assume_matrix_dot_observe_matrix(
     x=transpose([1.5 2.0;]), ::Type{TV}=Array{Float64}
 ) where {TV}
@@ -748,6 +804,7 @@ const MultivariateAssumeDemoModels = Union{
     Model{typeof(demo_dot_assume_observe_submodel)},
     Model{typeof(demo_dot_assume_dot_observe_matrix)},
     Model{typeof(demo_dot_assume_matrix_dot_observe_matrix)},
+    Model{typeof(demo_dot_assume_matrix_dot_observe_matrix2)},
 }
 function posterior_mean(model::MultivariateAssumeDemoModels)
     # Get some containers to fill.

test/pointwise_logdensities.jl

@@ -4,6 +4,7 @@
     mod_ctx = DynamicPPL.TestUtils.TestLogModifyingChildContext(1.2)
     mod_ctx2 = DynamicPPL.TestUtils.TestLogModifyingChildContext(1.4, mod_ctx)
     #m = DynamicPPL.TestUtils.DEMO_MODELS[12]
+    #m = model = DynamicPPL.TestUtils.demo_dot_assume_matrix_dot_observe_matrix2()
     @testset "$(m.f)" for (i, m) in enumerate(DynamicPPL.TestUtils.DEMO_MODELS)
         #@show i
         example_values = DynamicPPL.TestUtils.rand_prior_true(m)
@@ -25,13 +26,13 @@
         # Compute the pointwise loglikelihoods.
         lls = pointwise_logdensities(m, vi, likelihood_context)
         #lls2 = pointwise_loglikelihoods(m, vi)
-        loglikelihood = sum(sum, values(lls))
-        if loglikelihood ≈ 0.0 #isempty(lls)
+        loglikelihood_sum = sum(sum, values(lls))
+        if loglikelihood_sum ≈ 0.0 #isempty(lls)
             # One of the models with literal observations, so we just skip.
             # TODO: Think of better way to detect this special case 
             loglikelihood_true = 0.0
         end
-        @test loglikelihood ≈ loglikelihood_true
+        @test loglikelihood_sum ≈ loglikelihood_true
 
         # Compute the pointwise logdensities of the priors.
         lps_prior = pointwise_logdensities(m, vi, prior_context)