Use ParamsInit for predict; remove setval_and_resample! and friends

penelopeysm · penelopeysm · commit b24fc9679fc2 · 2025-07-19T23:37:22.000+01:00
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
@@ -28,7 +28,7 @@ end
 
 function _check_varname_indexing(c::MCMCChains.Chains)
     return DynamicPPL.supports_varname_indexing(c) ||
-           error("Chains do not support indexing using `VarName`s.")
+           error("This `Chains` object does not support indexing using `VarName`s.")
 end
 
 function DynamicPPL.getindex_varname(
@@ -42,6 +42,15 @@ function DynamicPPL.varnames(c::MCMCChains.Chains)
     return keys(c.info.varname_to_symbol)
 end
 
+function chain_sample_to_varname_dict(c::MCMCChains.Chains, sample_idx, chain_idx)
+    _check_varname_indexing(c)
+    d = Dict{DynamicPPL.VarName,Any}()
+    for vn in DynamicPPL.varnames(c)
+        d[vn] = DynamicPPL.getindex_varname(c, sample_idx, vn, chain_idx)
+    end
+    return d
+end
+
 """
     predict([rng::AbstractRNG,] model::Model, chain::MCMCChains.Chains; include_all=false)
 
@@ -114,9 +123,15 @@ function DynamicPPL.predict(
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     predictive_samples = map(iters) do (sample_idx, chain_idx)
-        DynamicPPL.setval_and_resample!(varinfo, parameter_only_chain, sample_idx, chain_idx)
-        varinfo = last(DynamicPPL.evaluate_and_sample!!(rng, model, varinfo))
-
+        # Extract values from the chain
+        values_dict = chain_sample_to_varname_dict(parameter_only_chain, sample_idx, chain_idx)
+        # Resample any variables that are not present in `values_dict`
+        _, varinfo = DynamicPPL.init!!(
+            rng,
+            model,
+            varinfo,
+            DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit()),
+        )
         vals = DynamicPPL.values_as_in_model(model, false, varinfo)
         varname_vals = mapreduce(
             collect,
@@ -248,13 +263,14 @@ function DynamicPPL.returned(model::DynamicPPL.Model, chain_full::MCMCChains.Cha
     varinfo = DynamicPPL.VarInfo(model)
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     return map(iters) do (sample_idx, chain_idx)
-        # TODO: Use `fix` once we've addressed https://github.com/TuringLang/DynamicPPL.jl/issues/702.
-        # Update the varinfo with the current sample and make variables not present in `chain`
-        # to be sampled.
-        DynamicPPL.setval_and_resample!(varinfo, chain, sample_idx, chain_idx)
-        # NOTE: Some of the varialbes can be a view into the `varinfo`, so we need to
-        # `deepcopy` the `varinfo` before passing it to the `model`.
-        model(deepcopy(varinfo))
+        # Extract values from the chain
+        values_dict = chain_sample_to_varname_dict(chain, sample_idx, chain_idx)
+        # Resample any variables that are not present in `values_dict`, and
+        # return the model's retval.
+        retval, _ = DynamicPPL.init!!(
+            model, varinfo, DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit())
+        )
+        retval
     end
 end
 
diff --git a/src/model.jl b/src/model.jl
@@ -1200,8 +1200,15 @@ function predict(
     varinfo = DynamicPPL.VarInfo(model)
     return map(chain) do params_varinfo
         vi = deepcopy(varinfo)
-        DynamicPPL.setval_and_resample!(vi, values_as(params_varinfo, NamedTuple))
-        model(rng, vi)
+        # TODO(penelopeysm): Requires two model evaluations, one to extract the
+        # parameters and one to set them. The reason why we need values_as_in_model
+        # is because `params_varinfo` may well have some weird combination of
+        # linked/unlinked, whereas `varinfo` is always unlinked since it is
+        # freshly constructed.
+        # This is quite inefficient. It would of course be alright if
+        # ValuesAsInModelAccumulator was a default acc.
+        values_nt = values_as_in_model(model, false, params_varinfo)
+        _, vi = DynamicPPL.init!!(rng, model, vi, ParamsInit(values_nt, PriorInit()))
         return vi
     end
 end
diff --git a/src/varinfo.jl b/src/varinfo.jl
@@ -1498,42 +1498,6 @@ function islinked(vi::VarInfo)
     return any(istrans(vi, vn) for vn in keys(vi))
 end
 
-function nested_setindex_maybe!(vi::UntypedVarInfo, val, vn::VarName)
-    return _nested_setindex_maybe!(vi, getmetadata(vi, vn), val, vn)
-end
-function nested_setindex_maybe!(
-    vi::VarInfo{<:NamedTuple{names}}, val, vn::VarName{sym}
-) where {names,sym}
-    return if sym in names
-        _nested_setindex_maybe!(vi, getmetadata(vi, vn), val, vn)
-    else
-        nothing
-    end
-end
-function _nested_setindex_maybe!(
-    vi::VarInfo, md::Union{Metadata,VarNamedVector}, val, vn::VarName
-)
-    # If `vn` is in `vns`, then we can just use the standard `setindex!`.
-    vns = Base.keys(md)
-    if vn in vns
-        setindex!(vi, val, vn)
-        return vn
-    end
-
-    # Otherwise, we need to check if either of the `vns` subsumes `vn`.
-    i = findfirst(Base.Fix2(subsumes, vn), vns)
-    i === nothing && return nothing
-
-    vn_parent = vns[i]
-    val_parent = getindex(vi, vn_parent)  # TODO: Ensure that we're working with a view here.
-    # Split the varname into its tail optic.
-    optic = remove_parent_optic(vn_parent, vn)
-    # Update the value for the parent.
-    val_parent_updated = set!!(val_parent, optic, val)
-    setindex!(vi, val_parent_updated, vn_parent)
-    return vn_parent
-end
-
 # The default getindex & setindex!() for get & set values
 # NOTE: vi[vn] will always transform the variable to its original space and Julia type
 function getindex(vi::VarInfo, vn::VarName)
@@ -1978,113 +1942,6 @@ function _setval_kernel!(vi::VarInfoOrThreadSafeVarInfo, vn::VarName, values, ke
     return indices
 end
 
-"""
-    setval_and_resample!(vi::VarInfo, x)
-    setval_and_resample!(vi::VarInfo, values, keys)
-    setval_and_resample!(vi::VarInfo, chains::AbstractChains, sample_idx, chain_idx)
-
-Set the values in `vi` to the provided values and those which are not present
-in `x` or `chains` to *be* resampled.
-
-Note that this does *not* resample the values not provided! It will call
-`setflag!(vi, vn, "del")` for variables `vn` for which no values are provided, which means
-that the next time we call `model(vi)` these variables will be resampled.
-
-## Note
-- This suffers from the same limitations as [`setval!`](@ref). See `setval!` for more info.
-
-## Example
-```jldoctest
-julia> using DynamicPPL, Distributions, StableRNGs
-
-julia> @model function demo(x)
-           m ~ Normal()
-           for i in eachindex(x)
-               x[i] ~ Normal(m, 1)
-           end
-       end;
-
-julia> rng = StableRNG(42);
-
-julia> m = demo([missing]);
-
-julia> var_info = DynamicPPL.VarInfo(rng, m);
-       # Checking the setting of "del" flags only makes sense for VarInfo{<:Metadata}. For VarInfo{<:VarNamedVector} the flag is effectively always set.
-
-julia> var_info[@varname(m)]
--0.6702516921145671
-
-julia> var_info[@varname(x[1])]
--0.22312984965118443
-
-julia> DynamicPPL.setval_and_resample!(var_info, (m = 100.0, )); # set `m` and ready `x[1]` for resampling
-
-julia> var_info[@varname(m)] # [✓] changed
-100.0
-
-julia> var_info[@varname(x[1])] # [✓] unchanged
--0.22312984965118443
-
-julia> m(rng, var_info); # sample `x[1]` conditioned on `m = 100.0`
-
-julia> var_info[@varname(m)] # [✓] unchanged
-100.0
-
-julia> var_info[@varname(x[1])] # [✓] changed
-101.37363069798343
-```
-
-## See also
-- [`setval!`](@ref)
-"""
-function setval_and_resample!(vi::VarInfoOrThreadSafeVarInfo, x)
-    return setval_and_resample!(vi, values(x), keys(x))
-end
-function setval_and_resample!(vi::VarInfoOrThreadSafeVarInfo, values, keys)
-    return _apply!(_setval_and_resample_kernel!, vi, values, keys)
-end
-function setval_and_resample!(
-    vi::VarInfoOrThreadSafeVarInfo, chains::AbstractChains, sample_idx::Int, chain_idx::Int
-)
-    if supports_varname_indexing(chains)
-        # First we need to set every variable to be resampled.
-        for vn in keys(vi)
-            set_flag!(vi, vn, "del")
-        end
-        # Then we set the variables in `varinfo` from `chain`.
-        for vn in varnames(chains)
-            vn_updated = nested_setindex_maybe!(
-                vi, getindex_varname(chains, sample_idx, vn, chain_idx), vn
-            )
-
-            # Unset the `del` flag if we found something.
-            if vn_updated !== nothing
-                # NOTE: This will be triggered even if only a subset of a variable has been set!
-                unset_flag!(vi, vn_updated, "del")
-            end
-        end
-    else
-        setval_and_resample!(vi, chains.value[sample_idx, :, chain_idx], keys(chains))
-    end
-end
-
-function _setval_and_resample_kernel!(
-    vi::VarInfoOrThreadSafeVarInfo, vn::VarName, values, keys
-)
-    indices = findall(Base.Fix1(subsumes_string, string(vn)), keys)
-    if !isempty(indices)
-        val = reduce(vcat, values[indices])
-        setval!(vi, val, vn)
-        settrans!!(vi, false, vn)
-    else
-        # Ensures that we'll resample the variable corresponding to `vn` if we run
-        # the model on `vi` again.
-        set_flag!(vi, vn, "del")
-    end
-
-    return indices
-end
-
 values_as(vi::VarInfo) = vi.metadata
 values_as(vi::VarInfo, ::Type{Vector}) = copy(getindex_internal(vi, Colon()))
 function values_as(vi::UntypedVarInfo, ::Type{NamedTuple})
diff --git a/test/ext/DynamicPPLMCMCChainsExt.jl b/test/ext/DynamicPPLMCMCChainsExt.jl
@@ -2,7 +2,12 @@
     @model demo() = x ~ Normal()
     model = demo()
 
-    chain = MCMCChains.Chains(randn(1000, 2, 1), [:x, :y], Dict(:internals => [:y]))
+    chain = MCMCChains.Chains(
+        randn(1000, 2, 1),
+        [:x, :y],
+        Dict(:internals => [:y]);
+        info=(; varname_to_symbol=Dict(@varname(x) => :x)),
+    )
     chain_generated = @test_nowarn returned(model, chain)
     @test size(chain_generated) == (1000, 1)
     @test mean(chain_generated) ≈ 0 atol = 0.1
diff --git a/test/model.jl b/test/model.jl
@@ -573,7 +573,7 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
             xs_train = 1:0.1:10
             ys_train = ground_truth_β .* xs_train + rand(Normal(0, 0.1), length(xs_train))
             m_lin_reg = linear_reg(xs_train, ys_train)
-            chain = [VarInfo(m_lin_reg) _ in 1:10000]
+            chain = [VarInfo(m_lin_reg) for _ in 1:10000]
 
             # chain is generated from the prior
             @test mean([chain[i][@varname(β)] for i in eachindex(chain)]) ≈ 1.0 atol = 0.1
diff --git a/test/test_util.jl b/test/test_util.jl
@@ -81,8 +81,10 @@ function make_chain_from_prior(rng::Random.AbstractRNG, model::Model, n_iters::I
     varnames = collect(varnames)
     # Construct matrix of values
     vals = [get(dict, vn, missing) for dict in dicts, vn in varnames]
+    # Construct dict of varnames -> symbol
+    vn_to_sym_dict = Dict(zip(varnames, map(Symbol, varnames)))
     # Construct and return the Chains object
-    return Chains(vals, varnames)
+    return Chains(vals, varnames; info=(; varname_to_symbol=vn_to_sym_dict))
 end
 function make_chain_from_prior(model::Model, n_iters::Int)
     return make_chain_from_prior(Random.default_rng(), model, n_iters)
diff --git a/test/varinfo.jl b/test/varinfo.jl
@@ -278,7 +278,7 @@ end
         @test typed_vi[vn_y] == 2.0
     end
 
-    @testset "setval! & setval_and_resample!" begin
+    @testset "setval!" begin
         @model function testmodel(x)
             n = length(x)
             s ~ truncated(Normal(); lower=0)
@@ -329,8 +329,8 @@ end
                 else
                     DynamicPPL.setval!(vicopy, (m=zeros(5),))
                 end
-                # Setting `m` fails for univariate due to limitations of `setval!`
-                # and `setval_and_resample!`. See docstring of `setval!` for more info.
+                # Setting `m` fails for univariate due to limitations of `setval!`.
+                # See docstring of `setval!` for more info.
                 if model == model_uv && vi in [vi_untyped, vi_typed]
                     @test_broken vicopy[m_vns] == zeros(5)
                 else
@@ -355,57 +355,6 @@ end
                 DynamicPPL.setval!(vicopy, (s=42,))
                 @test vicopy[m_vns] == 1:5
                 @test vicopy[s_vns] == 42
-
-                ### `setval_and_resample!` ###
-                if model == model_mv && vi == vi_untyped
-                    # Trying to re-run model with `MvNormal` on `vi_untyped` will call
-                    # `MvNormal(μ::Vector{Real}, Σ)` which causes `StackOverflowError`
-                    # so we skip this particular case.
-                    continue
-                end
-
-                if vi in [vi_vnv, vi_vnv_typed]
-                    # `setval_and_resample!` works differently for `VarNamedVector`: All
-                    # values will be resampled when model(vicopy) is called. Hence the below
-                    # tests are not applicable.
-                    continue
-                end
-
-                vicopy = deepcopy(vi)
-                DynamicPPL.setval_and_resample!(vicopy, (m=zeros(5),))
-                model(vicopy)
-                # Setting `m` fails for univariate due to limitations of `subsumes(::String, ::String)`
-                if model == model_uv
-                    @test_broken vicopy[m_vns] == zeros(5)
-                else
-                    @test vicopy[m_vns] == zeros(5)
-                end
-                @test vicopy[s_vns] != vi[s_vns]
-
-                # Ordering is NOT preserved.
-                DynamicPPL.setval_and_resample!(
-                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...)
-                )
-                model(vicopy)
-                if model == model_uv
-                    @test vicopy[m_vns] == 1:5
-                else
-                    @test vicopy[m_vns] == [1, 3, 5, 4, 2]
-                end
-                @test vicopy[s_vns] != vi[s_vns]
-
-                # Correct ordering.
-                DynamicPPL.setval_and_resample!(
-                    vicopy, (; (Symbol("m[$i]") => i for i in (1, 2, 3, 4, 5))...)
-                )
-                model(vicopy)
-                @test vicopy[m_vns] == 1:5
-                @test vicopy[s_vns] != vi[s_vns]
-
-                DynamicPPL.setval_and_resample!(vicopy, (s=42,))
-                model(vicopy)
-                @test vicopy[m_vns] != 1:5
-                @test vicopy[s_vns] == 42
             end
         end
 
@@ -419,9 +368,6 @@ end
         ks = [@varname(x[1, 1]), @varname(x[2, 1]), @varname(x[1, 2]), @varname(x[2, 2])]
         DynamicPPL.setval!(vi, vi.metadata.x.vals, ks)
         @test vals_prev == vi.metadata.x.vals
-
-        DynamicPPL.setval_and_resample!(vi, vi.metadata.x.vals, ks)
-        @test vals_prev == vi.metadata.x.vals
     end
 
     @testset "setval! on chain" begin
