Add tests; fix a bunch of nested submodel issues

Author: penelopeysm

src/context_implementations.jl

@@ -57,6 +57,7 @@ function tilde_assume(context::AbstractContext, args...)
     return tilde_assume(NodeTrait(tilde_assume, context), context, args...)
 end
 function tilde_assume(::IsLeaf, context::AbstractContext, right, vn, vi)
+    @show "isleaf", vn
     return assume(right, vn, vi)
 end
 function tilde_assume(::IsParent, context::AbstractContext, args...)
@@ -85,12 +86,25 @@ function tilde_assume(rng::Random.AbstractRNG, ::LikelihoodContext, sampler, rig
 end
 
 function tilde_assume(context::PrefixContext, right, vn, vi)
-    return tilde_assume(context.context, right, prefix(context, vn), vi)
+    # The slightly tricky thing about PrefixContext is that they are applied
+    # from the outside in, so `PrefixContext{:a}(PrefixContext{:b}(ctx))` means
+    # that variables get prefixed like `a.b.x`.
+    # This motivates the implementation shown here, where the function
+    # `prefix_and_strip_contexts` is responsible for not only adding the
+    # prefixes, but also removing the `PrefixContext`s from the context stack
+    # so that they don't get applied twice when recursing.
+    # TODO(penelopeysm): It would be nice to switch this round, but it's a very
+    # tricky task. Essentially it forces us to use a foldr inside
+    # `prefix_and_strip_contexts`, rather than a foldl which is what most of
+    # DynamicPPL uses.
+    new_vn, new_context = prefix_and_strip_contexts(context, vn)
+    return tilde_assume(new_context, right, new_vn, vi)
 end
 function tilde_assume(
     rng::Random.AbstractRNG, context::PrefixContext, sampler, right, vn, vi
 )
-    return tilde_assume(rng, context.context, sampler, right, prefix(context, vn), vi)
+    new_vn, new_context = prefix_and_strip_contexts(context, vn)
+    return tilde_assume(rng, new_context, sampler, right, new_vn, vi)
 end
 
 """

src/contexts.jl

@@ -276,6 +276,30 @@ function prefix(::IsParent, ctx::AbstractContext, vn::VarName)
     return prefix(childcontext(ctx), vn)
 end
 
+"""
+    prefix_and_strip_contexts(ctx::PrefixContext, vn::VarName)
+
+Same as `prefix`, but additionally returns a new context stack that has all the
+PrefixContexts removed.
+"""
+function prefix_and_strip_contexts(ctx::PrefixContext{Prefix}, vn::VarName) where {Prefix}
+    child_context = childcontext(ctx)
+    # vn_prefixed contains the prefixes from all lower levels
+    vn_prefixed, child_context_without_prefixes = prefix_and_strip_contexts(
+        child_context, vn
+    )
+    return AbstractPPL.prefix(vn_prefixed, VarName{Prefix}()),
+    child_context_without_prefixes
+end
+function prefix_and_strip_contexts(ctx::AbstractContext, vn::VarName)
+    return prefix_and_strip_contexts(NodeTrait(ctx), ctx, vn)
+end
+prefix_and_strip_contexts(::IsLeaf, ctx::AbstractContext, vn::VarName) = (vn, ctx)
+function prefix_and_strip_contexts(::IsParent, ctx::AbstractContext, vn::VarName)
+    vn, new_ctx = prefix_and_strip_contexts(childcontext(ctx), vn)
+    return vn, setchildcontext(ctx, new_ctx)
+end
+
 """
     prefix(model::Model, x)
 
@@ -351,6 +375,29 @@ NodeTrait(::ConditionContext) = IsParent()
 childcontext(context::ConditionContext) = context.context
 setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)
 
+"""
+    collapse_prefix_and_condition(context::AbstractContext)
+
+Apply `PrefixContext`s to any conditioned values inside them, and remove
+the `PrefixContext`s from the context stack.
+
+```jldoctest
+julia> using DynamicPPL: collapse_prefix_and_condition
+
+julia> c1 = PrefixContext({:a}(ConditionContext((x=1, )))
+```
+"""
+function collapse_prefix_and_condition(context::PrefixContext{Prefix}) where {Prefix}
+    # Collapse the child context (thus applying any inner prefixes first)
+    collapsed = collapse_prefix_and_condition(childcontext(context))
+    # Prefix any conditioned variables with the current prefix
+    # Note: prefix_conditioned_variables is O(N) in the depth of the context stack.
+    # So is this function. In the worst case scenario, this is O(N^2) in the
+    # depth of the context stack.
+    return prefix_conditioned_variables(collapsed, VarName{Prefix}())
+end
+collapse_prefix_and_condition(context::AbstractContext) = context
+
 """
     prefix_conditioned_variables(context::AbstractContext, prefix::VarName)
 
@@ -427,9 +474,7 @@ function hasconditioned_nested(::IsParent, context, vn)
     return hasconditioned(context, vn) || hasconditioned_nested(childcontext(context), vn)
 end
 function hasconditioned_nested(context::PrefixContext{Prefix}, vn) where {Prefix}
-    return hasconditioned_nested(
-        prefix_conditioned_variables(childcontext(context), VarName{Prefix}()), vn
-    )
+    return hasconditioned_nested(collapse_prefix_and_condition(context), vn)
 end
 
 """
@@ -447,9 +492,7 @@ function getconditioned_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
 function getconditioned_nested(context::PrefixContext{Prefix}, vn) where {Prefix}
-    return getconditioned_nested(
-        prefix_conditioned_variables(childcontext(context), VarName{Prefix}()), vn
-    )
+    return getconditioned_nested(collapse_prefix_and_condition(context), vn)
 end
 function getconditioned_nested(::IsParent, context, vn)
     return if hasconditioned(context, vn)

test/runtests.jl

@@ -67,6 +67,7 @@ include("test_util.jl")
         include("threadsafe.jl")
         include("debug_utils.jl")
         include("deprecated.jl")
+        include("submodels.jl")
     end
 
     if GROUP == "All" || GROUP == "Group2"

test/submodels.jl

@@ -0,0 +1,117 @@
+module DPPLSubmodelTests
+
+using DynamicPPL
+using Distributions
+using Test
+
+@testset "submodels.jl" begin
+    @testset "Conditioning variables" begin
+        @testset "Auto prefix" begin
+            @model function inner()
+                x ~ Normal()
+                return y ~ Normal()
+            end
+            @model function outer()
+                return a ~ to_submodel(inner())
+            end
+            inner_cond = inner() | (@varname(x) => 1.0)
+            with_outer_cond = outer() | (@varname(a.x) => 1.0)
+
+            # No conditioning
+            @test Set(keys(VarInfo(outer()))) == Set([@varname(a.x), @varname(a.y)])
+            # Conditioning from the outside
+            @test Set(keys(VarInfo(with_outer_cond))) == Set([@varname(a.y)])
+            # Conditioning from the inside
+            @model function outer2()
+                return a ~ to_submodel(inner_cond)
+            end
+            with_inner_cond = outer2()
+            @test Set(keys(VarInfo(with_inner_cond))) == Set([@varname(a.y)])
+        end
+
+        @testset "No prefix" begin
+            @model function inner()
+                x ~ Normal()
+                return y ~ Normal()
+            end
+            @model function outer()
+                return a ~ to_submodel(inner(), false)
+            end
+            inner_cond = inner() | (@varname(x) => 1.0)
+            with_outer_cond = outer() | (@varname(x) => 1.0)
+
+            # No conditioning
+            @test Set(keys(VarInfo(outer()))) == Set([@varname(x), @varname(y)])
+            # Conditioning from the outside
+            @test Set(keys(VarInfo(with_outer_cond))) == Set([@varname(y)])
+            # Conditioning from the inside
+            @model function outer2()
+                return a ~ to_submodel(inner_cond, false)
+            end
+            with_inner_cond = outer2()
+            @test Set(keys(VarInfo(with_inner_cond))) == Set([@varname(y)])
+        end
+
+        @testset "Manual prefix" begin
+            @model function inner()
+                x ~ Normal()
+                return y ~ Normal()
+            end
+            @model function outer()
+                return a ~ to_submodel(prefix(inner(), :b), false)
+            end
+            inner_cond = inner() | (@varname(x) => 1.0)
+            with_outer_cond = outer() | (@varname(b.x) => 1.0)
+
+            # No conditioning
+            @test Set(keys(VarInfo(outer()))) == Set([@varname(b.x), @varname(b.y)])
+            # Conditioning from the outside
+            @test Set(keys(VarInfo(with_outer_cond))) == Set([@varname(b.y)])
+            # Conditioning from the inside
+            @model function outer2()
+                return a ~ to_submodel(prefix(inner_cond, :b), false)
+            end
+            with_inner_cond = outer2()
+            @test Set(keys(VarInfo(with_inner_cond))) == Set([@varname(b.y)])
+        end
+
+        @testset "Nested submodels" begin
+            @model function f()
+                x ~ Normal()
+                return y ~ Normal()
+            end
+            @model function g()
+                return _unused ~ to_submodel(prefix(f(), :b), false)
+            end
+            @model function h()
+                return a ~ to_submodel(g())
+            end
+
+            # No conditioning
+            @test Set(keys(VarInfo(h()))) == Set([@varname(a.b.x), @varname(a.b.y)])
+
+            # Conditioning at the top level
+            condition_h = h() | (@varname(a.b.x) => 1.0)
+            @test Set(keys(VarInfo(condition_h))) == Set([@varname(a.b.y)])
+
+            # Conditioning at the second level
+            condition_g = g() | (@varname(b.x) => 1.0)
+            @model function h2()
+                return a ~ to_submodel(condition_g)
+            end
+            @test Set(keys(VarInfo(h2()))) == Set([@varname(a.b.y)])
+
+            # Conditioning at the very bottom
+            condition_f = f() | (@varname(x) => 1.0)
+            @model function g2()
+                return _unused ~ to_submodel(prefix(condition_f, :b), false)
+            end
+            @model function h3()
+                return a ~ to_submodel(g2())
+            end
+            @test Set(keys(VarInfo(h3()))) == Set([@varname(a.b.y)])
+        end
+    end
+end
+
+end