Use model field instead

Author: penelopeysm

src/DynamicPPL.jl

@@ -113,6 +113,7 @@ export AbstractVarInfo,
     pointwise_loglikelihoods,
     condition,
     decondition,
+    set_tracked_varnames,
     fix,
     unfix,
     predict,

src/model.jl

@@ -4,12 +4,20 @@
         args::NamedTuple{argnames,Targs}
         defaults::NamedTuple{defaultnames,Tdefaults}
         context::Ctx=DefaultContext()
+        tracked_varnames::Union{Nothing,Array{<:VarName}}
     end
 
 A `Model` struct with model evaluation function of type `F`, arguments of names `argnames`
 types `Targs`, default arguments of names `defaultnames` with types `Tdefaults`, missing
 arguments `missings`, and evaluation context of type `Ctx`.
 
+`tracked_varnames` is an array of VarNames that should be tracked during sampling. During
+model evaluation (with `DynamicPPL.evaluate!!`) all random variables are tracked; however,
+at the end of each iteration of MCMC sampling, `DynamicPPL.values_as_in_model` is used to
+extract the values of _only_ the tracked variables. This allows the user to control which
+variables are ultimately stored in the chain. This field can be set using the
+[`set_tracked_varnames`](@ref) function.
+
 Here `argnames`, `defaultargnames`, and `missings` are tuples of symbols, e.g. `(:a, :b)`.
 `context` is by default `DefaultContext()`.
 
@@ -23,14 +31,17 @@ different arguments.
 # Examples
 
 ```julia
+julia> f(x) = x + 1  # Dummy function
+f (generic function with 1 method)
+
 julia> Model(f, (x = 1.0, y = 2.0))
-Model{typeof(f),(:x, :y),(),(),Tuple{Float64,Float64},Tuple{}}(f, (x = 1.0, y = 2.0), NamedTuple())
+Model{typeof(f), (:x, :y), (), (), Tuple{Float64, Float64}, Tuple{}, DefaultContext}(f, (x = 1.0, y = 2.0), NamedTuple(), DefaultContext(), nothing)
 
 julia> Model(f, (x = 1.0, y = 2.0), (x = 42,))
-Model{typeof(f),(:x, :y),(:x,),(),Tuple{Float64,Float64},Tuple{Int64}}(f, (x = 1.0, y = 2.0), (x = 42,))
+Model{typeof(f), (:x, :y), (:x,), (), Tuple{Float64, Float64}, Tuple{Int64}, DefaultContext}(f, (x = 1.0, y = 2.0), (x = 42,), DefaultContext(), nothing)
 
 julia> Model{(:y,)}(f, (x = 1.0, y = 2.0), (x = 42,)) # with special definition of missings
-Model{typeof(f),(:x, :y),(:x,),(:y,),Tuple{Float64,Float64},Tuple{Int64}}(f, (x = 1.0, y = 2.0), (x = 42,))
+Model{typeof(f), (:x, :y), (:x,), (:y,), Tuple{Float64, Float64}, Tuple{Int64}, DefaultContext}(f, (x = 1.0, y = 2.0), (x = 42,), DefaultContext(), nothing)
 ```
 """
 struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractContext} <:
@@ -39,6 +50,7 @@ struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractConte
     args::NamedTuple{argnames,Targs}
     defaults::NamedTuple{defaultnames,Tdefaults}
     context::Ctx
+    tracked_varnames::Union{Nothing,Array{<:VarName}}
 
     @doc """
         Model{missings}(f, args::NamedTuple, defaults::NamedTuple)
@@ -51,9 +63,10 @@ struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx<:AbstractConte
         args::NamedTuple{argnames,Targs},
         defaults::NamedTuple{defaultnames,Tdefaults},
         context::Ctx=DefaultContext(),
+        tracked_varnames::Union{Nothing,Array{<:VarName}}=nothing,
     ) where {missings,F,argnames,Targs,defaultnames,Tdefaults,Ctx}
         return new{F,argnames,defaultnames,missings,Targs,Tdefaults,Ctx}(
-            f, args, defaults, context
+            f, args, defaults, context, tracked_varnames
         )
     end
 end
@@ -71,22 +84,44 @@ model with different arguments.
     args::NamedTuple{argnames,Targs},
     defaults::NamedTuple{kwargnames,Tkwargs},
     context::AbstractContext=DefaultContext(),
+    tracked_varnames::Union{Nothing,Array{<:VarName}}=nothing,
 ) where {F,argnames,Targs,kwargnames,Tkwargs}
     missing_args = Tuple(
         name for (name, typ) in zip(argnames, Targs.types) if typ <: Missing
     )
     missing_kwargs = Tuple(
         name for (name, typ) in zip(kwargnames, Tkwargs.types) if typ <: Missing
     )
-    return :(Model{$(missing_args..., missing_kwargs...)}(f, args, defaults, context))
+    return :(Model{$(missing_args..., missing_kwargs...)}(
+        f, args, defaults, context, tracked_varnames
+    ))
 end
 
-function Model(f, args::NamedTuple, context::AbstractContext=DefaultContext(); kwargs...)
-    return Model(f, args, NamedTuple(kwargs), context)
+function Model(
+    f,
+    args::NamedTuple,
+    context::AbstractContext=DefaultContext(),
+    tracked_varnames::Union{Nothing,Array{<:VarName}}=nothing;
+    kwargs...,
+)
+    return Model(f, args, NamedTuple(kwargs), context, tracked_varnames)
 end
 
 function contextualize(model::Model, context::AbstractContext)
-    return Model(model.f, model.args, model.defaults, context)
+    return Model(model.f, model.args, model.defaults, context, model.tracked_varnames)
+end
+
+"""
+    set_tracked_varnames(model::Model, varnames::Union{Nothing,Array{<:VarName}})
+
+Return a new `Model` which only tracks a subset of variables during sampling.
+
+If `varnames` is `nothing`, then all variables will be tracked. Otherwise, only
+the variables subsumed by `varnames` are tracked. For example, if `varnames =
+[@varname(x)]`, then any variable `x`, `x[1]`, `x.a`, ... will be tracked.
+"""
+function set_tracked_varnames(model::Model, varnames::Union{Nothing,Array{<:VarName}})
+    return Model(model.f, model.args, model.defaults, model.context, varnames)
 end
 
 """

src/values_as_in_model.jl

@@ -195,39 +195,12 @@ function values_as_in_model(
     model::Model,
     include_colon_eq::Bool,
     varinfo::AbstractVarInfo,
-    tracked_varnames=tracked_varnames(model),
+    tracked_varnames=model.tracked_varnames,
     context::AbstractContext=DefaultContext(),
 )
+    @show tracked_varnames
     tracked_varnames = isnothing(tracked_varnames) ? nothing : collect(tracked_varnames)
     context = ValuesAsInModelContext(include_colon_eq, tracked_varnames, context)
     evaluate!!(model, varinfo, context)
     return context.values
 end
-
-"""
-    tracked_varnames(model::Model)
-
-Returns a set of `VarName`s that the model should track.
-
-By default, this returns `nothing`, which means that all `VarName`s should be
-tracked.
-
-If you want to track only a subset of `VarName`s, you can override this method
-in your model definition:
-
-```julia
-@model function mymodel()
-    x ~ Normal()
-    y ~ Normal(x, 1)
-end
-
-DynamicPPL.tracked_varnames(::Model{typeof(mymodel)}) = [@varname(y)]
-```
-
-Then, when you sample from `mymodel()`, only the value of `y` will be tracked
-(and not `x`).
-
-Note that quantities on the left-hand side of `:=` are always tracked, and will
-ignore the varnames specified in this method.
-"""
-tracked_varnames(::Model) = nothing

test/contexts.jl

@@ -167,7 +167,7 @@ end
             ctx1 = PrefixContext{:a}(DefaultContext())
             ctx2 = SamplingContext(ctx1)
             ctx3 = PrefixContext{:b}(ctx2)
-            ctx4 = DynamicPPL.ValuesAsInModelContext(OrderedDict(), false, ctx3)
+            ctx4 = DynamicPPL.ValuesAsInModelContext(OrderedDict(), false, nothing, ctx3)
             vn_prefixed1 = prefix(ctx1, vn)
             vn_prefixed2 = prefix(ctx2, vn)
             vn_prefixed3 = prefix(ctx3, vn)

test/values_as_in_model.jl

@@ -82,9 +82,9 @@
         @test !haskey(values, @varname(y))
         @test haskey(values, @varname(z))  # := is always included
 
-        # Specify instead using `tracked_varnames` method
-        DynamicPPL.tracked_varnames(::Model{typeof(track_specific)}) = [@varname(y)]
-        values = values_as_in_model(model, true, vi)
+        # Specify instead using `set_tracked_varnames` method
+        model2 = DynamicPPL.set_tracked_varnames(model, [@varname(y)])
+        values = values_as_in_model(model2, true, vi)
         @test !haskey(values, @varname(x[1]))
         @test !haskey(values, @varname(x[2]))
         @test haskey(values, @varname(y))