Skip to content

InitContext, part 3 - Introduce InitContext #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: py/hasgetvalue
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions docs/src/api.md
Original file line number Diff line number Diff line change
Expand Up @@ -470,6 +470,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
Expand Down
9 changes: 8 additions & 1 deletion src/DynamicPPL.jl
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,12 @@ export AbstractVarInfo,
ConditionContext,
assume,
tilde_assume,
# Initialisation
InitContext,
AbstractInitStrategy,
PriorInit,
UniformInit,
ParamsInit,
# Pseudo distributions
NamedDist,
NoDist,
Expand Down Expand Up @@ -170,11 +176,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")
Expand Down
35 changes: 0 additions & 35 deletions src/contexts.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Comment on lines -283 to -287
Copy link
Member Author

@penelopeysm penelopeysm Jul 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This code was shifted verbatim to src/model.jl to avoid circular dependencies between files.

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}
Expand Down
176 changes: 176 additions & 0 deletions src/contexts/init.jl
Original file line number Diff line number Diff line change
@@ -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[]

Check warning on line 65 in src/contexts/init.jl

View check run for this annotation

Codecov / codecov/patch

src/contexts/init.jl#L65

Added line #L65 was not covered by tests
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)

Check warning on line 101 in src/contexts/init.jl

View check run for this annotation

Codecov / codecov/patch

src/contexts/init.jl#L101

Added line #L101 was not covered by tests
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)

Check warning on line 132 in src/contexts/init.jl

View check run for this annotation

Codecov / codecov/patch

src/contexts/init.jl#L131-L132

Added lines #L131 - L132 were not covered by tests
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)

Check warning on line 152 in src/contexts/init.jl

View check run for this annotation

Codecov / codecov/patch

src/contexts/init.jl#L152

Added line #L152 was not covered by tests
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
68 changes: 68 additions & 0 deletions src/model.jl
Original file line number Diff line number Diff line change
Expand Up @@ -799,6 +799,41 @@
"""
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))

Check warning on line 831 in src/model.jl

View check run for this annotation

Codecov / codecov/patch

src/model.jl#L829-L831

Added lines #L829 - L831 were not covered by tests
end
function prefix(model::Model, x)
return contextualize(model, PrefixContext(VarName{Symbol(x)}(), model.context))
end

"""
(model::Model)([rng, varinfo])

Expand Down Expand Up @@ -854,6 +889,39 @@
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)

Expand Down
Loading
Loading