Make PrefixContext contain a varname rather than symbol

Author: penelopeysm

docs/src/internals/submodel_condition.md

@@ -181,10 +181,10 @@ Putting all of the information so far together, what it means is that if we have
 using DynamicPPL: PrefixContext, ConditionContext, DefaultContext
 
 inner_ctx_with_outer_cond = ConditionContext(
-    Dict(@varname(a.x) => 1.0), PrefixContext{:a}(DefaultContext())
+    Dict(@varname(a.x) => 1.0), PrefixContext(@varname(a))
 )
-inner_ctx_with_inner_cond = PrefixContext{:a}(
-    ConditionContext(Dict(@varname(x) => 1.0), DefaultContext())
+inner_ctx_with_inner_cond = PrefixContext(
+    @varname(a), ConditionContext(Dict(@varname(x) => 1.0))
 )
 ```
 
@@ -252,10 +252,11 @@ The general strategy that we adopt is similar to above.
 Following the principle that `PrefixContext` should be nested inside the outer context, but outside the inner submodel's context, we can infer that the correct context inside `charlie` should be:
 
 ```@example
-big_ctx = PrefixContext{:a}(
+big_ctx = PrefixContext(
+    @varname(a),
     ConditionContext(
         Dict(@varname(b.y) => 1.0),
-        PrefixContext{:b}(ConditionContext(Dict(@varname(x) => 1.0))),
+        PrefixContext(@varname(b), ConditionContext(Dict(@varname(x) => 1.0))),
     ),
 )
 ```
@@ -280,9 +281,9 @@ end
 function myprefix(::IsParent, ctx::AbstractContext, vn::VarName)
     return myprefix(childcontext(ctx), vn)
 end
-function myprefix(ctx::DynamicPPL.PrefixContext{Prefix}, vn::VarName) where {Prefix}
+function myprefix(ctx::DynamicPPL.PrefixContext, vn::VarName)
     # The functionality to actually manipulate the VarNames is in AbstractPPL
-    new_vn = AbstractPPL.prefix(vn, VarName{Prefix}())
+    new_vn = AbstractPPL.prefix(vn, ctx.vn_prefix)
     # Then pass to the child context
     return myprefix(childcontext(ctx), new_vn)
 end
@@ -295,11 +296,11 @@ This implementation clearly is not correct, because it applies the _inner_ `Pref
 The right way to implement `myprefix` is to, essentially, reverse the order of two lines above:
 
 ```@example
-function myprefix(ctx::DynamicPPL.PrefixContext{Prefix}, vn::VarName) where {Prefix}
+function myprefix(ctx::DynamicPPL.PrefixContext, vn::VarName)
     # Pass to the child context first
     new_vn = myprefix(childcontext(ctx), vn)
     # Then apply this context's prefix
-    return AbstractPPL.prefix(new_vn, VarName{Prefix}())
+    return AbstractPPL.prefix(new_vn, ctx.vn_prefix)
 end
 
 myprefix(big_ctx, @varname(x))

src/context_implementations.jl

@@ -131,7 +131,8 @@ function tilde_assume!!(context, right, vn, vi)
         # change in the future.
         if should_auto_prefix(right)
             dppl_model = right.model.model # This isa DynamicPPL.Model
-            prefixed_submodel_context = PrefixContext{getsym(vn)}(dppl_model.context)
+            # TODO: This does _not_ correctly prefix varnames with non-identity lenses
+            prefixed_submodel_context = PrefixContext(vn, dppl_model.context)
             new_dppl_model = contextualize(dppl_model, prefixed_submodel_context)
             right = to_submodel(new_dppl_model, true)
         end

src/contexts.jl

@@ -237,36 +237,43 @@ function setchildcontext(parent::MiniBatchContext, child)
 end
 
 """
-    PrefixContext{Prefix}(context)
+    PrefixContext(vn::VarName[, context::AbstractContext])
+    PrefixContext(vn::Val{sym}[, context::AbstractContext]) where {sym}
 
 Create a context that allows you to use the wrapped `context` when running the model and
-adds the `Prefix` to all parameters.
+prefixes all parameters with the VarName `vn`.
+
+`PrefixContext(Val(:a), context)` is equivalent to `PrefixContext(@varname(a), context)`.
+If `context` is not provided, it defaults to `DefaultContext()`.
 
 This context is useful in nested models to ensure that the names of the parameters are
 unique.
 
 See also: [`to_submodel`](@ref)
 """
-struct PrefixContext{Prefix,C} <: AbstractContext
+struct PrefixContext{Tvn<:VarName,C<:AbstractContext} <: AbstractContext
+    vn_prefix::Tvn
     context::C
 end
-function PrefixContext{Prefix}(context::AbstractContext) where {Prefix}
-    return PrefixContext{Prefix,typeof(context)}(context)
+PrefixContext(vn::VarName) = PrefixContext(vn, DefaultContext())
+function PrefixContext(::Val{sym}, context::AbstractContext) where {sym}
+    return PrefixContext(VarName{sym}(), context)
 end
+PrefixContext(::Val{sym}) where {sym} = PrefixContext(VarName{sym}())
 
 NodeTrait(::PrefixContext) = IsParent()
 childcontext(context::PrefixContext) = context.context
-function setchildcontext(::PrefixContext{Prefix}, child) where {Prefix}
-    return PrefixContext{Prefix}(child)
+function setchildcontext(ctx::PrefixContext, child::AbstractContext)
+    return PrefixContext(ctx.vn_prefix, child)
 end
 
 """
     prefix(ctx::AbstractContext, vn::VarName)
 
 Apply the prefixes in the context `ctx` to the variable name `vn`.
 """
-function prefix(ctx::PrefixContext{Prefix}, vn::VarName) where {Prefix}
-    return AbstractPPL.prefix(prefix(childcontext(ctx), vn), VarName{Prefix}())
+function prefix(ctx::PrefixContext, vn::VarName)
+    return AbstractPPL.prefix(prefix(childcontext(ctx), vn), ctx.vn_prefix)
 end
 function prefix(ctx::AbstractContext, vn::VarName)
     return prefix(NodeTrait(ctx), ctx, vn)
@@ -295,14 +302,13 @@ not_ need to modify any inner `ConditionContext`s and `FixedContext`s. If you
 _do_ need to modify them, then you may need to use
 `prefix_cond_and_fixed_variables` instead.
 """
-function prefix_and_strip_contexts(ctx::PrefixContext{Prefix}, vn::VarName) where {Prefix}
+function prefix_and_strip_contexts(ctx::PrefixContext, vn::VarName)
     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
+    return AbstractPPL.prefix(vn_prefixed, ctx.vn_prefix), child_context_without_prefixes
 end
 function prefix_and_strip_contexts(ctx::AbstractContext, vn::VarName)
     return prefix_and_strip_contexts(NodeTrait(ctx), ctx, vn)
@@ -314,11 +320,15 @@ function prefix_and_strip_contexts(::IsParent, ctx::AbstractContext, vn::VarName
 end
 
 """
-    prefix(model::Model, x)
-
-Return `model` but with all random variables prefixed by `x`.
+    prefix(model::Model, x::VarName)
+    prefix(model::Model, x::Val{sym})
+    prefix(model::Model, x::Symbol)  # Not recommended
 
-If `x` is known at compile-time, use `Val{x}()` to avoid runtime overheads for prefixing.
+Return `model` but with all random variables prefixed by `x`, where `x` is either:
+- a `VarName` (e.g. `@varname(a)`),
+- a `Val{sym}` (e.g. `Val(:a)`), or
+- a `Symbol` (e.g. `:a`). This last method is not recommended as it
+  introduces runtime overheads.
 
 # Examples
 
@@ -328,17 +338,19 @@ julia> using DynamicPPL: prefix
 julia> @model demo() = x ~ Dirac(1)
 demo (generic function with 2 methods)
 
-julia> rand(prefix(demo(), :my_prefix))
+julia> rand(prefix(demo(), @my_prefix))
 (var"my_prefix.x" = 1,)
 
-julia> # One can also use `Val` to avoid runtime overheads.
-       rand(prefix(demo(), Val(:my_prefix)))
+julia> rand(prefix(demo(), Val(:my_prefix)))
 (var"my_prefix.x" = 1,)
 ```
 """
-prefix(model::Model, x) = contextualize(model, PrefixContext{Symbol(x)}(model.context))
-function prefix(model::Model, ::Val{x}) where {x}
-    return contextualize(model, PrefixContext{Symbol(x)}(model.context))
+prefix(model::Model, x::VarName) = contextualize(model, PrefixContext(x, model.context))
+function prefix(model::Model, x::Symbol)
+    return contextualize(model, PrefixContext(VarName{x}(), model.context))
+end
+function prefix(model::Model, x::Val{sym}) where {sym}
+    return contextualize(model, PrefixContext(VarName{sym}(), model.context))
 end
 
 """
@@ -426,7 +438,7 @@ hasconditioned_nested(::IsLeaf, context, vn) = hasconditioned(context, vn)
 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}
+function hasconditioned_nested(context::PrefixContext, vn)
     return hasconditioned_nested(collapse_prefix_stack(context), vn)
 end
 
@@ -444,7 +456,7 @@ end
 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}
+function getconditioned_nested(context::PrefixContext, vn)
     return getconditioned_nested(collapse_prefix_stack(context), vn)
 end
 function getconditioned_nested(::IsParent, context, vn)
@@ -715,13 +727,13 @@ which explains this in much more detail.
 ```jldoctest
 julia> using DynamicPPL: collapse_prefix_stack
 
-julia> c1 = PrefixContext{:a}(ConditionContext((x=1, )));
+julia> c1 = PrefixContext(@varname(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,)))));
+       c2 = PrefixContext(@varname(a), ConditionContext((x=1, ), PrefixContext(@varname(b), ConditionContext((y=2,)))));
 
 julia> collapsed = collapse_prefix_stack(c2);
 
@@ -733,14 +745,14 @@ julia> # `collapsed` really looks something like this:
 (1, 2)
 ```
 """
-function collapse_prefix_stack(context::PrefixContext{Prefix}) where {Prefix}
+function collapse_prefix_stack(context::PrefixContext)
     # 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}())
+    return prefix_cond_and_fixed_variables(collapsed, context.vn_prefix)
 end
 function collapse_prefix_stack(context::AbstractContext)
     return collapse_prefix_stack(NodeTrait(collapse_prefix_stack, context), context)

src/model.jl

@@ -429,7 +429,7 @@ julia> # Nested ones also work.
        # (Note that `PrefixContext` also prefixes the variables of any
        # ConditionContext that is _inside_ it; because of this, the type of the
        # container has to be broadened to a `Dict`.)
-       cm = condition(contextualize(m, PrefixContext{:a}(ConditionContext((m=1.0,)))), x=100.0);
+       cm = condition(contextualize(m, PrefixContext(@varname(a), ConditionContext((m=1.0,)))), x=100.0);
 
 julia> Set(keys(conditioned(cm))) == Set([@varname(a.m), @varname(x)])
 true
@@ -441,7 +441,7 @@ julia> # Since we conditioned on `a.m`, it is not treated as a random variable.
  a.x
 
 julia> # We can also condition on `a.m` _outside_ of the PrefixContext:
-       cm = condition(contextualize(m, PrefixContext{:a}(DefaultContext())), (@varname(a.m) => 1.0));
+       cm = condition(contextualize(m, PrefixContext(@varname(a))), (@varname(a.m) => 1.0));
 
 julia> conditioned(cm)
 Dict{VarName{:a, Accessors.PropertyLens{:m}}, Float64} with 1 entry:
@@ -769,7 +769,7 @@ julia> # Returns all the variables we have fixed on + their values.
 (x = 100.0, m = 1.0)
 
 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);
+       cm = fix(contextualize(m, PrefixContext(@varname(a), fix(m=1.0))), x=100.0);
 
 julia> Set(keys(fixed(cm))) == Set([@varname(a.m), @varname(x)])
 true
@@ -779,7 +779,7 @@ julia> keys(VarInfo(cm))
  a.x
 
 julia> # We can also condition on `a.m` _outside_ of the PrefixContext:
-       cm = fix(contextualize(m, PrefixContext{:a}(DefaultContext())), (@varname(a.m) => 1.0));
+       cm = fix(contextualize(m, PrefixContext(@varname(a))), (@varname(a.m) => 1.0));
 
 julia> fixed(cm)
 Dict{VarName{:a, Accessors.PropertyLens{:m}}, Float64} with 1 entry:

src/submodel_macro.jl

@@ -223,12 +223,12 @@ end
 prefix_submodel_context(prefix, left, ctx) = prefix_submodel_context(prefix, ctx)
 function prefix_submodel_context(prefix, ctx)
     # E.g. `prefix="asd[$i]"` or `prefix=asd` with `asd` to be evaluated.
-    return :($(PrefixContext){$(Symbol)($(esc(prefix)))}($ctx))
+    return :($(PrefixContext)($(Symbol)($(esc(prefix))), $ctx))
 end
 
 function prefix_submodel_context(prefix::Union{AbstractString,Symbol}, ctx)
     # E.g. `prefix="asd"`.
-    return :($(PrefixContext){$(esc(Meta.quot(Symbol(prefix))))}($ctx))
+    return :($(PrefixContext)($(esc(Meta.quot(Symbol(prefix)))), $ctx))
 end
 
 function prefix_submodel_context(prefix::Bool, ctx)