Fix fix as well

Author: penelopeysm

src/compiler.jl

@@ -105,7 +105,11 @@ function contextual_isassumption(context::AbstractContext, vn)
 end
 
 isfixed(expr, vn) = false
-isfixed(::Union{Symbol,Expr}, vn) = :($(DynamicPPL.contextual_isfixed)(__context__, $vn))
+function isfixed(::Union{Symbol,Expr}, vn)
+    return :($(DynamicPPL.contextual_isfixed)(
+        __context__, $(DynamicPPL.prefix)(__context__, $vn)
+    ))
+end
 
 """
     contextual_isfixed(context, vn)
@@ -443,7 +447,9 @@ function generate_tilde(left, right)
         )
         $isassumption = $(DynamicPPL.isassumption(left, vn))
         if $(DynamicPPL.isfixed(left, vn))
-            $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
+            $left = $(DynamicPPL.getfixed_nested)(
+                __context__, $(DynamicPPL.prefix)(__context__, $vn)
+            )
         elseif $isassumption
             $(generate_tilde_assume(left, dist, vn))
         else

src/context_implementations.jl

@@ -57,7 +57,6 @@ 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...)

src/contexts.jl

@@ -375,66 +375,6 @@ 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)
-
-Prefix all the conditioned variables in a given context with `prefix`.
-
-```jldoctest
-julia> using DynamicPPL: prefix_conditioned_variables, ConditionContext
-
-julia> c1 = ConditionContext((a=1, ))
-ConditionContext((a = 1,), DefaultContext())
-
-julia> prefix_conditioned_variables(c1, @varname(y))
-ConditionContext(Dict(y.a => 1), DefaultContext())
-```
-"""
-function prefix_conditioned_variables(ctx::ConditionContext, prefix::VarName)
-    # Replace the prefix of the conditioned variables
-    vn_dict = to_varname_dict(ctx.values)
-    prefixed_vn_dict = Dict(
-        AbstractPPL.prefix(vn, prefix) => value for (vn, value) in vn_dict
-    )
-    # Prefix the child context as well
-    prefixed_child_ctx = prefix_conditioned_variables(childcontext(ctx), prefix)
-    return ConditionContext(prefixed_vn_dict, prefixed_child_ctx)
-end
-function prefix_conditioned_variables(c::AbstractContext, prefix::VarName)
-    return prefix_conditioned_variables(
-        NodeTrait(prefix_conditioned_variables, c), c, prefix
-    )
-end
-prefix_conditioned_variables(::IsLeaf, context::AbstractContext, prefix::VarName) = context
-function prefix_conditioned_variables(::IsParent, context::AbstractContext, prefix::VarName)
-    return setchildcontext(
-        context, prefix_conditioned_variables(childcontext(context), prefix)
-    )
-end
-
 """
     hasconditioned(context::AbstractContext, vn::VarName)
 
@@ -474,7 +414,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(collapse_prefix_and_condition(context), vn)
+    return hasconditioned_nested(collapse_prefix_stack(context), vn)
 end
 
 """
@@ -492,7 +432,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(collapse_prefix_and_condition(context), vn)
+    return getconditioned_nested(collapse_prefix_stack(context), vn)
 end
 function getconditioned_nested(::IsParent, context, vn)
     return if hasconditioned(context, vn)
@@ -563,9 +503,7 @@ function conditioned(context::ConditionContext)
     return _merge(context.values, conditioned(childcontext(context)))
 end
 function conditioned(context::PrefixContext{Prefix}) where {Prefix}
-    return conditioned(
-        prefix_conditioned_variables(childcontext(context), VarName{Prefix}())
-    )
+    return conditioned(collapse_prefix_stack(context))
 end
 
 struct FixedContext{Values,Ctx<:AbstractContext} <: AbstractContext
@@ -630,7 +568,7 @@ function hasfixed_nested(::IsParent, context, vn)
     return hasfixed(context, vn) || hasfixed_nested(childcontext(context), vn)
 end
 function hasfixed_nested(context::PrefixContext, vn)
-    return hasfixed_nested(childcontext(context), prefix(context, vn))
+    return hasfixed_nested(collapse_prefix_stack(context), vn)
 end
 
 """
@@ -648,7 +586,7 @@ function getfixed_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
 function getfixed_nested(context::PrefixContext, vn)
-    return getfixed_nested(childcontext(context), prefix(context, vn))
+    return getfixed_nested(collapse_prefix_stack(context), vn)
 end
 function getfixed_nested(::IsParent, context, vn)
     return if hasfixed(context, vn)
@@ -743,3 +681,80 @@ function fixed(context::FixedContext)
     # precedence over decendants of `context`.
     return _merge(context.values, fixed(childcontext(context)))
 end
+
+"""
+    collapse_prefix_stack(context::AbstractContext)
+
+Apply `PrefixContext`s to any conditioned or fixed values inside them, and remove
+the `PrefixContext`s from the context stack.
+
+```jldoctest
+julia> using DynamicPPL: collapse_prefix_stack
+
+julia> c1 = PrefixContext({:a}(ConditionContext((x=1, )))
+```
+"""
+function collapse_prefix_stack(context::PrefixContext{Prefix}) where {Prefix}
+    # Collapse the child context (thus applying any inner prefixes first)
+    collapsed = collapse_prefix_stack(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_cond_and_fixed_variables(collapsed, VarName{Prefix}())
+end
+collapse_prefix_stack(context::AbstractContext) = context
+
+"""
+    prefix_cond_and_fixed_variables(context::AbstractContext, prefix::VarName)
+
+Prefix all the conditioned and fixed variables in a given context with a single
+`prefix`.
+
+```jldoctest
+julia> using DynamicPPL: prefix_cond_and_fixed_variables, ConditionContext
+
+julia> c1 = ConditionContext((a=1, ))
+ConditionContext((a = 1,), DefaultContext())
+
+julia> prefix_cond_and_fixed_variables(c1, @varname(y))
+ConditionContext(Dict(y.a => 1), DefaultContext())
+```
+"""
+function prefix_cond_and_fixed_variables(ctx::ConditionContext, prefix::VarName)
+    # Replace the prefix of the conditioned variables
+    vn_dict = to_varname_dict(ctx.values)
+    prefixed_vn_dict = Dict(
+        AbstractPPL.prefix(vn, prefix) => value for (vn, value) in vn_dict
+    )
+    # Prefix the child context as well
+    prefixed_child_ctx = prefix_cond_and_fixed_variables(childcontext(ctx), prefix)
+    return ConditionContext(prefixed_vn_dict, prefixed_child_ctx)
+end
+function prefix_cond_and_fixed_variables(ctx::FixedContext, prefix::VarName)
+    # Replace the prefix of the conditioned variables
+    vn_dict = to_varname_dict(ctx.values)
+    prefixed_vn_dict = Dict(
+        AbstractPPL.prefix(vn, prefix) => value for (vn, value) in vn_dict
+    )
+    # Prefix the child context as well
+    prefixed_child_ctx = prefix_cond_and_fixed_variables(childcontext(ctx), prefix)
+    return FixedContext(prefixed_vn_dict, prefixed_child_ctx)
+end
+function prefix_cond_and_fixed_variables(c::AbstractContext, prefix::VarName)
+    return prefix_cond_and_fixed_variables(
+        NodeTrait(prefix_cond_and_fixed_variables, c), c, prefix
+    )
+end
+function prefix_cond_and_fixed_variables(
+    ::IsLeaf, context::AbstractContext, prefix::VarName
+)
+    return context
+end
+function prefix_cond_and_fixed_variables(
+    ::IsParent, context::AbstractContext, prefix::VarName
+)
+    return setchildcontext(
+        context, prefix_cond_and_fixed_variables(childcontext(context), prefix)
+    )
+end

test/submodels.jl

@@ -5,74 +5,101 @@ using Distributions
 using Test
 
 @testset "submodels.jl" begin
-    @testset "Conditioning variables" begin
+    @testset "$op" for op in [condition, fix]
+        x_val = 1.0
+        x_logp = op == condition ? logpdf(Normal(), x_val) : 0.0
+
         @testset "Auto prefix" begin
             @model function inner()
                 x ~ Normal()
-                return y ~ Normal()
+                y ~ Normal()
+                return (x, y)
             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)
+            inner_op = op(inner(), (@varname(x) => x_val))
+            @model function outer2()
+                return a ~ to_submodel(inner_op)
+            end
+            with_inner_op = outer2()
+            with_outer_op = op(outer(), (@varname(a.x) => x_val))
 
-            # No conditioning
+            # No conditioning/fixing
             @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)
+
+            # With conditioning/fixing
+            @testset "$model" for model in [with_inner_op, with_outer_op]
+                # Test that the value was correctly set
+                @test model()[1] == x_val
+                # Test that the logp was correctly set
+                vi = VarInfo(model)
+                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(a.y)])
+                # Check the keys
+                @test Set(keys(VarInfo(model))) == Set([@varname(a.y)])
             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()
+                y ~ Normal()
+                return x
             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)
+            @model function outer2()
+                return a ~ to_submodel(inner_op, false)
+            end
+            with_inner_op = outer2()
+            inner_op = op(inner(), (@varname(x) => x_val))
+            with_outer_op = op(outer(), (@varname(x) => x_val))
 
-            # No conditioning
+            # No conditioning/fixing
             @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)
+
+            # With conditioning/fixing
+            @testset "$model" for model in [with_inner_op, with_outer_op]
+                # Test that the value was correctly set
+                @test model()[1] == x_val
+                # Test that the logp was correctly set
+                vi = VarInfo(model)
+                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(y)])
+                # Check the keys
+                @test Set(keys(VarInfo(model))) == Set([@varname(y)])
             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()
+                y ~ Normal()
+                return x
             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)
+            inner_op = op(inner(), (@varname(x) => x_val))
+            @model function outer2()
+                return a ~ to_submodel(prefix(inner_op, :b), false)
+            end
+            with_inner_op = outer2()
+            with_outer_op = op(outer(), (@varname(b.x) => x_val))
 
-            # No conditioning
+            # No conditioning/fixing
             @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)
+
+            # With conditioning/fixing
+            @testset "$model" for model in [with_inner_op, with_outer_op]
+                # Test that the value was correctly set
+                @test model()[1] == x_val
+                # Test that the logp was correctly set
+                vi = VarInfo(model)
+                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(b.y)])
+                # Check the keys
+                @test Set(keys(VarInfo(model))) == Set([@varname(b.y)])
             end
-            with_inner_cond = outer2()
-            @test Set(keys(VarInfo(with_inner_cond))) == Set([@varname(b.y)])
         end
 
         @testset "Nested submodels" begin
@@ -90,21 +117,21 @@ using Test
             # 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/fixing at the top level
+            op_h = op(h(), (@varname(a.b.x) => x_val))
+            @test Set(keys(VarInfo(op_h))) == Set([@varname(a.b.y)])
 
-            # Conditioning at the second level
-            condition_g = g() | (@varname(b.x) => 1.0)
+            # Conditioning/fixing at the second level
+            op_g = op(g(), (@varname(b.x) => x_val))
             @model function h2()
-                return a ~ to_submodel(condition_g)
+                return a ~ to_submodel(op_g)
             end
             @test Set(keys(VarInfo(h2()))) == Set([@varname(a.b.y)])
 
-            # Conditioning at the very bottom
-            condition_f = f() | (@varname(x) => 1.0)
+            # Conditioning/fixing at the very bottom
+            op_f = op(f(), (@varname(x) => x_val))
             @model function g2()
-                return _unused ~ to_submodel(prefix(condition_f, :b), false)
+                return _unused ~ to_submodel(prefix(op_f, :b), false)
             end
             @model function h3()
                 return a ~ to_submodel(g2())