Use ParamsInit for predict; remove setval_and_resample! and friends

penelopeysm · penelopeysm · commit 726588501260 · 2025-07-10T01:51:13.000+01:00
diff --git a/ext/DynamicPPLMCMCChainsExt.jl b/ext/DynamicPPLMCMCChainsExt.jl
@@ -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,17 @@ 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 = last(
+            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 +265,15 @@ 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 (`first`).
+        first(
+            DynamicPPL.init!!(
+                model, varinfo, DynamicPPL.ParamsInit(values_dict, DynamicPPL.PriorInit())
+            ),
+        )
     end
 end
 
diff --git a/src/model.jl b/src/model.jl
@@ -1233,8 +1233,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
@@ -2045,113 +2045,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/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
