Depreciate@submodel l ~ m in favour of l ~ to_submodel(m); rename generated_quantities to returned

docs/src/api.md

@@ -14,10 +14,18 @@ These statements are rewritten by `@model` as calls of [internal functions](@ref
 @model
 ```
 
-One can nest models and call another model inside the model function with [`@submodel`](@ref).
+One can nest models and call another model inside the model function with [`@submodel`](@ref) and [`@returned_quantities`](@ref).
 
 ```@docs
 @submodel
+@returned_quantities
+```
+
+In the context of nesting models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
+
+```@docs
+prefix(::DynamicPPL.Model, x)
+@prefix
 ```
 
 ### Type

src/DynamicPPL.jl

@@ -125,6 +125,9 @@ export AbstractVarInfo,
     # Convenience macros
     @addlogprob!,
     @submodel,
+    @returned_quantities,
+    prefix,
+    @prefix,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,

src/contexts.jl

@@ -294,6 +294,34 @@
     end
 end
 
+"""
+    prefix(model::Model, x)
+
+Return `model` but with all random variables prefixed by `x`.
+
+If `x` is known at compile-time, use `Val{x}()` to avoid runtime overheads for prefixing.
+"""
+prefix(model::Model, x) = contextualize(model, PrefixContext{Symbol(x)}(model.context))
+function prefix(model::Model, ::Val{x}) where {x}
+    return contextualize(model, PrefixContext{Symbol(x)}(model.context))
+end
+
+"""
+    @prefix(model, prefix_expr)
+
+Return `model` but with all random variables prefixed by `prefix_expr`.
+
+The result of `prefix_expr` must will be converted to a `Symbol` and used as the prefix.
+
+!!! note
+    This is effectively just a convenience macro for the method [`prefix(::Model, x)`](@ref),
+    which automatically converts the result of `prefix_expr` into a `Val` to avoid runtime overheads
+    for static prefixes. For more control over the prefixing, use the method directly.
+"""
+macro prefix(model, prefix_expr)
+    return :($prefix($(esc(model)), $Val{$Symbol($(esc(prefix_expr)))}()))
+end
+
 struct ConditionContext{Values,Ctx<:AbstractContext} <: AbstractContext
     values::Values
     context::Ctx

src/model.jl

@@ -1051,7 +1051,9 @@ function Base.rand(rng::Random.AbstractRNG, ::Type{T}, model::Model) where {T}
         evaluate!!(
             model,
             SimpleVarInfo{Float64}(OrderedDict()),
-            SamplingContext(rng, SampleFromPrior(), model.context),
+            # NOTE: Use `leafcontext` here so we a) avoid overriding the leaf context of `model`,
+            # and b) avoid double-stacking the parent contexts.
+            SamplingContext(rng, SampleFromPrior(), leafcontext(model.context)),
         ),
     )
     return values_as(x, T)

src/submodel_macro.jl

@@ -248,3 +248,177 @@
         end
     end
 end
+
+"""
+    @returned_quantities model
+
+Run `model` nested inside of another model and return the return-values of the `model`.
+
+!!! warning
+    It's generally recommended to use [`prefix(::Model, x)`](@ref) or
+    [`@prefix(model, prefix_expr)`](@ref) in combination with `@returned_quantities`
+    to ensure that the variables in `model` are unique and do not clash with other variables in the
+    parent model or in other submodels.
+
+# Examples
+
+## Simple example
+```jldoctest submodel-returned-quantities; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y)
+            a = @returned_quantities(demo1(x))
+            return y ~ Uniform(0, a)
+       end;
+```
+
+When we sample from the model `demo2(missing, 0.4)` random variable `x` will be sampled:
+```jldoctest submodel-returned-quantities
+julia> vi = VarInfo(demo2(missing, 0.4));
+
+julia> @varname(x) in keys(vi)
+true
+```
+
+Variable `a` is not tracked since it can be computed from the random variable `x` that was
+tracked when running `demo1`:
+```jldoctest submodel-returned-quantities
+julia> @varname(a) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-returned-quantities
+julia> x = vi[@varname(x)];
+
+julia> getlogp(vi) ≈ logpdf(Normal(), x) + logpdf(Uniform(0, 1 + abs(x)), 0.4)
+true
+```
+
+## With prefixing
+```jldoctest submodel-returned-quantities-prefix; setup=:(using Distributions)
+julia> @model function demo1(x)
+           x ~ Normal()
+           return 1 + abs(x)
+       end;
+
+julia> @model function demo2(x, y, z)
+            a = @returned_quantities @prefix(demo1(x), :sub1)
+            b = @returned_quantities @prefix(demo1(y), :sub2)
+            return z ~ Uniform(-a, b)
+       end;
+```
+
+When we sample from the model `demo2(missing, missing, 0.4)` random variables `sub1.x` and
+`sub2.x` will be sampled:
+```jldoctest submodel-returned-quantities-prefix
+julia> vi = VarInfo(demo2(missing, missing, 0.4));
+
+julia> @varname(var"sub1.x") in keys(vi)
+true
+
+julia> @varname(var"sub2.x") in keys(vi)
+true
+```
+
+Variables `a` and `b` are not tracked since they can be computed from the random variables `sub1.x` and
+`sub2.x` that were tracked when running `demo1`:
+```jldoctest submodel-returned-quantities-prefix
+julia> @varname(a) in keys(vi)
+false
+
+julia> @varname(b) in keys(vi)
+false
+```
+
+We can check that the log joint probability of the model accumulated in `vi` is correct:
+
+```jldoctest submodel-returned-quantities-prefix
+julia> sub1_x = vi[@varname(var"sub1.x")];
+
+julia> sub2_x = vi[@varname(var"sub2.x")];
+
+julia> logprior = logpdf(Normal(), sub1_x) + logpdf(Normal(), sub2_x);
+
+julia> loglikelihood = logpdf(Uniform(-1 - abs(sub1_x), 1 + abs(sub2_x)), 0.4);
+
+julia> getlogp(vi) ≈ logprior + loglikelihood
+true
+```
+
+## Different ways of setting the prefix
+```jldoctest submodel-returned-quantities-prefix-alts; setup=:(using DynamicPPL, Distributions)
+julia> @model inner() = x ~ Normal()
+inner (generic function with 2 methods)
+
+julia> # When `prefix` is unspecified, no prefix is used.
+       @model submodel_noprefix() = a = @returned_quantities inner()
+submodel_noprefix (generic function with 2 methods)
+
+julia> @varname(x) in keys(VarInfo(submodel_noprefix()))
+true
+
+julia> # Using a static string.
+       @model submodel_prefix_string() = a = @returned_quantities @prefix inner() "my prefix"
+submodel_prefix_string (generic function with 2 methods)
+
+julia> @varname(var"my prefix.x") in keys(VarInfo(submodel_prefix_string()))
+true
+
+julia> # Using string interpolation.
+       @model submodel_prefix_interpolation() = a = @returned_quantities @prefix inner() "\$(nameof(inner()))"
+submodel_prefix_interpolation (generic function with 2 methods)
+
+julia> @varname(var"inner.x") in keys(VarInfo(submodel_prefix_interpolation()))
+true
+
+julia> # Or using some arbitrary expression.
+       @model submodel_prefix_expr() = a = @returned_quantities prefix=1 + 2 inner()
+submodel_prefix_expr (generic function with 2 methods)
+
+julia> @varname(var"3.x") in keys(VarInfo(submodel_prefix_expr()))
+true
+```
+"""
+macro returned_quantities(expr)
+    return returned_quantities_expr(:(prefix = false), expr)
+end
+
+macro returned_quantities(prefix_expr, expr)
+    return returned_quantities_expr(prefix_expr, expr)
+end
+
+"""
+    @returned_quantities_expr model
+
+Returns an expression that captures the return-values of a model in addition to the varinfo.
+"""
+function returned_quantities_expr(prefix_expr, expr, ctx=esc(:__context__))
+    prefix_left, prefix = getargs_assignment(prefix_expr)
+    if prefix_left !== :prefix
+        error("$(prefix_left) is not a valid kwarg")
+    end
+
+    # The user expects `@returned_quantities ...` to return the
+    # return-value of the `...`, hence we need to capture
+    # the return-value and handle it correctly.
+    @gensym retval
+
+    # Prefix.
+    if prefix !== nothing
+        ctx = prefix_submodel_context(prefix, ctx)
+    end
+    return quote
+        # Evaluate the model and capture the return values + varinfo.
+        $retval, $(esc(:__varinfo__)) = $(_evaluate!!)(
+            $(esc(expr)), $(esc(:__varinfo__)), $(ctx)
+        )
+
+        # Return the return-value of the model.
+        $retval
+    end
+end