Replace evaluate_and_sample!! -> init!!

Author: penelopeysm

docs/src/api.md

@@ -456,11 +456,6 @@ AbstractPPL.evaluate!!
 
 This method mutates the `varinfo` used for execution.
 By default, it does not perform any actual sampling: it only evaluates the model using the values of the variables that are already in the `varinfo`.
-To perform sampling, you can either wrap `model.context` in a `SamplingContext`, or use this convenience method:
-
-```@docs
-DynamicPPL.evaluate_and_sample!!
-```
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
@@ -475,7 +470,12 @@ InitContext
 
 ### VarInfo initialisation
 
-`InitContext` is used to initialise, or overwrite, values in a VarInfo.
+The function `init!!` is used to initialise, or overwrite, values in a VarInfo.
+It is really a thin wrapper around using `evaluate!!` with an `InitContext`.
+
+```@docs
+DynamicPPL.init!!
+```
 
 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:

src/extract_priors.jl

@@ -117,7 +117,7 @@ extract_priors(args::Union{Model,AbstractVarInfo}...) =
 function extract_priors(rng::Random.AbstractRNG, model::Model)
     varinfo = VarInfo()
     varinfo = setaccs!!(varinfo, (PriorDistributionAccumulator(),))
-    varinfo = last(evaluate_and_sample!!(rng, model, varinfo))
+    varinfo = last(init!!(rng, model, varinfo))
     return getacc(varinfo, Val(:PriorDistributionAccumulator)).priors
 end
 

src/model.jl

@@ -850,7 +850,7 @@ end
 # ^ Weird Documenter.jl bug means that we have to write the two above separately
 # as it can only detect the `function`-less syntax.
 function (model::Model)(rng::Random.AbstractRNG, varinfo::AbstractVarInfo=VarInfo())
-    return first(evaluate_and_sample!!(rng, model, varinfo))
+    return first(init!!(rng, model, varinfo))
 end
 
 """
@@ -863,46 +863,19 @@ function use_threadsafe_eval(context::AbstractContext, varinfo::AbstractVarInfo)
     return Threads.nthreads() > 1
 end
 
-"""
-    evaluate_and_sample!!([rng::Random.AbstractRNG, ]model::Model, varinfo[, sampler])
-
-Evaluate the `model` with the given `varinfo`, but perform sampling during the
-evaluation using the given `sampler` by wrapping the model's context in a
-`SamplingContext`.
-
-If `sampler` is not provided, defaults to [`SampleFromPrior`](@ref).
-
-Returns a tuple of the model's return value, plus the updated `varinfo` object.
-"""
-function evaluate_and_sample!!(
-    rng::Random.AbstractRNG,
-    model::Model,
-    varinfo::AbstractVarInfo,
-    sampler::AbstractSampler=SampleFromPrior(),
-)
-    sampling_model = contextualize(model, SamplingContext(rng, sampler, model.context))
-    return evaluate!!(sampling_model, varinfo)
-end
-function evaluate_and_sample!!(
-    model::Model, varinfo::AbstractVarInfo, sampler::AbstractSampler=SampleFromPrior()
-)
-    return evaluate_and_sample!!(Random.default_rng(), model, varinfo, sampler)
-end
-
 """
     init!!(
-        [rng::Random.AbstractRNG,]
+        [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().
+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 set, they will be added.
+using a specified initialisation strategy. If `init_strategy` is not provided,
+defaults to PriorInit().
 
 Returns a tuple of the model's return value, plus the updated `varinfo` object.
 """
@@ -1049,11 +1022,7 @@ Base.nameof(model::Model{<:Function}) = nameof(model.f)
 Generate a sample of type `T` from the prior distribution of the `model`.
 """
 function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
-    x = last(
-        evaluate_and_sample!!(
-            rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())
-        ),
-    )
+    x = last(init!!(rng, model, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}())))
     return values_as(x, T)
 end
 

src/sampler.jl

@@ -58,7 +58,8 @@ function AbstractMCMC.step(
     kwargs...,
 )
     vi = VarInfo()
-    DynamicPPL.evaluate_and_sample!!(rng, model, vi, sampler)
+    strategy = sampler isa SampleFromPrior ? PriorInit() : UniformInit()
+    DynamicPPL.init!!(rng, model, vi, strategy)
     return vi, nothing
 end
 

src/simple_varinfo.jl

@@ -39,7 +39,7 @@ julia> rng = StableRNG(42);
 julia> # In the `NamedTuple` version we need to provide the place-holder values for
        # the variables which are using "containers", e.g. `Array`.
        # In this case, this means that we need to specify `x` but not `m`.
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo((x = ones(2), )));
+       _, vi = DynamicPPL.init!!(rng, m, SimpleVarInfo((x = ones(2), )));
 
 julia> # (✓) Vroom, vroom! FAST!!!
        vi[@varname(x[1])]
@@ -57,12 +57,12 @@ julia> vi[@varname(x[1:2])]
  1.3736306979834252
 
 julia> # (×) If we don't provide the container...
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo()); vi
+       _, vi = DynamicPPL.init!!(rng, m, SimpleVarInfo()); vi
 ERROR: type NamedTuple has no field x
 [...]
 
 julia> # If one does not know the varnames, we can use a `OrderedDict` instead.
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}()));
+       _, vi = DynamicPPL.init!!(rng, m, SimpleVarInfo{Float64}(OrderedDict{VarName,Any}()));
 
 julia> # (✓) Sort of fast, but only possible at runtime.
        vi[@varname(x[1])]
@@ -91,28 +91,28 @@ demo_constrained (generic function with 2 methods)
 
 julia> m = demo_constrained();
 
-julia> _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, SimpleVarInfo());
+julia> _, vi = DynamicPPL.init!!(rng, m, SimpleVarInfo());
 
 julia> vi[@varname(x)] # (✓) 0 ≤ x < ∞
 1.8632965762164932
 
-julia> _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true));
+julia> _, vi = DynamicPPL.init!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true));
 
 julia> vi[@varname(x)] # (✓) -∞ < x < ∞
 -0.21080155351918753
 
-julia> xs = [last(DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true)))[@varname(x)] for i = 1:10];
+julia> xs = [last(DynamicPPL.init!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true)))[@varname(x)] for i = 1:10];
 
 julia> any(xs .< 0)  # (✓) Positive probability mass on negative numbers!
 true
 
 julia> # And with `OrderedDict` of course!
-       _, vi = DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(OrderedDict{VarName,Any}()), true));
+       _, vi = DynamicPPL.init!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(OrderedDict{VarName,Any}()), true));
 
 julia> vi[@varname(x)] # (✓) -∞ < x < ∞
 0.6225185067787314
 
-julia> xs = [last(DynamicPPL.evaluate_and_sample!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true)))[@varname(x)] for i = 1:10];
+julia> xs = [last(DynamicPPL.init!!(rng, m, DynamicPPL.settrans!!(SimpleVarInfo(), true)))[@varname(x)] for i = 1:10];
 
 julia> any(xs .< 0) # (✓) Positive probability mass on negative numbers!
 true
@@ -232,24 +232,25 @@ end
 
 # Constructor from `Model`.
 function SimpleVarInfo{T}(
-    rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler=SampleFromPrior()
+    rng::Random.AbstractRNG, model::Model, init_strategy::AbstractInitStrategy=PriorInit()
 ) where {T<:Real}
-    new_model = contextualize(model, SamplingContext(rng, sampler, model.context))
+    new_context = setleafcontext(model.context, InitContext(rng, init_strategy))
+    new_model = contextualize(model, new_context)
     return last(evaluate!!(new_model, SimpleVarInfo{T}()))
 end
 function SimpleVarInfo{T}(
-    model::Model, sampler::AbstractSampler=SampleFromPrior()
+    model::Model, init_strategy::AbstractInitStrategy=PriorInit()
 ) where {T<:Real}
-    return SimpleVarInfo{T}(Random.default_rng(), model, sampler)
+    return SimpleVarInfo{T}(Random.default_rng(), model, init_strategy)
 end
 # Constructors without type param
 function SimpleVarInfo(
-    rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler=SampleFromPrior()
+    rng::Random.AbstractRNG, model::Model, init_strategy::AbstractInitStrategy=PriorInit()
 )
-    return SimpleVarInfo{LogProbType}(rng, model, sampler)
+    return SimpleVarInfo{LogProbType}(rng, model, init_strategy)
 end
-function SimpleVarInfo(model::Model, sampler::AbstractSampler=SampleFromPrior())
-    return SimpleVarInfo{LogProbType}(Random.default_rng(), model, sampler)
+function SimpleVarInfo(model::Model, init_strategy::AbstractInitStrategy=PriorInit())
+    return SimpleVarInfo{LogProbType}(Random.default_rng(), model, init_strategy)
 end
 
 # Constructor from `VarInfo`.
@@ -265,12 +266,12 @@ end
 
 function untyped_simple_varinfo(model::Model)
     varinfo = SimpleVarInfo(OrderedDict{VarName,Any}())
-    return last(evaluate_and_sample!!(model, varinfo))
+    return last(init!!(model, varinfo))
 end
 
 function typed_simple_varinfo(model::Model)
     varinfo = SimpleVarInfo{Float64}()
-    return last(evaluate_and_sample!!(model, varinfo))
+    return last(init!!(model, varinfo))
 end
 
 function unflatten(svi::SimpleVarInfo, x::AbstractVector)
@@ -480,7 +481,6 @@ function assume(
     return value, vi
 end
 
-# NOTE: We don't implement `settrans!!(vi, trans, vn)`.
 function settrans!!(vi::SimpleVarInfo, trans)
     return settrans!!(vi, trans ? DynamicTransformation() : NoTransformation())
 end
@@ -490,6 +490,15 @@ end
 function settrans!!(vi::ThreadSafeVarInfo{<:SimpleVarInfo}, trans)
     return Accessors.@set vi.varinfo = settrans!!(vi.varinfo, trans)
 end
+function settrans!!(vi::SimpleOrThreadSafeSimple, trans::Bool, ::VarName)
+    # We keep this method around just to obey the AbstractVarInfo interface; however,
+    # this is only a valid operation if it would be a no-op.
+    if trans != istrans(vi)
+        error(
+            "Individual variables in SimpleVarInfo cannot have different `settrans` statuses.",
+        )
+    end
+end
 
 istrans(vi::SimpleVarInfo) = !(vi.transformation isa NoTransformation)
 istrans(vi::SimpleVarInfo, ::VarName) = istrans(vi)

src/test_utils/contexts.jl

@@ -29,21 +29,45 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
     node_trait = DynamicPPL.NodeTrait(context)
     # Throw error immediately if it it's missing a `NodeTrait` implementation.
     node_trait isa Union{DynamicPPL.IsLeaf,DynamicPPL.IsParent} ||
-        throw(ValueError("Invalid NodeTrait: $node_trait"))
+        error("Invalid NodeTrait: $node_trait")
 
-    # To see change, let's make sure we're using a different leaf context than the current.
-    leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
-        DynamicPPL.DynamicTransformationContext{false}()
+    if node_trait isa DynamicPPL.IsLeaf
+        test_leaf_context(context, model)
     else
-        DefaultContext()
+        test_parent_context(context, model)
     end
-    @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
-        leafcontext_new
+end
+
+function test_leaf_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
+    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsLeaf
+
+    # Note that for a leaf context we can't assume that it will work with an
+    # empty VarInfo. Thus we only test evaluation (i.e., assuming that the
+    # varinfo already contains all necessary variables).
+    @testset "evaluation" begin
+        # Generate a new filled untyped varinfo
+        _, untyped_vi = DynamicPPL.init!!(model, DynamicPPL.VarInfo())
+        typed_vi = DynamicPPL.typed_varinfo(untyped_vi)
+        new_model = contextualize(model, context)
+        for vi in [untyped_vi, typed_vi]
+            _, vi = DynamicPPL.evaluate!!(new_model, vi)
+            @test vi isa DynamicPPL.VarInfo
+        end
+    end
+end
+
+function test_parent_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Model)
+    @test DynamicPPL.NodeTrait(context) isa DynamicPPL.IsParent
 
-    # The interface methods.
-    if node_trait isa DynamicPPL.IsParent
-        # `childcontext` and `setchildcontext`
-        # With new child context
+    @testset "{set,}{leaf,child}context" begin
+        # Ensure we're using a different leaf context than the current.
+        leafcontext_new = if DynamicPPL.leafcontext(context) isa DefaultContext
+            DynamicPPL.DynamicTransformationContext{false}()
+        else
+            DefaultContext()
+        end
+        @test DynamicPPL.leafcontext(DynamicPPL.setleafcontext(context, leafcontext_new)) ==
+            leafcontext_new
         childcontext_new = TestParentContext()
         @test DynamicPPL.childcontext(
             DynamicPPL.setchildcontext(context, childcontext_new)
@@ -56,19 +80,15 @@ function test_context(context::DynamicPPL.AbstractContext, model::DynamicPPL.Mod
             leafcontext_new
     end
 
-    # Make sure that the we can evaluate the model with the context (i.e. that none of the tilde-functions are incorrectly overloaded).
-    # The tilde-pipeline contains two different paths: with `SamplingContext` as a parent, and without it.
-    # NOTE(torfjelde): Need to sample with the untyped varinfo _using_ the context, since the
-    # context might alter which variables are present, their names, etc., e.g. `PrefixContext`.
-    # TODO(torfjelde): Make the `varinfo` used for testing a kwarg once it makes sense for other varinfos.
-    # Untyped varinfo.
-    varinfo_untyped = DynamicPPL.VarInfo()
-    model_with_spl = contextualize(model, SamplingContext(context))
-    model_without_spl = contextualize(model, context)
-    @test DynamicPPL.evaluate!!(model_with_spl, varinfo_untyped) isa Any
-    @test DynamicPPL.evaluate!!(model_without_spl, varinfo_untyped) isa Any
-    # Typed varinfo.
-    varinfo_typed = DynamicPPL.typed_varinfo(varinfo_untyped)
-    @test DynamicPPL.evaluate!!(model_with_spl, varinfo_typed) isa Any
-    @test DynamicPPL.evaluate!!(model_without_spl, varinfo_typed) isa Any
+    @testset "initialisation and evaluation" begin
+        new_model = contextualize(model, context)
+        for vi in [DynamicPPL.VarInfo(), DynamicPPL.typed_varinfo(DynamicPPL.VarInfo())]
+            # Initialisation
+            _, vi = DynamicPPL.init!!(new_model, DynamicPPL.VarInfo())
+            @test vi isa DynamicPPL.VarInfo
+            # Evaluation
+            _, vi = DynamicPPL.evaluate!!(new_model, vi)
+            @test vi isa DynamicPPL.VarInfo
+        end
+    end
 end

src/test_utils/model_interface.jl

@@ -92,9 +92,7 @@ Even though it is recommended to implement this by hand for a particular `Model`
 a default implementation using [`SimpleVarInfo{<:Dict}`](@ref) is provided.
 """
 function varnames(model::Model)
-    return collect(
-        keys(last(DynamicPPL.evaluate_and_sample!!(model, SimpleVarInfo(Dict()))))
-    )
+    return collect(keys(last(DynamicPPL.init!!(model, SimpleVarInfo(Dict())))))
 end
 
 """