AbstractPPL 0.11; change prefixing behaviour

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,54 @@
 # DynamicPPL Changelog
 
+## 0.36.0
+
+**Breaking changes**
+
+### VarName prefixing behaviour
+
+The way in which VarNames in submodels are prefixed has been changed.
+This is best explained through an example.
+Consider this model and submodel:
+
+```julia
+using DynamicPPL, Distributions
+@model inner() = x ~ Normal()
+@model outer() = a ~ to_submodel(inner())
+```
+
+In previous versions, the inner variable `x` would be saved as `a.x`.
+However, this was represented as a single symbol `Symbol("a.x")`:
+
+```julia
+julia> dump(keys(VarInfo(outer()))[1])
+VarName{Symbol("a.x"), typeof(identity)}
+  optic: identity (function of type typeof(identity))
+```
+
+Now, the inner variable is stored as a field `x` on the VarName `a`:
+
+```julia
+julia> dump(keys(VarInfo(outer()))[1])
+VarName{:a, Accessors.PropertyLens{:x}}
+  optic: Accessors.PropertyLens{:x} (@o _.x)
+```
+
+In practice, this means that if you are trying to condition a variable in the submodel, you now need to use
+
+```julia
+outer() | (@varname(a.x) => 1.0,)
+```
+
+instead of either of these (which would have worked previously)
+
+```julia
+outer() | (@varname(var"a.x") => 1.0,)
+outer() | (a.x = 1.0,)
+```
+
+If you are sampling from a model with submodels, this doesn't affect the way you interact with the `MCMCChains.Chains` object, because VarNames are converted into Symbols when stored in the chain.
+(This behaviour will likely be changed in the future, in that Chains should be indexable by VarNames and not just Symbols, but that has not been implemented yet.)
+
 ## 0.35.5
 
 Several internal methods have been removed:

Project.toml

@@ -44,7 +44,7 @@ DynamicPPLMooncakeExt = ["Mooncake"]
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.10.1"
+AbstractPPL = "0.11"
 Accessors = "0.1"
 BangBang = "0.4.1"
 Bijectors = "0.13.18, 0.14, 0.15"

src/DynamicPPL.jl

@@ -21,6 +21,9 @@ using DocStringExtensions
 
 using Random: Random
 
+# For extending
+import AbstractPPL: predict, prefix
+
 # TODO: Remove these when it's possible.
 import Bijectors: link, invlink
 
@@ -39,8 +42,6 @@ import Base:
     keys,
     haskey
 
-import AbstractPPL: predict
-
 # VarInfo
 export AbstractVarInfo,
     VarInfo,

src/compiler.jl

@@ -113,7 +113,7 @@ function contextual_isassumption(context::ConditionContext, vn)
     return contextual_isassumption(childcontext(context), vn)
 end
 function contextual_isassumption(context::PrefixContext, vn)
-    return contextual_isassumption(childcontext(context), prefix(context, vn))
+    return contextual_isassumption(childcontext(context), prefix_with_context(context, vn))
 end
 
 isfixed(expr, vn) = false
@@ -132,7 +132,7 @@ function contextual_isfixed(context::AbstractContext, vn)
     return contextual_isfixed(NodeTrait(context), context, vn)
 end
 function contextual_isfixed(context::PrefixContext, vn)
-    return contextual_isfixed(childcontext(context), prefix(context, vn))
+    return contextual_isfixed(childcontext(context), prefix_with_context(context, vn))
 end
 function contextual_isfixed(context::FixedContext, vn)
     if hasfixed(context, vn)

src/context_implementations.jl

@@ -85,12 +85,14 @@ function tilde_assume(rng::Random.AbstractRNG, ::LikelihoodContext, sampler, rig
 end
 
 function tilde_assume(context::PrefixContext, right, vn, vi)
-    return tilde_assume(context.context, right, prefix(context, vn), vi)
+    return tilde_assume(context.context, right, prefix_with_context(context, vn), vi)
 end
 function tilde_assume(
     rng::Random.AbstractRNG, context::PrefixContext, sampler, right, vn, vi
 )
-    return tilde_assume(rng, context.context, sampler, right, prefix(context, vn), vi)
+    return tilde_assume(
+        rng, context.context, sampler, right, prefix_with_context(context, vn), vi
+    )
 end
 
 """

src/contexts.jl

@@ -260,25 +260,25 @@ function setchildcontext(::PrefixContext{Prefix}, child) where {Prefix}
     return PrefixContext{Prefix}(child)
 end
 
-const PREFIX_SEPARATOR = Symbol(".")
-
-@generated function PrefixContext{PrefixOuter}(
-    context::PrefixContext{PrefixInner}
-) where {PrefixOuter,PrefixInner}
-    return :(PrefixContext{$(QuoteNode(Symbol(PrefixOuter, PREFIX_SEPARATOR, PrefixInner)))}(
-        context.context
-    ))
-end
+"""
+    prefix_with_context(ctx::AbstractContext, vn::VarName)
 
-function prefix(ctx::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
-    vn_prefixed_inner = prefix(childcontext(ctx), vn)
-    return VarName{Symbol(Prefix, PREFIX_SEPARATOR, getsym(vn_prefixed_inner))}(
-        getoptic(vn_prefixed_inner)
+Apply the prefixes in the context `ctx` to the variable name `vn`.
+"""
+function prefix_with_context(
+    ctx::PrefixContext{Prefix}, vn::VarName{Sym}
+) where {Prefix,Sym}
+    return AbstractPPL.prefix(
+        prefix_with_context(childcontext(ctx), vn), VarName{Symbol(Prefix)}()
     )
 end
-prefix(ctx::AbstractContext, vn::VarName) = prefix(NodeTrait(ctx), ctx, vn)
-prefix(::IsLeaf, ::AbstractContext, vn::VarName) = vn
-prefix(::IsParent, ctx::AbstractContext, vn::VarName) = prefix(childcontext(ctx), vn)
+function prefix_with_context(ctx::AbstractContext, vn::VarName)
+    return prefix_with_context(NodeTrait(ctx), ctx, vn)
+end
+prefix_with_context(::IsLeaf, ::AbstractContext, vn::VarName) = vn
+function prefix_with_context(::IsParent, ctx::AbstractContext, vn::VarName)
+    return prefix_with_context(childcontext(ctx), vn)
+end
 
 """
     prefix(model::Model, x)
@@ -392,7 +392,7 @@ function hasconditioned_nested(::IsParent, context, vn)
     return hasconditioned(context, vn) || hasconditioned_nested(childcontext(context), vn)
 end
 function hasconditioned_nested(context::PrefixContext, vn)
-    return hasconditioned_nested(childcontext(context), prefix(context, vn))
+    return hasconditioned_nested(childcontext(context), prefix_with_context(context, vn))
 end
 
 """
@@ -410,7 +410,7 @@ function getconditioned_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
 function getconditioned_nested(context::PrefixContext, vn)
-    return getconditioned_nested(childcontext(context), prefix(context, vn))
+    return getconditioned_nested(childcontext(context), prefix_with_context(context, vn))
 end
 function getconditioned_nested(::IsParent, context, vn)
     return if hasconditioned(context, vn)
@@ -543,7 +543,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(childcontext(context), prefix_with_context(context, vn))
 end
 
 """
@@ -561,7 +561,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(childcontext(context), prefix_with_context(context, vn))
 end
 function getfixed_nested(::IsParent, context, vn)
     return if hasfixed(context, vn)

src/debug_utils.jl

@@ -183,7 +183,7 @@ function DynamicPPL.setchildcontext(context::DebugContext, child)
 end
 
 function record_varname!(context::DebugContext, varname::VarName, dist)
-    prefixed_varname = prefix(context, varname)
+    prefixed_varname = DynamicPPL.prefix_with_context(context, varname)
     if haskey(context.varnames_seen, prefixed_varname)
         if context.error_on_failure
             error("varname $prefixed_varname used multiple times in model")

src/model.jl

@@ -243,7 +243,7 @@ julia> model() ≠ 1.0
 true
 
 julia> # To condition the variable inside `demo_inner` we need to refer to it as `inner.m`.
-       conditioned_model = model | (var"inner.m" = 1.0, );
+       conditioned_model = model | (@varname(inner.m) => 1.0, );
 
 julia> conditioned_model()
 1.0
@@ -255,15 +255,6 @@ julia> conditioned_model_fail()
 ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
 [...]
 ```
-
-And similarly when using `Dict`:
-
-```jldoctest condition
-julia> conditioned_model_dict = model | (@varname(var"inner.m") => 1.0);
-
-julia> conditioned_model_dict()
-1.0
-```
 """
 function AbstractPPL.condition(model::Model, values...)
     # Positional arguments - need to handle cases carefully
@@ -443,16 +434,16 @@ julia> conditioned(cm)
 julia> # Since we conditioned 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{Symbol("a.m"), typeof(identity)}}:
+1-element Vector{VarName{:a, Accessors.PropertyLens{:m}}}:
  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}(ConditionContext((var"a.m"=1.0,)))), x=100.0);
+       cm = condition(contextualize(m, PrefixContext{:a}(ConditionContext(Dict(@varname(a.m) => 1.0)))), x=100.0);
 
-julia> conditioned(cm).x
+julia> conditioned(cm)[@varname(x)]
 100.0
 
-julia> conditioned(cm).var"a.m"
+julia> conditioned(cm)[@varname(a.m)]
 1.0
 
 julia> keys(VarInfo(cm)) # No variables are sampled
@@ -583,7 +574,7 @@ julia> model = demo_outer();
 julia> model() ≠ 1.0
 true
 
-julia> fixed_model = fix(model, var"inner.m" = 1.0, );
+julia> fixed_model = fix(model, (@varname(inner.m) => 1.0, ));
 
 julia> fixed_model()
 1.0
@@ -599,24 +590,9 @@ julia> fixed_model()
 2.0
 ```
 
-And similarly when using `Dict`:
-
-```jldoctest fix
-julia> fixed_model_dict = fix(model, @varname(var"inner.m") => 1.0);
-
-julia> fixed_model_dict()
-1.0
-
-julia> fixed_model_dict = fix(model, @varname(inner) => 2.0);
-
-julia> fixed_model_dict()
-2.0
-```
-
 ## Difference from `condition`
 
-A very similar functionality is also provided by [`condition`](@ref) which,
-not surprisingly, _conditions_ variables instead of fixing them. The only
+A very similar functionality is also provided by [`condition`](@ref). The only
 difference between fixing and conditioning is as follows:
 - `condition`ed variables are considered to be observations, and are thus
   included in the computation [`logjoint`](@ref) and [`loglikelihood`](@ref),
@@ -798,16 +774,16 @@ julia> fixed(cm)
 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{Symbol("a.m"), typeof(identity)}}:
+1-element Vector{VarName{:a, Accessors.PropertyLens{:m}}}:
  a.m
 
 julia> # If we instead fix on `a.m`, `m` in the model will be considered an observation.
-       cm = fix(contextualize(m, PrefixContext{:a}(fix(var"a.m"=1.0))), x=100.0);
+       cm = fix(contextualize(m, PrefixContext{:a}(fix(@varname(a.m) => 1.0,))), x=100.0);
 
-julia> fixed(cm).x
+julia> fixed(cm)[@varname(x)]
 100.0
 
-julia> fixed(cm).var"a.m"
+julia> fixed(cm)[@varname(a.m)]
 1.0
 
 julia> keys(VarInfo(cm)) # <= no variables are sampled
@@ -1365,7 +1341,7 @@ When we sample from the model `demo2(missing, 0.4)` random variable `x` will be
 ```jldoctest submodel-to_submodel
 julia> vi = VarInfo(demo2(missing, 0.4));
 
-julia> @varname(var\"a.x\") in keys(vi)
+julia> @varname(a.x) in keys(vi)
 true
 ```
 
@@ -1379,7 +1355,7 @@ false
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
 ```jldoctest submodel-to_submodel
-julia> x = vi[@varname(var\"a.x\")];
+julia> x = vi[@varname(a.x)];
 
 julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
 true
@@ -1417,10 +1393,10 @@ julia> @model function demo2(x, y, z)
 
 julia> vi = VarInfo(demo2(missing, missing, 0.4));
 
-julia> @varname(var"sub1.x") in keys(vi)
+julia> @varname(sub1.x) in keys(vi)
 true
 
-julia> @varname(var"sub2.x") in keys(vi)
+julia> @varname(sub2.x) in keys(vi)
 true
 ```
 
@@ -1437,9 +1413,9 @@ false
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
 ```jldoctest submodel-to_submodel-prefix
-julia> sub1_x = vi[@varname(var"sub1.x")];
+julia> sub1_x = vi[@varname(sub1.x)];
 
-julia> sub2_x = vi[@varname(var"sub2.x")];
+julia> sub2_x = vi[@varname(sub2.x)];
 
 julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);