Fix remaining bugs e.g. in nested decondition_context

Author: penelopeysm

src/contexts.jl

@@ -447,9 +447,17 @@ function decondition_context(context::ConditionContext, sym, syms...)
         # No more values left, can unwrap
         decondition_context(childcontext(context), syms...)
     else
-        ConditionContext(new_values, decondition_context(childcontext(context), syms...))
+        ConditionContext(
+            new_values, decondition_context(childcontext(context), sym, syms...)
+        )
     end
 end
+function decondition_context(context::NamedConditionContext, vn::VarName{sym}) where {sym}
+    return ConditionContext(
+        BangBang.delete!!(context.values, sym),
+        decondition_context(childcontext(context), vn),
+    )
+end
 
 """
     conditioned(context::AbstractContext)

src/model.jl

@@ -96,7 +96,7 @@ Return a `Model` which now treats variables on the right-hand side as observatio
 
 See [`condition`](@ref) for more information and examples.
 """
-Base.:|(model::Model, values::Union{Tuple,NamedTuple,AbstractDict{<:VarName}}) =
+Base.:|(model::Model, values::Union{Pair,Tuple,NamedTuple,AbstractDict{<:VarName}}) =
     condition(model, values)
 
 """
@@ -281,14 +281,14 @@ end
 
 Convert different types of input to either a `NamedTuple` or `AbstractDict` of
 conditioning values, suitable for storage in a `ConditionContext`.
+
+This handles all the cases where `vals` is either already a NamedTuple or
+AbstractDict (e.g. `model | (x=1, y=2)`), as well as if they are splatted (e.g.
+`condition(model, x=1, y=2)`).
 """
-# Case 1: Already in the right format, e.g. condition(model, (x=1, y=2))
 _make_conditioning_values(values::Union{NamedTuple,AbstractDict}) = values
-# Case 2: condition(model, (@varname(x) => 1, @varname(y) => 2))
 _make_conditioning_values(values::Tuple{Pair{<:VarName}}) = Dict(values)
-# Case 3: Case 1 but splatted, e.g. condition(model, x=1, y=2)
 _make_conditioning_values(v::Pair{<:Symbol}, vs::Pair{<:Symbol}...) = NamedTuple(v, vs...)
-# Case 4: Case 2 but splatted, e.g. condition(model, @varname(x) => 1, @varname(y) => 2)
 _make_conditioning_values(v::Pair{<:VarName}, vs::Pair{<:VarName}...) = Dict(v, vs...)
 
 """
@@ -401,7 +401,7 @@ true
 ```
 """
 function AbstractPPL.decondition(model::Model, syms...)
-    return contextualize(model, decondition(model.context, syms...))
+    return contextualize(model, decondition_context(model.context, syms...))
 end
 
 """
@@ -435,7 +435,7 @@ julia> # Returns all the variables we have conditioned on + their values.
 (x = 100.0, m = 1.0)
 
 julia> # Nested ones also work (note that `PrefixContext` does nothing to the result).
-       cm = condition(contextualize(m, PrefixContext{:a}(condition(m=1.0))), x=100.0);
+       cm = condition(contextualize(m, PrefixContext{:a}(ConditionContext((m=1.0,)))), x=100.0);
 
 julia> conditioned(cm)
 (x = 100.0, m = 1.0)
@@ -447,15 +447,15 @@ julia> # Since we conditioned on `m`, not `a.m` as it will appear after prefixed
  a.m
 
 julia> # If we instead condition on `a.m`, `m` in the model will be considered an observation.
-       cm = condition(contextualize(m, PrefixContext{:a}(condition(var"a.m"=1.0))), x=100.0);
+       cm = condition(contextualize(m, PrefixContext{:a}(ConditionContext((var"a.m"=1.0,)))), x=100.0);
 
 julia> conditioned(cm).x
 100.0
 
 julia> conditioned(cm).var"a.m"
 1.0
 
-julia> keys(VarInfo(cm)) # <= no variables are sampled
+julia> keys(VarInfo(cm)) # No variables are sampled
 VarName[]
 ```
 """

test/contexts.jl

@@ -259,6 +259,23 @@ end
                 @test decondition_context(ctx, @varname(x), @varname(y), @varname(z)) isa
                     DefaultContext
             end
+
+            @testset "Nesting" begin
+                ctx = ConditionContext(
+                    (x=1, y=2), ConditionContext(Dict(@varname(a) => 3, @varname(b) => 4))
+                )
+                # Decondition an outer variable
+                dctx = decondition_context(ctx, :x)
+                @test dctx.values == (y=2,)
+                @test childcontext(dctx).values == Dict(@varname(a) => 3, @varname(b) => 4)
+                # Decondition an inner variable
+                dctx = decondition_context(ctx, @varname(a))
+                @test dctx.values == (x=1, y=2)
+                @test childcontext(dctx).values == Dict(@varname(b) => 4)
+                # Try deconditioning everything
+                dctx = decondition_context(ctx)
+                @test dctx isa DefaultContext
+            end
         end
     end
 

test/model.jl

@@ -100,25 +100,36 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
         end
     end
 
-    @testset "model conditioning with various arguments" begin
+    @testset "model de/conditioning" begin
         @model function demo_condition()
             x ~ Normal()
             return y ~ Normal(x)
         end
         model = demo_condition()
+
         # Test that different syntaxes work and give the same underlying ConditionContext
-        @testset "NamedTuple ConditionContext" begin
+        @testset "conditioning NamedTuple" begin
             expected_values = (y=2,)
             @test condition(model, (y=2,)).context.values == expected_values
             @test condition(model; y=2).context.values == expected_values
             @test condition(model; y=2).context.values == expected_values
             @test (model | (y=2,)).context.values == expected_values
+            conditioned_model = condition(model, (y=2,))
+            @test keys(VarInfo(conditioned_model)) == [@varname(x)]
         end
-        @testset "AbstractDict ConditionContext" begin
+        @testset "conditioning AbstractDict" begin
             expected_values = Dict(@varname(y) => 2)
             @test condition(model, Dict(@varname(y) => 2)).context.values == expected_values
             @test condition(model, @varname(y) => 2).context.values == expected_values
             @test (model | (@varname(y) => 2,)).context.values == expected_values
+            conditioned_model = condition(model, Dict(@varname(y) => 2))
+            @test keys(VarInfo(conditioned_model)) == [@varname(x)]
+        end
+
+        @testset "deconditioning" begin
+            conditioned_model = condition(model, (y=2,))
+            deconditioned_model = decondition(conditioned_model)
+            @test keys(VarInfo(deconditioned_model)) == [@varname(x), @varname(y)]
         end
     end