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,36 @@ 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/runtests.jl

@@ -45,58 +45,58 @@ include("test_util.jl")
     # groups are chosen to make both groups take roughly the same amount of
     # time, but beyond that there is no particular reason for the split.
     if GROUP == "All" || GROUP == "Group1"
-        include("utils.jl")
-        include("compiler.jl")
-        include("varnamedvector.jl")
-        include("varinfo.jl")
-        include("simple_varinfo.jl")
-        include("model.jl")
+        # include("utils.jl")
+        # include("compiler.jl")
+        # include("varnamedvector.jl")
+        # include("varinfo.jl")
+        # include("simple_varinfo.jl")
+        # include("model.jl")
         include("values_as_in_model.jl")
-        include("sampler.jl")
-        include("independence.jl")
-        include("distribution_wrappers.jl")
-        include("logdensityfunction.jl")
-        include("linking.jl")
-        include("serialization.jl")
-        include("pointwise_logdensities.jl")
-        include("lkj.jl")
-        include("deprecated.jl")
+        # include("sampler.jl")
+        # include("independence.jl")
+        # include("distribution_wrappers.jl")
+        # include("logdensityfunction.jl")
+        # include("linking.jl")
+        # include("serialization.jl")
+        # include("pointwise_logdensities.jl")
+        # include("lkj.jl")
+        # include("deprecated.jl")
     end
 
     if GROUP == "All" || GROUP == "Group2"
-        include("contexts.jl")
-        include("context_implementations.jl")
-        include("threadsafe.jl")
-        include("debug_utils.jl")
-        @testset "compat" begin
-            include(joinpath("compat", "ad.jl"))
-        end
-        @testset "extensions" begin
-            include("ext/DynamicPPLMCMCChainsExt.jl")
-            include("ext/DynamicPPLJETExt.jl")
-        end
-        @testset "ad" begin
-            include("ext/DynamicPPLForwardDiffExt.jl")
-            include("ext/DynamicPPLMooncakeExt.jl")
-            include("ad.jl")
-        end
-        @testset "prob and logprob macro" begin
-            @test_throws ErrorException prob"..."
-            @test_throws ErrorException logprob"..."
-        end
-        @testset "doctests" begin
-            DocMeta.setdocmeta!(
-                DynamicPPL,
-                :DocTestSetup,
-                :(using DynamicPPL, Distributions);
-                recursive=true,
-            )
-            doctestfilters = [
-                # Ignore the source of a warning in the doctest output, since this is dependent on host.
-                # This is a line that starts with "└ @ " and ends with the line number.
-                r"└ @ .+:[0-9]+",
-            ]
-            doctest(DynamicPPL; manual=false, doctestfilters=doctestfilters)
-        end
+        # include("contexts.jl")
+        # include("context_implementations.jl")
+        # include("threadsafe.jl")
+        # include("debug_utils.jl")
+        # @testset "compat" begin
+        #     include(joinpath("compat", "ad.jl"))
+        # end
+        # @testset "extensions" begin
+        #     include("ext/DynamicPPLMCMCChainsExt.jl")
+        #     include("ext/DynamicPPLJETExt.jl")
+        # end
+        # @testset "ad" begin
+        #     include("ext/DynamicPPLForwardDiffExt.jl")
+        #     include("ext/DynamicPPLMooncakeExt.jl")
+        #     include("ad.jl")
+        # end
+        # @testset "prob and logprob macro" begin
+        #     @test_throws ErrorException prob"..."
+        #     @test_throws ErrorException logprob"..."
+        # end
+        # @testset "doctests" begin
+        #     DocMeta.setdocmeta!(
+        #         DynamicPPL,
+        #         :DocTestSetup,
+        #         :(using DynamicPPL, Distributions);
+        #         recursive=true,
+        #     )
+        #     doctestfilters = [
+        #         # Ignore the source of a warning in the doctest output, since this is dependent on host.
+        #         # This is a line that starts with "└ @ " and ends with the line number.
+        #         r"└ @ .+:[0-9]+",
+        #     ]
+        #     doctest(DynamicPPL; manual=false, doctestfilters=doctestfilters)
+        # end
     end
 end

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))