InitContext, part 3 - Introduce InitContext

docs/src/api.md

@@ -463,6 +463,27 @@ SamplingContext
 DefaultContext
 PrefixContext
 ConditionContext
+InitContext
+```
+
+### VarInfo initialisation
+
+`InitContext` is used to initialise, or overwrite, values in a VarInfo.
+
+To accomplish this, an initialisation _strategy_ is required, which defines how new values are to be obtained.
+There are three concrete strategies provided in DynamicPPL:
+
+```@docs
+PriorInit
+UniformInit
+ParamsInit
+```
+
+If you wish to write your own, you have to subtype [`DynamicPPL.AbstractInitStrategy`](@ref) and implement the `init` method.
+
+```@docs
+DynamicPPL.AbstractInitStrategy
+DynamicPPL.init
 ```
 
 ### Samplers

src/DynamicPPL.jl

@@ -108,6 +108,12 @@ export AbstractVarInfo,
     ConditionContext,
     assume,
     tilde_assume,
+    # Initialisation
+    InitContext,
+    AbstractInitStrategy,
+    PriorInit,
+    UniformInit,
+    ParamsInit,
     # Pseudo distributions
     NamedDist,
     NoDist,
@@ -169,11 +175,12 @@ abstract type AbstractVarInfo <: AbstractModelTrace end
 # Necessary forward declarations
 include("utils.jl")
 include("chains.jl")
+include("contexts.jl")
+include("contexts/init.jl")
 include("model.jl")
 include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")
-include("contexts.jl")
 include("submodel.jl")
 include("varnamedvector.jl")
 include("accumulators.jl")

src/contexts.jl

@@ -280,41 +280,6 @@ function prefix_and_strip_contexts(::IsParent, ctx::AbstractContext, vn::VarName
     return vn, setchildcontext(ctx, new_ctx)
 end
 
-"""
-    prefix(model::Model, x::VarName)
-    prefix(model::Model, x::Val{sym})
-    prefix(model::Model, x::Any)
-
-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
-- for any other type, `x` is converted to a Symbol and then to a `VarName`. Note that
-  this will introduce runtime overheads so is not recommended unless absolutely
-  necessary.
-
-# Examples
-
-```jldoctest
-julia> using DynamicPPL: prefix
-
-julia> @model demo() = x ~ Dirac(1)
-demo (generic function with 2 methods)
-
-julia> rand(prefix(demo(), @varname(my_prefix)))
-(var"my_prefix.x" = 1,)
-
-julia> rand(prefix(demo(), Val(:my_prefix)))
-(var"my_prefix.x" = 1,)
-```
-"""
-prefix(model::Model, x::VarName) = contextualize(model, PrefixContext(x, model.context))
-function prefix(model::Model, x::Val{sym}) where {sym}
-    return contextualize(model, PrefixContext(VarName{sym}(), model.context))
-end
-function prefix(model::Model, x)
-    return contextualize(model, PrefixContext(VarName{Symbol(x)}(), model.context))
-end
-
 """
 
     ConditionContext{Values<:Union{NamedTuple,AbstractDict},Ctx<:AbstractContext}

src/contexts/init.jl

@@ -0,0 +1,176 @@
+"""
+    AbstractInitStrategy
+
+Abstract type representing the possible ways of initialising new values for
+the random variables in a model (e.g., when creating a new VarInfo).
+"""
+abstract type AbstractInitStrategy end
+
+"""
+    init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, strategy::AbstractInitStrategy)
+
+Generate a new value for a random variable with the given distribution.
+
+!!! warning "Values must be unlinked"
+    The values returned by `init` are always in the untransformed space, i.e.,
+    they must be within the support of the original distribution. That means that,
+    for example, `init(rng, dist, u::UniformInit)` will in general return values that
+    are outside the range [u.lower, u.upper].
+"""
+function init end
+
+"""
+    PriorInit()
+
+Obtain new values by sampling from the prior distribution.
+"""
+struct PriorInit <: AbstractInitStrategy end
+init(rng::Random.AbstractRNG, ::VarName, dist::Distribution, ::PriorInit) = rand(rng, dist)
+
+"""
+    UniformInit()
+    UniformInit(lower, upper)
+
+Obtain new values by first transforming the distribution of the random variable
+to unconstrained space, then sampling a value uniformly between `lower` and
+`upper`, and transforming that value back to the original space.
+
+If `lower` and `upper` are unspecified, they default to `(-2, 2)`, which mimics
+Stan's default initialisation strategy.
+
+Requires that `lower <= upper`.
+
+# References
+
+[Stan reference manual page on initialization](https://mc-stan.org/docs/reference-manual/execution.html#initialization)
+"""
+struct UniformInit{T<:AbstractFloat} <: AbstractInitStrategy
+    lower::T
+    upper::T
+    function UniformInit(lower::T, upper::T) where {T<:AbstractFloat}
+        lower > upper &&
+            throw(ArgumentError("`lower` must be less than or equal to `upper`"))
+        return new{T}(lower, upper)
+    end
+    UniformInit() = UniformInit(-2.0, 2.0)
+end
+function init(rng::Random.AbstractRNG, ::VarName, dist::Distribution, u::UniformInit)
+    b = Bijectors.bijector(dist)
+    sz = Bijectors.output_size(b, size(dist))
+    y = u.lower .+ ((u.upper - u.lower) .* rand(rng, sz...))
+    b_inv = Bijectors.inverse(b)
+    x = b_inv(y)
+    # 0-dim arrays: https://github.com/TuringLang/Bijectors.jl/issues/398
+    if x isa Array{<:Any,0}
+        x = x[]
+    end
+    return x
+end
+
+"""
+    ParamsInit(params::AbstractDict{<:VarName}, default::AbstractInitStrategy=PriorInit())
+    ParamsInit(params::NamedTuple, default::AbstractInitStrategy=PriorInit())
+
+Obtain new values by extracting them from the given dictionary or NamedTuple.
+The parameter `default` specifies how new values are to be obtained if they
+cannot be found in `params`, or they are specified as `missing`. The default
+for `default` is `PriorInit()`.
+
+!!! note
+    These values must be provided in the space of the untransformed distribution.
+"""
+struct ParamsInit{P,S<:AbstractInitStrategy} <: AbstractInitStrategy
+    params::P
+    default::S
+    function ParamsInit(params::AbstractDict{<:VarName}, default::AbstractInitStrategy)
+        return new{typeof(params),typeof(default)}(params, default)
+    end
+    ParamsInit(params::AbstractDict{<:VarName}) = ParamsInit(params, PriorInit())
+    function ParamsInit(params::NamedTuple, default::AbstractInitStrategy=PriorInit())
+        return ParamsInit(to_varname_dict(params), default)
+    end
+end
+function init(rng::Random.AbstractRNG, vn::VarName, dist::Distribution, p::ParamsInit)
+    # TODO(penelopeysm): It would be nice to do a check to make sure that all
+    # of the parameters in `p.params` were actually used, and either warn or
+    # error if they aren't. This is actually quite non-trivial though because
+    # the structure of Dicts in particular can have arbitrary nesting.
+    return if hasvalue(p.params, vn, dist)
+        x = getvalue(p.params, vn, dist)
+        if x === missing
+            init(rng, vn, dist, p.default)
+        else
+            # TODO(penelopeysm): Since x is user-supplied, maybe we could also
+            # check here that the type / size of x matches the dist?
+            x
+        end
+    else
+        init(rng, vn, dist, p.default)
+    end
+end
+
+"""
+    InitContext(
+            [rng::Random.AbstractRNG=Random.default_rng()],
+            [strategy::AbstractInitStrategy=PriorInit()],
+    )
+
+A leaf context that indicates that new values for random variables are
+currently being obtained through sampling. Used e.g. when initialising a fresh
+VarInfo. Note that, if `leafcontext(model.context) isa InitContext`, then
+`evaluate!!(model, varinfo)` will override all values in the VarInfo.
+"""
+struct InitContext{R<:Random.AbstractRNG,S<:AbstractInitStrategy} <: AbstractContext
+    rng::R
+    strategy::S
+    function InitContext(
+        rng::Random.AbstractRNG, strategy::AbstractInitStrategy=PriorInit()
+    )
+        return new{typeof(rng),typeof(strategy)}(rng, strategy)
+    end
+    function InitContext(strategy::AbstractInitStrategy=PriorInit())
+        return InitContext(Random.default_rng(), strategy)
+    end
+end
+NodeTrait(::InitContext) = IsLeaf()
+
+function tilde_assume(
+    ctx::InitContext, dist::Distribution, vn::VarName, vi::AbstractVarInfo
+)
+    in_varinfo = haskey(vi, vn)
+    # `init()` always returns values in original space, i.e. possibly
+    # constrained
+    x = init(ctx.rng, vn, dist, ctx.strategy)
+    # Determine whether to insert a transformed value into the VarInfo.
+    # If the VarInfo alrady had a value for this variable, we will
+    # keep the same linked status as in the original VarInfo. If not, we
+    # check the rest of the VarInfo to see if other variables are linked.
+    # istrans(vi) returns true if vi is nonempty and all variables in vi
+    # are linked.
+    insert_transformed_value = in_varinfo ? istrans(vi, vn) : istrans(vi)
+    f = if insert_transformed_value
+        link_transform(dist)
+    else
+        identity
+    end
+    y, logjac = with_logabsdet_jacobian(f, x)
+    # Add the new value to the VarInfo. `push!!` errors if the value already
+    # exists, hence the need for setindex!!.
+    if in_varinfo
+        vi = setindex!!(vi, y, vn)
+    else
+        vi = push!!(vi, vn, y, dist)
+    end
+    # Neither of these set the `trans` flag so we have to do it manually if
+    # necessary.
+    insert_transformed_value && settrans!!(vi, true, vn)
+    # `accumulate_assume!!` wants untransformed values as the second argument.
+    vi = accumulate_assume!!(vi, x, logjac, vn, dist)
+    # We always return the untransformed value here, as that will determine
+    # what the lhs of the tilde-statement is set to.
+    return x, vi
+end
+
+function tilde_observe!!(::InitContext, right, left, vn, vi)
+    return tilde_observe!!(DefaultContext(), right, left, vn, vi)
+end

src/model.jl

@@ -799,6 +799,41 @@ julia> # Now `a.x` will be sampled.
 """
 fixed(model::Model) = fixed(model.context)
 
+"""
+    prefix(model::Model, x::VarName)
+    prefix(model::Model, x::Val{sym})
+    prefix(model::Model, x::Any)
+
+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
+- for any other type, `x` is converted to a Symbol and then to a `VarName`. Note that
+  this will introduce runtime overheads so is not recommended unless absolutely
+  necessary.
+
+# Examples
+
+```jldoctest
+julia> using DynamicPPL: prefix
+
+julia> @model demo() = x ~ Dirac(1)
+demo (generic function with 2 methods)
+
+julia> rand(prefix(demo(), @varname(my_prefix)))
+(var"my_prefix.x" = 1,)
+
+julia> rand(prefix(demo(), Val(:my_prefix)))
+(var"my_prefix.x" = 1,)
+```
+"""
+prefix(model::Model, x::VarName) = contextualize(model, PrefixContext(x, model.context))
+function prefix(model::Model, x::Val{sym}) where {sym}
+    return contextualize(model, PrefixContext(VarName{sym}(), model.context))
+end
+function prefix(model::Model, x)
+    return contextualize(model, PrefixContext(VarName{Symbol(x)}(), model.context))
+end
+
 """
     (model::Model)([rng, varinfo])
 
@@ -854,6 +889,39 @@ function evaluate_and_sample!!(
     return evaluate_and_sample!!(Random.default_rng(), model, varinfo, sampler)
 end
 
+"""
+    init!!(
+        [rng::Random.AbstractRNG,]
+        model::Model,
+        varinfo::AbstractVarInfo,
+        [init_strategy::AbstractInitStrategy=PriorInit()]
+    )
+
+Evaluate the `model` and replace the values of the model's random variables in
+the given `varinfo` with new values using a specified initialisation strategy.
+If the values in `varinfo` are not already present, they will be added using
+that same strategy.
+
+If `init_strategy` is not provided, defaults to PriorInit().
+
+Returns a tuple of the model's return value, plus the updated `varinfo` object.
+"""
+function init!!(
+    rng::Random.AbstractRNG,
+    model::Model,
+    varinfo::AbstractVarInfo,
+    init_strategy::AbstractInitStrategy=PriorInit(),
+)
+    new_context = setleafcontext(model.context, InitContext(rng, init_strategy))
+    new_model = contextualize(model, new_context)
+    return evaluate!!(new_model, varinfo)
+end
+function init!!(
+    model::Model, varinfo::AbstractVarInfo, init_strategy::AbstractInitStrategy=PriorInit()
+)
+    return init!!(Random.default_rng(), model, varinfo, init_strategy)
+end
+
 """
     evaluate!!(model::Model, varinfo)
 

src/varnamedvector.jl

@@ -766,6 +766,11 @@ function update_internal!(
     return nothing
 end
 
+function BangBang.push!(vnv::VarNamedVector, vn, val, dist)
+    f = from_vec_transform(dist)
+    return setindex_internal!(vnv, tovec(val), vn, f)
+end
+
 # BangBang versions of the above functions.
 # The only difference is that update_internal!! and insert_internal!! check whether the
 # container types of the VarNamedVector vector need to be expanded to accommodate the new