Fix doctests

Author: penelopeysm

src/contexts.jl

@@ -502,7 +502,7 @@ function conditioned(context::ConditionContext)
     # precedence over decendants of `context`.
     return _merge(context.values, conditioned(childcontext(context)))
 end
-function conditioned(context::PrefixContext{Prefix}) where {Prefix}
+function conditioned(context::PrefixContext)
     return conditioned(collapse_prefix_stack(context))
 end
 
@@ -681,6 +681,9 @@ function fixed(context::FixedContext)
     # precedence over decendants of `context`.
     return _merge(context.values, fixed(childcontext(context)))
 end
+function fixed(context::PrefixContext)
+    return fixed(collapse_prefix_stack(context))
+end
 
 """
     collapse_prefix_stack(context::AbstractContext)
@@ -691,7 +694,22 @@ the `PrefixContext`s from the context stack.
 ```jldoctest
 julia> using DynamicPPL: collapse_prefix_stack
 
-julia> c1 = PrefixContext({:a}(ConditionContext((x=1, )))
+julia> c1 = PrefixContext{:a}(ConditionContext((x=1, )));
+
+julia> collapse_prefix_stack(c1)
+ConditionContext(Dict(a.x => 1), DefaultContext())
+
+julia> # Here, `x` gets prefixed only with `a`, whereas `y` is prefixed with both.
+       c2 = PrefixContext{:a}(ConditionContext((x=1, ), PrefixContext{:b}(ConditionContext((y=2,)))));
+
+julia> collapsed = collapse_prefix_stack(c2);
+
+julia> # `collapsed` really looks something like this:
+       # ConditionContext(Dict{VarName{:a}, Int64}(a.b.y => 2, a.x => 1), DefaultContext())
+       # To avoid fragility arising from the order of the keys in the doctest, we test
+       # this indirectly:
+       collapsed.values[@varname(a.x)], collapsed.values[@varname(a.b.y)]
+(1, 2)
 ```
 """
 function collapse_prefix_stack(context::PrefixContext{Prefix}) where {Prefix}
@@ -703,7 +721,14 @@ function collapse_prefix_stack(context::PrefixContext{Prefix}) where {Prefix}
     # depth of the context stack.
     return prefix_cond_and_fixed_variables(collapsed, VarName{Prefix}())
 end
-collapse_prefix_stack(context::AbstractContext) = context
+function collapse_prefix_stack(context::AbstractContext)
+    return collapse_prefix_stack(NodeTrait(collapse_prefix_stack, context), context)
+end
+collapse_prefix_stack(::IsLeaf, context) = context
+function collapse_prefix_stack(::IsParent, context)
+    new_child_context = collapse_prefix_stack(childcontext(context))
+    return setchildcontext(context, new_child_context)
+end
 
 """
     prefix_cond_and_fixed_variables(context::AbstractContext, prefix::VarName)

src/model.jl

@@ -431,10 +431,8 @@ julia> # Nested ones also work.
        # container has to be broadened to a `Dict`.)
        cm = condition(contextualize(m, PrefixContext{:a}(ConditionContext((m=1.0,)))), x=100.0);
 
-julia> conditioned(cm)
-Dict{VarName, Any} with 2 entries:
-  a.m => 1.0
-  x   => 100.0
+julia> Set(keys(conditioned(cm))) == Set([@varname(a.m), @varname(x)])
+true
 
 julia> # Since we conditioned on `a.m`, it is not treated as a random variable.
        # However, `a.x` will still be a random variable.
@@ -770,29 +768,27 @@ julia> # Returns all the variables we have fixed on + their values.
        fixed(fix(m, x=100.0, m=1.0))
 (x = 100.0, m = 1.0)
 
-julia> # Nested ones also work (note that `PrefixContext` does nothing to the result).
+julia> # The rest of this is the same as the `condition` example above.
        cm = fix(contextualize(m, PrefixContext{:a}(fix(m=1.0))), x=100.0);
 
-julia> fixed(cm)
-(x = 100.0, m = 1.0)
-
-julia> # Since we fixed on `m`, not `a.m` as it will appear after prefixed,
-       # `a.m` is treated as a random variable.
-       keys(VarInfo(cm))
-1-element Vector{VarName{:a, Accessors.PropertyLens{:m}}}:
- a.m
+julia> Set(keys(fixed(cm))) == Set([@varname(a.m), @varname(x)])
+true
 
-julia> # If we instead fix on `a.m`, `m` in the model will be considered an observation.
-       cm = fix(contextualize(m, PrefixContext{:a}(fix(@varname(a.m) => 1.0,))), x=100.0);
+julia> keys(VarInfo(cm))
+1-element Vector{VarName{:a, Accessors.PropertyLens{:x}}}:
+ a.x
 
-julia> fixed(cm)[@varname(x)]
-100.0
+julia> # We can also condition on `a.m` _outside_ of the PrefixContext:
+       cm = fix(contextualize(m, PrefixContext{:a}(DefaultContext())), (@varname(a.m) => 1.0));
 
-julia> fixed(cm)[@varname(a.m)]
-1.0
+julia> fixed(cm)
+Dict{VarName{:a, Accessors.PropertyLens{:m}}, Float64} with 1 entry:
+  a.m => 1.0
 
-julia> keys(VarInfo(cm)) # <= no variables are sampled
-VarName[]
+julia> # Now `a.x` will be sampled.
+       keys(VarInfo(cm))
+1-element Vector{VarName{:a, Accessors.PropertyLens{:x}}}:
+ a.x
 ```
 """
 fixed(model::Model) = fixed(model.context)