Add tests for ConditionContext/decondition_context

Author: penelopeysm

src/contexts.jl

@@ -331,13 +331,18 @@ const NamedConditionContext{Names} = ConditionContext{<:NamedTuple{Names}}
 const DictConditionContext = ConditionContext{<:AbstractDict}
 
 # Use DefaultContext as the default base context
-ConditionContext(values::Union{NamedTuple,AbstractDict}) = ConditionContext(values, DefaultContext())
+function ConditionContext(values::Union{NamedTuple,AbstractDict})
+    return ConditionContext(values, DefaultContext())
+end
 # Optimisation when there are no values to condition on
 ConditionContext(::NamedTuple{()}, context::AbstractContext) = context
-# Try to avoid nested `ConditionContext`.
+# Collapse consecutive levels of `ConditionContext`. Note that this overrides
+# values inside the child context, thus giving precedence to the outermost
+# `ConditionContext`.
 function ConditionContext(values::NamedTuple, context::NamedConditionContext)
-    # Note that this potentially overrides values from `context`, thus giving
-    # precedence to the outmost `ConditionContext`.
+    return ConditionContext(merge(context.values, values), childcontext(context))
+end
+function ConditionContext(values::AbstractDict{<:VarName}, context::DictConditionContext)
     return ConditionContext(merge(context.values, values), childcontext(context))
 end
 
@@ -433,32 +438,17 @@ end
 function decondition_context(context::ConditionContext)
     return decondition_context(childcontext(context))
 end
-function decondition_context(context::ConditionContext, sym)
-    return ConditionContext(
-        decondition_context(childcontext(context), sym), BangBang.delete!!(context.values, sym)
-    )
-end
 function decondition_context(context::ConditionContext, sym, syms...)
-    return decondition_context(
-        ConditionContext(
-            decondition_context(childcontext(context), syms...),
-            BangBang.delete!!(context.values, sym),
-        ),
-        syms...,
-    )
-end
-
-function decondition_context(
-    context::NamedConditionContext, vn::VarName{sym}
-) where {sym}
-    return ConditionContext(
-        decondition_context(childcontext(context), vn), BangBang.delete!!(context.values, sym)
-    )
-end
-function decondition_context(context::ConditionContext, vn::VarName)
-    return ConditionContext(
-        decondition_context(childcontext(context), vn), BangBang.delete!!(context.values, vn)
-    )
+    new_values = deepcopy(context.values)
+    for s in (sym, syms...)
+        new_values = BangBang.delete!!(new_values, s)
+    end
+    return if length(new_values) == 0
+        # No more values left, can unwrap
+        decondition_context(childcontext(context), syms...)
+    else
+        ConditionContext(new_values, decondition_context(childcontext(context), syms...))
+    end
 end
 
 """

test/contexts.jl

@@ -11,6 +11,7 @@ using DynamicPPL:
     PointwiseLogdensityContext,
     contextual_isassumption,
     ConditionContext,
+    decondition_context,
     hasconditioned,
     getconditioned,
     hasconditioned_nested,
@@ -196,6 +197,68 @@ end
         @test EnzymeCore.EnzymeRules.inactive_type(typeof(context))
     end
 
+    @testset "ConditionContext" begin
+        @testset "Nesting" begin
+            @testset "NamedTuple" begin
+                n1 = (x=1, y=2)
+                n2 = (x=3,)
+                # Values from outer context should override inner one
+                ctx1 = ConditionContext(n1, ConditionContext(n2))
+                @test ctx1.values == (x=1, y=2)
+                # Check that the two ConditionContexts are collapsed 
+                @test childcontext(ctx1) isa DefaultContext
+                # Then test the nesting the other way round
+                ctx2 = ConditionContext(n2, ConditionContext(n1))
+                @test ctx2.values == (x=3, y=2)
+                @test childcontext(ctx2) isa DefaultContext
+            end
+
+            @testset "Dict" begin
+                # Same tests as NamedTuple above
+                d1 = Dict(@varname(x) => 1, @varname(y) => 2)
+                d2 = Dict(@varname(x) => 3)
+                ctx1 = ConditionContext(d1, ConditionContext(d2))
+                @test ctx1.values == Dict(@varname(x) => 1, @varname(y) => 2)
+                @test childcontext(ctx1) isa DefaultContext
+                ctx2 = ConditionContext(d2, ConditionContext(d1))
+                @test ctx2.values == Dict(@varname(x) => 3, @varname(y) => 2)
+                @test childcontext(ctx2) isa DefaultContext
+            end
+        end
+
+        @testset "decondition_context" begin
+            @testset "NamedTuple" begin
+                ctx = ConditionContext((x=1, y=2, z=3))
+                # Decondition all variables
+                @test decondition_context(ctx) isa DefaultContext
+                # Decondition only some variables
+                dctx = decondition_context(ctx, :x)
+                @test dctx isa ConditionContext
+                @test dctx.values == (y=2, z=3)
+                dctx = decondition_context(ctx, :y, :z)
+                @test dctx isa ConditionContext
+                @test dctx.values == (x=1,)
+                # Decondition all variables manually
+                @test decondition_context(ctx, :x, :y, :z) isa DefaultContext
+            end
+
+            @testset "Dict" begin
+                ctx = ConditionContext(Dict(@varname(x) => 1, @varname(y) => 2, @varname(z) => 3))
+                # Decondition all variables
+                @test decondition_context(ctx) isa DefaultContext
+                # Decondition only some variables
+                dctx = decondition_context(ctx, @varname(x))
+                @test dctx isa ConditionContext
+                @test dctx.values == Dict(@varname(y) => 2, @varname(z) => 3)
+                dctx = decondition_context(ctx, @varname(y), @varname(z))
+                @test dctx isa ConditionContext
+                @test dctx.values == Dict(@varname(x) => 1)
+                # Decondition all variables manually
+                @test decondition_context(ctx, @varname(x), @varname(y), @varname(z)) isa DefaultContext
+            end
+        end
+    end
+
     @testset "FixedContext" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
             retval = model()