AbstractPPL 0.11 + change prefixing behaviour (#830)

* AbstractPPL 0.11; change prefixing behaviour

* Use DynamicPPL.prefix rather than overloading

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"

docs/Project.toml

@@ -4,6 +4,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177"
+DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"

docs/src/api.md

@@ -149,7 +149,7 @@ In the past, one would instead embed sub-models using [`@submodel`](@ref), which
 In the context of including models within models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
 
 ```@docs
-prefix
+DynamicPPL.prefix
 ```
 
 Under the hood, [`to_submodel`](@ref) makes use of the following method to indicate that the model it's wrapping is a model over its return-values rather than something else

src/DynamicPPL.jl

@@ -21,6 +21,9 @@ using DocStringExtensions
 
 using Random: Random
 
+# For extending
+import AbstractPPL: predict
+
 # 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/contexts.jl

@@ -260,25 +260,21 @@ 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(ctx::AbstractContext, vn::VarName)
 
+Apply the prefixes in the context `ctx` to the variable name `vn`.
+"""
 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)
-    )
+    return AbstractPPL.prefix(prefix(childcontext(ctx), vn), VarName{Symbol(Prefix)}())
+end
+function prefix(ctx::AbstractContext, vn::VarName)
+    return prefix(NodeTrait(ctx), ctx, vn)
 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(::IsParent, ctx::AbstractContext, vn::VarName)
+    return prefix(childcontext(ctx), vn)
+end
 
 """
     prefix(model::Model, x)

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(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);
 

src/submodel_macro.jl

@@ -96,10 +96,10 @@ julia> vi = VarInfo(demo2(missing, missing, 0.4));
 │   caller = ip:0x0
 └ @ Core :-1
 
-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
 ```
 
@@ -116,9 +116,9 @@ false
 We can check that the log joint probability of the model accumulated in `vi` is correct:
 
 ```jldoctest submodelprefix
-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);
 
@@ -157,7 +157,7 @@ julia> # Automatically determined from `a`.
        @model submodel_prefix_true() = @submodel prefix=true a = inner()
 submodel_prefix_true (generic function with 2 methods)
 
-julia> @varname(var"a.x") in keys(VarInfo(submodel_prefix_true()))
+julia> @varname(a.x) in keys(VarInfo(submodel_prefix_true()))
 ┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
@@ -167,7 +167,7 @@ julia> # Using a static string.
        @model submodel_prefix_string() = @submodel prefix="my prefix" a = inner()
 submodel_prefix_string (generic function with 2 methods)
 
-julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+julia> @varname(var"my prefix".x) in keys(VarInfo(submodel_prefix_string()))
 ┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
@@ -177,7 +177,7 @@ julia> # Using string interpolation.
        @model submodel_prefix_interpolation() = @submodel prefix="\$(nameof(inner()))" a = inner()
 submodel_prefix_interpolation (generic function with 2 methods)
 
-julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+julia> @varname(inner.x) in keys(VarInfo(submodel_prefix_interpolation()))
 ┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1
@@ -187,7 +187,7 @@ julia> # Or using some arbitrary expression.
        @model submodel_prefix_expr() = @submodel prefix=1 + 2 a = inner()
 submodel_prefix_expr (generic function with 2 methods)
 
-julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+julia> @varname(var"3".x) in keys(VarInfo(submodel_prefix_expr()))
 ┌ Warning: `@submodel model` and `@submodel prefix=... model` are deprecated; see `to_submodel` for the up-to-date syntax.
 │   caller = ip:0x0
 └ @ Core :-1

src/utils.jl

@@ -1285,14 +1285,18 @@ broadcast_safe(x) = x
 broadcast_safe(x::Distribution) = (x,)
 broadcast_safe(x::AbstractContext) = (x,)
 
+# Convert (x=1,) to Dict(@varname(x) => 1)
+_nt_to_varname_dict(nt) = Dict(VarName{k}() => v for (k, v) in pairs(nt))
 # Version of `merge` used by `conditioned` and `fixed` to handle
 # the scenario where we might try to merge a dict with an empty
 # tuple.
 # TODO: Maybe replace the default of returning `NamedTuple` with `nothing`?
 _merge(left::NamedTuple, right::NamedTuple) = merge(left, right)
 _merge(left::AbstractDict, right::AbstractDict) = merge(left, right)
-_merge(left::AbstractDict, right::NamedTuple{()}) = left
-_merge(left::NamedTuple{()}, right::AbstractDict) = right
+_merge(left::AbstractDict, ::NamedTuple{()}) = left
+_merge(left::AbstractDict, right::NamedTuple) = merge(left, _nt_to_varname_dict(right))
+_merge(::NamedTuple{()}, right::AbstractDict) = right
+_merge(left::NamedTuple, right::AbstractDict) = merge(_nt_to_varname_dict(left), right)
 
 """
     unique_syms(vns::T) where {T<:NTuple{N,VarName}}

test/Project.toml

@@ -32,7 +32,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.10.1"
+AbstractPPL = "0.11"
 Accessors = "0.1"
 Aqua = "0.8"
 Bijectors = "0.15.1"