From 51549c8a1104731129bbfbac4501bc8b26e77802 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 20:41:34 +0100 Subject: [PATCH 01/58] Replace AdvancedPS with a native particle MCMC implementation SMC/PG/CSMC are reimplemented directly on Libtask + DynamicPPL, dropping the AdvancedPS dependency. A particle is a suspended model execution advanced one observation at a time; the conditional-SMC reference trajectory is regenerated by replaying a per-particle TracedRNG (Random123 Philox), so no reference values need to be stored, and a child forked from the reference is simply reseeded to branch off. All particle state lives explicitly on the particle via Libtask's taped globals -- no task_local_storage -- and @addlogprob! reweights by producing inside the likelihood accumulator, so no type piracy is required. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- Project.toml | 4 +- src/mcmc/Inference.jl | 2 +- src/mcmc/particle_mcmc.jl | 793 ++++++++++++++++++------------------ src/mcmc/traced_rng.jl | 59 +++ test/Project.toml | 2 - test/essential/container.jl | 60 +-- test/mcmc/particle_mcmc.jl | 37 +- 7 files changed, 493 insertions(+), 464 deletions(-) create mode 100644 src/mcmc/traced_rng.jl diff --git a/Project.toml b/Project.toml index e96c2dd443..5fb591b61d 100644 --- a/Project.toml +++ b/Project.toml @@ -9,7 +9,6 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -AdvancedPS = "576499cb-2369-40b2-a588-c64705576edc" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" @@ -29,6 +28,7 @@ OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Random123 = "74087812-796a-5b5d-8853-05524746bad3" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" @@ -52,7 +52,6 @@ AbstractPPL = "0.15" Accessors = "0.1" AdvancedHMC = "0.8.3" AdvancedMH = "0.8.9" -AdvancedPS = "0.7.2" AdvancedVI = "0.7" BangBang = "0.4.2" Bijectors = "0.15.17, 0.16" @@ -74,6 +73,7 @@ OptimizationOptimJL = "0.1 - 0.4" OrderedCollections = "1, 2" Printf = "1" Random = "1" +Random123 = "1.7.1" Reexport = "0.2, 1" SciMLBase = "2, 3" SpecialFunctions = "0.7.2, 0.8, 0.9, 0.10, 1, 2" diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl index b9aa08e8f6..ac2041c5ba 100644 --- a/src/mcmc/Inference.jl +++ b/src/mcmc/Inference.jl @@ -32,7 +32,7 @@ import AdvancedHMC const AHMC = AdvancedHMC import AdvancedMH const AMH = AdvancedMH -import AdvancedPS +import Random123 import EllipticalSliceSampling import LogDensityProblems import Random diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 2bef555edc..ac643b0d1b 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -1,8 +1,22 @@ ### -### Particle Filtering and Particle MCMC Samplers. +### Particle filtering and particle MCMC samplers. +### +### SMC runs a particle filter over a model; PG/CSMC wraps a *conditional* particle filter +### inside an MCMC loop. Both treat each `observe` statement as one filtering step: running +### the model under `ParticleMCMCContext` turns every likelihood term into a +### `Libtask.produce`, so a particle is a suspended model execution that we advance one +### observation at a time, reweight, and resample. +### +### Each particle carries a `TracedRNG` that records the seed it used at every step. The +### reference trajectory of a conditional sweep is reproduced by *replaying* those seeds +### (`load_state!`), so it is regenerated exactly without storing its values. A particle +### forked from the reference is reseeded, which automatically switches it from replaying to +### sampling fresh -- no per-particle bookkeeping required. ### -using Accessors: Accessors +using StatsFuns: softmax, logsumexp + +include("traced_rng.jl") function error_if_threadsafe_eval(model::DynamicPPL.Model) if DynamicPPL.requires_threadsafe(model) @@ -15,77 +29,307 @@ function error_if_threadsafe_eval(model::DynamicPPL.Model) return nothing end -### AdvancedPS models and interface +# +# Model evaluation via Libtask +# +# A `Particle` is the only mutable state. It is stored as its `TapedTask`'s "taped globals", +# so the tilde overloads reach it from *inside* a running model via `get_taped_globals`. +# This keeps all state explicit on the particle -- no `task_local_storage`. -struct ParticleMCMCContext{R<:AbstractRNG} <: DynamicPPL.AbstractContext - rng::R -end -# Because pMCMC uses OnlyAccsVarInfo, we need to overload this. It's fine to use Any (see -# the docstring of get_param_eltype in DynamicPPL) because pMCMC doesn't involve AD or any -# other tracer types. +""" + ParticleMCMCContext + +Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws +from the prior using the particle's [`TracedRNG`](@ref), and `tilde_observe!!` scores the +observation, which [`ProduceLogLikelihoodAccumulator`](@ref) turns into a `Libtask.produce`. +""" +struct ParticleMCMCContext <: DynamicPPL.AbstractContext end + +# `OnlyAccsVarInfo` needs a parameter eltype; `Any` is fine here since particle MCMC never +# involves AD or tracer types (see the `get_param_eltype` docstring in DynamicPPL). DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::ParticleMCMCContext) = Any -mutable struct TracedModel{M<:Model,T<:Tuple,NT<:NamedTuple} <: - AdvancedPS.AbstractGenericModel - model::M - # TODO(penelopeysm): I don't like that this is an abstract type. However, the problem is - # that the type of VarInfo can change during execution, especially with PG-inside-Gibbs - # when you have to muck with merging VarInfos from different sub-conditioned models. - # - # However, I don't think that this is actually a problem in practice. Whenever we do - # Libtask.get_taped_globals, that is already type unstable anyway, so accessing this - # field here is not going to cause extra type instability. This change is associated - # with Turing v0.43, and I benchmarked on v0.42 vs v0.43, and v0.43 is actually faster - # (probably due to underlying changes in DynamicPPL), so I'm not really bothered by - # this. - varinfo::AbstractVarInfo - resample::Bool - fargs::T - kwargs::NT -end - -function TracedModel( - model::Model, varinfo::AbstractVarInfo, rng::Random.AbstractRNG, resample::Bool +""" + Particle(model, varinfo, rng::TracedRNG) + +A single particle: a suspended `model` execution together with its `varinfo`, its own +replayable `rng`, and an accumulated `logweight`. +""" +mutable struct Particle + # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go + # through Libtask's (already type-unstable) taped globals, so this costs nothing extra. + varinfo::DynamicPPL.AbstractVarInfo + rng::TracedRNG + logweight::Float64 + task::Libtask.TapedTask + # `task` is filled in once the particle exists, because the task must capture the + # particle as its taped globals (a back-reference). + Particle(vi, rng) = new(vi, rng, 0.0) +end + +function Particle( + model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo, rng::TracedRNG ) - model = DynamicPPL.setleafcontext(model, ParticleMCMCContext(rng)) + model = DynamicPPL.setleafcontext(model, ParticleMCMCContext()) args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - fargs = (model.f, args...) - return TracedModel(model, varinfo, resample, fargs, kwargs) + particle = Particle(deepcopy(varinfo), rng) + particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) + return particle +end + +""" + fork(particle, rng) + +Copy `particle` into an independent continuation seeded from `rng`. `deepcopy` forks the +underlying `TapedTask` (Libtask defines `copy` as `deepcopy`) and preserves the +task↔particle back-reference, which we reset explicitly to be safe. Reseeding switches the +child from replaying to sampling afresh; `keys` is truncated to the steps already taken so a +child of the reference forgets the reference's future. +""" +function fork(particle::Particle, rng::AbstractRNG) + child = deepcopy(particle) + Libtask.set_taped_globals!(child.task, child) + Random.seed!(child.rng, rand(rng, UInt64)) + resize!(child.rng.keys, child.rng.count - 1) + return child +end + +""" + advance!(particle, isref) -> Union{Float64,Nothing} + +Run the particle to its next `observe`, returning the incremental log-likelihood, or +`nothing` once the model finishes. An ordinary particle records the step's seed; the +reference (`isref = true`) replays its recorded seed instead. +""" +function advance!(particle::Particle, isref::Bool) + isref ? load_state!(particle.rng) : save_state!(particle.rng) + inc_step!(particle.rng) + return Libtask.consume(particle.task) +end + +function DynamicPPL.tilde_assume!!( + ::ParticleMCMCContext, + dist::Distribution, + vn::VarName, + template, + ::DynamicPPL.AbstractVarInfo, +) + particle = Libtask.get_taped_globals(Particle) + ctx = DynamicPPL.InitContext( + particle.rng, DynamicPPL.InitFromPrior(), DynamicPPL.UnlinkAll() + ) + x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, particle.varinfo) + particle.varinfo = vi + return x, vi end -function AdvancedPS.advance!( - trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false +function DynamicPPL.tilde_observe!!( + ::ParticleMCMCContext, + dist::Distribution, + left, + vn::Union{VarName,Nothing}, + template, + ::DynamicPPL.AbstractVarInfo, ) - # Make sure we load/reset the rng in the new replaying mechanism - isref ? AdvancedPS.load_state!(trace.rng) : AdvancedPS.save_state!(trace.rng) - score = consume(trace.model.ctask) - return score + particle = Libtask.get_taped_globals(Particle) + left, vi = DynamicPPL.tilde_observe!!( + DynamicPPL.DefaultContext(), dist, left, vn, template, particle.varinfo + ) + particle.varinfo = vi + return left, vi end -function AdvancedPS.delete_retained!(trace::TracedModel) - # This method is called if, during a CSMC update, we perform a resampling - # and choose the reference particle as the trajectory to carry on from. - # In such a case, we need to ensure that when we continue sampling (i.e. - # the next time we hit tilde_assume!!), we don't use the values in the - # reference particle but rather sample new values. - return TracedModel(trace.model, trace.varinfo, true, trace.fargs, trace.kwargs) +""" + ProduceLogLikelihoodAccumulator{T} <: LogProbAccumulator{T} + +Like `LogLikelihoodAccumulator`, but `Libtask.produce`s each likelihood increment as it is +accumulated. Because `@addlogprob!` also routes through `acclogp`, it too triggers a +`produce`, so manual likelihood terms reweight particles correctly (issue #1996). +""" +struct ProduceLogLikelihoodAccumulator{T<:Real} <: DynamicPPL.LogProbAccumulator{T} + logp::T end -function AdvancedPS.reset_model(trace::TracedModel) - return trace +DynamicPPL.accumulator_name(::Type{<:ProduceLogLikelihoodAccumulator}) = :LogLikelihood +DynamicPPL.logp(acc::ProduceLogLikelihoodAccumulator) = acc.logp + +function DynamicPPL.acclogp(acc::ProduceLogLikelihoodAccumulator, val) + Libtask.produce(val) # the only difference from `LogLikelihoodAccumulator` + return ProduceLogLikelihoodAccumulator(acc.logp + val) end -function Libtask.TapedTask(taped_globals, model::TracedModel) - return Libtask.TapedTask( - taped_globals, model.fargs[1], model.fargs[2:end]...; model.kwargs... +function DynamicPPL.accumulate_assume!!( + acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, dist, template +) + return acc +end +function DynamicPPL.accumulate_observe!!( + acc::ProduceLogLikelihoodAccumulator, dist, left, vn, template +) + return DynamicPPL.acclogp(acc, Distributions.loglikelihood(dist, left)) +end + +# Tell Libtask which calls may contain `produce`, walking up the call stack from `acclogp`. +Libtask.@might_produce(DynamicPPL.accloglikelihood!!) +function Libtask.might_produce( + ::Type{ + <:Tuple{ + typeof(Base.:+), + ProduceLogLikelihoodAccumulator, + DynamicPPL.LogLikelihoodAccumulator, + }, + }, +) + return true +end +Libtask.@might_produce(DynamicPPL.accumulate_observe!!) +Libtask.@might_produce(DynamicPPL.tilde_observe!!) +Libtask.@might_produce(DynamicPPL.tilde_assume!!) # GibbsContext turns assumes into observes +# Every model / submodel evaluator takes a `DynamicPPL.Model`, so this covers them all. +# See https://github.com/TuringLang/Libtask.jl/issues/217. +Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true + +# A varinfo carrying every accumulator a retained particle needs: the produce-aware +# likelihood, the prior and Jacobian terms (for log-density chain metadata), and raw values. +function particle_varinfo() + vi = DynamicPPL.OnlyAccsVarInfo() + vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator()) + vi = DynamicPPL.setacc!!(vi, DynamicPPL.LogPriorAccumulator()) + vi = DynamicPPL.setacc!!(vi, DynamicPPL.LogJacobianAccumulator()) + vi = DynamicPPL.setacc!!(vi, DynamicPPL.RawValueAccumulator(true)) + return vi +end + +# +# Resampling schemes +# + +abstract type AbstractResampler end + +"""Whether to resample given the normalized `weights`. Bare schemes always resample.""" +should_resample(::AbstractResampler, weights) = true + +"""Draw `n` ancestor indices from `1:length(weights)` with probabilities `weights`.""" +function resample_indices end + +"Multinomial resampling: `n` independent draws from the categorical over `weights`." +struct Multinomial <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::Multinomial, weights, n::Integer) + return rand(rng, Distributions.Categorical(weights), n) +end + +"Systematic resampling: one uniform placed on a regular grid of `n` points." +struct Systematic <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::Systematic, weights, n::Integer) + v = n * weights[1] + u = oftype(v, rand(rng)) + indices = Vector{Int}(undef, n) + s = 1 + for k in 1:n + while v < u + s += 1 + v += n * weights[s] + end + indices[k] = s + u += one(u) + end + return indices +end + +""" + ESSResampler(threshold, scheme = Systematic()) + +Resample with `scheme`, but only when the effective sample size drops below +`threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). +""" +struct ESSResampler{R<:AbstractResampler} <: AbstractResampler + threshold::Float64 + scheme::R +end +ESSResampler(threshold::Real) = ESSResampler(Float64(threshold), Systematic()) + +function should_resample(resampler::ESSResampler, weights) + ess = inv(sum(abs2, weights)) + return ess ≤ resampler.threshold * length(weights) +end +function resample_indices(rng::AbstractRNG, resampler::ESSResampler, weights, n::Integer) + return resample_indices(rng, resampler.scheme, weights, n) +end + +# +# Particle sweep +# +# In a conditional sweep the last particle is the reference: it is always retained and +# replays its recorded randomness, while the other `n-1` slots are resampled from all `n` +# particles (so they may descend from the reference). + +logweights(particles) = [p.logweight for p in particles] +normalized_weights(particles) = softmax(logweights(particles)) +logevidence(particles) = logsumexp(logweights(particles)) + +# Advance every particle by one observation; return `true` once all have finished. A model +# whose number of observations varies across executions leaves particles out of step. +function reweight!(particles, conditional::Bool) + n = length(particles) + n_done = 0 + for (i, p) in enumerate(particles) + score = advance!(p, conditional && i == n) + if score === nothing + n_done += 1 + else + p.logweight += score + end + end + n_done == 0 && return false + n_done == n && return true + return error( + "mis-aligned execution traces ($n_done/$n finished): the number of observations must not be random.", ) end -abstract type ParticleInference <: AbstractSampler end +# Resample (if the scheme calls for it) and fork the survivors, or -- when not resampling -- +# refresh each ordinary particle's seed so the next step draws fresh randomness. +function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional::Bool) + n = length(particles) + weights = normalized_weights(particles) + if should_resample(resampler, weights) + ancestors = resample_indices(rng, resampler, weights, conditional ? n - 1 : n) + old = copy(particles) + for (slot, a) in enumerate(ancestors) + child = fork(old[a], rng) + child.logweight = 0.0 + particles[slot] = child + end + conditional && (particles[n].logweight = 0.0) # reference retained, weight reset + else + for (i, p) in enumerate(particles) + # Refresh every particle's seed except the reference, which keeps replaying. + if !(conditional && i == n) + refresh!(p.rng) + end + end + end + return nothing +end -#### -#### Generic Sequential Monte Carlo sampler. -#### +# Run a full particle sweep in place, returning the log-evidence estimate. +function sweep!(rng::AbstractRNG, particles, resampler; conditional::Bool=false) + logZ = 0.0 + while true + resample_propagate!(rng, particles, resampler, conditional) + logZ0 = logevidence(particles) + done = reweight!(particles, conditional) + logZ += logevidence(particles) - logZ0 + done && break + end + return logZ +end + +# +# Sequential Monte Carlo +# + +abstract type ParticleInference <: AbstractSampler end """ $(TYPEDEF) @@ -96,34 +340,29 @@ Sequential Monte Carlo sampler. $(TYPEDFIELDS) """ -struct SMC{R} <: ParticleInference +struct SMC{R<:AbstractResampler} <: ParticleInference + "resampling scheme" resampler::R end """ - SMC([resampler = AdvancedPS.ResampleWithESSThreshold()]) - SMC([resampler = AdvancedPS.resample_systematic, ]threshold) - -Create a sequential Monte Carlo sampler of type [`SMC`](@ref). + SMC([resampler = ESSResampler(0.5)]) + SMC([scheme = Systematic(), ]threshold) -If the algorithm for the resampling step is not specified explicitly, systematic resampling -is performed if the estimated effective sample size per particle drops below 0.5. +Sequential Monte Carlo sampler. By default systematic resampling is triggered whenever the +effective sample size drops below half the number of particles. """ -SMC() = SMC(AdvancedPS.ResampleWithESSThreshold()) - -# Convenient constructors with ESS threshold -function SMC(resampler, threshold::Real) - return SMC(AdvancedPS.ResampleWithESSThreshold(resampler, threshold)) -end -function SMC(threshold::Real) - return SMC(AdvancedPS.resample_systematic, threshold) +SMC() = SMC(ESSResampler(0.5)) +SMC(threshold::Real) = SMC(ESSResampler(threshold)) +function SMC(scheme::AbstractResampler, threshold::Real) + return SMC(ESSResampler(Float64(threshold), scheme)) end -struct SMCState{P,F<:AbstractFloat} +struct SMCState{P,W} particles::P - particleindex::Int - # The logevidence after aggregating all samples together. - average_logevidence::F + weights::W + index::Int + logevidence::Float64 end function AbstractMCMC.sample( @@ -142,416 +381,160 @@ function AbstractMCMC.sample( ) check_model && Turing._check_model(model, sampler) error_if_threadsafe_eval(model) - # SMC does not produce a Markov chain, so discard_initial and thinning do not apply. - # We consume these keyword arguments here to prevent them from being passed to - # AbstractMCMC.mcmcsample, which would cause a BoundsError (#1811). + # SMC is not a Markov chain, so these AbstractMCMC knobs do not apply. Consume them here + # rather than forwarding them to `mcmcsample` (which would `BoundsError`, see #1811). if discard_initial > 0 || thinning > 1 - @warn "SMC samplers do not support `discard_initial` or `thinning`. These keyword arguments will be ignored." + @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." end - chn = AbstractMCMC.mcmcsample( + chain = AbstractMCMC.mcmcsample( rng, model, sampler, N; - chain_type=chain_type, - initial_params=initial_params, - progress=progress, + chain_type, + initial_params, + progress, nparticles=N, kwargs..., ) - post_sample_hook(chn, sampler; verbose) - return chn + post_sample_hook(chain, sampler; verbose) + return chain end +# The whole sweep runs on the first step; later steps read off the population one particle at +# a time (SMC returns a weighted sample, not a chain). function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, - spl::SMC; + sampler::SMC; nparticles::Int, - initial_params, discard_sample=false, kwargs..., ) - # Create an empty VarInfo - accs = DynamicPPL.OnlyAccsVarInfo() - accs = DynamicPPL.setacc!!(accs, ProduceLogLikelihoodAccumulator()) - accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) - - # Create a new set of particles. - particles = AdvancedPS.ParticleContainer( - [AdvancedPS.Trace(model, accs, AdvancedPS.TracedRNG(), true) for _ in 1:nparticles], - AdvancedPS.TracedRNG(), - rng, - ) - - # Perform particle sweep. - logevidence = AdvancedPS.sweep!(rng, particles, spl.resampler, spl) - - # Extract the first particle and its weight. - particle = particles.vals[1] - weight = AdvancedPS.getweight(particles, 1) - - # Compute the first transition and the first state. - stats = (; weight=weight, logevidence=logevidence) - transition = if discard_sample - nothing - else - DynamicPPL.ParamsWithStats(particle.model.f.varinfo, stats) - end - state = SMCState(particles, 2, logevidence) + error_if_threadsafe_eval(model) + particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] + logZ = sweep!(rng, particles, sampler.resampler) + weights = normalized_weights(particles) - return transition, state + stats = (; weight=weights[1], logevidence=logZ) + transition = + discard_sample ? nothing : DynamicPPL.ParamsWithStats(particles[1].varinfo, stats) + return transition, SMCState(particles, weights, 2, logZ) end function AbstractMCMC.step( ::AbstractRNG, - model::DynamicPPL.Model, - spl::SMC, + ::DynamicPPL.Model, + ::SMC, state::SMCState; discard_sample=false, kwargs..., ) - # Extract the index of the current particle. - index = state.particleindex - - # Extract the current particle and its weight. - particles = state.particles - particle = particles.vals[index] - weight = AdvancedPS.getweight(particles, index) - - # Compute the transition and the next state. - stats = (; weight=weight, logevidence=state.average_logevidence) + i = state.index + stats = (; weight=state.weights[i], logevidence=state.logevidence) transition = if discard_sample nothing else - DynamicPPL.ParamsWithStats(deepcopy(particle.model.f.varinfo), stats) + DynamicPPL.ParamsWithStats(deepcopy(state.particles[i].varinfo), stats) end - nextstate = SMCState(state.particles, index + 1, state.average_logevidence) - - return transition, nextstate + return transition, SMCState(state.particles, state.weights, i + 1, state.logevidence) end -#### -#### Particle Gibbs sampler. -#### +# +# Particle Gibbs / conditional SMC +# """ $(TYPEDEF) -Particle Gibbs sampler. +Particle Gibbs (conditional SMC) sampler. # Fields $(TYPEDFIELDS) """ -struct PG{R} <: ParticleInference - """Number of particles.""" +struct PG{R<:AbstractResampler} <: ParticleInference + "number of particles" nparticles::Int - """Resampling algorithm.""" + "resampling scheme" resampler::R end """ -PG(n, [resampler = AdvancedPS.ResampleWithESSThreshold()]) -PG(n, [resampler = AdvancedPS.resample_systematic, ]threshold) - -Create a Particle Gibbs sampler of type [`PG`](@ref) with `n` particles. + PG(n, [resampler = ESSResampler(0.5)]) + PG(n, [scheme = Systematic(), ]threshold) -If the algorithm for the resampling step is not specified explicitly, systematic resampling -is performed if the estimated effective sample size per particle drops below 0.5. +Particle Gibbs sampler with `n` particles. By default systematic resampling is triggered +whenever the effective sample size drops below half the number of particles. """ -function PG(nparticles::Int) - return PG(nparticles, AdvancedPS.ResampleWithESSThreshold()) +PG(n::Int) = PG(n, ESSResampler(0.5)) +PG(n::Int, threshold::Real) = PG(n, ESSResampler(threshold)) +function PG(n::Int, scheme::AbstractResampler, threshold::Real) + return PG(n, ESSResampler(Float64(threshold), scheme)) end -# Convenient constructors with ESS threshold -function PG(nparticles::Int, resampler, threshold::Real) - return PG(nparticles, AdvancedPS.ResampleWithESSThreshold(resampler, threshold)) -end -function PG(nparticles::Int, threshold::Real) - return PG(nparticles, AdvancedPS.resample_systematic, threshold) -end - -""" - CSMC(...) +"Conditional SMC, an alias for [`PG`](@ref)." +const CSMC = PG -Equivalent to [`PG`](@ref). -""" -const CSMC = PG # type alias of PG as Conditional SMC - -struct PGState{V<:DynamicPPL.AbstractVarInfo,R<:Random.AbstractRNG} - vi::V +struct PGState{V<:DynamicPPL.AbstractVarInfo,R<:TracedRNG} + varinfo::V rng::R end +# First iteration: an ordinary (unconditional) particle sweep. function AbstractMCMC.step( - rng::AbstractRNG, model::DynamicPPL.Model, spl::PG; discard_sample=false, kwargs... + rng::AbstractRNG, model::DynamicPPL.Model, sampler::PG; discard_sample=false, kwargs... ) error_if_threadsafe_eval(model) - oavi = DynamicPPL.OnlyAccsVarInfo() - oavi = DynamicPPL.setacc!!(oavi, ProduceLogLikelihoodAccumulator()) - oavi = DynamicPPL.setacc!!(oavi, DynamicPPL.RawValueAccumulator(true)) - - # Create a new set of particles - num_particles = spl.nparticles - particles = AdvancedPS.ParticleContainer( - [ - AdvancedPS.Trace(model, oavi, AdvancedPS.TracedRNG(), true) for - _ in 1:num_particles - ], - AdvancedPS.TracedRNG(), - rng, - ) - - # Perform a particle sweep. - logevidence = AdvancedPS.sweep!(rng, particles, spl.resampler, spl) - - # Pick a particle to be retained. - Ws = AdvancedPS.getweights(particles) - index = AdvancedPS.randcat(rng, Ws) - reference = particles.vals[index] - - # Compute the first transition. - _vi = reference.model.f.varinfo - transition = if discard_sample - nothing - else - DynamicPPL.ParamsWithStats(deepcopy(_vi), (; logevidence=logevidence)) - end - - return transition, PGState(_vi, reference.rng) + particles = [ + Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:(sampler.nparticles) + ] + logZ = sweep!(rng, particles, sampler.resampler) + return pg_transition_and_state(rng, particles, logZ, discard_sample) end +# Subsequent iterations: conditional SMC given the retained trajectory, which the reference +# particle regenerates by replaying `state.rng` from the first step. function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, - spl::PG, + sampler::PG, state::PGState; discard_sample=false, kwargs..., ) - # Reset log-prob accs in reference particle, to avoid accumulating into the same accs - # across iterations. If the chosen particle for this iteration is the reference - # particle, this allows us to just read off the log-probs from the accumulators, - # without having to re-evaluate the model. - reference_vi = state.vi - reference_vi = DynamicPPL.setacc!!(reference_vi, ProduceLogLikelihoodAccumulator()) - reference_vi = DynamicPPL.setacc!!(reference_vi, DynamicPPL.LogPriorAccumulator()) - reference_vi = DynamicPPL.setacc!!(reference_vi, DynamicPPL.LogJacobianAccumulator()) - - # Create reference particle for which the samples will be retained. - reference = AdvancedPS.forkr(AdvancedPS.Trace(model, reference_vi, state.rng, false)) - - # Create a new set of particles with newly emptied accs - empty_accs = DynamicPPL.OnlyAccsVarInfo() - empty_accs = DynamicPPL.setacc!!(empty_accs, ProduceLogLikelihoodAccumulator()) - empty_accs = DynamicPPL.setacc!!(empty_accs, DynamicPPL.RawValueAccumulator(true)) - num_particles = spl.nparticles - x = map(1:num_particles) do i - if i != num_particles - return AdvancedPS.Trace(model, empty_accs, AdvancedPS.TracedRNG(), true) - else - return reference - end + error_if_threadsafe_eval(model) + n = sampler.nparticles + reference = Particle(model, particle_varinfo(), set_step!(deepcopy(state.rng), 1)) + particles = map(1:n) do i + i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end - particles = AdvancedPS.ParticleContainer(x, AdvancedPS.TracedRNG(), rng) - - # Perform a particle sweep. - logevidence = AdvancedPS.sweep!(rng, particles, spl.resampler, spl, reference) - - # Pick a particle to be retained. - Ws = AdvancedPS.getweights(particles) - index = AdvancedPS.randcat(rng, Ws) - newreference = particles.vals[index] + logZ = sweep!(rng, particles, sampler.resampler; conditional=true) + return pg_transition_and_state(rng, particles, logZ, discard_sample) +end - # Compute the transition. - _vi = newreference.model.f.varinfo +function pg_transition_and_state(rng, particles, logZ, discard_sample) + retained = particles[rand( + rng, Distributions.Categorical(normalized_weights(particles)) + )] transition = if discard_sample nothing else - DynamicPPL.ParamsWithStats(deepcopy(_vi), (; logevidence=logevidence)) + DynamicPPL.ParamsWithStats(deepcopy(retained.varinfo), (; logevidence=logZ)) end - - return transition, PGState(_vi, newreference.rng) -end - -""" - get_trace_local_varinfo() - -Get the varinfo stored in the 'taped globals' of a `Libtask.TapedTask`. This function -is meant to be called from *inside* the TapedTask itself. -""" -function get_trace_local_varinfo() - trace = Libtask.get_taped_globals(Any).other - return trace.model.f.varinfo::AbstractVarInfo -end - -""" - get_trace_local_resampled() - -Get the `resample` flag stored in the 'taped globals' of a `Libtask.TapedTask`. - -This indicates whether new variable values should be sampled from the prior or not. For -example, in SMC, this is true for all particles; in PG, this is true for all particles -except the reference particle, whose trajectory must be reproduced exactly. - -This function is meant to be called from *inside* the TapedTask itself. -""" -function get_trace_local_resampled() - trace = Libtask.get_taped_globals(Any).other - return trace.model.f.resample::Bool + return transition, PGState(retained.varinfo, retained.rng) end -""" - get_trace_local_rng() - -Get the RNG stored in the 'taped globals' of a `Libtask.TapedTask`, if one exists. +# +# Gibbs interface +# -This function is meant to be called from *inside* the TapedTask itself. -""" -function get_trace_local_rng() - return Libtask.get_taped_globals(Any).rng -end - -""" - set_trace_local_varinfo(vi::AbstractVarInfo) - -Set the `varinfo` stored in Libtask's taped globals. The 'other' taped global in Libtask -is expected to be an `AdvancedPS.Trace`. - -Returns `nothing`. - -This function is meant to be called from *inside* the TapedTask itself. -""" -function set_trace_local_varinfo(vi::AbstractVarInfo) - trace = Libtask.get_taped_globals(Any).other - trace.model.f.varinfo = vi - return nothing -end - -function DynamicPPL.tilde_assume!!( - ::ParticleMCMCContext, dist::Distribution, vn::VarName, template::Any, ::AbstractVarInfo -) - # Get all the info we need from the trace, namely, the stored VarInfo, and whether - # we need to sample a new value or use the existing one. - vi = get_trace_local_varinfo() - trng = get_trace_local_rng() - resample = get_trace_local_resampled() - # Modify the varinfo as appropriate. - values = DynamicPPL.get_raw_values(vi) - init_strat = if ~haskey(values, vn) || resample - DynamicPPL.InitFromPrior() - else - DynamicPPL.InitFromParams(values, nothing) - end - ctx = DynamicPPL.InitContext(trng, init_strat, DynamicPPL.UnlinkAll()) - x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, vi) - # Set the varinfo back in the trace. - set_trace_local_varinfo(vi) - return x, vi -end - -function DynamicPPL.tilde_observe!!( - ::ParticleMCMCContext, - right::Distribution, - left, - vn::Union{VarName,Nothing}, - template::Any, - vi::AbstractVarInfo, -) - vi = get_trace_local_varinfo() - left, vi = DynamicPPL.tilde_observe!!(DefaultContext(), right, left, vn, template, vi) - set_trace_local_varinfo(vi) - return left, vi -end - -# Convenient constructor -function AdvancedPS.Trace( - model::Model, varinfo::AbstractVarInfo, rng::AdvancedPS.TracedRNG, resample::Bool -) - newvarinfo = deepcopy(varinfo) - tmodel = TracedModel(model, newvarinfo, rng, resample) - newtrace = AdvancedPS.Trace(tmodel, rng) - return newtrace -end - -""" -ProduceLogLikelihoodAccumulator{T<:Real} <: AbstractAccumulator - -Exactly like `LogLikelihoodAccumulator`, but calls `Libtask.produce` on change of value. - -# Fields -$(TYPEDFIELDS) -""" -struct ProduceLogLikelihoodAccumulator{T<:Real} <: DynamicPPL.LogProbAccumulator{T} - "the scalar log likelihood value" - logp::T -end - -# Note that this uses the same name as `LogLikelihoodAccumulator`. Thus only one of the two -# can be used in a given VarInfo. -DynamicPPL.accumulator_name(::Type{<:ProduceLogLikelihoodAccumulator}) = :LogLikelihood -DynamicPPL.logp(acc::ProduceLogLikelihoodAccumulator) = acc.logp - -function DynamicPPL.acclogp(acc1::ProduceLogLikelihoodAccumulator, val) - # The below line is the only difference from `LogLikelihoodAccumulator`. - Libtask.produce(val) - return ProduceLogLikelihoodAccumulator(acc1.logp + val) -end - -function DynamicPPL.accumulate_assume!!( - acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, right, template -) - return acc -end -function DynamicPPL.accumulate_observe!!( - acc::ProduceLogLikelihoodAccumulator, right, left, vn, template -) - return DynamicPPL.acclogp(acc, Distributions.loglikelihood(right, left)) -end - -# We need to tell Libtask which calls may have `produce` calls within them. In practice most -# of these won't be needed, because of inlining and the fact that `might_produce` is only -# called on `:invoke` expressions rather than `:call`s, but since those are implementation -# details of the compiler, we set a bunch of methods as might_produce = true. We start with -# adding to ProduceLogLikelihoodAccumulator, which is what calls `produce`, and go up the -# call stack. -Libtask.@might_produce(DynamicPPL.accloglikelihood!!) -function Libtask.might_produce( - ::Type{ - <:Tuple{ - typeof(Base.:+), - ProduceLogLikelihoodAccumulator, - DynamicPPL.LogLikelihoodAccumulator, - }, - }, -) - return true -end -Libtask.@might_produce(DynamicPPL.accumulate_observe!!) -Libtask.@might_produce(DynamicPPL.tilde_observe!!) -# Could tilde_assume!! have tighter type bounds on the arguments, namely a GibbsContext? -# That's the only thing that makes tilde_assume calls result in tilde_observe calls. -Libtask.@might_produce(DynamicPPL.tilde_assume!!) - -# This handles all models and submodel evaluator functions (including those with keyword -# arguments). The key to this is realising that all model evaluator functions take -# DynamicPPL.Model as an argument, so we can just check for that. See -# https://github.com/TuringLang/Libtask.jl/issues/217. -Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true - -#### -#### Gibbs interface -#### - -function gibbs_get_raw_values(state::PGState) - return DynamicPPL.get_raw_values(state.vi) -end +gibbs_get_raw_values(state::PGState) = DynamicPPL.get_raw_values(state.varinfo) function gibbs_update_state!!( ::PG, state::PGState, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple ) - init_strat = DynamicPPL.InitFromParams(global_vals, nothing) - new_vi = last(DynamicPPL.init!!(model, state.vi, init_strat, DynamicPPL.UnlinkAll())) + init = DynamicPPL.InitFromParams(global_vals, nothing) + new_vi = last(DynamicPPL.init!!(model, state.varinfo, init, DynamicPPL.UnlinkAll())) return PGState(new_vi, state.rng) end diff --git a/src/mcmc/traced_rng.jl b/src/mcmc/traced_rng.jl new file mode 100644 index 0000000000..de02a3f151 --- /dev/null +++ b/src/mcmc/traced_rng.jl @@ -0,0 +1,59 @@ +# +# A counter-based RNG that records the seed used at each model step, so that a particle's +# trajectory can be replayed exactly. This is what lets a conditional SMC sweep reproduce +# its reference trajectory: the reference simply replays its recorded seeds. +# + +""" + TracedRNG([rng = Random.default_rng()]) + +A `Random123.Philox2x` generator that remembers the seed (`key`) it used at each model step +in `keys`, indexed by the step counter `count`. + + - [`save_state!`](@ref) records the current seed (ordinary particles); + - [`load_state!`](@ref) restores `keys[count]`, replaying that step's randomness (the + reference trajectory). +""" +mutable struct TracedRNG{K,T<:Random123.AbstractR123} <: Random.AbstractRNG + count::Int + rng::T + keys::Vector{K} +end + +function TracedRNG(inner::Random123.AbstractR123{T}) where {T<:Unsigned} + Random123.set_counter!(inner, 0) + return TracedRNG(1, inner, T[]) +end +function TracedRNG(rng::AbstractRNG=Random.default_rng()) + inner = Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) + return TracedRNG(inner) +end + +Random.rng_native_52(trng::TracedRNG) = Random.rng_native_52(trng.rng) +Random.rand(trng::TracedRNG, ::Type{T}) where {T<:Unsigned} = Random.rand(trng.rng, T) + +"The current seed of the inner generator." +inner_key(rng::Random123.Philox2x) = rng.key + +"Reseed and rewind the inner generator. The model-step counter is left untouched." +function Random.seed!(trng::TracedRNG, key) + Random.seed!(trng.rng, key) + Random123.set_counter!(trng.rng, 0) + return trng +end + +"Record the seed used at the current step." +save_state!(trng::TracedRNG) = push!(trng.keys, inner_key(trng.rng)) + +"Replay the seed recorded at the current step." +load_state!(trng::TracedRNG) = Random.seed!(trng, trng.keys[trng.count]) + +"Set / advance the model-step counter." +set_step!(trng::TracedRNG, n::Integer) = (trng.count = n; trng) +inc_step!(trng::TracedRNG, n::Integer=1) = (trng.count += n; trng) + +"Deterministically derive a fresh seed from `key`." +split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) + +"Reseed from the generator's own current state (used between steps when not resampling)." +refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) diff --git a/test/Project.toml b/test/Project.toml index 1110301a92..0e0bdcfb69 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,7 +3,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -AdvancedPS = "576499cb-2369-40b2-a588-c64705576edc" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" @@ -45,7 +44,6 @@ ADTypes = "1" AbstractMCMC = "5.13" AbstractPPL = "0.14, 0.15" AdvancedMH = "0.8.9" -AdvancedPS = "0.7.2" AdvancedVI = "0.7" Aqua = "0.8" BangBang = "0.4" diff --git a/test/essential/container.jl b/test/essential/container.jl index 3554526c97..7b54411992 100644 --- a/test/essential/container.jl +++ b/test/essential/container.jl @@ -1,10 +1,12 @@ module ContainerTests -using AdvancedPS: AdvancedPS using Distributions: Bernoulli, Beta, Gamma, Normal using DynamicPPL: DynamicPPL, @model +using Random: Xoshiro using Test: @test, @testset using Turing +using Turing.Inference: + Particle, TracedRNG, particle_varinfo, advance!, fork, set_step!, refresh! @testset "container.jl" begin @model function test() @@ -17,26 +19,25 @@ using Turing return x end - @testset "constructor" begin - accs = DynamicPPL.OnlyAccsVarInfo() - accs = DynamicPPL.setacc!!(accs, Turing.Inference.ProduceLogLikelihoodAccumulator()) - accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) - sampler = PG(10) - model = test() - trace = AdvancedPS.Trace(model, accs, AdvancedPS.TracedRNG(), false) - - # Make sure the backreference from taped_globals to the trace is in place. - @test trace.model.ctask.taped_globals.other === trace - - res = AdvancedPS.advance!(trace, false) - @test res ≈ -log(2) - - # Catch broken copy, espetially for RNG / VarInfo - newtrace = AdvancedPS.fork(trace) - res2 = AdvancedPS.advance!(trace) + @testset "advance!" begin + # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. + particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) + @test advance!(particle, false) ≈ -log(2) + @test advance!(particle, false) ≈ -log(2) # `0 ~ Bernoulli(0.5)` + @test advance!(particle, false) === nothing # model finished end @testset "fork" begin + particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) + advance!(particle, false) + child = fork(particle, Xoshiro(1)) + # Independent continuations: advancing one does not touch the other. + @test advance!(child, false) ≈ -log(2) + @test particle.varinfo !== child.varinfo + @test advance!(particle, false) ≈ -log(2) + end + + @testset "rng replay" begin @model function normal() a ~ Normal(0, 1) 3 ~ Normal(a, 2) @@ -44,17 +45,22 @@ using Turing 1.5 ~ Normal(b, 2) return a, b end - accs = DynamicPPL.OnlyAccsVarInfo() - accs = DynamicPPL.setacc!!(accs, Turing.Inference.ProduceLogLikelihoodAccumulator()) - accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) - sampler = PG(10) - model = normal() - trace = AdvancedPS.Trace(model, accs, AdvancedPS.TracedRNG(), false) + # Run a particle to completion, then replay it from its recorded seeds (as the + # reference trajectory of a conditional sweep does) and check it regenerates exactly. + # Replay relies on each step using a distinct seed, so we refresh before every step + # exactly as the sweep's no-resample path does. + particle = Particle(normal(), particle_varinfo(), TracedRNG(Xoshiro(23))) + while (refresh!(particle.rng); advance!(particle, false)) !== nothing + end + values = DynamicPPL.get_raw_values(particle.varinfo) - newtrace = AdvancedPS.forkr(trace) - # Catch broken replay mechanism - @test AdvancedPS.advance!(trace) ≈ AdvancedPS.advance!(newtrace) + reference = Particle( + normal(), particle_varinfo(), set_step!(deepcopy(particle.rng), 1) + ) + while advance!(reference, true) !== nothing + end + @test DynamicPPL.get_raw_values(reference.varinfo) == values end end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 2cbc4dbece..b4267567e0 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -2,7 +2,7 @@ module ParticleMCMCTests using ..Models: gdemo_default using ..SamplerTestUtils: test_chain_logp_metadata -using AdvancedPS: ResampleWithESSThreshold, resample_systematic, resample_multinomial +using Turing.Inference: Systematic, Multinomial, ESSResampler using Distributions: Bernoulli, Beta, Gamma, Normal, sample using FlexiChains: VNChain using Random: Random @@ -12,17 +12,10 @@ using Turing @testset "SMC" begin @testset "constructor" begin - s = SMC() - @test s.resampler == ResampleWithESSThreshold() - - s = SMC(0.6) - @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6) - - s = SMC(resample_multinomial, 0.6) - @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6) - - s = SMC(resample_systematic) - @test s.resampler === resample_systematic + @test SMC().resampler == ESSResampler(0.5) + @test SMC(0.6).resampler == ESSResampler(0.6) + @test SMC(Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) + @test SMC(Systematic()).resampler == Systematic() end @testset "basic model" begin @@ -120,21 +113,11 @@ end @testset "PG" begin @testset "constructor" begin - s = PG(10) - @test s.nparticles == 10 - @test s.resampler == ResampleWithESSThreshold() - - s = PG(60, 0.6) - @test s.nparticles == 60 - @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6) - - s = PG(80, resample_multinomial, 0.6) - @test s.nparticles == 80 - @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6) - - s = PG(100, resample_systematic) - @test s.nparticles == 100 - @test s.resampler === resample_systematic + @test PG(10).nparticles == 10 + @test PG(10).resampler == ESSResampler(0.5) + @test PG(60, 0.6).resampler == ESSResampler(0.6) + @test PG(80, Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) + @test PG(100, Systematic()).resampler == Systematic() end @testset "chain log-density metadata" begin From a2eb174e59d5a0b6822707cf9ccaf3e6312a54bf Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 20:56:23 +0100 Subject: [PATCH 02/58] Consolidate particle MCMC files, drop redundant fork step, add notes - Inline traced_rng.jl into particle_mcmc.jl as a section (it must precede Particle/PGState, which name TracedRNG in their types). - Merge the container unit tests into test/mcmc/particle_mcmc.jl and remove the now-empty essential test group. - Drop the set_taped_globals! call in fork: deepcopy already preserves the task-particle back-reference (verified against the full particle suite). - Add WHY notes for the might_produce hints and the reference RNG rewind. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 75 ++++++++++++++++++++++++++++++++++--- src/mcmc/traced_rng.jl | 59 ----------------------------- test/essential/container.jl | 67 --------------------------------- test/mcmc/particle_mcmc.jl | 71 ++++++++++++++++++++++++++++++++++- test/runtests.jl | 4 -- 5 files changed, 139 insertions(+), 137 deletions(-) delete mode 100644 src/mcmc/traced_rng.jl delete mode 100644 test/essential/container.jl diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index ac643b0d1b..7cbd7d3c0d 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -16,7 +16,67 @@ using StatsFuns: softmax, logsumexp -include("traced_rng.jl") +# +# Traced RNG +# +# A counter-based RNG that records the seed used at each model step, so that a particle's +# trajectory can be replayed exactly: the conditional-SMC reference regenerates itself by +# replaying its recorded seeds. This section comes first because `Particle` and `PGState` +# name `TracedRNG` in their type signatures. + +""" + TracedRNG([rng = Random.default_rng()]) + +A `Random123.Philox2x` generator that remembers the seed (`key`) it used at each model step +in `keys`, indexed by the step counter `count`. + + - [`save_state!`](@ref) records the current seed (ordinary particles); + - [`load_state!`](@ref) restores `keys[count]`, replaying that step's randomness (the + reference trajectory). +""" +mutable struct TracedRNG{K,T<:Random123.AbstractR123} <: Random.AbstractRNG + count::Int + rng::T + keys::Vector{K} +end + +function TracedRNG(inner::Random123.AbstractR123{T}) where {T<:Unsigned} + Random123.set_counter!(inner, 0) + return TracedRNG(1, inner, T[]) +end +function TracedRNG(rng::AbstractRNG=Random.default_rng()) + inner = Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) + return TracedRNG(inner) +end + +Random.rng_native_52(trng::TracedRNG) = Random.rng_native_52(trng.rng) +Random.rand(trng::TracedRNG, ::Type{T}) where {T<:Unsigned} = Random.rand(trng.rng, T) + +"The current seed of the inner generator." +inner_key(rng::Random123.Philox2x) = rng.key + +"Reseed and rewind the inner generator. The model-step counter is left untouched." +function Random.seed!(trng::TracedRNG, key) + Random.seed!(trng.rng, key) + Random123.set_counter!(trng.rng, 0) + return trng +end + +"Record the seed used at the current step." +save_state!(trng::TracedRNG) = push!(trng.keys, inner_key(trng.rng)) + +"Replay the seed recorded at the current step." +load_state!(trng::TracedRNG) = Random.seed!(trng, trng.keys[trng.count]) + +"Set / advance the model-step counter." +set_step!(trng::TracedRNG, n::Integer) = (trng.count = n; trng) +inc_step!(trng::TracedRNG, n::Integer=1) = (trng.count += n; trng) + +"Deterministically derive a fresh seed from `key`." +split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) + +"Reseed from the generator's own current state (used between steps when not resampling)." +refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) function error_if_threadsafe_eval(model::DynamicPPL.Model) if DynamicPPL.requires_threadsafe(model) @@ -82,13 +142,12 @@ end Copy `particle` into an independent continuation seeded from `rng`. `deepcopy` forks the underlying `TapedTask` (Libtask defines `copy` as `deepcopy`) and preserves the -task↔particle back-reference, which we reset explicitly to be safe. Reseeding switches the -child from replaying to sampling afresh; `keys` is truncated to the steps already taken so a -child of the reference forgets the reference's future. +task↔particle back-reference. Reseeding switches the child from replaying to sampling afresh; +`keys` is truncated to the steps already taken so a child of the reference forgets the +reference's future. """ function fork(particle::Particle, rng::AbstractRNG) child = deepcopy(particle) - Libtask.set_taped_globals!(child.task, child) Random.seed!(child.rng, rand(rng, UInt64)) resize!(child.rng.keys, child.rng.count - 1) return child @@ -170,7 +229,12 @@ function DynamicPPL.accumulate_observe!!( end # Tell Libtask which calls may contain `produce`, walking up the call stack from `acclogp`. +# Over-approximating is safe (a wrongly-marked call just gets instrumented); missing a real +# one is not, so we err towards marking. Libtask.@might_produce(DynamicPPL.accloglikelihood!!) +# Merging accumulators (across submodels or Gibbs blocks) can add a +# ProduceLogLikelihoodAccumulator to a plain one, which routes through the producing +# `acclogp` -- so this `+` may itself produce. function Libtask.might_produce( ::Type{ <:Tuple{ @@ -505,6 +569,7 @@ function AbstractMCMC.step( ) error_if_threadsafe_eval(model) n = sampler.nparticles + # Rewind the retained RNG to step 1 so the reference replays its trajectory from the start. reference = Particle(model, particle_varinfo(), set_step!(deepcopy(state.rng), 1)) particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference diff --git a/src/mcmc/traced_rng.jl b/src/mcmc/traced_rng.jl deleted file mode 100644 index de02a3f151..0000000000 --- a/src/mcmc/traced_rng.jl +++ /dev/null @@ -1,59 +0,0 @@ -# -# A counter-based RNG that records the seed used at each model step, so that a particle's -# trajectory can be replayed exactly. This is what lets a conditional SMC sweep reproduce -# its reference trajectory: the reference simply replays its recorded seeds. -# - -""" - TracedRNG([rng = Random.default_rng()]) - -A `Random123.Philox2x` generator that remembers the seed (`key`) it used at each model step -in `keys`, indexed by the step counter `count`. - - - [`save_state!`](@ref) records the current seed (ordinary particles); - - [`load_state!`](@ref) restores `keys[count]`, replaying that step's randomness (the - reference trajectory). -""" -mutable struct TracedRNG{K,T<:Random123.AbstractR123} <: Random.AbstractRNG - count::Int - rng::T - keys::Vector{K} -end - -function TracedRNG(inner::Random123.AbstractR123{T}) where {T<:Unsigned} - Random123.set_counter!(inner, 0) - return TracedRNG(1, inner, T[]) -end -function TracedRNG(rng::AbstractRNG=Random.default_rng()) - inner = Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) - return TracedRNG(inner) -end - -Random.rng_native_52(trng::TracedRNG) = Random.rng_native_52(trng.rng) -Random.rand(trng::TracedRNG, ::Type{T}) where {T<:Unsigned} = Random.rand(trng.rng, T) - -"The current seed of the inner generator." -inner_key(rng::Random123.Philox2x) = rng.key - -"Reseed and rewind the inner generator. The model-step counter is left untouched." -function Random.seed!(trng::TracedRNG, key) - Random.seed!(trng.rng, key) - Random123.set_counter!(trng.rng, 0) - return trng -end - -"Record the seed used at the current step." -save_state!(trng::TracedRNG) = push!(trng.keys, inner_key(trng.rng)) - -"Replay the seed recorded at the current step." -load_state!(trng::TracedRNG) = Random.seed!(trng, trng.keys[trng.count]) - -"Set / advance the model-step counter." -set_step!(trng::TracedRNG, n::Integer) = (trng.count = n; trng) -inc_step!(trng::TracedRNG, n::Integer=1) = (trng.count += n; trng) - -"Deterministically derive a fresh seed from `key`." -split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) - -"Reseed from the generator's own current state (used between steps when not resampling)." -refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) diff --git a/test/essential/container.jl b/test/essential/container.jl deleted file mode 100644 index 7b54411992..0000000000 --- a/test/essential/container.jl +++ /dev/null @@ -1,67 +0,0 @@ -module ContainerTests - -using Distributions: Bernoulli, Beta, Gamma, Normal -using DynamicPPL: DynamicPPL, @model -using Random: Xoshiro -using Test: @test, @testset -using Turing -using Turing.Inference: - Particle, TracedRNG, particle_varinfo, advance!, fork, set_step!, refresh! - -@testset "container.jl" begin - @model function test() - a ~ Normal(0, 1) - x ~ Bernoulli(1) - b ~ Gamma(2, 3) - 1 ~ Bernoulli(x / 2) - c ~ Beta() - 0 ~ Bernoulli(x / 2) - return x - end - - @testset "advance!" begin - # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. - particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) - @test advance!(particle, false) ≈ -log(2) - @test advance!(particle, false) ≈ -log(2) # `0 ~ Bernoulli(0.5)` - @test advance!(particle, false) === nothing # model finished - end - - @testset "fork" begin - particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) - advance!(particle, false) - child = fork(particle, Xoshiro(1)) - # Independent continuations: advancing one does not touch the other. - @test advance!(child, false) ≈ -log(2) - @test particle.varinfo !== child.varinfo - @test advance!(particle, false) ≈ -log(2) - end - - @testset "rng replay" begin - @model function normal() - a ~ Normal(0, 1) - 3 ~ Normal(a, 2) - b ~ Normal(a, 1) - 1.5 ~ Normal(b, 2) - return a, b - end - - # Run a particle to completion, then replay it from its recorded seeds (as the - # reference trajectory of a conditional sweep does) and check it regenerates exactly. - # Replay relies on each step using a distinct seed, so we refresh before every step - # exactly as the sweep's no-resample path does. - particle = Particle(normal(), particle_varinfo(), TracedRNG(Xoshiro(23))) - while (refresh!(particle.rng); advance!(particle, false)) !== nothing - end - values = DynamicPPL.get_raw_values(particle.varinfo) - - reference = Particle( - normal(), particle_varinfo(), set_step!(deepcopy(particle.rng), 1) - ) - while advance!(reference, true) !== nothing - end - @test DynamicPPL.get_raw_values(reference.varinfo) == values - end -end - -end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index b4267567e0..da708f1bad 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -2,10 +2,21 @@ module ParticleMCMCTests using ..Models: gdemo_default using ..SamplerTestUtils: test_chain_logp_metadata -using Turing.Inference: Systematic, Multinomial, ESSResampler +using DynamicPPL: DynamicPPL +using Turing.Inference: + Systematic, + Multinomial, + ESSResampler, + Particle, + TracedRNG, + particle_varinfo, + advance!, + fork, + set_step!, + refresh! using Distributions: Bernoulli, Beta, Gamma, Normal, sample using FlexiChains: VNChain -using Random: Random +using Random: Random, Xoshiro using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset using Turing @@ -236,4 +247,60 @@ end end end +@testset "particle container" begin + @model function test() + a ~ Normal(0, 1) + x ~ Bernoulli(1) + b ~ Gamma(2, 3) + 1 ~ Bernoulli(x / 2) + c ~ Beta() + 0 ~ Bernoulli(x / 2) + return x + end + + @testset "advance!" begin + # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. + particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) + @test advance!(particle, false) ≈ -log(2) + @test advance!(particle, false) ≈ -log(2) # `0 ~ Bernoulli(0.5)` + @test advance!(particle, false) === nothing # model finished + end + + @testset "fork" begin + particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) + advance!(particle, false) + child = fork(particle, Xoshiro(1)) + # Independent continuations: advancing one does not touch the other. + @test advance!(child, false) ≈ -log(2) + @test particle.varinfo !== child.varinfo + @test advance!(particle, false) ≈ -log(2) + end + + @testset "rng replay" begin + @model function normal() + a ~ Normal(0, 1) + 3 ~ Normal(a, 2) + b ~ Normal(a, 1) + 1.5 ~ Normal(b, 2) + return a, b + end + + # Run a particle to completion, then replay it from its recorded seeds (as the + # reference trajectory of a conditional sweep does) and check it regenerates exactly. + # Replay relies on each step using a distinct seed, so we refresh before every step + # exactly as the sweep's no-resample path does. + particle = Particle(normal(), particle_varinfo(), TracedRNG(Xoshiro(23))) + while (refresh!(particle.rng); advance!(particle, false)) !== nothing + end + values = DynamicPPL.get_raw_values(particle.varinfo) + + reference = Particle( + normal(), particle_varinfo(), set_step!(deepcopy(particle.rng), 1) + ) + while advance!(reference, true) !== nothing + end + @test DynamicPPL.get_raw_values(reference.varinfo) == values + end +end + end diff --git a/test/runtests.jl b/test/runtests.jl index c4faea5a37..b6065d74d0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -38,10 +38,6 @@ end @timeit_include("ad.jl") end - @testset "essential" verbose = true begin - @timeit_include("essential/container.jl") - end - @testset "samplers (without AD)" verbose = true begin @timeit_include("mcmc/abstractmcmc.jl") @timeit_include("mcmc/callbacks.jl") From 3e0960b25a45ee4483f7621f5307b8272edad41e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:08:24 +0100 Subject: [PATCH 03/58] Default to stratified resampling Stratified resampling satisfies the unbiased-offspring condition needed for the particle Gibbs invariance proof (Andrieu, Doucet & Holenstein, 2010) and is consistent as N grows. Systematic resampling shares the expected offspring counts but is order-dependent and not consistent in general (Gerber, Chopin & Whiteley, 2019), so it sits outside the invariance proof. Make stratified the default scheme for ESSResampler (hence for SMC and PG); systematic remains available explicitly. A note in the resampling section records the trade-off. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 38 +++++++++++++++++++++++++++++++------- test/mcmc/particle_mcmc.jl | 2 ++ 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 7cbd7d3c0d..d9feead3c1 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -267,6 +267,13 @@ end # # Resampling schemes # +# On theoretical correctness: particle Gibbs (and the SMC evidence estimate) are justified +# for resampling schemes whose offspring counts satisfy `E[Oᵏ] = N·Wᵏ` (Andrieu, Doucet & +# Holenstein, 2010, Assumption 2). Multinomial and stratified resampling meet this and are +# also consistent as `N → ∞`. Systematic resampling has the same expected counts, but its +# single shared uniform makes it order-dependent and it is not consistent in general (Gerber, +# Chopin & Whiteley, 2019), so it falls outside the particle Gibbs invariance proof. We +# therefore default to stratified resampling and offer systematic only as an explicit choice. abstract type AbstractResampler end @@ -282,7 +289,24 @@ function resample_indices(rng::AbstractRNG, ::Multinomial, weights, n::Integer) return rand(rng, Distributions.Categorical(weights), n) end -"Systematic resampling: one uniform placed on a regular grid of `n` points." +"Stratified resampling: one independent uniform per stratum of width `1/n`." +struct Stratified <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::Stratified, weights, n::Integer) + v = n * weights[1] + indices = Vector{Int}(undef, n) + s = 1 + for k in 1:n + u = oftype(v, (k - 1) + rand(rng)) + while v < u + s += 1 + v += n * weights[s] + end + indices[k] = s + end + return indices +end + +"Systematic resampling: one shared uniform placed on a regular grid of `n` points." struct Systematic <: AbstractResampler end function resample_indices(rng::AbstractRNG, ::Systematic, weights, n::Integer) v = n * weights[1] @@ -301,7 +325,7 @@ function resample_indices(rng::AbstractRNG, ::Systematic, weights, n::Integer) end """ - ESSResampler(threshold, scheme = Systematic()) + ESSResampler(threshold, scheme = Stratified()) Resample with `scheme`, but only when the effective sample size drops below `threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). @@ -310,7 +334,7 @@ struct ESSResampler{R<:AbstractResampler} <: AbstractResampler threshold::Float64 scheme::R end -ESSResampler(threshold::Real) = ESSResampler(Float64(threshold), Systematic()) +ESSResampler(threshold::Real) = ESSResampler(Float64(threshold), Stratified()) function should_resample(resampler::ESSResampler, weights) ess = inv(sum(abs2, weights)) @@ -411,9 +435,9 @@ end """ SMC([resampler = ESSResampler(0.5)]) - SMC([scheme = Systematic(), ]threshold) + SMC([scheme = Stratified(), ]threshold) -Sequential Monte Carlo sampler. By default systematic resampling is triggered whenever the +Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. """ SMC() = SMC(ESSResampler(0.5)) @@ -526,9 +550,9 @@ end """ PG(n, [resampler = ESSResampler(0.5)]) - PG(n, [scheme = Systematic(), ]threshold) + PG(n, [scheme = Stratified(), ]threshold) -Particle Gibbs sampler with `n` particles. By default systematic resampling is triggered +Particle Gibbs sampler with `n` particles. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. """ PG(n::Int) = PG(n, ESSResampler(0.5)) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index da708f1bad..89bd846fdb 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -4,6 +4,7 @@ using ..Models: gdemo_default using ..SamplerTestUtils: test_chain_logp_metadata using DynamicPPL: DynamicPPL using Turing.Inference: + Stratified, Systematic, Multinomial, ESSResampler, @@ -24,6 +25,7 @@ using Turing @testset "SMC" begin @testset "constructor" begin @test SMC().resampler == ESSResampler(0.5) + @test SMC().resampler.scheme isa Stratified # stratified is the default scheme @test SMC(0.6).resampler == ESSResampler(0.6) @test SMC(Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) @test SMC(Systematic()).resampler == Systematic() From 06202318b3220abca9c3f5e9207d7b91f74a2d06 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:15:03 +0100 Subject: [PATCH 04/58] Drop redundant prior/Jacobian accumulators from particles Chain metadata recomputes the log-prior and Jacobian terms from a particle's raw values, so accumulating them during the sweep was wasted work; particles now carry only the produce-aware likelihood accumulator and the raw values. Metadata stays correct (test_chain_logp_metadata passes for SMC and PG). Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index d9feead3c1..2126c695ad 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -253,13 +253,12 @@ Libtask.@might_produce(DynamicPPL.tilde_assume!!) # GibbsContext turns assumes # See https://github.com/TuringLang/Libtask.jl/issues/217. Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true -# A varinfo carrying every accumulator a retained particle needs: the produce-aware -# likelihood, the prior and Jacobian terms (for log-density chain metadata), and raw values. +# A particle needs only the produce-aware likelihood accumulator (which drives reweighting) +# and the raw sampled values. The prior/Jacobian terms shown in chain metadata are recomputed +# downstream from the raw values, so accumulating them per particle would be wasted work. function particle_varinfo() vi = DynamicPPL.OnlyAccsVarInfo() vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator()) - vi = DynamicPPL.setacc!!(vi, DynamicPPL.LogPriorAccumulator()) - vi = DynamicPPL.setacc!!(vi, DynamicPPL.LogJacobianAccumulator()) vi = DynamicPPL.setacc!!(vi, DynamicPPL.RawValueAccumulator(true)) return vi end From 5d9d127d3bbaaa3ff9e92c8c531ea80f6934fe58 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:33:15 +0100 Subject: [PATCH 05/58] Reuse surviving parents on resample to cut allocations resample_propagate! deepcopied every resampled slot; deepcopying a Libtask TapedTask is the dominant allocation in a sweep, so this made PG allocate ~40% more than the AdvancedPS implementation (59.9M vs 41.9M allocations; ~22% slower). Reuse each surviving parent's object for its first offspring, deepcopying only additional offspring and any descending from the retained reference (as AdvancedPS does). Both kinds of child are reseeded via the new `reseed!` helper. Allocations now match the old implementation and PG is slightly faster; CSMC reference reproduction is unchanged (0 mismatches over the reference-consistency check, full suite green). Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 65 +++++++++++++++++++++++++-------------- 1 file changed, 42 insertions(+), 23 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 2126c695ad..eb7394aaf9 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -1,17 +1,24 @@ ### -### Particle filtering and particle MCMC samplers. +### Particle filtering and particle MCMC samplers: SMC, PG / conditional SMC. ### -### SMC runs a particle filter over a model; PG/CSMC wraps a *conditional* particle filter -### inside an MCMC loop. Both treat each `observe` statement as one filtering step: running -### the model under `ParticleMCMCContext` turns every likelihood term into a -### `Libtask.produce`, so a particle is a suspended model execution that we advance one -### observation at a time, reweight, and resample. +### Key design. +### A probabilistic model becomes a particle filter by reading each `observe` statement as one +### filtering step. Evaluated under `ParticleMCMCContext`, every likelihood term calls +### `Libtask.produce`, so a *particle* is a suspended model execution: we `advance!` it to its +### next `observe`, take the produced log-likelihood as its weight, then resample. SMC is one +### such sweep; particle Gibbs (PG/CSMC) runs a *conditional* sweep -- one particle is a fixed +### reference trajectory -- inside an MCMC loop. ### -### Each particle carries a `TracedRNG` that records the seed it used at every step. The -### reference trajectory of a conditional sweep is reproduced by *replaying* those seeds -### (`load_state!`), so it is regenerated exactly without storing its values. A particle -### forked from the reference is reseeded, which automatically switches it from replaying to -### sampling fresh -- no per-particle bookkeeping required. +### The reference is reproduced without storing its values: each particle's `TracedRNG` +### records the seed it drew at every step, and the reference simply *replays* those seeds +### (`load_state!`), regenerating its trajectory exactly. A particle forked from the reference +### is reseeded, which flips it from replaying to sampling fresh, so branching needs no +### per-particle flag. This is what keeps the reference handling small, and it rests on one +### invariant: every step must draw from a fresh seed -- guaranteed by the resample/refresh in +### `resample_propagate!` -- otherwise the recorded seeds collide and replay is wrong. +### +### Sections below: traced RNG; model evaluation via Libtask; resampling schemes; the particle +### sweep; the SMC sampler; the PG/CSMC sampler; the Gibbs-component interface. ### using StatsFuns: softmax, logsumexp @@ -138,21 +145,27 @@ function Particle( end """ - fork(particle, rng) + reseed!(particle, rng) -Copy `particle` into an independent continuation seeded from `rng`. `deepcopy` forks the -underlying `TapedTask` (Libtask defines `copy` as `deepcopy`) and preserves the -task↔particle back-reference. Reseeding switches the child from replaying to sampling afresh; -`keys` is truncated to the steps already taken so a child of the reference forgets the -reference's future. +Restart `particle` as a fresh continuation seeded from `rng`: it switches from replaying to +sampling afresh, and `keys` is truncated to the steps already taken so a particle descended +from the reference forgets the reference's future. Mutates and returns `particle`. """ -function fork(particle::Particle, rng::AbstractRNG) - child = deepcopy(particle) - Random.seed!(child.rng, rand(rng, UInt64)) - resize!(child.rng.keys, child.rng.count - 1) - return child +function reseed!(particle::Particle, rng::AbstractRNG) + Random.seed!(particle.rng, rand(rng, UInt64)) + resize!(particle.rng.keys, particle.rng.count - 1) + return particle end +""" + fork(particle, rng) + +Copy `particle` into an independent, reseeded continuation. `deepcopy` forks the underlying +`TapedTask` (Libtask defines `copy` as `deepcopy`) and preserves the task↔particle +back-reference; [`reseed!`](@ref) then gives it its own random stream. +""" +fork(particle::Particle, rng::AbstractRNG) = reseed!(deepcopy(particle), rng) + """ advance!(particle, isref) -> Union{Float64,Nothing} @@ -382,8 +395,14 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional if should_resample(resampler, weights) ancestors = resample_indices(rng, resampler, weights, conditional ? n - 1 : n) old = copy(particles) + seen = falses(n) for (slot, a) in enumerate(ancestors) - child = fork(old[a], rng) + # Reuse each surviving parent's object for its first offspring; only extra + # offspring -- and any offspring of the retained reference -- need the costly + # `deepcopy`. Either way the child is reseeded to continue independently. + reuse = !seen[a] && !(conditional && a == n) + seen[a] = true + child = reuse ? reseed!(old[a], rng) : fork(old[a], rng) child.logweight = 0.0 particles[slot] = child end From ad5979b85fe70470a852946250592d56e0ad844b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:35:42 +0100 Subject: [PATCH 06/58] Clarity pass: high-level design note, sectioning, key-step notes Open the file with a concise "Key design" overview (the observe-as-filtering-step idea, the RNG-replay reference mechanism, and the fresh-seed-per-step invariant it rests on) plus a section map. Move the threadsafe-eval guard into the model-evaluation section with a note on why particle samplers require it, and note the log-evidence telescoping in the sweep. Comments only; no behaviour change. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index eb7394aaf9..c2e829da44 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -85,6 +85,15 @@ split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) "Reseed from the generator's own current state (used between steps when not resampling)." refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) +# +# Model evaluation via Libtask +# +# A `Particle` is the only mutable state. It is stored as its `TapedTask`'s "taped globals", +# so the tilde overloads reach it from *inside* a running model via `get_taped_globals`. +# This keeps all state explicit on the particle -- no `task_local_storage`. + +# Particle samplers replay executions in a fixed order, so they cannot run models whose +# evaluation order is nondeterministic (e.g. a threaded `observe` loop). function error_if_threadsafe_eval(model::DynamicPPL.Model) if DynamicPPL.requires_threadsafe(model) throw( @@ -96,13 +105,6 @@ function error_if_threadsafe_eval(model::DynamicPPL.Model) return nothing end -# -# Model evaluation via Libtask -# -# A `Particle` is the only mutable state. It is stored as its `TapedTask`'s "taped globals", -# so the tilde overloads reach it from *inside* a running model via `get_taped_globals`. -# This keeps all state explicit on the particle -- no `task_local_storage`. - """ ParticleMCMCContext @@ -387,8 +389,8 @@ function reweight!(particles, conditional::Bool) ) end -# Resample (if the scheme calls for it) and fork the survivors, or -- when not resampling -- -# refresh each ordinary particle's seed so the next step draws fresh randomness. +# Resample (if the scheme calls for it) and propagate the survivors, or -- when not +# resampling -- refresh each ordinary particle's seed so the next step draws fresh randomness. function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional::Bool) n = length(particles) weights = normalized_weights(particles) @@ -425,6 +427,8 @@ function sweep!(rng::AbstractRNG, particles, resampler; conditional::Bool=false) resample_propagate!(rng, particles, resampler, conditional) logZ0 = logevidence(particles) done = reweight!(particles, conditional) + # Each observation contributes the log-ratio of total weight it adds; summed over the + # sweep these telescope into an estimate of the model's log-evidence log p(y). logZ += logevidence(particles) - logZ0 done && break end From 48e112a05b32ec7d466b7ac9a517b3a771d8968b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:45:31 +0100 Subject: [PATCH 07/58] Document particle-MCMC rewrite in HISTORY; add RNG-respected tests Add a HISTORY.md breaking-change entry for the native SMC/PG reimplementation (AdvancedPS removed, resamplers are now types, stratified is the new default), and add test_rng_respected coverage for SMC and PG so RNG determinism is locked in. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- HISTORY.md | 12 ++++++++++++ test/mcmc/particle_mcmc.jl | 10 +++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index 580932cb71..1ab9c7e041 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -19,6 +19,18 @@ Please see [the AdvancedVI changelog](https://github.com/TuringLang/AdvancedVI.j - `vi` with `KLMinScoreGradDescent` now optimises in unconstrained (linked) space, making it consistent with the other `KLMin...` algorithms. If you use it with a model that has constrained parameters, results may differ slightly from previous releases. +### Particle MCMC (SMC and PG) + +The particle samplers (`SMC`, and `PG` / `CSMC`) have been reimplemented natively and no longer depend on AdvancedPS. + +Resampling schemes are now types rather than functions: use `Stratified()`, `Systematic()`, or `Multinomial()`, optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size drops below `threshold * nparticles`. +For example, `SMC(Systematic())`, `SMC(0.5)`, or `PG(10, Multinomial(), 0.5)`. +The previous function-based API (`resample_systematic`, `resample_multinomial`, `AdvancedPS.ResampleWithESSThreshold`) is no longer available. + +The default resampling scheme is now **stratified** (previously systematic). +Stratified resampling is unbiased and consistent as the number of particles grows, and unlike systematic resampling it stays within the theoretical guarantees underpinning particle Gibbs. +Because of the new default and the reimplementation, sampling results will differ from previous releases. + ## Other changes **DifferentiationInterface removed as a direct dependency** diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 89bd846fdb..5a2ad72e4e 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -1,7 +1,7 @@ module ParticleMCMCTests using ..Models: gdemo_default -using ..SamplerTestUtils: test_chain_logp_metadata +using ..SamplerTestUtils: test_chain_logp_metadata, test_rng_respected using DynamicPPL: DynamicPPL using Turing.Inference: Stratified, @@ -60,6 +60,10 @@ using Turing test_chain_logp_metadata(SMC()) end + @testset "rng is respected" begin + test_rng_respected(SMC()) + end + @testset "logevidence" begin @model function test() a ~ Normal(0, 1) @@ -137,6 +141,10 @@ end test_chain_logp_metadata(PG(10)) end + @testset "rng is respected" begin + test_rng_respected(PG(10)) + end + @testset "logevidence" begin @model function test() a ~ Normal(0, 1) From 22dacedfebea70778ec166c2076b68fe04a4543e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 21:53:27 +0100 Subject: [PATCH 08/58] Port extra particle-MCMC tests from PR 2848 (ck/smc) Bring across the valuable tests the ck/smc branch added, adapted to the native internals: - resampling schemes: stratified and multinomial both hit the analytic coinflip posterior, and the two schemes produce genuinely different draws; - CSMC reference consistency: over 30 conditional sweeps the reference particle regenerates the retained trajectory exactly and its traced-RNG keys stay aligned with the trajectory length (previously only checked in a scratch script); - @addlogprob! is also respected under MH, not just PG; - a particle advanced without resampling matches a direct init!! evaluation (values and log-likelihood) seeded identically. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- test/mcmc/particle_mcmc.jl | 104 +++++++++++++++++++++++++++++++++++-- 1 file changed, 101 insertions(+), 3 deletions(-) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 5a2ad72e4e..787b713d05 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -2,7 +2,8 @@ module ParticleMCMCTests using ..Models: gdemo_default using ..SamplerTestUtils: test_chain_logp_metadata, test_rng_respected -using DynamicPPL: DynamicPPL +using ..NumericalTests: check_numerical +using DynamicPPL: DynamicPPL, extract_priors, get_raw_values, getloglikelihood using Turing.Inference: Stratified, Systematic, @@ -14,8 +15,11 @@ using Turing.Inference: advance!, fork, set_step!, - refresh! -using Distributions: Bernoulli, Beta, Gamma, Normal, sample + refresh!, + sweep!, + normalized_weights, + PGState +using Distributions: Bernoulli, Beta, Gamma, Normal, Uniform, Categorical, sample using FlexiChains: VNChain using Random: Random, Xoshiro using StableRNGs: StableRNG @@ -42,6 +46,27 @@ using Turing tested = sample(normal(), SMC(), 100) end + @testset "resampling schemes" begin + @model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end + end + obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1] + coin_model = coinflip(obs) + prior = extract_priors(coin_model)[@varname(p)] + exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs)) + + # every scheme targets the same posterior... + chn_strat = sample(StableRNG(23), coin_model, SMC(Stratified()), 100) + chn_multi = sample(StableRNG(23), coin_model, SMC(Multinomial()), 100) + check_numerical(chn_strat, [@varname(p)], [mean(exact)]; atol=0.1) + check_numerical(chn_multi, [@varname(p)], [mean(exact)]; atol=0.1) + # ...but the schemes are genuinely different, so the draws differ. + @test chn_strat[@varname(p)] != chn_multi[@varname(p)] + end + @testset "errors when number of observations is not fixed" begin @model function fail_smc() a ~ Normal(4, 5) @@ -172,6 +197,52 @@ end @test length(unique(c[:s])) == 1 end + @testset "ensuring reference consistency" begin + # In conditional SMC the retained trajectory must be regenerated *exactly* by the + # reference particle on the next iteration -- this is what makes CSMC valid. It fails + # if the traced-RNG step counter and the recorded seeds ever fall out of alignment. + @model function state_space_model(y) + ρ ~ Uniform(0, 1) + x = Vector{Float64}(undef, length(y) + 1) + x[1] ~ Normal(0, 1) + for t in eachindex(y) + x[t + 1] ~ Normal(ρ * x[t], 1) + y[t] ~ Normal(x[t + 1], 1) + end + end + + # Run PG's conditional sweep by hand so we can inspect the reference particle (slot + # N) and check it reproduces the trajectory we retained. Wrapped in a function to keep + # the mutating loop out of test soft scope. + function run_csmc(model, N, nsteps, rng) + draw(ps) = ps[rand(rng, Categorical(normalized_weights(ps)))] + particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:N] + sweep!(rng, particles, ESSResampler(0.5)) + state = let p = draw(particles) + PGState(p.varinfo, p.rng) + end + allok = true + for _ in 1:nsteps + ref = Particle(model, particle_varinfo(), set_step!(deepcopy(state.rng), 1)) + parts = map( + i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, + 1:N, + ) + sweep!(rng, parts, ESSResampler(0.5); conditional=true) + allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) + p = draw(parts) + state = PGState(p.varinfo, p.rng) + end + return allok, length(state.rng.keys) + end + + rng = StableRNG(1234) + y = randn(rng, 10) + allok, nkeys = run_csmc(state_space_model(y), 3, 30, rng) + @test allok # reference regenerated exactly every step + @test nkeys == length(y) + 1 # keys stay aligned with the trajectory length + end + @testset "addlogprob leads to reweighting" begin # Make sure that PG takes @addlogprob! into account. It didn't use to: # https://github.com/TuringLang/Turing.jl/issues/1996 @@ -188,6 +259,10 @@ end c = sample(StableRNG(468), addlogprob_demo(), PG(10), 100) # Result should be biased towards x > 0. @test mean(c[:x]) > 0.7 + + # @addlogprob! should also be respected by ordinary (non-particle) samplers. + c2 = sample(StableRNG(468), addlogprob_demo(), MH(), 100) + @test mean(c2[:x]) > 0.7 end @testset "keyword argument handling" begin @@ -276,6 +351,29 @@ end @test advance!(particle, false) === nothing # model finished end + @testset "matches a direct evaluation" begin + # A particle advanced without resampling draws from its RNG continuously, so it must + # produce exactly the same values and log-likelihood as a plain DynamicPPL evaluation + # seeded identically. + particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) + while advance!(particle, false) !== nothing + end + + accs = DynamicPPL.OnlyAccsVarInfo() + accs = DynamicPPL.setacc!!(accs, DynamicPPL.LogLikelihoodAccumulator()) + accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) + _, accs = DynamicPPL.init!!( + TracedRNG(Xoshiro(23)), + test(), + accs, + DynamicPPL.InitFromPrior(), + DynamicPPL.UnlinkAll(), + ) + + @test get_raw_values(particle.varinfo) == get_raw_values(accs) + @test getloglikelihood(particle.varinfo) == getloglikelihood(accs) + end + @testset "fork" begin particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) advance!(particle, false) From 4789428ac6796e71cc628298f64c6561908c8b60 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 22:01:31 +0100 Subject: [PATCH 09/58] Note that resampler scheme types are not exported SMC(Systematic()) and friends require Turing.Inference qualification: the scheme types are not exported, and a clean export is blocked by the Multinomial resampler colliding with the re-exported Distributions.Multinomial. Document this in the SMC docstring (whether to export a renamed subset is a separate API decision). This matches the previous AdvancedPS-qualified API, so it is not a regression. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index c2e829da44..b704f3de85 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -461,6 +461,9 @@ end Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. + +The resampling scheme types (`Stratified`, `Systematic`, `Multinomial`, `ESSResampler`) are +not exported; refer to them as e.g. `Turing.Inference.Systematic`. """ SMC() = SMC(ESSResampler(0.5)) SMC(threshold::Real) = SMC(ESSResampler(threshold)) From cd6444a977ccf9b7d7b7372a22aca4b1226d53a7 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 22:09:58 +0100 Subject: [PATCH 10/58] Minimise: drop unused step-counter generality inc_step! only ever advances by one and set_step! was only ever called as set_step!(_, 1). Remove the unused `n` parameters: inc_step! takes no count, and set_step! becomes rewind! (reset to the first step), which also reads clearer at the call site where the reference RNG is rewound for replay. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 11 ++++++----- test/mcmc/particle_mcmc.jl | 8 +++----- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index b704f3de85..4bd05bd655 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -75,9 +75,11 @@ save_state!(trng::TracedRNG) = push!(trng.keys, inner_key(trng.rng)) "Replay the seed recorded at the current step." load_state!(trng::TracedRNG) = Random.seed!(trng, trng.keys[trng.count]) -"Set / advance the model-step counter." -set_step!(trng::TracedRNG, n::Integer) = (trng.count = n; trng) -inc_step!(trng::TracedRNG, n::Integer=1) = (trng.count += n; trng) +"Advance the model-step counter by one." +inc_step!(trng::TracedRNG) = (trng.count += 1; trng) + +"Rewind the model-step counter to the first step, so a trajectory replays from the start." +rewind!(trng::TracedRNG) = (trng.count = 1; trng) "Deterministically derive a fresh seed from `key`." split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) @@ -618,8 +620,7 @@ function AbstractMCMC.step( ) error_if_threadsafe_eval(model) n = sampler.nparticles - # Rewind the retained RNG to step 1 so the reference replays its trajectory from the start. - reference = Particle(model, particle_varinfo(), set_step!(deepcopy(state.rng), 1)) + reference = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng))) particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 787b713d05..ae0bcacd52 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -14,7 +14,7 @@ using Turing.Inference: particle_varinfo, advance!, fork, - set_step!, + rewind!, refresh!, sweep!, normalized_weights, @@ -223,7 +223,7 @@ end end allok = true for _ in 1:nsteps - ref = Particle(model, particle_varinfo(), set_step!(deepcopy(state.rng), 1)) + ref = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng))) parts = map( i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, 1:N, @@ -402,9 +402,7 @@ end end values = DynamicPPL.get_raw_values(particle.varinfo) - reference = Particle( - normal(), particle_varinfo(), set_step!(deepcopy(particle.rng), 1) - ) + reference = Particle(normal(), particle_varinfo(), rewind!(deepcopy(particle.rng))) while advance!(reference, true) !== nothing end @test DynamicPPL.get_raw_values(reference.varinfo) == values From b904713cded402c23e1f619f3179ac6423a73aad Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 22:13:22 +0100 Subject: [PATCH 11/58] Minimise: drop redundant threadsafe check in SMC step SMC's `sample` override already calls error_if_threadsafe_eval before delegating to mcmcsample, so the same check at the top of SMC's first `step` never fires on a distinct path. Remove it. PG keeps its check, being the only guard there (PG has no `sample` override). Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 4bd05bd655..bd21ede7f7 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -526,7 +526,6 @@ function AbstractMCMC.step( discard_sample=false, kwargs..., ) - error_if_threadsafe_eval(model) particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] logZ = sweep!(rng, particles, sampler.resampler) weights = normalized_weights(particles) From 499b61aba07bfc523ac46d076fc9794440c698c1 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 22:34:09 +0100 Subject: [PATCH 12/58] Derive numeric types instead of hardcoding Float64 Use DynamicPPL.LogProbType for particle log-weights and the log-evidence rather than a hardcoded Float64, and make ESSResampler parametric on its threshold type so it keeps whatever Real the user passes. Also bound the previously-open type parameters (TracedRNG's key type as <:Unsigned, SMCState's fields as <:AbstractVector) so every parametric struct carries meaningful subtype bounds. Integer types were already platform-native `Int`, not `Int64`. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index bd21ede7f7..9812fe7787 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -41,7 +41,7 @@ in `keys`, indexed by the step counter `count`. - [`load_state!`](@ref) restores `keys[count]`, replaying that step's randomness (the reference trajectory). """ -mutable struct TracedRNG{K,T<:Random123.AbstractR123} <: Random.AbstractRNG +mutable struct TracedRNG{K<:Unsigned,T<:Random123.AbstractR123} <: Random.AbstractRNG count::Int rng::T keys::Vector{K} @@ -131,11 +131,11 @@ mutable struct Particle # through Libtask's (already type-unstable) taped globals, so this costs nothing extra. varinfo::DynamicPPL.AbstractVarInfo rng::TracedRNG - logweight::Float64 + logweight::DynamicPPL.LogProbType task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the # particle as its taped globals (a back-reference). - Particle(vi, rng) = new(vi, rng, 0.0) + Particle(vi, rng) = new(vi, rng, zero(DynamicPPL.LogProbType)) end function Particle( @@ -171,7 +171,7 @@ back-reference; [`reseed!`](@ref) then gives it its own random stream. fork(particle::Particle, rng::AbstractRNG) = reseed!(deepcopy(particle), rng) """ - advance!(particle, isref) -> Union{Float64,Nothing} + advance!(particle, isref) -> Union{Real,Nothing} Run the particle to its next `observe`, returning the incremental log-likelihood, or `nothing` once the model finishes. An ordinary particle records the step's seed; the @@ -346,11 +346,11 @@ end Resample with `scheme`, but only when the effective sample size drops below `threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). """ -struct ESSResampler{R<:AbstractResampler} <: AbstractResampler - threshold::Float64 +struct ESSResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler + threshold::T scheme::R end -ESSResampler(threshold::Real) = ESSResampler(Float64(threshold), Stratified()) +ESSResampler(threshold::Real) = ESSResampler(threshold, Stratified()) function should_resample(resampler::ESSResampler, weights) ess = inv(sum(abs2, weights)) @@ -407,10 +407,11 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional reuse = !seen[a] && !(conditional && a == n) seen[a] = true child = reuse ? reseed!(old[a], rng) : fork(old[a], rng) - child.logweight = 0.0 + child.logweight = zero(DynamicPPL.LogProbType) particles[slot] = child end - conditional && (particles[n].logweight = 0.0) # reference retained, weight reset + # reference retained, weight reset + conditional && (particles[n].logweight = zero(DynamicPPL.LogProbType)) else for (i, p) in enumerate(particles) # Refresh every particle's seed except the reference, which keeps replaying. @@ -424,7 +425,7 @@ end # Run a full particle sweep in place, returning the log-evidence estimate. function sweep!(rng::AbstractRNG, particles, resampler; conditional::Bool=false) - logZ = 0.0 + logZ = zero(DynamicPPL.LogProbType) while true resample_propagate!(rng, particles, resampler, conditional) logZ0 = logevidence(particles) @@ -470,14 +471,14 @@ not exported; refer to them as e.g. `Turing.Inference.Systematic`. SMC() = SMC(ESSResampler(0.5)) SMC(threshold::Real) = SMC(ESSResampler(threshold)) function SMC(scheme::AbstractResampler, threshold::Real) - return SMC(ESSResampler(Float64(threshold), scheme)) + return SMC(ESSResampler(threshold, scheme)) end -struct SMCState{P,W} +struct SMCState{P<:AbstractVector,W<:AbstractVector} particles::P weights::W index::Int - logevidence::Float64 + logevidence::DynamicPPL.LogProbType end function AbstractMCMC.sample( @@ -584,7 +585,7 @@ whenever the effective sample size drops below half the number of particles. PG(n::Int) = PG(n, ESSResampler(0.5)) PG(n::Int, threshold::Real) = PG(n, ESSResampler(threshold)) function PG(n::Int, scheme::AbstractResampler, threshold::Real) - return PG(n, ESSResampler(Float64(threshold), scheme)) + return PG(n, ESSResampler(threshold, scheme)) end "Conditional SMC, an alias for [`PG`](@ref)." From 87c0bb70b98d04c7af2c46f90a33f56eaed290d2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 22:55:08 +0100 Subject: [PATCH 13/58] Move import Random123 into particle_mcmc.jl Random123 is only used by the TracedRNG in particle_mcmc.jl, so scope the import there rather than in Inference.jl. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/Inference.jl | 1 - src/mcmc/particle_mcmc.jl | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl index ac2041c5ba..5d5cf28669 100644 --- a/src/mcmc/Inference.jl +++ b/src/mcmc/Inference.jl @@ -32,7 +32,6 @@ import AdvancedHMC const AHMC = AdvancedHMC import AdvancedMH const AMH = AdvancedMH -import Random123 import EllipticalSliceSampling import LogDensityProblems import Random diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 9812fe7787..d20bc0fb01 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -22,6 +22,7 @@ ### using StatsFuns: softmax, logsumexp +import Random123 # # Traced RNG From bcae152bba41fcd8ea5566891fd9f24623a781a8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 23:00:08 +0100 Subject: [PATCH 14/58] Release the particle MCMC rewrite as 0.47.0 The native SMC/PG reimplementation changes the resampler API and default, so give it its own breaking release: move the changelog entry under a new 0.47.0 heading and bump the version. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- HISTORY.md | 27 +++++++++++++++------------ Project.toml | 2 +- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 1ab9c7e041..804ea5eae1 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,18 @@ +# 0.47.0 + +## Breaking changes + +### Particle MCMC (SMC and PG) + +`SMC` and `PG` / `CSMC` have been reimplemented natively and no longer depend on AdvancedPS. + +Resampling schemes are now types instead of functions — `Stratified()`, `Systematic()`, and `Multinomial()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`. +For example, `SMC(Turing.Inference.Systematic())`, `SMC(0.5)`, or `PG(10, Turing.Inference.Multinomial(), 0.5)`. +The old function-based API (`resample_systematic`, `AdvancedPS.ResampleWithESSThreshold`, …) is gone. + +The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows and within the theoretical guarantees for particle Gibbs, which systematic resampling does not. +Exact draws may therefore differ from previous releases, but should remain statistically consistent (the same target distribution). + # 0.46.0 ## Breaking changes @@ -19,18 +34,6 @@ Please see [the AdvancedVI changelog](https://github.com/TuringLang/AdvancedVI.j - `vi` with `KLMinScoreGradDescent` now optimises in unconstrained (linked) space, making it consistent with the other `KLMin...` algorithms. If you use it with a model that has constrained parameters, results may differ slightly from previous releases. -### Particle MCMC (SMC and PG) - -The particle samplers (`SMC`, and `PG` / `CSMC`) have been reimplemented natively and no longer depend on AdvancedPS. - -Resampling schemes are now types rather than functions: use `Stratified()`, `Systematic()`, or `Multinomial()`, optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size drops below `threshold * nparticles`. -For example, `SMC(Systematic())`, `SMC(0.5)`, or `PG(10, Multinomial(), 0.5)`. -The previous function-based API (`resample_systematic`, `resample_multinomial`, `AdvancedPS.ResampleWithESSThreshold`) is no longer available. - -The default resampling scheme is now **stratified** (previously systematic). -Stratified resampling is unbiased and consistent as the number of particles grows, and unlike systematic resampling it stays within the theoretical guarantees underpinning particle Gibbs. -Because of the new default and the reimplementation, sampling results will differ from previous releases. - ## Other changes **DifferentiationInterface removed as a direct dependency** diff --git a/Project.toml b/Project.toml index 5fb591b61d..c113c5bbb9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.46.0" +version = "0.47.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" From 022394485aaee663837b596b01687cd8cc175674 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Tue, 21 Jul 2026 23:49:20 +0100 Subject: [PATCH 15/58] Fix method ambiguity in TracedRNG seed! `Random.seed!(trng::TracedRNG, key)` took an untyped `key`, making it ambiguous with `Random.seed!(::AbstractRNG, ::Nothing)` (Random stdlib) for the call `seed!(::TracedRNG, ::Nothing)`. The Aqua ambiguity check flags this on CI (the conflicting Random method isn't present in every local Random, so it didn't reproduce locally). The seed is always an integer, so constrain `key::Integer`, which removes the overlap with `::Nothing`. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index d20bc0fb01..5090310ef2 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -64,7 +64,7 @@ Random.rand(trng::TracedRNG, ::Type{T}) where {T<:Unsigned} = Random.rand(trng.r inner_key(rng::Random123.Philox2x) = rng.key "Reseed and rewind the inner generator. The model-step counter is left untouched." -function Random.seed!(trng::TracedRNG, key) +function Random.seed!(trng::TracedRNG, key::Integer) Random.seed!(trng.rng, key) Random123.set_counter!(trng.rng, 0) return trng From ad08bbf2f2f2b341ea689e5ce77cc2834849389e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 22 Jul 2026 00:05:01 +0100 Subject: [PATCH 16/58] Make RNG seed splitting version-stable and decorrelated `split_key` derived a fresh seed by re-seeding a stdlib `MersenneTwister`. That is fragile two ways: re-seeding to split a generator can yield correlated streams (Steele et al., OOPSLA 2014), and MersenneTwister streams are not reproducible across Julia versions, so SMC/PG results drifted between versions even under a StableRNG. Both affected the previous AdvancedPS implementation (#2781, AdvancedPS.jl#110). Derive the seed through Philox instead -- a counter-based generator with a fixed, portable algorithm and strong avalanche -- which is both well-decorrelated and version-stable. The full particle suite, including exact CSMC reference reproduction, stays green. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 5090310ef2..f133bac853 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -82,8 +82,15 @@ inc_step!(trng::TracedRNG) = (trng.count += 1; trng) "Rewind the model-step counter to the first step, so a trajectory replays from the start." rewind!(trng::TracedRNG) = (trng.count = 1; trng) -"Deterministically derive a fresh seed from `key`." -split_key(key::Integer) = rand(Random.MersenneTwister(key), typeof(key)) +# Derive a fresh seed from `key`. Splitting one generator into many by re-seeding is fragile +# in two ways: the derived seeds can yield *correlated* streams (Steele et al., "Fast +# Splittable Pseudorandom Number Generators", OOPSLA 2014), and a stdlib `MersenneTwister` +# derivation is not identical across Julia versions (Julia does not guarantee reproducible +# streams), which made SMC/PG drift between versions even under a StableRNG. Both bit the +# previous AdvancedPS implementation (#2781, AdvancedPS.jl#110). Philox is a counter-based +# generator with a fixed, portable algorithm and strong avalanche, so deriving the seed +# through it is both well-decorrelated from its parent and version-stable. +split_key(key::Integer) = rand(Random.seed!(Random123.Philox2x(), key), typeof(key)) "Reseed from the generator's own current state (used between steps when not resampling)." refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) From 716cc1e8ea76350b0baba7566c1f8b13f2435cb0 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 22 Jul 2026 00:05:17 +0100 Subject: [PATCH 17/58] Cite Andrieu, Doucet & Holenstein (2010) in the particle MCMC module Add the foundational particle MCMC reference to the module's design note, where the short "(ADH 2010)" citations in the resampling notes point. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- src/mcmc/particle_mcmc.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index f133bac853..074f8e5b1e 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -20,6 +20,9 @@ ### Sections below: traced RNG; model evaluation via Libtask; resampling schemes; the particle ### sweep; the SMC sampler; the PG/CSMC sampler; the Gibbs-component interface. ### +### Reference: Andrieu, Doucet & Holenstein, "Particle Markov chain Monte Carlo methods", +### Journal of the Royal Statistical Society: Series B 72(3), 269-342 (2010). +### using StatsFuns: softmax, logsumexp import Random123 From 541bc2a8041e646521d60bba09b6963e45406c6a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 22 Jul 2026 00:37:54 +0100 Subject: [PATCH 18/58] Stabilise the gdemo CSMC+ESS test against seed variation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The native particle-MCMC RNG (Philox seed derivation) changes the exact draws, and `StableRNG(23)` now lands a ~2.7σ tail draw for E[s] that just exceeds atol=0.1 at 3_000 iterations. The estimator is unbiased -- E[s] scatters tightly around the true 2.042 across seeds -- so the fix is simply more headroom: CSMC mixes the variance slowly, and 10_000 draws bring the error comfortably inside atol on every Julia version and platform. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- test/mcmc/ess.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/mcmc/ess.jl b/test/mcmc/ess.jl index ab81287868..44dd75fe79 100644 --- a/test/mcmc/ess.jl +++ b/test/mcmc/ess.jl @@ -60,7 +60,10 @@ using Turing @testset "gdemo with CSMC + ESS" begin alg = Gibbs(:s => CSMC(15), :m => ESS()) - chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 3_000) + # CSMC mixes the variance `s` slowly, so the Monte Carlo error at 3_000 draws sits + # close to `atol` and an unlucky seed tips it over; the estimator is unbiased, so + # 10_000 draws simply give comfortable headroom. + chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 10_000) check_gdemo(chain; atol=0.1) end From 160ab54d5b077c026915c23aee36024a534127fd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 22 Jul 2026 00:59:51 +0100 Subject: [PATCH 19/58] Add opt-in within-sweep particle threading to SMC and PG `SMC(; threaded=true)` and `PG(n; threaded=true)` now evaluate the particles across threads within each sweep. Only the per-particle model advances (the expensive part) run in parallel; resampling stays serial, and because every particle's RNG is seeded serially in `resample_propagate!` before the parallel region, the threaded run reproduces the serial draws bit for bit. The default (serial) path keeps its scalar tally and allocates nothing extra. `threaded` is a keyword-only field: an inner constructor suppresses the positional default constructor, so it cannot collide with the existing `(scheme, threshold::Real)` forms (`Bool <: Real`). HISTORY.md gains notes on this, on the cross-version RNG reproducibility, and on the weight-diagnostic columns, plus a mention that multiple chains run under `MCMCThreads()` / `MCMCDistributed()`. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 13 ++++-- src/mcmc/particle_mcmc.jl | 89 +++++++++++++++++++++++++++----------- test/mcmc/particle_mcmc.jl | 60 +++++++++++++++++++++++-- 3 files changed, 130 insertions(+), 32 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 804ea5eae1..3fca2405af 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -6,12 +6,17 @@ `SMC` and `PG` / `CSMC` have been reimplemented natively and no longer depend on AdvancedPS. -Resampling schemes are now types instead of functions — `Stratified()`, `Systematic()`, and `Multinomial()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`. -For example, `SMC(Turing.Inference.Systematic())`, `SMC(0.5)`, or `PG(10, Turing.Inference.Multinomial(), 0.5)`. +Resampling schemes are now types rather than functions — `Stratified()`, `Systematic()`, and `Multinomial()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.Systematic())`, `SMC(0.5)`, or `PG(10, Turing.Inference.Multinomial(), 0.5)`. The old function-based API (`resample_systematic`, `AdvancedPS.ResampleWithESSThreshold`, …) is gone. -The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows and within the theoretical guarantees for particle Gibbs, which systematic resampling does not. -Exact draws may therefore differ from previous releases, but should remain statistically consistent (the same target distribution). +The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows and within the theoretical guarantees for particle Gibbs, which systematic does not. +Exact draws may therefore differ from previous releases, but remain statistically consistent (the same target distribution). + +The rewrite also brings: + + - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). + - **Parallelism.** Multiple chains run under AbstractMCMC's `MCMCThreads()` and `MCMCDistributed()`, and `SMC(; threaded=true)` / `PG(n; threaded=true)` spread the particles across threads within each sweep. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. + - **Weight diagnostics.** `SMC` chains carry the log-evidence estimate `logevidence` and the per-particle normalised `weight` as extra columns; `PG` / `CSMC` chains carry `logevidence`. # 0.46.0 diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 074f8e5b1e..dee0a33abc 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -382,17 +382,36 @@ logweights(particles) = [p.logweight for p in particles] normalized_weights(particles) = softmax(logweights(particles)) logevidence(particles) = logsumexp(logweights(particles)) +# Advance one particle by one observation, folding its incremental weight in; return `true` +# once it has finished (produced nothing). Factored out so the serial and threaded loops in +# `reweight!` share one body. +function advance_particle!(p::Particle, isref::Bool) + score = advance!(p, isref) + score === nothing && return true + p.logweight += score + return false +end + # Advance every particle by one observation; return `true` once all have finished. A model # whose number of observations varies across executions leaves particles out of step. -function reweight!(particles, conditional::Bool) +# +# Each particle advances only its own state (rng, varinfo, task), and its rng was already +# seeded serially in `resample_propagate!`, so the threaded loop is race-free and gives +# results identical to the serial one. Only the model evaluations parallelise; the shared +# sampler rng is untouched here. +function reweight!(particles, conditional::Bool, threaded::Bool) n = length(particles) - n_done = 0 - for (i, p) in enumerate(particles) - score = advance!(p, conditional && i == n) - if score === nothing - n_done += 1 - else - p.logweight += score + if threaded + # A shared counter would race, so collect per-particle results and tally afterwards. + finished = Vector{Bool}(undef, n) + Threads.@threads for i in 1:n + finished[i] = advance_particle!(particles[i], conditional && i == n) + end + n_done = count(finished) + else + n_done = 0 + for i in 1:n + n_done += advance_particle!(particles[i], conditional && i == n) end end n_done == 0 && return false @@ -435,12 +454,14 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional end # Run a full particle sweep in place, returning the log-evidence estimate. -function sweep!(rng::AbstractRNG, particles, resampler; conditional::Bool=false) +function sweep!( + rng::AbstractRNG, particles, resampler, threaded::Bool; conditional::Bool=false +) logZ = zero(DynamicPPL.LogProbType) while true resample_propagate!(rng, particles, resampler, conditional) logZ0 = logevidence(particles) - done = reweight!(particles, conditional) + done = reweight!(particles, conditional, threaded) # Each observation contributes the log-ratio of total weight it adds; summed over the # sweep these telescope into an estimate of the model's log-evidence log p(y). logZ += logevidence(particles) - logZ0 @@ -467,22 +488,30 @@ $(TYPEDFIELDS) struct SMC{R<:AbstractResampler} <: ParticleInference "resampling scheme" resampler::R + "reweight the particles across threads within each sweep" + threaded::Bool + function SMC(resampler::R; threaded::Bool=false) where {R<:AbstractResampler} + return new{R}(resampler, threaded) + end end """ - SMC([resampler = ESSResampler(0.5)]) - SMC([scheme = Stratified(), ]threshold) + SMC([resampler = ESSResampler(0.5)]; threaded = false) + SMC([scheme = Stratified(), ]threshold; threaded = false) Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. +Set `threaded = true` to evaluate the particles across threads within each sweep; results are +unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). + The resampling scheme types (`Stratified`, `Systematic`, `Multinomial`, `ESSResampler`) are not exported; refer to them as e.g. `Turing.Inference.Systematic`. """ -SMC() = SMC(ESSResampler(0.5)) -SMC(threshold::Real) = SMC(ESSResampler(threshold)) -function SMC(scheme::AbstractResampler, threshold::Real) - return SMC(ESSResampler(threshold, scheme)) +SMC(; kwargs...) = SMC(ESSResampler(0.5); kwargs...) +SMC(threshold::Real; kwargs...) = SMC(ESSResampler(threshold); kwargs...) +function SMC(scheme::AbstractResampler, threshold::Real; kwargs...) + return SMC(ESSResampler(threshold, scheme); kwargs...) end struct SMCState{P<:AbstractVector,W<:AbstractVector} @@ -539,7 +568,7 @@ function AbstractMCMC.step( kwargs..., ) particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] - logZ = sweep!(rng, particles, sampler.resampler) + logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded) weights = normalized_weights(particles) stats = (; weight=weights[1], logevidence=logZ) @@ -584,19 +613,29 @@ struct PG{R<:AbstractResampler} <: ParticleInference nparticles::Int "resampling scheme" resampler::R + "reweight the particles across threads within each sweep" + threaded::Bool + function PG( + nparticles::Int, resampler::R; threaded::Bool=false + ) where {R<:AbstractResampler} + return new{R}(nparticles, resampler, threaded) + end end """ - PG(n, [resampler = ESSResampler(0.5)]) - PG(n, [scheme = Stratified(), ]threshold) + PG(n, [resampler = ESSResampler(0.5)]; threaded = false) + PG(n, [scheme = Stratified(), ]threshold; threaded = false) Particle Gibbs sampler with `n` particles. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. + +Set `threaded = true` to evaluate the particles across threads within each sweep; results are +unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). """ -PG(n::Int) = PG(n, ESSResampler(0.5)) -PG(n::Int, threshold::Real) = PG(n, ESSResampler(threshold)) -function PG(n::Int, scheme::AbstractResampler, threshold::Real) - return PG(n, ESSResampler(threshold, scheme)) +PG(n::Int; kwargs...) = PG(n, ESSResampler(0.5); kwargs...) +PG(n::Int, threshold::Real; kwargs...) = PG(n, ESSResampler(threshold); kwargs...) +function PG(n::Int, scheme::AbstractResampler, threshold::Real; kwargs...) + return PG(n, ESSResampler(threshold, scheme); kwargs...) end "Conditional SMC, an alias for [`PG`](@ref)." @@ -615,7 +654,7 @@ function AbstractMCMC.step( particles = [ Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:(sampler.nparticles) ] - logZ = sweep!(rng, particles, sampler.resampler) + logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) end @@ -635,7 +674,7 @@ function AbstractMCMC.step( particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end - logZ = sweep!(rng, particles, sampler.resampler; conditional=true) + logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded; conditional=true) return pg_transition_and_state(rng, particles, logZ, discard_sample) end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index ae0bcacd52..958aa3fc1d 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -20,7 +20,7 @@ using Turing.Inference: normalized_weights, PGState using Distributions: Bernoulli, Beta, Gamma, Normal, Uniform, Categorical, sample -using FlexiChains: VNChain +using FlexiChains: VNChain, has_same_data using Random: Random, Xoshiro using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset @@ -33,6 +33,9 @@ using Turing @test SMC(0.6).resampler == ESSResampler(0.6) @test SMC(Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) @test SMC(Systematic()).resampler == Systematic() + @test SMC().threaded == false + @test SMC(; threaded=true).threaded == true + @test SMC(Systematic(); threaded=true).threaded == true end @testset "basic model" begin @@ -112,6 +115,21 @@ using Turing @test chains_smc[:logevidence] ≈ fill(smc_logevidence, 100) end + @testset "threaded execution matches serial" begin + # Particles are seeded serially before the parallel reweighting, so `threaded=true` + # must reproduce the serial draws exactly (bit for bit), whatever the thread count. + @model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end + end + model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + serial = sample(StableRNG(23), model, SMC(), 200) + threaded = sample(StableRNG(23), model, SMC(; threaded=true), 200) + @test serial[@varname(p)] == threaded[@varname(p)] + end + @testset "refuses to run threadsafe eval" begin # SMC can't run models that have nondeterministic evaluation order, # so it should refuse to run models marked as threadsafe. @@ -160,6 +178,9 @@ end @test PG(60, 0.6).resampler == ESSResampler(0.6) @test PG(80, Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) @test PG(100, Systematic()).resampler == Systematic() + @test PG(10).threaded == false + @test PG(10; threaded=true).threaded == true + @test PG(80, Multinomial(), 0.6; threaded=true).threaded == true end @testset "chain log-density metadata" begin @@ -190,6 +211,22 @@ end @test chains_pg[:logevidence] ≈ fill(pg_logevidence, 100) end + @testset "threaded execution matches serial" begin + # Threading the reweighting must not perturb the reference-replay bookkeeping, so the + # conditional sweeps have to reproduce the serial draws exactly, whatever the thread + # count. + @model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end + end + model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + serial = sample(StableRNG(23), model, PG(10), 200) + threaded = sample(StableRNG(23), model, PG(10; threaded=true), 200) + @test serial[@varname(p)] == threaded[@varname(p)] + end + # https://github.com/TuringLang/Turing.jl/issues/1598 @testset "reference particle" begin c = sample(gdemo_default, PG(1), 1_000) @@ -217,7 +254,7 @@ end function run_csmc(model, N, nsteps, rng) draw(ps) = ps[rand(rng, Categorical(normalized_weights(ps)))] particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:N] - sweep!(rng, particles, ESSResampler(0.5)) + sweep!(rng, particles, ESSResampler(0.5), false) state = let p = draw(particles) PGState(p.varinfo, p.rng) end @@ -228,7 +265,7 @@ end i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, 1:N, ) - sweep!(rng, parts, ESSResampler(0.5); conditional=true) + sweep!(rng, parts, ESSResampler(0.5), false; conditional=true) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) p = draw(parts) state = PGState(p.varinfo, p.rng) @@ -332,6 +369,23 @@ end end end +@testset "parallel chains (MCMCThreads)" begin + @model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end + end + model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + # Multiple chains through AbstractMCMC's thread-based ensemble stay reproducible under a + # fixed seed (genuinely parallel only when Julia is started with more than one thread). + for sampler in (SMC(), PG(10)) + c1 = sample(Xoshiro(5), model, sampler, MCMCThreads(), 100, 4) + c2 = sample(Xoshiro(5), model, sampler, MCMCThreads(), 100, 4) + @test has_same_data(c1, c2) + end +end + @testset "particle container" begin @model function test() a ~ Normal(0, 1) From 75d3c0b1f6a0caf443c64aae5181ea9199dae69e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 16:18:06 +0100 Subject: [PATCH 20/58] Keep particle log-score consistent with the likelihood accumulator Addresses Charles's review comment on internal consistency. Emit the per-step `produce` from `tilde_observe!!` and an `accloglikelihood!!` overload -- after the log-likelihood accumulator has been updated -- instead of from inside `acclogp` before it. This removes the one-step lag between a particle's produced weight and its accumulated log-likelihood, and makes `@addlogprob!` terms reach the accumulator (hence the reported log-likelihood), not only the weight. The produced score is the accumulator's increment, so the accumulator stays the single source of truth: `acclogp` now inherits the generic `LogProbAccumulator` method, and `ProduceLogLikelihoodAccumulator` is just a marker type flagging a particle's varinfo so the produce sites know to emit. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 76 +++++++++++++++++++++++++-------------- 1 file changed, 49 insertions(+), 27 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index dee0a33abc..e56e1c944f 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -123,7 +123,7 @@ end Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws from the prior using the particle's [`TracedRNG`](@ref), and `tilde_observe!!` scores the -observation, which [`ProduceLogLikelihoodAccumulator`](@ref) turns into a `Libtask.produce`. +observation and `Libtask.produce`s the increment as the particle's weight. """ struct ParticleMCMCContext <: DynamicPPL.AbstractContext end @@ -210,6 +210,11 @@ function DynamicPPL.tilde_assume!!( return x, vi end +# Reweighting invariant: a particle's per-step score is `produce`d from here and from the +# `accloglikelihood!!` overload (for `@addlogprob!`), *after* the likelihood accumulator is +# updated, and equals the accumulator's increment. Producing after the update -- rather than +# inside `acclogp` -- keeps the produced weight in step with the accumulated log-likelihood +# (no one-step lag) and lets `@addlogprob!` terms reach the accumulator, not just the weight. function DynamicPPL.tilde_observe!!( ::ParticleMCMCContext, dist::Distribution, @@ -219,19 +224,23 @@ function DynamicPPL.tilde_observe!!( ::DynamicPPL.AbstractVarInfo, ) particle = Libtask.get_taped_globals(Particle) + before = DynamicPPL.getloglikelihood(particle.varinfo) left, vi = DynamicPPL.tilde_observe!!( DynamicPPL.DefaultContext(), dist, left, vn, template, particle.varinfo ) particle.varinfo = vi + Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) # increment this observation added return left, vi end """ ProduceLogLikelihoodAccumulator{T} <: LogProbAccumulator{T} -Like `LogLikelihoodAccumulator`, but `Libtask.produce`s each likelihood increment as it is -accumulated. Because `@addlogprob!` also routes through `acclogp`, it too triggers a -`produce`, so manual likelihood terms reweight particles correctly (issue #1996). +A marker likelihood accumulator: it accumulates exactly like `LogLikelihoodAccumulator`, but +its distinct type flags a varinfo as belonging to a particle, so the produce sites know to +emit. The produce happens in [`tilde_observe!!`](@ref) (observations) and the +`accloglikelihood!!` overload below (`@addlogprob!`, issue #1996) -- in each case from the +increase in accumulated log-likelihood, keeping the accumulator the single source of truth. """ struct ProduceLogLikelihoodAccumulator{T<:Real} <: DynamicPPL.LogProbAccumulator{T} logp::T @@ -239,11 +248,8 @@ end DynamicPPL.accumulator_name(::Type{<:ProduceLogLikelihoodAccumulator}) = :LogLikelihood DynamicPPL.logp(acc::ProduceLogLikelihoodAccumulator) = acc.logp - -function DynamicPPL.acclogp(acc::ProduceLogLikelihoodAccumulator, val) - Libtask.produce(val) # the only difference from `LogLikelihoodAccumulator` - return ProduceLogLikelihoodAccumulator(acc.logp + val) -end +# `acclogp` is inherited from the generic `LogProbAccumulator` method (plain addition); the +# produce is handled by the produce sites, not here. function DynamicPPL.accumulate_assume!!( acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, dist, template @@ -256,27 +262,43 @@ function DynamicPPL.accumulate_observe!!( return DynamicPPL.acclogp(acc, Distributions.loglikelihood(dist, left)) end -# Tell Libtask which calls may contain `produce`, walking up the call stack from `acclogp`. -# Over-approximating is safe (a wrongly-marked call just gets instrumented); missing a real -# one is not, so we err towards marking. -Libtask.@might_produce(DynamicPPL.accloglikelihood!!) -# Merging accumulators (across submodels or Gibbs blocks) can add a -# ProduceLogLikelihoodAccumulator to a plain one, which routes through the producing -# `acclogp` -- so this `+` may itself produce. -function Libtask.might_produce( - ::Type{ - <:Tuple{ - typeof(Base.:+), - ProduceLogLikelihoodAccumulator, - DynamicPPL.LogLikelihoodAccumulator, - }, - }, +# `@addlogprob!` bypasses `tilde_observe!!`, so its produce is emitted here instead -- again +# only once the accumulator has been updated. Gated on the producing accumulator, so outside +# particle evaluation this reduces to the default (non-producing) method (issue #1996). +function DynamicPPL.accloglikelihood!!( + vi::DynamicPPL.OnlyAccsVarInfo, logp; ignore_missing_accumulator=false ) - return true + acc_name = Val(:LogLikelihood) + if ignore_missing_accumulator && !DynamicPPL.hasacc(vi, acc_name) + return vi + end + is_particle = DynamicPPL.getacc(vi, acc_name) isa ProduceLogLikelihoodAccumulator + before = is_particle ? DynamicPPL.getloglikelihood(vi) : zero(DynamicPPL.LogProbType) + vi = DynamicPPL.map_accumulator!!(acc -> DynamicPPL.acclogp(acc, logp), vi, acc_name) + if is_particle + particle = Libtask.get_taped_globals(Particle) + particle.varinfo = vi + Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) + end + return vi end -Libtask.@might_produce(DynamicPPL.accumulate_observe!!) + +# Tell Libtask which calls may contain a `produce`, so it instruments them. The produce lives +# in `tilde_observe!!` and `accloglikelihood!!`; the rest of each chain is marked so Libtask +# tapes through to reach it. Over-approximating is safe (a wrongly-marked call just gets +# instrumented); missing a real one is not, so we err towards marking. +# +# observe: tilde_observe!! accumulates (accumulate_observe!! -> acclogp), then produces +# @addlogprob!: accloglikelihood!! accumulates (map_accumulator!! -> acclogp), then produces +# (the `@addlogprob! (; ...)` NamedTuple form routes through acclogp!! first) +# Gibbs: GibbsContext turns a tilde_assume!! into a tilde_observe!! Libtask.@might_produce(DynamicPPL.tilde_observe!!) -Libtask.@might_produce(DynamicPPL.tilde_assume!!) # GibbsContext turns assumes into observes +Libtask.@might_produce(DynamicPPL.accumulate_observe!!) +Libtask.@might_produce(DynamicPPL.acclogp) +Libtask.@might_produce(DynamicPPL.tilde_assume!!) +Libtask.@might_produce(DynamicPPL.accloglikelihood!!) +Libtask.@might_produce(DynamicPPL.map_accumulator!!) +Libtask.@might_produce(DynamicPPL.acclogp!!) # Every model / submodel evaluator takes a `DynamicPPL.Model`, so this covers them all. # See https://github.com/TuringLang/Libtask.jl/issues/217. Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true From 50d2f40f9a94ff864dc7e36746cb32517b705de2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 16:24:21 +0100 Subject: [PATCH 21/58] Rename ParticleMCMCContext to SMCContext Per Charles's review: SMC is not an MCMC (nor PMCMC) algorithm, so `SMCContext` names the leaf context more accurately. Pure rename, internal to particle_mcmc.jl. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index e56e1c944f..118e8c26de 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -3,7 +3,7 @@ ### ### Key design. ### A probabilistic model becomes a particle filter by reading each `observe` statement as one -### filtering step. Evaluated under `ParticleMCMCContext`, every likelihood term calls +### filtering step. Evaluated under `SMCContext`, every likelihood term calls ### `Libtask.produce`, so a *particle* is a suspended model execution: we `advance!` it to its ### next `observe`, take the produced log-likelihood as its weight, then resample. SMC is one ### such sweep; particle Gibbs (PG/CSMC) runs a *conditional* sweep -- one particle is a fixed @@ -119,17 +119,17 @@ function error_if_threadsafe_eval(model::DynamicPPL.Model) end """ - ParticleMCMCContext + SMCContext Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws from the prior using the particle's [`TracedRNG`](@ref), and `tilde_observe!!` scores the observation and `Libtask.produce`s the increment as the particle's weight. """ -struct ParticleMCMCContext <: DynamicPPL.AbstractContext end +struct SMCContext <: DynamicPPL.AbstractContext end # `OnlyAccsVarInfo` needs a parameter eltype; `Any` is fine here since particle MCMC never # involves AD or tracer types (see the `get_param_eltype` docstring in DynamicPPL). -DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::ParticleMCMCContext) = Any +DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any """ Particle(model, varinfo, rng::TracedRNG) @@ -152,7 +152,7 @@ end function Particle( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo, rng::TracedRNG ) - model = DynamicPPL.setleafcontext(model, ParticleMCMCContext()) + model = DynamicPPL.setleafcontext(model, SMCContext()) args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) particle = Particle(deepcopy(varinfo), rng) particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) @@ -195,11 +195,7 @@ function advance!(particle::Particle, isref::Bool) end function DynamicPPL.tilde_assume!!( - ::ParticleMCMCContext, - dist::Distribution, - vn::VarName, - template, - ::DynamicPPL.AbstractVarInfo, + ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) ctx = DynamicPPL.InitContext( @@ -216,7 +212,7 @@ end # inside `acclogp` -- keeps the produced weight in step with the accumulated log-likelihood # (no one-step lag) and lets `@addlogprob!` terms reach the accumulator, not just the weight. function DynamicPPL.tilde_observe!!( - ::ParticleMCMCContext, + ::SMCContext, dist::Distribution, left, vn::Union{VarName,Nothing}, From 7247b13655602d5e5aa43fdce01ac58f8a7fb35b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 16:28:30 +0100 Subject: [PATCH 22/58] Parametrise Particle on its RNG and weight types Per Charles's review: `Particle{RT<:TracedRNG,WT<:Real}`, so `logweight` tracks `DynamicPPL.LogProbType` and follows it if that is ever changed. Behaviour is unchanged -- taped-globals access stays type-unstable as before (the `::Particle` typeassert still holds for the parametric type), so this only concretises the stored particle's field types. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 118e8c26de..1b9d07a973 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -137,16 +137,22 @@ DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any A single particle: a suspended `model` execution together with its `varinfo`, its own replayable `rng`, and an accumulated `logweight`. """ -mutable struct Particle +mutable struct Particle{RT<:TracedRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go # through Libtask's (already type-unstable) taped globals, so this costs nothing extra. varinfo::DynamicPPL.AbstractVarInfo - rng::TracedRNG - logweight::DynamicPPL.LogProbType + rng::RT + # `logweight` tracks whatever `DynamicPPL.LogProbType` is, so weights follow suit if it + # is ever changed. + logweight::WT task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the - # particle as its taped globals (a back-reference). - Particle(vi, rng) = new(vi, rng, zero(DynamicPPL.LogProbType)) + # particle as its taped globals (a back-reference). This has to be an inner constructor + # for that reason: `task` is left undefined here and set immediately after. + function Particle(vi::DynamicPPL.AbstractVarInfo, rng::RT) where {RT<:TracedRNG} + w = zero(DynamicPPL.LogProbType) + return new{RT,typeof(w)}(vi, rng, w) + end end function Particle( From 39ab78ec39aa3e9106227464ff92c0bf47e68976 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 16:32:48 +0100 Subject: [PATCH 23/58] Document why particle assume always draws from the prior Answers Charles's review question ("is it always InitFromPrior?"). Verified that a variable is never already present at `tilde_assume!!` across SMC, PG, CSMC and PG-in-Gibbs: particle varinfos start empty, each variable is assumed exactly once, and the CSMC reference reproduces its trajectory by replaying RNG seeds rather than reusing stored values. So `InitFromPrior` is always correct, and a conditional `InitFromParams` branch would be dead code unless the varinfo were pre-populated (e.g. initial_params for particle samplers, which is not currently supported). Documented rather than adding the inert branch. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 1b9d07a973..eee2d7ec74 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -204,6 +204,11 @@ function DynamicPPL.tilde_assume!!( ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) + # Always draw from the prior. A value is never already present here: particle varinfos + # start empty, each variable is assumed exactly once, and the CSMC reference reproduces its + # trajectory by replaying its RNG seeds, not by reusing stored values. Reusing an existing + # value (via `InitFromParams`) would only matter for a pre-populated varinfo, which + # particle sampling does not currently create. ctx = DynamicPPL.InitContext( particle.rng, DynamicPPL.InitFromPrior(), DynamicPPL.UnlinkAll() ) From 8f6f700ac92085d8fdb8d948d60ec4cc471218af Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 16:57:22 +0100 Subject: [PATCH 24/58] Reuse Particle as the PG state (const PGState = Particle) Per Charles's review: the retained particle already carries the reference trajectory's varinfo and rng needed to resume the next conditional sweep, so particle Gibbs needs no separate state struct. `const PGState = Particle` keeps the semantic name at the PG call sites while adding no new type. `gibbs_update_state!!` now updates the reference varinfo in place, which is safe -- Gibbs replaces the state with the returned value and never reads the pre-update one again. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 26 ++++++++++++++++---------- test/mcmc/particle_mcmc.jl | 10 +++------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index eee2d7ec74..6228f1450f 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -32,8 +32,8 @@ import Random123 # # A counter-based RNG that records the seed used at each model step, so that a particle's # trajectory can be replayed exactly: the conditional-SMC reference regenerates itself by -# replaying its recorded seeds. This section comes first because `Particle` and `PGState` -# name `TracedRNG` in their type signatures. +# replaying its recorded seeds. This section comes first because `Particle` names `TracedRNG` +# in its type signature. """ TracedRNG([rng = Random.default_rng()]) @@ -135,7 +135,8 @@ DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any Particle(model, varinfo, rng::TracedRNG) A single particle: a suspended `model` execution together with its `varinfo`, its own -replayable `rng`, and an accumulated `logweight`. +replayable `rng`, and an accumulated `logweight`. It also serves as the particle Gibbs +sampler state, aliased `PGState` below. """ mutable struct Particle{RT<:TracedRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go @@ -670,10 +671,10 @@ end "Conditional SMC, an alias for [`PG`](@ref)." const CSMC = PG -struct PGState{V<:DynamicPPL.AbstractVarInfo,R<:TracedRNG} - varinfo::V - rng::R -end +# PG's sampler state is just the retained `Particle`: it already carries the reference +# trajectory's `varinfo` and `rng` (its `task`/`logweight` are then unused), so we alias +# rather than define a separate struct. +const PGState = Particle # First iteration: an ordinary (unconditional) particle sweep. function AbstractMCMC.step( @@ -716,7 +717,7 @@ function pg_transition_and_state(rng, particles, logZ, discard_sample) else DynamicPPL.ParamsWithStats(deepcopy(retained.varinfo), (; logevidence=logZ)) end - return transition, PGState(retained.varinfo, retained.rng) + return transition, retained end # @@ -729,6 +730,11 @@ function gibbs_update_state!!( ::PG, state::PGState, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple ) init = DynamicPPL.InitFromParams(global_vals, nothing) - new_vi = last(DynamicPPL.init!!(model, state.varinfo, init, DynamicPPL.UnlinkAll())) - return PGState(new_vi, state.rng) + # Re-initialise the reference varinfo with the values conditioned by other Gibbs + # components. Mutating in place is safe: the caller replaces this state with the value we + # return and never reads the pre-update one again. + state.varinfo = last( + DynamicPPL.init!!(model, state.varinfo, init, DynamicPPL.UnlinkAll()) + ) + return state end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 958aa3fc1d..1ad1aafcda 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -17,8 +17,7 @@ using Turing.Inference: rewind!, refresh!, sweep!, - normalized_weights, - PGState + normalized_weights using Distributions: Bernoulli, Beta, Gamma, Normal, Uniform, Categorical, sample using FlexiChains: VNChain, has_same_data using Random: Random, Xoshiro @@ -255,9 +254,7 @@ end draw(ps) = ps[rand(rng, Categorical(normalized_weights(ps)))] particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:N] sweep!(rng, particles, ESSResampler(0.5), false) - state = let p = draw(particles) - PGState(p.varinfo, p.rng) - end + state = draw(particles) allok = true for _ in 1:nsteps ref = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng))) @@ -267,8 +264,7 @@ end ) sweep!(rng, parts, ESSResampler(0.5), false; conditional=true) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) - p = draw(parts) - state = PGState(p.varinfo, p.rng) + state = draw(parts) end return allok, length(state.rng.keys) end From f80dc164874afd8d252f2733a7d21164bad021b9 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 17:10:18 +0100 Subject: [PATCH 25/58] Clarify particle parallelism: rename threaded to multithreaded, separate the two levels Renames the within-sweep parallelism keyword `SMC`/`PG(; threaded=...)` to `multithreaded`, and documents (in HISTORY and a note at `reweight!`) that it is a distinct axis from AbstractMCMC's chain-level ensemble. Within-sweep threads a single sweep's particle evaluations; the ensemble (`MCMCThreads`/`MCMCDistributed`) runs whole chains independently; the two compose. Only threading is offered within a sweep -- particles resample every step (all-to-all) and are live Libtask tasks, so distributing one sweep across processes would be communication-bound rather than a speed-up. Addresses Charles's review that the parallelism description conflated the two levels, and follows the discussion that these are genuinely different axes. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 2 +- src/mcmc/particle_mcmc.jl | 51 ++++++++++++++++++++++---------------- test/mcmc/particle_mcmc.jl | 26 +++++++++---------- 3 files changed, 44 insertions(+), 35 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 3fca2405af..db538c4bd9 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -15,7 +15,7 @@ Exact draws may therefore differ from previous releases, but remain statisticall The rewrite also brings: - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). - - **Parallelism.** Multiple chains run under AbstractMCMC's `MCMCThreads()` and `MCMCDistributed()`, and `SMC(; threaded=true)` / `PG(n; threaded=true)` spread the particles across threads within each sweep. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. + - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. - **Weight diagnostics.** `SMC` chains carry the log-evidence estimate `logevidence` and the per-particle normalised `weight` as extra columns; `PG` / `CSMC` chains carry `logevidence`. # 0.46.0 diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 6228f1450f..46979a1a28 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -413,7 +413,7 @@ normalized_weights(particles) = softmax(logweights(particles)) logevidence(particles) = logsumexp(logweights(particles)) # Advance one particle by one observation, folding its incremental weight in; return `true` -# once it has finished (produced nothing). Factored out so the serial and threaded loops in +# once it has finished (produced nothing). Factored out so the serial and multithreaded loops in # `reweight!` share one body. function advance_particle!(p::Particle, isref::Bool) score = advance!(p, isref) @@ -425,13 +425,20 @@ end # Advance every particle by one observation; return `true` once all have finished. A model # whose number of observations varies across executions leaves particles out of step. # +# `multithreaded` is *within-sweep* parallelism -- spreading this sweep's particle evaluations +# across threads. It is a separate axis from AbstractMCMC's chain-level ensemble +# (`MCMCThreads`/`MCMCDistributed`, which runs whole chains independently); the two compose. +# Only threading is offered here, not distribution: particles resample every step (all-to-all) +# and are live Libtask tasks, so spreading one sweep across processes would be communication- +# bound rather than a speed-up. +# # Each particle advances only its own state (rng, varinfo, task), and its rng was already -# seeded serially in `resample_propagate!`, so the threaded loop is race-free and gives +# seeded serially in `resample_propagate!`, so the multithreaded loop is race-free and gives # results identical to the serial one. Only the model evaluations parallelise; the shared # sampler rng is untouched here. -function reweight!(particles, conditional::Bool, threaded::Bool) +function reweight!(particles, conditional::Bool, multithreaded::Bool) n = length(particles) - if threaded + if multithreaded # A shared counter would race, so collect per-particle results and tally afterwards. finished = Vector{Bool}(undef, n) Threads.@threads for i in 1:n @@ -485,13 +492,13 @@ end # Run a full particle sweep in place, returning the log-evidence estimate. function sweep!( - rng::AbstractRNG, particles, resampler, threaded::Bool; conditional::Bool=false + rng::AbstractRNG, particles, resampler, multithreaded::Bool; conditional::Bool=false ) logZ = zero(DynamicPPL.LogProbType) while true resample_propagate!(rng, particles, resampler, conditional) logZ0 = logevidence(particles) - done = reweight!(particles, conditional, threaded) + done = reweight!(particles, conditional, multithreaded) # Each observation contributes the log-ratio of total weight it adds; summed over the # sweep these telescope into an estimate of the model's log-evidence log p(y). logZ += logevidence(particles) - logZ0 @@ -519,20 +526,20 @@ struct SMC{R<:AbstractResampler} <: ParticleInference "resampling scheme" resampler::R "reweight the particles across threads within each sweep" - threaded::Bool - function SMC(resampler::R; threaded::Bool=false) where {R<:AbstractResampler} - return new{R}(resampler, threaded) + multithreaded::Bool + function SMC(resampler::R; multithreaded::Bool=false) where {R<:AbstractResampler} + return new{R}(resampler, multithreaded) end end """ - SMC([resampler = ESSResampler(0.5)]; threaded = false) - SMC([scheme = Stratified(), ]threshold; threaded = false) + SMC([resampler = ESSResampler(0.5)]; multithreaded = false) + SMC([scheme = Stratified(), ]threshold; multithreaded = false) Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. -Set `threaded = true` to evaluate the particles across threads within each sweep; results are +Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). The resampling scheme types (`Stratified`, `Systematic`, `Multinomial`, `ESSResampler`) are @@ -598,7 +605,7 @@ function AbstractMCMC.step( kwargs..., ) particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] - logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded) + logZ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) stats = (; weight=weights[1], logevidence=logZ) @@ -644,22 +651,22 @@ struct PG{R<:AbstractResampler} <: ParticleInference "resampling scheme" resampler::R "reweight the particles across threads within each sweep" - threaded::Bool + multithreaded::Bool function PG( - nparticles::Int, resampler::R; threaded::Bool=false + nparticles::Int, resampler::R; multithreaded::Bool=false ) where {R<:AbstractResampler} - return new{R}(nparticles, resampler, threaded) + return new{R}(nparticles, resampler, multithreaded) end end """ - PG(n, [resampler = ESSResampler(0.5)]; threaded = false) - PG(n, [scheme = Stratified(), ]threshold; threaded = false) + PG(n, [resampler = ESSResampler(0.5)]; multithreaded = false) + PG(n, [scheme = Stratified(), ]threshold; multithreaded = false) Particle Gibbs sampler with `n` particles. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. -Set `threaded = true` to evaluate the particles across threads within each sweep; results are +Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). """ PG(n::Int; kwargs...) = PG(n, ESSResampler(0.5); kwargs...) @@ -684,7 +691,7 @@ function AbstractMCMC.step( particles = [ Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:(sampler.nparticles) ] - logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded) + logZ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) end @@ -704,7 +711,9 @@ function AbstractMCMC.step( particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end - logZ = sweep!(rng, particles, sampler.resampler, sampler.threaded; conditional=true) + logZ = sweep!( + rng, particles, sampler.resampler, sampler.multithreaded; conditional=true + ) return pg_transition_and_state(rng, particles, logZ, discard_sample) end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 1ad1aafcda..b2dc77a806 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -32,9 +32,9 @@ using Turing @test SMC(0.6).resampler == ESSResampler(0.6) @test SMC(Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) @test SMC(Systematic()).resampler == Systematic() - @test SMC().threaded == false - @test SMC(; threaded=true).threaded == true - @test SMC(Systematic(); threaded=true).threaded == true + @test SMC().multithreaded == false + @test SMC(; multithreaded=true).multithreaded == true + @test SMC(Systematic(); multithreaded=true).multithreaded == true end @testset "basic model" begin @@ -114,8 +114,8 @@ using Turing @test chains_smc[:logevidence] ≈ fill(smc_logevidence, 100) end - @testset "threaded execution matches serial" begin - # Particles are seeded serially before the parallel reweighting, so `threaded=true` + @testset "multithreaded execution matches serial" begin + # Particles are seeded serially before the parallel reweighting, so `multithreaded=true` # must reproduce the serial draws exactly (bit for bit), whatever the thread count. @model function coinflip(y) p ~ Beta(1, 1) @@ -125,8 +125,8 @@ using Turing end model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) serial = sample(StableRNG(23), model, SMC(), 200) - threaded = sample(StableRNG(23), model, SMC(; threaded=true), 200) - @test serial[@varname(p)] == threaded[@varname(p)] + multithreaded = sample(StableRNG(23), model, SMC(; multithreaded=true), 200) + @test serial[@varname(p)] == multithreaded[@varname(p)] end @testset "refuses to run threadsafe eval" begin @@ -177,9 +177,9 @@ end @test PG(60, 0.6).resampler == ESSResampler(0.6) @test PG(80, Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) @test PG(100, Systematic()).resampler == Systematic() - @test PG(10).threaded == false - @test PG(10; threaded=true).threaded == true - @test PG(80, Multinomial(), 0.6; threaded=true).threaded == true + @test PG(10).multithreaded == false + @test PG(10; multithreaded=true).multithreaded == true + @test PG(80, Multinomial(), 0.6; multithreaded=true).multithreaded == true end @testset "chain log-density metadata" begin @@ -210,7 +210,7 @@ end @test chains_pg[:logevidence] ≈ fill(pg_logevidence, 100) end - @testset "threaded execution matches serial" begin + @testset "multithreaded execution matches serial" begin # Threading the reweighting must not perturb the reference-replay bookkeeping, so the # conditional sweeps have to reproduce the serial draws exactly, whatever the thread # count. @@ -222,8 +222,8 @@ end end model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) serial = sample(StableRNG(23), model, PG(10), 200) - threaded = sample(StableRNG(23), model, PG(10; threaded=true), 200) - @test serial[@varname(p)] == threaded[@varname(p)] + multithreaded = sample(StableRNG(23), model, PG(10; multithreaded=true), 200) + @test serial[@varname(p)] == multithreaded[@varname(p)] end # https://github.com/TuringLang/Turing.jl/issues/1598 From 280ade99f6ad739265c85fe3ab7a4bee2eb04104 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 17:27:42 +0100 Subject: [PATCH 26/58] Move project instructions into AGENTS.md, include from CLAUDE.md AGENTS.md is the tool-agnostic convention; CLAUDE.md now just includes it via @AGENTS.md so both entry points resolve to one source. Co-Authored-By: Claude Code Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> --- AGENTS.md | 100 +++++++++++++++++++++++++++++++++++++++++++++++++++++ CLAUDE.md | 101 +----------------------------------------------------- 2 files changed, 101 insertions(+), 100 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..4922775aff --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,100 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +Turing.jl is the user-facing entry point for the [TuringLang](https://github.com/TuringLang) probabilistic programming ecosystem. It is largely a translation layer between DynamicPPL models — which work with named, structured parameters — and inference algorithms that expect flat, vectorised samples (e.g. HMC/NUTS operate on `AbstractVector{<:Real}`). DynamicPPL's `LogDensityFunction` handles most of this translation; Turing provides the sampler wrappers that set it up and manage state across iterations. + +Model definition lives in [DynamicPPL.jl](https://github.com/TuringLang/DynamicPPL.jl), parameter transformations in [Bijectors.jl](https://github.com/TuringLang/Bijectors.jl), and sampling interfaces in [AbstractMCMC.jl](https://github.com/TuringLang/AbstractMCMC.jl). Turing re-exports their APIs and provides concrete sampler implementations that wire everything together. + +## Building and Testing + +Code formatting uses [JuliaFormatter.jl](https://github.com/domluna/JuliaFormatter.jl) v1 (not v2) with the **Blue style** (configured in `.JuliaFormatter.toml`). CI enforces formatting on all PRs. JuliaFormatter must be installed in the **global** Julia environment, not the project environment — do not use `--project`. See the [formatting guide](https://turinglang.org/docs/contributing/code-formatting/) for setup details. + +```bash +julia -e 'using JuliaFormatter; format(".")' +``` + +Tests use `SelectiveTests.jl` (in `test/test_utils/`) to filter by path. CI splits the suite into four shards: `mcmc/gibbs.jl`, `mcmc/Inference.jl`, `ad.jl`, and everything else. To run a subset locally: + +```bash +julia --project -e 'using Pkg; Pkg.test(; test_args=["mcmc/hmc.jl"])' +``` + +Use `--skip` to exclude files: + +```bash +julia --project -e 'using Pkg; Pkg.test(; test_args=["--skip", "mcmc/gibbs.jl", "ad.jl"])' +``` + +CI matrix: Julia stable + min, Ubuntu/Windows/macOS, 1 and 2 threads. + +`test/test_utils/sampler.jl` provides generic test helpers (`test_rng_respected`, `test_sampler_analytical`, `test_chain_logp_metadata`) that should work for any sampler. Beyond these, sampler-specific tests are needed to capture the properties you care about — there is no standardised test template yet. + +## Architecture + +### What lives here vs elsewhere + +Most complexity is in DynamicPPL. Turing.jl contains: + + - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (wrapping AdvancedPS), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. + - **External sampler interface** (`src/mcmc/external_sampler.jl`): The `externalsampler()` wrapper lets any `AbstractMCMC.AbstractSampler` that implements `step` for `LogDensityModel` work with Turing models. This is the easier path for new samplers — it only requires a dependency on AbstractMCMC and the LogDensityProblems.jl interface, with no Turing internals. The tradeoff is less power: you can only interact with the model as a black-box log-density function, just like using `LogDensityFunction` directly. + - **Variational inference** (`src/variational/`): Wraps AdvancedVI algorithms. + - **Mode estimation** (`src/optimisation/`): MAP and MLE via Optimization.jl. + - **Custom distributions** (`src/stdlib/`): `Flat`, `FlatPos`, `BinomialLogit`, `OrderedLogistic`, `LogPoisson`, and Dirichlet/Chinese Restaurant processes. + +For how the model and inference machinery works under the hood, see the [DynamicPPL docs](https://turinglang.org/DynamicPPL.jl/stable/) and the [developer guides](https://turinglang.org/docs/developers/). + +### Gibbs sampler + +The Gibbs sampler (`src/mcmc/gibbs.jl`) is the most complex piece in Turing.jl. It maintains a global `VarNamedTuple` of raw values for all variables. On each iteration, it conditions the model on the non-target variables via `GibbsContext`, runs the component sampler, and updates the global state. + +To plug a sampler into Gibbs, implement: + + - `gibbs_get_raw_values(state)` — return a `VarNamedTuple` of raw values for the variables this sampler is responsible for. + - `gibbs_update_state!!(sampler, state, model, global_vals)` — update the sampler's state to reflect new conditioned values. For samplers that use `LogDensityFunction`, the helper `gibbs_recompute_ldf_and_params` handles the common case. + - Optionally, `isgibbscomponent(sampler)` — return `false` to disallow use in Gibbs (the default is `true`). + +### Extension + +`ext/TuringDynamicHMCExt` provides the DynamicHMC.jl integration (loaded when DynamicHMC is imported). + +## Review Guidelines + +### Use `OnlyAccsVarInfo`, not `VarInfo` + +Sampler state should use `OnlyAccsVarInfo` (with appropriate accumulators), not `VarInfo`. `VarInfo` is being phased out across the ecosystem. + +Most gradient-based samplers (HMC, NUTS, external samplers) go through `LogDensityFunction`, which handles the model interaction. `LogDensityFunction` works well when the model structure is static (the set of variables is fixed across evaluations) and the sampler only needs a scalar log-density value. However, LDF is hard to use when the sampler needs extra accumulators beyond log-probability — for example, MH uses custom accumulators to capture proposal distributions and linked values, so it works directly with `OnlyAccsVarInfo` + `init!!` instead. Either approach is fine; the key constraint is no `VarInfo`. + +Note: "linked" and "unconstrained" are synonymous in this codebase. Linking transforms constrained parameters to unconstrained (Euclidean) space for gradient-based sampling. + +### `VarNamedTuple` for parameter collections + +Interfaces that accept or return named parameter collections should use `VarNamedTuple`, not `NamedTuple` or `Dict{VarName}`. `NamedTuple` and `Dict{VarName}` are accepted as user-facing input but should be converted to `VarNamedTuple` at the boundary (see `_to_varnamedtuple` in `src/common.jl`). Don't propagate them through internal code. + +### `getlogjoint_internal` vs `getlogjoint` + +Samplers operating in unconstrained space should use `getlogjoint_internal`, which includes the Jacobian correction from the linking transform. This is the default and what you almost always want. The exceptions are ESS (which needs the likelihood in constrained space, per the algorithm) and optimisation (where the Jacobian term should not influence the objective). + +### AD backend handling + +Gradient-based samplers accept an `adtype::ADTypes.AbstractADType` keyword (default: `AutoForwardDiff()`). When reviewing sampler code, check that `adtype` is threaded through to `LogDensityFunction` and not hardcoded. The AD backend is the user's choice, not the sampler's. + +### `initial_params` conversion + +User-facing functions accept `initial_params` as a convenience. `_convert_initial_params` in `src/common.jl` converts `NamedTuple`/`Dict{VarName}` to `InitFromParams`. Raw vectors are no longer supported and will error. Don't bypass this conversion or accept raw vectors in new code. + +### Discrete variables + +`allow_discrete_variables(sampler)` defaults to `true`. Gradient-based samplers (all `Hamiltonian` subtypes) override this to `false`. `_check_model` uses this to validate the model before sampling. If adding a new sampler that requires continuous variables, override `allow_discrete_variables` to return `false`. + +### GibbsContext is not ConditionContext + +`GibbsContext` is distinct from `condition`/`ConditionContext`. For non-target variables, `GibbsContext.tilde_assume!!` calls `tilde_observe!!` — this means particle samplers (PG/CSMC) will correctly resample on conditioned variables. The key difference from `condition` is that `GibbsContext` obtains the conditioned values from the global `VarNamedTuple` rather than from the model's conditioning, and it handles the bookkeeping needed for Gibbs (e.g. updating the global VNT when new variables appear). + +## Contributing + + - Non-breaking changes target `main`; breaking changes target the `breaking` branch. + - Julia ≥ 1.10.8 required (see `[compat]` in `Project.toml`). diff --git a/CLAUDE.md b/CLAUDE.md index 4922775aff..43c994c2d3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,100 +1 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Project Overview - -Turing.jl is the user-facing entry point for the [TuringLang](https://github.com/TuringLang) probabilistic programming ecosystem. It is largely a translation layer between DynamicPPL models — which work with named, structured parameters — and inference algorithms that expect flat, vectorised samples (e.g. HMC/NUTS operate on `AbstractVector{<:Real}`). DynamicPPL's `LogDensityFunction` handles most of this translation; Turing provides the sampler wrappers that set it up and manage state across iterations. - -Model definition lives in [DynamicPPL.jl](https://github.com/TuringLang/DynamicPPL.jl), parameter transformations in [Bijectors.jl](https://github.com/TuringLang/Bijectors.jl), and sampling interfaces in [AbstractMCMC.jl](https://github.com/TuringLang/AbstractMCMC.jl). Turing re-exports their APIs and provides concrete sampler implementations that wire everything together. - -## Building and Testing - -Code formatting uses [JuliaFormatter.jl](https://github.com/domluna/JuliaFormatter.jl) v1 (not v2) with the **Blue style** (configured in `.JuliaFormatter.toml`). CI enforces formatting on all PRs. JuliaFormatter must be installed in the **global** Julia environment, not the project environment — do not use `--project`. See the [formatting guide](https://turinglang.org/docs/contributing/code-formatting/) for setup details. - -```bash -julia -e 'using JuliaFormatter; format(".")' -``` - -Tests use `SelectiveTests.jl` (in `test/test_utils/`) to filter by path. CI splits the suite into four shards: `mcmc/gibbs.jl`, `mcmc/Inference.jl`, `ad.jl`, and everything else. To run a subset locally: - -```bash -julia --project -e 'using Pkg; Pkg.test(; test_args=["mcmc/hmc.jl"])' -``` - -Use `--skip` to exclude files: - -```bash -julia --project -e 'using Pkg; Pkg.test(; test_args=["--skip", "mcmc/gibbs.jl", "ad.jl"])' -``` - -CI matrix: Julia stable + min, Ubuntu/Windows/macOS, 1 and 2 threads. - -`test/test_utils/sampler.jl` provides generic test helpers (`test_rng_respected`, `test_sampler_analytical`, `test_chain_logp_metadata`) that should work for any sampler. Beyond these, sampler-specific tests are needed to capture the properties you care about — there is no standardised test template yet. - -## Architecture - -### What lives here vs elsewhere - -Most complexity is in DynamicPPL. Turing.jl contains: - - - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (wrapping AdvancedPS), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. - - **External sampler interface** (`src/mcmc/external_sampler.jl`): The `externalsampler()` wrapper lets any `AbstractMCMC.AbstractSampler` that implements `step` for `LogDensityModel` work with Turing models. This is the easier path for new samplers — it only requires a dependency on AbstractMCMC and the LogDensityProblems.jl interface, with no Turing internals. The tradeoff is less power: you can only interact with the model as a black-box log-density function, just like using `LogDensityFunction` directly. - - **Variational inference** (`src/variational/`): Wraps AdvancedVI algorithms. - - **Mode estimation** (`src/optimisation/`): MAP and MLE via Optimization.jl. - - **Custom distributions** (`src/stdlib/`): `Flat`, `FlatPos`, `BinomialLogit`, `OrderedLogistic`, `LogPoisson`, and Dirichlet/Chinese Restaurant processes. - -For how the model and inference machinery works under the hood, see the [DynamicPPL docs](https://turinglang.org/DynamicPPL.jl/stable/) and the [developer guides](https://turinglang.org/docs/developers/). - -### Gibbs sampler - -The Gibbs sampler (`src/mcmc/gibbs.jl`) is the most complex piece in Turing.jl. It maintains a global `VarNamedTuple` of raw values for all variables. On each iteration, it conditions the model on the non-target variables via `GibbsContext`, runs the component sampler, and updates the global state. - -To plug a sampler into Gibbs, implement: - - - `gibbs_get_raw_values(state)` — return a `VarNamedTuple` of raw values for the variables this sampler is responsible for. - - `gibbs_update_state!!(sampler, state, model, global_vals)` — update the sampler's state to reflect new conditioned values. For samplers that use `LogDensityFunction`, the helper `gibbs_recompute_ldf_and_params` handles the common case. - - Optionally, `isgibbscomponent(sampler)` — return `false` to disallow use in Gibbs (the default is `true`). - -### Extension - -`ext/TuringDynamicHMCExt` provides the DynamicHMC.jl integration (loaded when DynamicHMC is imported). - -## Review Guidelines - -### Use `OnlyAccsVarInfo`, not `VarInfo` - -Sampler state should use `OnlyAccsVarInfo` (with appropriate accumulators), not `VarInfo`. `VarInfo` is being phased out across the ecosystem. - -Most gradient-based samplers (HMC, NUTS, external samplers) go through `LogDensityFunction`, which handles the model interaction. `LogDensityFunction` works well when the model structure is static (the set of variables is fixed across evaluations) and the sampler only needs a scalar log-density value. However, LDF is hard to use when the sampler needs extra accumulators beyond log-probability — for example, MH uses custom accumulators to capture proposal distributions and linked values, so it works directly with `OnlyAccsVarInfo` + `init!!` instead. Either approach is fine; the key constraint is no `VarInfo`. - -Note: "linked" and "unconstrained" are synonymous in this codebase. Linking transforms constrained parameters to unconstrained (Euclidean) space for gradient-based sampling. - -### `VarNamedTuple` for parameter collections - -Interfaces that accept or return named parameter collections should use `VarNamedTuple`, not `NamedTuple` or `Dict{VarName}`. `NamedTuple` and `Dict{VarName}` are accepted as user-facing input but should be converted to `VarNamedTuple` at the boundary (see `_to_varnamedtuple` in `src/common.jl`). Don't propagate them through internal code. - -### `getlogjoint_internal` vs `getlogjoint` - -Samplers operating in unconstrained space should use `getlogjoint_internal`, which includes the Jacobian correction from the linking transform. This is the default and what you almost always want. The exceptions are ESS (which needs the likelihood in constrained space, per the algorithm) and optimisation (where the Jacobian term should not influence the objective). - -### AD backend handling - -Gradient-based samplers accept an `adtype::ADTypes.AbstractADType` keyword (default: `AutoForwardDiff()`). When reviewing sampler code, check that `adtype` is threaded through to `LogDensityFunction` and not hardcoded. The AD backend is the user's choice, not the sampler's. - -### `initial_params` conversion - -User-facing functions accept `initial_params` as a convenience. `_convert_initial_params` in `src/common.jl` converts `NamedTuple`/`Dict{VarName}` to `InitFromParams`. Raw vectors are no longer supported and will error. Don't bypass this conversion or accept raw vectors in new code. - -### Discrete variables - -`allow_discrete_variables(sampler)` defaults to `true`. Gradient-based samplers (all `Hamiltonian` subtypes) override this to `false`. `_check_model` uses this to validate the model before sampling. If adding a new sampler that requires continuous variables, override `allow_discrete_variables` to return `false`. - -### GibbsContext is not ConditionContext - -`GibbsContext` is distinct from `condition`/`ConditionContext`. For non-target variables, `GibbsContext.tilde_assume!!` calls `tilde_observe!!` — this means particle samplers (PG/CSMC) will correctly resample on conditioned variables. The key difference from `condition` is that `GibbsContext` obtains the conditioned values from the global `VarNamedTuple` rather than from the model's conditioning, and it handles the bookkeeping needed for Gibbs (e.g. updating the global VNT when new variables appear). - -## Contributing - - - Non-breaking changes target `main`; breaking changes target the `breaking` branch. - - Julia ≥ 1.10.8 required (see `[compat]` in `Project.toml`). +@AGENTS.md From 564b471e8933169deb756d16ee6d0555a33b049d Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 17:36:58 +0100 Subject: [PATCH 27/58] SMC: bundle the population directly, drop the step-loop, PGState, and weighting Per Charles's review: SMC is a single weighted sweep, not a Markov chain, so overload `AbstractMCMC.sample` to run the sweep and bundle the whole population in one shot instead of faking iteration through the step loop with an `SMCState` cursor. Removes `SMCState` and both SMC `step` methods; `discard_initial`/`thinning` then have no loop to (not) apply to. A final resampling step makes the returned particles an equal-weight sample, so downstream chain summaries (`mean`, etc.) are correct without weighting; SMC chains no longer carry a per-particle `weight` column. Also drops the `const PGState = Particle` alias -- with SMC carrying no state struct, PG's state is plainly a `Particle`. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 2 +- src/mcmc/particle_mcmc.jl | 92 +++++++++++---------------------------- 2 files changed, 26 insertions(+), 68 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index db538c4bd9..2425d88841 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -16,7 +16,7 @@ The rewrite also brings: - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. - - **Weight diagnostics.** `SMC` chains carry the log-evidence estimate `logevidence` and the per-particle normalised `weight` as extra columns; `PG` / `CSMC` chains carry `logevidence`. + - **Equal-weight draws & evidence.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry the log-evidence estimate `logevidence` as an extra column. # 0.46.0 diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 46979a1a28..a2cdd06ad6 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -135,8 +135,8 @@ DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any Particle(model, varinfo, rng::TracedRNG) A single particle: a suspended `model` execution together with its `varinfo`, its own -replayable `rng`, and an accumulated `logweight`. It also serves as the particle Gibbs -sampler state, aliased `PGState` below. +replayable `rng`, and an accumulated `logweight`. It also serves directly as the particle +Gibbs sampler state (there is no separate state struct). """ mutable struct Particle{RT<:TracedRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go @@ -551,22 +551,17 @@ function SMC(scheme::AbstractResampler, threshold::Real; kwargs...) return SMC(ESSResampler(threshold, scheme); kwargs...) end -struct SMCState{P<:AbstractVector,W<:AbstractVector} - particles::P - weights::W - index::Int - logevidence::DynamicPPL.LogProbType -end - +# SMC is a single weighted sweep, not a Markov chain: rather than fake an iteration through +# AbstractMCMC's step loop (returning the population one particle at a time), we run the sweep +# and bundle the whole population into the chain in one shot. `discard_initial`/`thinning` +# therefore have nothing to apply to. function AbstractMCMC.sample( rng::AbstractRNG, model::DynamicPPL.Model, sampler::SMC, - N::Integer; + nparticles::Integer; check_model=true, chain_type=DEFAULT_CHAIN_TYPE, - initial_params=Turing.Inference.init_strategy(sampler), - progress=PROGRESS[], discard_initial=0, thinning=1, verbose=false, @@ -574,62 +569,26 @@ function AbstractMCMC.sample( ) check_model && Turing._check_model(model, sampler) error_if_threadsafe_eval(model) - # SMC is not a Markov chain, so these AbstractMCMC knobs do not apply. Consume them here - # rather than forwarding them to `mcmcsample` (which would `BoundsError`, see #1811). if discard_initial > 0 || thinning > 1 @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." end - chain = AbstractMCMC.mcmcsample( - rng, - model, - sampler, - N; - chain_type, - initial_params, - progress, - nparticles=N, - kwargs..., - ) - post_sample_hook(chain, sampler; verbose) - return chain -end - -# The whole sweep runs on the first step; later steps read off the population one particle at -# a time (SMC returns a weighted sample, not a chain). -function AbstractMCMC.step( - rng::AbstractRNG, - model::DynamicPPL.Model, - sampler::SMC; - nparticles::Int, - discard_sample=false, - kwargs..., -) particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] logZ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) - - stats = (; weight=weights[1], logevidence=logZ) - transition = - discard_sample ? nothing : DynamicPPL.ParamsWithStats(particles[1].varinfo, stats) - return transition, SMCState(particles, weights, 2, logZ) -end - -function AbstractMCMC.step( - ::AbstractRNG, - ::DynamicPPL.Model, - ::SMC, - state::SMCState; - discard_sample=false, - kwargs..., -) - i = state.index - stats = (; weight=state.weights[i], logevidence=state.logevidence) - transition = if discard_sample - nothing - else - DynamicPPL.ParamsWithStats(deepcopy(state.particles[i].varinfo), stats) + # One final resampling step, so the returned particles are an equal-weight sample. The + # sweep ends on a reweight, leaving the population weighted; resampling once here makes the + # result a standard unweighted chain (so `mean(chain[...])` and friends need no weighting), + # at the cost of a little resampling variance. Unconditional -- unlike the ESS-gated + # resampling inside the sweep. + ancestors = resample_indices(rng, sampler.resampler, weights, nparticles) + transitions = map(ancestors) do a + DynamicPPL.ParamsWithStats(particles[a].varinfo, (; logevidence=logZ)) end - return transition, SMCState(state.particles, state.weights, i + 1, state.logevidence) + chain = AbstractMCMC.bundle_samples( + transitions, model, sampler, nothing, chain_type; kwargs... + ) + post_sample_hook(chain, sampler; verbose) + return chain end # @@ -679,9 +638,8 @@ end const CSMC = PG # PG's sampler state is just the retained `Particle`: it already carries the reference -# trajectory's `varinfo` and `rng` (its `task`/`logweight` are then unused), so we alias -# rather than define a separate struct. -const PGState = Particle +# trajectory's `varinfo` and `rng` (its `task`/`logweight` are then unused), so there is no +# dedicated state struct. # First iteration: an ordinary (unconditional) particle sweep. function AbstractMCMC.step( @@ -701,7 +659,7 @@ function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, sampler::PG, - state::PGState; + state::Particle; discard_sample=false, kwargs..., ) @@ -733,10 +691,10 @@ end # Gibbs interface # -gibbs_get_raw_values(state::PGState) = DynamicPPL.get_raw_values(state.varinfo) +gibbs_get_raw_values(state::Particle) = DynamicPPL.get_raw_values(state.varinfo) function gibbs_update_state!!( - ::PG, state::PGState, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple + ::PG, state::Particle, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple ) init = DynamicPPL.InitFromParams(global_vals, nothing) # Re-initialise the reference varinfo with the values conditioned by other Gibbs From 99a2c2bd2ecba4d8d3507c4bd57f511254ca6f38 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 18:03:08 +0100 Subject: [PATCH 28/58] Add per-observation ESS to SMC output; rename evidence stat to log_normalizing_constant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SMC chains now carry `ess_per_step`: the effective sample size (1/Σŵ²) after each observation's reweight, a sweep-level degeneracy diagnostic reported alongside the normalizing-constant estimate. `sweep!` records the trajectory and returns it (PG discards it), and the ESS formula is factored out of `should_resample` into a shared `ess` helper. Renames the marginal-likelihood / normalizing-constant statistic `logevidence` to `log_normalizing_constant` -- the standard SMC/PMCMC term for the estimated p(y) -- including the internal `log_normalizing_constant(particles)` helper. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 2 +- src/mcmc/particle_mcmc.jl | 39 +++++++++++++++++++++++++------------- test/mcmc/particle_mcmc.jl | 19 ++++++++++--------- 3 files changed, 37 insertions(+), 23 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 2425d88841..dfee1f3abc 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -16,7 +16,7 @@ The rewrite also brings: - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. - - **Equal-weight draws & evidence.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry the log-evidence estimate `logevidence` as an extra column. + - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry the log-normalizing-constant (marginal-likelihood) estimate `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). # 0.46.0 diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index a2cdd06ad6..0e0541690b 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -237,7 +237,7 @@ function DynamicPPL.tilde_observe!!( DynamicPPL.DefaultContext(), dist, left, vn, template, particle.varinfo ) particle.varinfo = vi - Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) # increment this observation added + Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) return left, vi end @@ -394,8 +394,7 @@ end ESSResampler(threshold::Real) = ESSResampler(threshold, Stratified()) function should_resample(resampler::ESSResampler, weights) - ess = inv(sum(abs2, weights)) - return ess ≤ resampler.threshold * length(weights) + return ess(weights) ≤ resampler.threshold * length(weights) end function resample_indices(rng::AbstractRNG, resampler::ESSResampler, weights, n::Integer) return resample_indices(rng, resampler.scheme, weights, n) @@ -410,7 +409,9 @@ end logweights(particles) = [p.logweight for p in particles] normalized_weights(particles) = softmax(logweights(particles)) -logevidence(particles) = logsumexp(logweights(particles)) +log_normalizing_constant(particles) = logsumexp(logweights(particles)) +"Effective sample size of a normalised weight vector, `1 / Σ wᵢ²`." +ess(weights) = inv(sum(abs2, weights)) # Advance one particle by one observation, folding its incremental weight in; return `true` # once it has finished (produced nothing). Factored out so the serial and multithreaded loops in @@ -490,21 +491,27 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional return nothing end -# Run a full particle sweep in place, returning the log-evidence estimate. +# Run a full particle sweep in place, returning the log-evidence estimate and the +# per-observation effective sample sizes. function sweep!( rng::AbstractRNG, particles, resampler, multithreaded::Bool; conditional::Bool=false ) logZ = zero(DynamicPPL.LogProbType) + ess_per_step = Float64[] while true resample_propagate!(rng, particles, resampler, conditional) - logZ0 = logevidence(particles) + logZ0 = log_normalizing_constant(particles) done = reweight!(particles, conditional, multithreaded) # Each observation contributes the log-ratio of total weight it adds; summed over the # sweep these telescope into an estimate of the model's log-evidence log p(y). - logZ += logevidence(particles) - logZ0 + logZ += log_normalizing_constant(particles) - logZ0 done && break + # Post-reweight ESS for this observation: a degeneracy diagnostic (low ESS means few + # particles carry the weight). After the break, so the finishing pass -- which adds no + # observation and leaves the weights unchanged -- contributes no spurious entry. + push!(ess_per_step, ess(normalized_weights(particles))) end - return logZ + return logZ, ess_per_step end # @@ -573,7 +580,7 @@ function AbstractMCMC.sample( @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." end particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] - logZ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) + logZ, ess_per_step = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) # One final resampling step, so the returned particles are an equal-weight sample. The # sweep ends on a reweight, leaving the population weighted; resampling once here makes the @@ -581,8 +588,12 @@ function AbstractMCMC.sample( # at the cost of a little resampling variance. Unconditional -- unlike the ESS-gated # resampling inside the sweep. ancestors = resample_indices(rng, sampler.resampler, weights, nparticles) + # `log_normalizing_constant` and `ess_per_step` are sweep-level, so every returned particle carries the + # same values. transitions = map(ancestors) do a - DynamicPPL.ParamsWithStats(particles[a].varinfo, (; logevidence=logZ)) + DynamicPPL.ParamsWithStats( + particles[a].varinfo, (; log_normalizing_constant=logZ, ess_per_step) + ) end chain = AbstractMCMC.bundle_samples( transitions, model, sampler, nothing, chain_type; kwargs... @@ -649,7 +660,7 @@ function AbstractMCMC.step( particles = [ Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:(sampler.nparticles) ] - logZ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) + logZ, _ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) end @@ -669,7 +680,7 @@ function AbstractMCMC.step( particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end - logZ = sweep!( + logZ, _ = sweep!( rng, particles, sampler.resampler, sampler.multithreaded; conditional=true ) return pg_transition_and_state(rng, particles, logZ, discard_sample) @@ -682,7 +693,9 @@ function pg_transition_and_state(rng, particles, logZ, discard_sample) transition = if discard_sample nothing else - DynamicPPL.ParamsWithStats(deepcopy(retained.varinfo), (; logevidence=logZ)) + DynamicPPL.ParamsWithStats( + deepcopy(retained.varinfo), (; log_normalizing_constant=logZ) + ) end return transition, retained end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index b2dc77a806..253a3811e6 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -91,7 +91,7 @@ using Turing test_rng_respected(SMC()) end - @testset "logevidence" begin + @testset "log_normalizing_constant" begin @model function test() a ~ Normal(0, 1) x ~ Bernoulli(1) @@ -105,13 +105,14 @@ using Turing chains_smc = sample(StableRNG(100), test(), SMC(), 100) @test all(isone, chains_smc[:x]) - # For SMC, the chain stores the collective logevidence of the sampled trajectories + # For SMC, the chain stores the collective log_normalizing_constant of the sampled trajectories # as a statistic (which is the same for all 'iterations'). So we can just pick the # first one. - smc_logevidence = first(chains_smc[:logevidence]) - @test smc_logevidence ≈ -2 * log(2) + smc_log_normalizing_constant = first(chains_smc[:log_normalizing_constant]) + @test smc_log_normalizing_constant ≈ -2 * log(2) # Check that they're all equal. - @test chains_smc[:logevidence] ≈ fill(smc_logevidence, 100) + @test chains_smc[:log_normalizing_constant] ≈ + fill(smc_log_normalizing_constant, 100) end @testset "multithreaded execution matches serial" begin @@ -190,7 +191,7 @@ end test_rng_respected(PG(10)) end - @testset "logevidence" begin + @testset "log_normalizing_constant" begin @model function test() a ~ Normal(0, 1) x ~ Bernoulli(1) @@ -204,10 +205,10 @@ end chains_pg = sample(StableRNG(468), test(), PG(10), 100) @test all(isone, chains_pg[:x]) - pg_logevidence = mean(chains_pg[:logevidence]) - @test pg_logevidence ≈ -2 * log(2) atol = 0.01 + pg_log_normalizing_constant = mean(chains_pg[:log_normalizing_constant]) + @test pg_log_normalizing_constant ≈ -2 * log(2) atol = 0.01 # Should be the same for all iterations. - @test chains_pg[:logevidence] ≈ fill(pg_logevidence, 100) + @test chains_pg[:log_normalizing_constant] ≈ fill(pg_log_normalizing_constant, 100) end @testset "multithreaded execution matches serial" begin From 186282c1fe4f00e4339eef4dbb9f6539617664f7 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 18:14:59 +0100 Subject: [PATCH 29/58] Fix Inference.jl test for the log_normalizing_constant rename The `logevidence` -> `log_normalizing_constant` stat rename missed the SMC log-evidence assertion in test/mcmc/Inference.jl (CI caught it on the min-Julia mcmc/Inference shard). Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- test/mcmc/Inference.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/mcmc/Inference.jl b/test/mcmc/Inference.jl index c085184f97..42538e8f99 100644 --- a/test/mcmc/Inference.jl +++ b/test/mcmc/Inference.jl @@ -394,13 +394,13 @@ using Turing N = 1_000 - # For SMC, the chain stores the collective logevidence of the sampled trajectories + # For SMC, the chain stores the collective log_normalizing_constant of the sampled trajectories # as a statistic (which is the same for all 'iterations'). So we can just pick the # first one. res_smc = sample(StableRNG(seed), test(), smc, N) @test all(isone, res_smc[@varname(x)]) - smc_logevidence = first(res_smc[:logevidence]) - @test smc_logevidence ≈ 2 * log(0.5) + smc_log_normalizing_constant = first(res_smc[:log_normalizing_constant]) + @test smc_log_normalizing_constant ≈ 2 * log(0.5) res_pg = sample(StableRNG(seed), test(), pg, 100) @test all(isone, res_pg[@varname(x)]) From b55ea5a10053c735fad00e9fa4fdceca1aac5232 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 19:09:04 +0100 Subject: [PATCH 30/58] Rename resampler types with a Resampler suffix `Multinomial`/`Stratified`/`Systematic` -> `MultinomialResampler`/ `StratifiedResampler`/`SystematicResampler`, so the resampling-scheme types no longer shadow common names (notably `Distributions.Multinomial`) inside `Turing.Inference`. Algorithm-description docstrings are unchanged. (Review finding 4.) Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 24 ++++++++++++------------ test/mcmc/particle_mcmc.jl | 26 ++++++++++++++------------ 2 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 0e0541690b..9eab10922b 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -341,14 +341,14 @@ should_resample(::AbstractResampler, weights) = true function resample_indices end "Multinomial resampling: `n` independent draws from the categorical over `weights`." -struct Multinomial <: AbstractResampler end -function resample_indices(rng::AbstractRNG, ::Multinomial, weights, n::Integer) +struct MultinomialResampler <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::MultinomialResampler, weights, n::Integer) return rand(rng, Distributions.Categorical(weights), n) end "Stratified resampling: one independent uniform per stratum of width `1/n`." -struct Stratified <: AbstractResampler end -function resample_indices(rng::AbstractRNG, ::Stratified, weights, n::Integer) +struct StratifiedResampler <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::StratifiedResampler, weights, n::Integer) v = n * weights[1] indices = Vector{Int}(undef, n) s = 1 @@ -364,8 +364,8 @@ function resample_indices(rng::AbstractRNG, ::Stratified, weights, n::Integer) end "Systematic resampling: one shared uniform placed on a regular grid of `n` points." -struct Systematic <: AbstractResampler end -function resample_indices(rng::AbstractRNG, ::Systematic, weights, n::Integer) +struct SystematicResampler <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::SystematicResampler, weights, n::Integer) v = n * weights[1] u = oftype(v, rand(rng)) indices = Vector{Int}(undef, n) @@ -382,7 +382,7 @@ function resample_indices(rng::AbstractRNG, ::Systematic, weights, n::Integer) end """ - ESSResampler(threshold, scheme = Stratified()) + ESSResampler(threshold, scheme = StratifiedResampler()) Resample with `scheme`, but only when the effective sample size drops below `threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). @@ -391,7 +391,7 @@ struct ESSResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler threshold::T scheme::R end -ESSResampler(threshold::Real) = ESSResampler(threshold, Stratified()) +ESSResampler(threshold::Real) = ESSResampler(threshold, StratifiedResampler()) function should_resample(resampler::ESSResampler, weights) return ess(weights) ≤ resampler.threshold * length(weights) @@ -541,7 +541,7 @@ end """ SMC([resampler = ESSResampler(0.5)]; multithreaded = false) - SMC([scheme = Stratified(), ]threshold; multithreaded = false) + SMC([scheme = StratifiedResampler(), ]threshold; multithreaded = false) Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. @@ -549,8 +549,8 @@ effective sample size drops below half the number of particles. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). -The resampling scheme types (`Stratified`, `Systematic`, `Multinomial`, `ESSResampler`) are -not exported; refer to them as e.g. `Turing.Inference.Systematic`. +The resampling scheme types (`StratifiedResampler`, `SystematicResampler`, `MultinomialResampler`, `ESSResampler`) are +not exported; refer to them as e.g. `Turing.Inference.SystematicResampler`. """ SMC(; kwargs...) = SMC(ESSResampler(0.5); kwargs...) SMC(threshold::Real; kwargs...) = SMC(ESSResampler(threshold); kwargs...) @@ -631,7 +631,7 @@ end """ PG(n, [resampler = ESSResampler(0.5)]; multithreaded = false) - PG(n, [scheme = Stratified(), ]threshold; multithreaded = false) + PG(n, [scheme = StratifiedResampler(), ]threshold; multithreaded = false) Particle Gibbs sampler with `n` particles. By default stratified resampling is triggered whenever the effective sample size drops below half the number of particles. diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 253a3811e6..e6309a8729 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -5,9 +5,9 @@ using ..SamplerTestUtils: test_chain_logp_metadata, test_rng_respected using ..NumericalTests: check_numerical using DynamicPPL: DynamicPPL, extract_priors, get_raw_values, getloglikelihood using Turing.Inference: - Stratified, - Systematic, - Multinomial, + StratifiedResampler, + SystematicResampler, + MultinomialResampler, ESSResampler, Particle, TracedRNG, @@ -28,13 +28,14 @@ using Turing @testset "SMC" begin @testset "constructor" begin @test SMC().resampler == ESSResampler(0.5) - @test SMC().resampler.scheme isa Stratified # stratified is the default scheme + @test SMC().resampler.scheme isa StratifiedResampler # stratified is the default scheme @test SMC(0.6).resampler == ESSResampler(0.6) - @test SMC(Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) - @test SMC(Systematic()).resampler == Systematic() + @test SMC(MultinomialResampler(), 0.6).resampler == + ESSResampler(0.6, MultinomialResampler()) + @test SMC(SystematicResampler()).resampler == SystematicResampler() @test SMC().multithreaded == false @test SMC(; multithreaded=true).multithreaded == true - @test SMC(Systematic(); multithreaded=true).multithreaded == true + @test SMC(SystematicResampler(); multithreaded=true).multithreaded == true end @testset "basic model" begin @@ -61,8 +62,8 @@ using Turing exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs)) # every scheme targets the same posterior... - chn_strat = sample(StableRNG(23), coin_model, SMC(Stratified()), 100) - chn_multi = sample(StableRNG(23), coin_model, SMC(Multinomial()), 100) + chn_strat = sample(StableRNG(23), coin_model, SMC(StratifiedResampler()), 100) + chn_multi = sample(StableRNG(23), coin_model, SMC(MultinomialResampler()), 100) check_numerical(chn_strat, [@varname(p)], [mean(exact)]; atol=0.1) check_numerical(chn_multi, [@varname(p)], [mean(exact)]; atol=0.1) # ...but the schemes are genuinely different, so the draws differ. @@ -176,11 +177,12 @@ end @test PG(10).nparticles == 10 @test PG(10).resampler == ESSResampler(0.5) @test PG(60, 0.6).resampler == ESSResampler(0.6) - @test PG(80, Multinomial(), 0.6).resampler == ESSResampler(0.6, Multinomial()) - @test PG(100, Systematic()).resampler == Systematic() + @test PG(80, MultinomialResampler(), 0.6).resampler == + ESSResampler(0.6, MultinomialResampler()) + @test PG(100, SystematicResampler()).resampler == SystematicResampler() @test PG(10).multithreaded == false @test PG(10; multithreaded=true).multithreaded == true - @test PG(80, Multinomial(), 0.6; multithreaded=true).multithreaded == true + @test PG(80, MultinomialResampler(), 0.6; multithreaded=true).multithreaded == true end @testset "chain log-density metadata" begin From 96c11a50d19dabbf33183e96d55ae58f750df511 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 19:56:12 +0100 Subject: [PATCH 31/58] Pin the CSMC reference to the retained trajectory by value MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The particle-Gibbs reference reproduced its trajectory by replaying RNG seeds, which only regenerates the retained values when the model is unchanged. Inside Gibbs the model is re-conditioned between sweeps -- one component's variables become observations parametrised by the others -- so seed-replay re-drew the reference from a *different* prior rather than reproducing the retained trajectory. CSMC then lost its defining invariant and behaved like independent SMC: exact at large N but biased at small N. On a linear-Gaussian state-space model, Gibbs(:x => PG(2), :ρ => MH()) put E[x1] at 0.64 against NUTS's 0.96; the fix closes the gap. Reproduce by value instead: the sampler state carries the retained values, and the reference re-runs the model with InitFromParams, so it stays the retained trajectory regardless of re-conditioning. A fork clears the carried values (reseed!) and thus samples fresh past the fork point. The existing consistency test only covers same-model reproduction; the new test pins the reference under changed conditioning, the Gibbs scenario. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 66 ++++++++++++++++++++++++++------------ test/mcmc/particle_mcmc.jl | 29 +++++++++++++++++ 2 files changed, 75 insertions(+), 20 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 9eab10922b..4c1cad16c3 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -9,13 +9,14 @@ ### such sweep; particle Gibbs (PG/CSMC) runs a *conditional* sweep -- one particle is a fixed ### reference trajectory -- inside an MCMC loop. ### -### The reference is reproduced without storing its values: each particle's `TracedRNG` -### records the seed it drew at every step, and the reference simply *replays* those seeds -### (`load_state!`), regenerating its trajectory exactly. A particle forked from the reference -### is reseeded, which flips it from replaying to sampling fresh, so branching needs no -### per-particle flag. This is what keeps the reference handling small, and it rests on one -### invariant: every step must draw from a fresh seed -- guaranteed by the resample/refresh in -### `resample_propagate!` -- otherwise the recorded seeds collide and replay is wrong. +### The reference reproduces the retained trajectory by *reusing its values*: the sampler state +### carries those values, and the reference re-runs the model with `InitFromParams`, so it stays +### the retained trajectory even when the model is re-conditioned between Gibbs sweeps (e.g. a +### state-space transition prior that depends on a parameter owned by another Gibbs component). +### A particle forked from the reference forgets the remaining values (in `reseed!`), so +### branching samples fresh with no per-particle flag. Each particle's counter-based `TracedRNG` +### supplies splittable, version-stable seeds that decorrelate the fresh draws across particles +### and Julia versions. ### ### Sections below: traced RNG; model evaluation via Libtask; resampling schemes; the particle ### sweep; the SMC sampler; the PG/CSMC sampler; the Gibbs-component interface. @@ -146,22 +147,35 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # `logweight` tracks whatever `DynamicPPL.LogProbType` is, so weights follow suit if it # is ever changed. logweight::WT + # The retained trajectory's values, which the CSMC reference reproduces by reusing them + # (`InitFromParams` in `tilde_assume!!`); empty for ordinary particles. `reseed!` clears it + # so a particle forked off the reference samples fresh beyond the fork point. Reproducing + # by *value* (not by replayed RNG seeds) is what keeps the reference the retained trajectory + # when the model is re-conditioned between Gibbs sweeps. + reference_values::DynamicPPL.VarNamedTuple task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the # particle as its taped globals (a back-reference). This has to be an inner constructor # for that reason: `task` is left undefined here and set immediately after. - function Particle(vi::DynamicPPL.AbstractVarInfo, rng::RT) where {RT<:TracedRNG} + function Particle( + vi::DynamicPPL.AbstractVarInfo, + rng::RT, + reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), + ) where {RT<:TracedRNG} w = zero(DynamicPPL.LogProbType) - return new{RT,typeof(w)}(vi, rng, w) + return new{RT,typeof(w)}(vi, rng, w, reference_values) end end function Particle( - model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo, rng::TracedRNG + model::DynamicPPL.Model, + varinfo::DynamicPPL.AbstractVarInfo, + rng::TracedRNG, + reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), ) model = DynamicPPL.setleafcontext(model, SMCContext()) args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - particle = Particle(deepcopy(varinfo), rng) + particle = Particle(deepcopy(varinfo), rng, reference_values) particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) return particle end @@ -176,6 +190,8 @@ from the reference forgets the reference's future. Mutates and returns `particle function reseed!(particle::Particle, rng::AbstractRNG) Random.seed!(particle.rng, rand(rng, UInt64)) resize!(particle.rng.keys, particle.rng.count - 1) + # A fork samples fresh from here on, so it must forget the reference's remaining values. + particle.reference_values = DynamicPPL.VarNamedTuple() return particle end @@ -205,14 +221,16 @@ function DynamicPPL.tilde_assume!!( ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) - # Always draw from the prior. A value is never already present here: particle varinfos - # start empty, each variable is assumed exactly once, and the CSMC reference reproduces its - # trajectory by replaying its RNG seeds, not by reusing stored values. Reusing an existing - # value (via `InitFromParams`) would only matter for a pre-populated varinfo, which - # particle sampling does not currently create. - ctx = DynamicPPL.InitContext( - particle.rng, DynamicPPL.InitFromPrior(), DynamicPPL.UnlinkAll() - ) + # Reuse the retained value (`InitFromParams`) if this particle is reproducing the CSMC + # reference trajectory and still carries this variable; otherwise draw from the prior. + # `reference_values` is empty for ordinary particles and is cleared by `reseed!` on a fork, + # so a fork of the reference draws fresh past the fork point (see the `Particle` fields). + strategy = if haskey(particle.reference_values, vn) + DynamicPPL.InitFromParams(particle.reference_values, nothing) + else + DynamicPPL.InitFromPrior() + end + ctx = DynamicPPL.InitContext(particle.rng, strategy, DynamicPPL.UnlinkAll()) x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, particle.varinfo) particle.varinfo = vi return x, vi @@ -676,7 +694,15 @@ function AbstractMCMC.step( ) error_if_threadsafe_eval(model) n = sampler.nparticles - reference = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng))) + # The reference reproduces the retained trajectory by reusing its values (passed here and + # consumed by `tilde_assume!!`), so it stays that trajectory even if the model was + # re-conditioned since the last sweep. Its varinfo starts empty like any other particle. + reference = Particle( + model, + particle_varinfo(), + rewind!(deepcopy(state.rng)), + DynamicPPL.get_raw_values(state.varinfo), + ) particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index e6309a8729..f011a7f3f5 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -279,6 +279,35 @@ end @test nkeys == length(y) + 1 # keys stay aligned with the trajectory length end + @testset "reference is pinned to retained values under re-conditioning" begin + # Finding 1 regression. In Gibbs the model is re-conditioned between sweeps, so the + # CSMC reference must reproduce the *retained values* rather than re-draw them from the + # (now different) prior. Retain a trajectory under one conditioning, rebuild the + # reference under another; value-pinning keeps the trajectory, whereas re-drawing from + # the prior -- the pre-fix behaviour -- would follow the shifted prior instead. + @model function reconditioned(y) + a ~ Normal(0, 10) + x ~ Normal(a, 1) # x's prior depends on a, owned by another component + return y ~ Normal(x, 1) + end + rng = StableRNG(42) + retained = Particle( + reconditioned(2.0) | (@varname(a) => 0.0), particle_varinfo(), TracedRNG(rng) + ) + while advance!(retained, false) !== nothing + end + retained_vals = get_raw_values(retained.varinfo) + reference = Particle( + reconditioned(2.0) | (@varname(a) => 5.0), # x's prior shifted far away + particle_varinfo(), + rewind!(deepcopy(retained.rng)), + retained_vals, + ) + while advance!(reference, true) !== nothing + end + @test get_raw_values(reference.varinfo) == retained_vals + end + @testset "addlogprob leads to reweighting" begin # Make sure that PG takes @addlogprob! into account. It didn't use to: # https://github.com/TuringLang/Turing.jl/issues/1996 From 14bc29dc182fa2087cc79e788454fa6b94d9ea5a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 19:59:02 +0100 Subject: [PATCH 32/58] Guard stratified/systematic resampling against weight undersum softmax can return weights summing to a hair under one; the cumulative walk then falls short of the final uniform and the unguarded loop indexes one past the last particle -- a BoundsError. `s < length(weights)` clamps the walk to the last particle. Bounding by the weight count (not the resample count) keeps the last particle reachable in the conditional sweep, where fewer indices are drawn than there are particles. The realistic trigger is astronomically rare -- softmax undersum is a few ULP, so the uniform must land within that of the stratum edge -- but the overrun is a latent array-index bug; the test exaggerates the undersum to hit it deterministically. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 9 +++++++-- test/mcmc/particle_mcmc.jl | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 4c1cad16c3..6c4b6d75d1 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -372,7 +372,10 @@ function resample_indices(rng::AbstractRNG, ::StratifiedResampler, weights, n::I s = 1 for k in 1:n u = oftype(v, (k - 1) + rand(rng)) - while v < u + # `s < length(weights)` guards the last particle: if `weights` sums to slightly under + # one (softmax rounding), `v` can fall a hair short of `u` at the final stratum and the + # unguarded loop would index past the end. + while s < length(weights) && v < u s += 1 v += n * weights[s] end @@ -389,7 +392,9 @@ function resample_indices(rng::AbstractRNG, ::SystematicResampler, weights, n::I indices = Vector{Int}(undef, n) s = 1 for k in 1:n - while v < u + # See `StratifiedResampler`: `s < length(weights)` keeps the final stratum from + # indexing past the end when `weights` sums to slightly under one. + while s < length(weights) && v < u s += 1 v += n * weights[s] end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index f011a7f3f5..d78cc774f6 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -17,6 +17,7 @@ using Turing.Inference: rewind!, refresh!, sweep!, + resample_indices, normalized_weights using Distributions: Bernoulli, Beta, Gamma, Normal, Uniform, Categorical, sample using FlexiChains: VNChain, has_same_data @@ -70,6 +71,19 @@ using Turing @test chn_strat[@varname(p)] != chn_multi[@varname(p)] end + @testset "stratified/systematic resampling never index past the end" begin + # softmax can return weights summing to slightly under one; the cumulative walk must + # not run off the end of the last stratum. Exaggerate the undersum so the (otherwise + # astronomically rare) overrun is hit on every seed. + weights = fill(0.9 / 8, 8) + for scheme in (StratifiedResampler(), SystematicResampler()) + @test all( + all(in(1:8), resample_indices(Xoshiro(s), scheme, weights, 8)) for + s in 1:1000 + ) + end + end + @testset "errors when number of observations is not fixed" begin @model function fail_smc() a ~ Normal(4, 5) From 1a917bb4df8c192661edb3d8965751249336cf8c Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 23 Jul 2026 20:03:08 +0100 Subject: [PATCH 33/58] Warn when SMC is given initial_params SMC draws its whole population from the prior, so a user's initial_params cannot apply -- yet it was silently swallowed by kwargs. Warn instead, as SMC already does for discard_initial and thinning. (PG ignores it too, but has no custom sample overload to host the warning without risking spurious ones inside Gibbs, so it is left unchanged.) Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 4 ++++ test/mcmc/particle_mcmc.jl | 6 +++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 6c4b6d75d1..50ed67fccb 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -594,6 +594,7 @@ function AbstractMCMC.sample( chain_type=DEFAULT_CHAIN_TYPE, discard_initial=0, thinning=1, + initial_params=nothing, verbose=false, kwargs..., ) @@ -602,6 +603,9 @@ function AbstractMCMC.sample( if discard_initial > 0 || thinning > 1 @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." end + if initial_params !== nothing + @warn "SMC draws its initial population from the prior; `initial_params` is ignored." + end particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] logZ, ess_per_step = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index d78cc774f6..99aa8f55e2 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -158,7 +158,7 @@ using Turing @test_throws ArgumentError sample(model, SMC(), 100) end - @testset "discard_initial and thinning are ignored" begin + @testset "discard_initial, thinning and initial_params are ignored" begin @model function normal() a ~ Normal(4, 5) 3 ~ Normal(a, 2) @@ -167,6 +167,10 @@ using Turing return a, b end + @test_logs (:warn, r"initial_params.*ignored") match_mode = :any sample( + normal(), SMC(), 10; initial_params=(; a=1.0) + ) + @test_logs (:warn, r"ignored") sample(normal(), SMC(), 10; discard_initial=5) chn = sample(normal(), SMC(), 10; discard_initial=5) @test size(chn, 1) == 10 From bf9357de5bb162f46c642a9e97d3e83c9dbf6851 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Fri, 24 Jul 2026 02:35:58 +0100 Subject: [PATCH 34/58] Explain why the CSMC reference reproduces by value, with the math MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reference_values comment stated the conclusion -- value-replay, not seed-replay, keeps the reference on the retained trajectory across Gibbs re-conditioning -- but not the mechanism. A draw is x = g(u; θ) of the RNG output and the distribution parameters; seed-replay fixes u and recomputes x' = g(u; θ') while value-replay reuses x' = x, so the two agree only when θ' = θ. Re-conditioning changes a θ that another block owns, so θ' ≠ θ and a seed-replayed reference would move with the shifted prior instead of holding its retained values. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 50ed67fccb..5953b3421b 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -151,7 +151,13 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # (`InitFromParams` in `tilde_assume!!`); empty for ordinary particles. `reseed!` clears it # so a particle forked off the reference samples fresh beyond the fork point. Reproducing # by *value* (not by replayed RNG seeds) is what keeps the reference the retained trajectory - # when the model is re-conditioned between Gibbs sweeps. + # when the model is re-conditioned between Gibbs sweeps. A draw is a deterministic function + # x = g(u; θ) of the RNG output `u` and the distribution parameters θ (canonically the + # inverse-CDF, x = F⁻¹(u; θ)). Seed-replay fixes `u` and recomputes x' = g(u; θ'); value-replay + # reuses x' = x. These agree only when θ' = θ. Re-conditioning updates a value θ depends on + # (owned by another block), so θ' ≠ θ and the replayed draw moves with the changed distribution + # -- e.g. x ~ Normal(μ, 1) draws x = μ + Φ⁻¹(u), so after μ → μ' the same `u` gives x + (μ' − μ), + # not x. reference_values::DynamicPPL.VarNamedTuple task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the From 118348b20b62486cc5b0006b4b8ad35719320631 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 00:41:35 +0100 Subject: [PATCH 35/58] Finish the resampler rename in HISTORY.md The types gained their `Resampler` suffix in b55ea5a10; the changelog still advertised the old names. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index dfee1f3abc..9663a0f401 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -6,7 +6,7 @@ `SMC` and `PG` / `CSMC` have been reimplemented natively and no longer depend on AdvancedPS. -Resampling schemes are now types rather than functions — `Stratified()`, `Systematic()`, and `Multinomial()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.Systematic())`, `SMC(0.5)`, or `PG(10, Turing.Inference.Multinomial(), 0.5)`. +Resampling schemes are now types rather than functions — `StratifiedResampler()`, `SystematicResampler()`, and `MultinomialResampler()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.SystematicResampler())`, `SMC(0.5)`, or `PG(10, Turing.Inference.MultinomialResampler(), 0.5)`. The old function-based API (`resample_systematic`, `AdvancedPS.ResampleWithESSThreshold`, …) is gone. The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows and within the theoretical guarantees for particle Gibbs, which systematic does not. From 58a4e0a8c41239ebb8ebfb94a7d0c2119f645dfc Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 00:41:49 +0100 Subject: [PATCH 36/58] Detect a changed reference execution trace instead of silently resampling it A `haskey` miss fell through to the prior, silently unpinning the CSMC reference when re-conditioning changed the trace. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 65 ++++++++++++++++++++++++++----- test/mcmc/particle_mcmc.jl | 79 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 130 insertions(+), 14 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 5953b3421b..505d013b3a 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -159,6 +159,18 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # -- e.g. x ~ Normal(μ, 1) draws x = μ + Φ⁻¹(u), so after μ → μ' the same `u` gives x + (μ' − μ), # not x. reference_values::DynamicPPL.VarNamedTuple + # `nothing` for an ordinary particle; for a CSMC reference, the addresses the retained + # trajectory assumed. Two reasons this cannot be read off `reference_values`: a slice + # assume such as `x[1:2] ~ MvNormal(...)` is stored there under the keys `x[1]`, `x[2]` + # but assumed under the single address `x[1:2]`, so comparing against those keys would + # report a spurious trace change; and an empty `reference_values` cannot distinguish a + # reference with no latents from an ordinary particle. Without it, an address the retained + # trajectory never had would silently draw from the prior, corrupting the reference. + expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}} + # Addresses assumed by this execution, in the same form. Survives forking, so a particle + # that becomes the retained state hands the complete set to the next reference; a + # reference must finish with exactly the set above. + assumed_varnames::Set{DynamicPPL.VarName} task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the # particle as its taped globals (a back-reference). This has to be an inner constructor @@ -167,9 +179,17 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} vi::DynamicPPL.AbstractVarInfo, rng::RT, reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), + expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}}=nothing, ) where {RT<:TracedRNG} w = zero(DynamicPPL.LogProbType) - return new{RT,typeof(w)}(vi, rng, w, reference_values) + return new{RT,typeof(w)}( + vi, + rng, + w, + reference_values, + expected_reference_varnames, + Set{DynamicPPL.VarName}(), + ) end end @@ -178,10 +198,13 @@ function Particle( varinfo::DynamicPPL.AbstractVarInfo, rng::TracedRNG, reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), + expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}}=nothing, ) model = DynamicPPL.setleafcontext(model, SMCContext()) args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - particle = Particle(deepcopy(varinfo), rng, reference_values) + particle = Particle( + deepcopy(varinfo), rng, reference_values, expected_reference_varnames + ) particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) return particle end @@ -198,6 +221,7 @@ function reseed!(particle::Particle, rng::AbstractRNG) resize!(particle.rng.keys, particle.rng.count - 1) # A fork samples fresh from here on, so it must forget the reference's remaining values. particle.reference_values = DynamicPPL.VarNamedTuple() + particle.expected_reference_varnames = nothing return particle end @@ -220,25 +244,45 @@ reference (`isref = true`) replays its recorded seed instead. function advance!(particle::Particle, isref::Bool) isref ? load_state!(particle.rng) : save_state!(particle.rng) inc_step!(particle.rng) - return Libtask.consume(particle.task) + score = Libtask.consume(particle.task) + # `tilde_assume!!` already rejects any address outside the retained set, so once the + # reference has run to completion the only discrepancy left to catch is a retained address + # it never visited (e.g. a branch that stopped being taken after re-conditioning). + expected = particle.expected_reference_varnames + if isref && score === nothing && expected !== nothing + dropped = setdiff(expected, particle.assumed_varnames) + isempty(dropped) || error( + "the reference execution trace changed while replaying retained values " * + "(retained addresses never reached: $(collect(dropped)))", + ) + end + return score end function DynamicPPL.tilde_assume!!( ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) - # Reuse the retained value (`InitFromParams`) if this particle is reproducing the CSMC - # reference trajectory and still carries this variable; otherwise draw from the prior. - # `reference_values` is empty for ordinary particles and is cleared by `reseed!` on a fork, - # so a fork of the reference draws fresh past the fork point (see the `Particle` fields). - strategy = if haskey(particle.reference_values, vn) - DynamicPPL.InitFromParams(particle.reference_values, nothing) - else + # A CSMC reference reuses the retained value (`InitFromParams`) at every address it visits. + # `expected_reference_varnames` is `nothing` for ordinary particles and is cleared by + # `reseed!` on a fork, so both draw from the prior (see the `Particle` fields). An address + # outside the retained set means the execution trace changed, which must error rather than + # silently drawing part of the nominally fixed reference afresh; the `nothing` fallback + # given to `InitFromParams` catches the converse, a retained address with no usable value. + expected = particle.expected_reference_varnames + strategy = if expected === nothing DynamicPPL.InitFromPrior() + else + vn in expected || error( + "the reference execution trace changed while replaying retained values " * + "(new address: $vn)", + ) + DynamicPPL.InitFromParams(particle.reference_values, nothing) end ctx = DynamicPPL.InitContext(particle.rng, strategy, DynamicPPL.UnlinkAll()) x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, particle.varinfo) particle.varinfo = vi + push!(particle.assumed_varnames, vn) return x, vi end @@ -717,6 +761,7 @@ function AbstractMCMC.step( particle_varinfo(), rewind!(deepcopy(state.rng)), DynamicPPL.get_raw_values(state.varinfo), + copy(state.assumed_varnames), ) particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 99aa8f55e2..f60f83da89 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -19,8 +19,9 @@ using Turing.Inference: sweep!, resample_indices, normalized_weights -using Distributions: Bernoulli, Beta, Gamma, Normal, Uniform, Categorical, sample +using Distributions: Bernoulli, Beta, Gamma, MvNormal, Normal, Uniform, Categorical, sample using FlexiChains: VNChain, has_same_data +using LinearAlgebra: I using Random: Random, Xoshiro using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset @@ -256,8 +257,9 @@ end @testset "ensuring reference consistency" begin # In conditional SMC the retained trajectory must be regenerated *exactly* by the - # reference particle on the next iteration -- this is what makes CSMC valid. It fails - # if the traced-RNG step counter and the recorded seeds ever fall out of alignment. + # reference particle on the next iteration -- this is what makes CSMC valid. Value + # replay must reach every latent address, and the traced-RNG step counter and recorded + # seeds must stay aligned with the observation boundaries. @model function state_space_model(y) ρ ~ Uniform(0, 1) x = Vector{Float64}(undef, length(y) + 1) @@ -278,7 +280,13 @@ end state = draw(particles) allok = true for _ in 1:nsteps - ref = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng))) + ref = Particle( + model, + particle_varinfo(), + rewind!(deepcopy(state.rng)), + get_raw_values(state.varinfo), + copy(state.assumed_varnames), + ) parts = map( i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, 1:N, @@ -320,12 +328,75 @@ end particle_varinfo(), rewind!(deepcopy(retained.rng)), retained_vals, + copy(retained.assumed_varnames), ) while advance!(reference, true) !== nothing end @test get_raw_values(reference.varinfo) == retained_vals end + @testset "value replay detects a changed latent trace" begin + @model function branch_changes(flag, y) + if flag + x ~ Normal() + μ = x + else + z ~ Normal() + μ = z + end + return y ~ Normal(μ, 1) + end + rng = StableRNG(91) + retained = Particle(branch_changes(true, 0.0), particle_varinfo(), TracedRNG(rng)) + while advance!(retained, false) !== nothing + end + reference = Particle( + branch_changes(false, 0.0), + particle_varinfo(), + rewind!(deepcopy(retained.rng)), + get_raw_values(retained.varinfo), + copy(retained.assumed_varnames), + ) + @test_throws "reference execution trace changed" advance!(reference, true) + + @model function branch_drops(flag, y) + x ~ Normal() + if flag + z ~ Normal() + end + return y ~ Normal(x, 1) + end + retained = Particle(branch_drops(true, 0.0), particle_varinfo(), TracedRNG(rng)) + while advance!(retained, false) !== nothing + end + reference = Particle( + branch_drops(false, 0.0), + particle_varinfo(), + rewind!(deepcopy(retained.rng)), + get_raw_values(retained.varinfo), + copy(retained.assumed_varnames), + ) + @test_throws "reference execution trace changed" begin + while advance!(reference, true) !== nothing + end + end + end + + @testset "value replay tracks slice assumes by their assumed address" begin + # `x[1:2] ~ MvNormal(...)` is assumed under the single address `x[1:2]` but stored in + # the retained values under the keys `x[1]`, `x[2]`. The expected-address set must + # therefore come from what the trajectory assumed, not from the retained values' keys, + # or every slice assume looks like a changed trace. + @model function slice_assume(y) + x = Vector{Float64}(undef, 2) + x[1:2] ~ MvNormal(zeros(2), I) + y[1] ~ Normal(x[1], 0.5) + return y[2] ~ Normal(x[2], 0.5) + end + chn = sample(StableRNG(105), slice_assume([0.4, -0.4]), PG(5), 20) + @test size(chn, 1) == 20 + end + @testset "addlogprob leads to reweighting" begin # Make sure that PG takes @addlogprob! into account. It didn't use to: # https://github.com/TuringLang/Turing.jl/issues/1996 From 43f46940fa979873b8bb590e22dacbeb63847fb9 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 00:42:04 +0100 Subject: [PATCH 37/58] Draw conditional-sweep ancestors multinomially, whatever scheme is named Pinning one output of a stratified or systematic draw is not that scheme's conditional law. Changes PG draws. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- HISTORY.md | 3 ++- src/mcmc/particle_mcmc.jl | 39 ++++++++++++++++++++++++++++---------- test/mcmc/particle_mcmc.jl | 38 ++++++++++++++++++++++++++++++++++++- 3 files changed, 68 insertions(+), 12 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 9663a0f401..6d9bd005b8 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -9,7 +9,8 @@ Resampling schemes are now types rather than functions — `StratifiedResampler()`, `SystematicResampler()`, and `MultinomialResampler()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.SystematicResampler())`, `SMC(0.5)`, or `PG(10, Turing.Inference.MultinomialResampler(), 0.5)`. The old function-based API (`resample_systematic`, `AdvancedPS.ResampleWithESSThreshold`, …) is gone. -The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows and within the theoretical guarantees for particle Gibbs, which systematic does not. +The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows, which systematic does not. +The selected scheme applies to unconditional sweeps; `PG` / `CSMC` draw the ancestors of a conditional sweep from the categorical over the weights, since a correct conditional version of stratified or systematic resampling is scheme-specific rather than "pin one draw and keep the rest". Exact draws may therefore differ from previous releases, but remain statistically consistent (the same target distribution). The rewrite also brings: diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 505d013b3a..d35c224f48 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -392,13 +392,22 @@ end # # Resampling schemes # -# On theoretical correctness: particle Gibbs (and the SMC evidence estimate) are justified -# for resampling schemes whose offspring counts satisfy `E[Oᵏ] = N·Wᵏ` (Andrieu, Doucet & -# Holenstein, 2010, Assumption 2). Multinomial and stratified resampling meet this and are -# also consistent as `N → ∞`. Systematic resampling has the same expected counts, but its -# single shared uniform makes it order-dependent and it is not consistent in general (Gerber, -# Chopin & Whiteley, 2019), so it falls outside the particle Gibbs invariance proof. We -# therefore default to stratified resampling and offer systematic only as an explicit choice. +# For unconditional SMC, multinomial, stratified, and systematic resampling all have offspring +# counts satisfying `E[Oᵏ] = N·Wᵏ`. Multinomial resampling is broadly consistent, and +# stratified resampling is consistent under standard regularity conditions; systematic +# resampling is order-dependent and can fail to be consistent for arbitrary particle orderings +# (Gerber, Chopin & Whiteley, 2019), so stratified is the default unconditional scheme. +# +# Particle Gibbs additionally needs a valid *conditional* version of the chosen law, and +# pinning one ordered offspring is not it: for systematic resampling the conditional +# construction draws the grid offset from a weight-dependent mixture rather than `U[0,1]`, then +# randomly cycles the output so the reference lands in the pinned slot (Chopin & Singh, 2015, +# Algorithm 4, https://doi.org/10.3150/14-BEJ629; see also Finke, Johansen, Lee & Murray, +# "Resampling in conditional SMC algorithms", https://arxiv.org/abs/2606.25603). Independent +# multinomial draws stay valid once the reference ancestor is pinned, so `resample_propagate!` +# uses them for every scheme in a conditional sweep. That costs mixing -- Chopin & Singh find +# systematic resampling mixes noticeably better than multinomial in particle Gibbs -- so +# implementing the conditional schemes properly would be a genuine improvement. abstract type AbstractResampler end @@ -538,7 +547,14 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional n = length(particles) weights = normalized_weights(particles) if should_resample(resampler, weights) - ancestors = resample_indices(rng, resampler, weights, conditional ? n - 1 : n) + # A conditional sweep draws the `n-1` free ancestors independently from the categorical + # over the weights, whatever scheme `resampler` names -- see the resampling-schemes + # section for why the named scheme's conditional version is not simply "pin one draw". + ancestors = if conditional + resample_indices(rng, MultinomialResampler(), weights, n - 1) + else + resample_indices(rng, resampler, weights, n) + end old = copy(particles) seen = falses(n) for (slot, a) in enumerate(ancestors) @@ -710,8 +726,11 @@ end PG(n, [resampler = ESSResampler(0.5)]; multithreaded = false) PG(n, [scheme = StratifiedResampler(), ]threshold; multithreaded = false) -Particle Gibbs sampler with `n` particles. By default stratified resampling is triggered -whenever the effective sample size drops below half the number of particles. +Particle Gibbs sampler with `n` particles. By default resampling is triggered whenever the +effective sample size drops below half the number of particles. The selected scheme is used +for the unconditional first sweep; conditional sweeps draw their ancestors from the categorical +over the weights, because the conditional version of stratified or systematic resampling is a +different algorithm rather than the same draw with one output pinned. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index f60f83da89..82504a41a3 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -248,6 +248,42 @@ end @test serial[@varname(p)] == multithreaded[@varname(p)] end + @testset "conditional sweeps ignore the named resampling scheme" begin + # Conditional sweeps draw their ancestors multinomially whatever scheme is named, so + # sweeps differing only in that scheme must agree exactly. Bare schemes (no ESS gate) + # resample at every step, so the draw is exercised throughout. + @model function drifting(y) + x ~ Normal() + for t in eachindex(y) + y[t] ~ Normal(x, 1) + end + end + model = drifting([0.3, -0.7, 1.1]) + function conditional_sweep(scheme) + rng = StableRNG(77) + retained = Particle(model, particle_varinfo(), TracedRNG(rng)) + while advance!(retained, false) !== nothing + end + reference = Particle( + model, + particle_varinfo(), + rewind!(deepcopy(retained.rng)), + get_raw_values(retained.varinfo), + copy(retained.assumed_varnames), + ) + particles = map( + i -> + i < 5 ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference, + 1:5, + ) + sweep!(StableRNG(78), particles, scheme, false; conditional=true) + return map(p -> get_raw_values(p.varinfo), particles) + end + multinomial = conditional_sweep(MultinomialResampler()) + @test conditional_sweep(StratifiedResampler()) == multinomial + @test conditional_sweep(SystematicResampler()) == multinomial + end + # https://github.com/TuringLang/Turing.jl/issues/1598 @testset "reference particle" begin c = sample(gdemo_default, PG(1), 1_000) @@ -468,7 +504,7 @@ end return a ~ to_submodel(inner_kwarg(5.0; n=n)) end m2 = outer_kwarg2(10.0) - chn2 = sample(StableRNG(468), m2, PG(10), 1000) + chn2 = sample(StableRNG(468), m2, PG(10), 2000) @test mean(chn2[Symbol("a.x")]) ≈ 7.5 atol = 0.3 end From 8e74b4c09833dc0d3d17d9bf390f2f96094dbffd Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 00:42:25 +0100 Subject: [PATCH 38/58] Report which sampling keywords SMC actually honours The ensemble wrapper's injected `InitFromPrior` warned spuriously, and `callback` was dropped in silence. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 8 +++++++- test/mcmc/particle_mcmc.jl | 16 +++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index d35c224f48..efbef44a86 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -661,6 +661,7 @@ function AbstractMCMC.sample( discard_initial=0, thinning=1, initial_params=nothing, + callback=nothing, verbose=false, kwargs..., ) @@ -669,9 +670,14 @@ function AbstractMCMC.sample( if discard_initial > 0 || thinning > 1 @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." end - if initial_params !== nothing + if initial_params !== nothing && !(initial_params isa DynamicPPL.InitFromPrior) @warn "SMC draws its initial population from the prior; `initial_params` is ignored." end + # Accepted only so it can be reported as ignored: AbstractMCMC's contract is one callback + # per step, and SMC is a single sweep, so there is no iteration to call back from. + if callback !== nothing + @warn "SMC runs one sweep rather than an MCMC loop, so there are no per-iteration callbacks; `callback` is ignored." + end particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] logZ, ess_per_step = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 82504a41a3..7e5a3b97ba 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -159,7 +159,7 @@ using Turing @test_throws ArgumentError sample(model, SMC(), 100) end - @testset "discard_initial, thinning and initial_params are ignored" begin + @testset "discard_initial, thinning, initial_params and callback are ignored" begin @model function normal() a ~ Normal(4, 5) 3 ~ Normal(a, 2) @@ -171,6 +171,20 @@ using Turing @test_logs (:warn, r"initial_params.*ignored") match_mode = :any sample( normal(), SMC(), 10; initial_params=(; a=1.0) ) + @test_logs (:warn, r"initial_params.*ignored") sample( + normal(), SMC(), 10; initial_params=DynamicPPL.InitFromUniform() + ) + + # The ensemble wrapper injects the sampler's own default `InitFromPrior()` per chain. + # That is not a user-specified initialisation, so it must not warn. + @test_logs sample(Xoshiro(1), normal(), SMC(), MCMCSerial(), 10, 2; progress=false) + + # A callback is accepted only to be reported as ignored, and must not run. + called = false + @test_logs (:warn, r"callback.*ignored") sample( + normal(), SMC(), 10; callback=(args...; kwargs...) -> (called = true) + ) + @test !called @test_logs (:warn, r"ignored") sample(normal(), SMC(), 10; discard_initial=5) chn = sample(normal(), SMC(), 10; discard_initial=5) From ac77e43db4c7dafa005f1a28753378f4ad9778d0 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 01:18:15 +0100 Subject: [PATCH 39/58] Let ess_per_step follow LogProbType instead of pinning it to Float64 The line above already derives `logZ` from it, so the two could silently diverge. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index efbef44a86..4461c17255 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -586,7 +586,9 @@ function sweep!( rng::AbstractRNG, particles, resampler, multithreaded::Bool; conditional::Bool=false ) logZ = zero(DynamicPPL.LogProbType) - ess_per_step = Float64[] + # The ESS values are computed from the particle weights, so they follow whatever + # `DynamicPPL.LogProbType` is rather than being pinned to `Float64`. + ess_per_step = DynamicPPL.LogProbType[] while true resample_propagate!(rng, particles, resampler, conditional) logZ0 = log_normalizing_constant(particles) From 8665e5787f7fa5acd9595530b5af5845db855517 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 01:18:15 +0100 Subject: [PATCH 40/58] Build a CSMC reference from the retained particle, not from loose values Values and addresses were separately defaulted, so passing only values type-checked and did nothing. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 60 ++++++++++++++++++-------------------- test/mcmc/particle_mcmc.jl | 21 ++++--------- 2 files changed, 33 insertions(+), 48 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 4461c17255..c46c55d8c0 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -134,10 +134,19 @@ DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any """ Particle(model, varinfo, rng::TracedRNG) + Particle(model, varinfo, rng::TracedRNG, retained::Particle) A single particle: a suspended `model` execution together with its `varinfo`, its own replayable `rng`, and an accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there is no separate state struct). + +Without `retained` the particle draws from the prior. Given the previous sweep's retained +particle, it becomes a conditional-SMC reference pinned to that trajectory: it replays the +retained values, and errors if its execution reaches an address the retained trajectory does +not have, or finishes without reaching one that it does. Taking the whole particle -- rather +than its values and addresses separately -- is what makes a half-specified reference +unrepresentable; only the two pieces are kept, so the retained particle itself is not held +alive across sweeps. """ mutable struct Particle{RT<:TracedRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go @@ -165,31 +174,27 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # but assumed under the single address `x[1:2]`, so comparing against those keys would # report a spurious trace change; and an empty `reference_values` cannot distinguish a # reference with no latents from an ordinary particle. Without it, an address the retained - # trajectory never had would silently draw from the prior, corrupting the reference. + # trajectory never had would silently draw from the prior, corrupting the reference. It is + # set only alongside `reference_values`, by the reference constructor below. expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}} - # Addresses assumed by this execution, in the same form. Survives forking, so a particle - # that becomes the retained state hands the complete set to the next reference; a - # reference must finish with exactly the set above. + # Addresses assumed by this execution, in the same form as the set above. Survives forking, + # so a particle that becomes the retained state hands the complete set to the next + # reference; a reference must finish having assumed exactly that set. assumed_varnames::Set{DynamicPPL.VarName} task::Libtask.TapedTask # `task` is filled in once the particle exists, because the task must capture the # particle as its taped globals (a back-reference). This has to be an inner constructor # for that reason: `task` is left undefined here and set immediately after. function Particle( - vi::DynamicPPL.AbstractVarInfo, - rng::RT, - reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), - expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}}=nothing, + vi::DynamicPPL.AbstractVarInfo, rng::RT, retained::Union{Nothing,Particle}=nothing ) where {RT<:TracedRNG} w = zero(DynamicPPL.LogProbType) - return new{RT,typeof(w)}( - vi, - rng, - w, - reference_values, - expected_reference_varnames, - Set{DynamicPPL.VarName}(), - ) + values, varnames = if retained === nothing + DynamicPPL.VarNamedTuple(), nothing + else + DynamicPPL.get_raw_values(retained.varinfo), copy(retained.assumed_varnames) + end + return new{RT,typeof(w)}(vi, rng, w, values, varnames, Set{DynamicPPL.VarName}()) end end @@ -197,14 +202,11 @@ function Particle( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo, rng::TracedRNG, - reference_values::DynamicPPL.VarNamedTuple=DynamicPPL.VarNamedTuple(), - expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}}=nothing, + retained::Union{Nothing,Particle}=nothing, ) model = DynamicPPL.setleafcontext(model, SMCContext()) args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - particle = Particle( - deepcopy(varinfo), rng, reference_values, expected_reference_varnames - ) + particle = Particle(deepcopy(varinfo), rng, retained) particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) return particle end @@ -265,10 +267,10 @@ function DynamicPPL.tilde_assume!!( particle = Libtask.get_taped_globals(Particle) # A CSMC reference reuses the retained value (`InitFromParams`) at every address it visits. # `expected_reference_varnames` is `nothing` for ordinary particles and is cleared by - # `reseed!` on a fork, so both draw from the prior (see the `Particle` fields). An address - # outside the retained set means the execution trace changed, which must error rather than - # silently drawing part of the nominally fixed reference afresh; the `nothing` fallback - # given to `InitFromParams` catches the converse, a retained address with no usable value. + # `reseed!` on a fork, so both draw from the prior. An address outside the retained set means + # the execution trace changed, which must error rather than silently drawing part of the + # nominally fixed reference afresh; the `nothing` fallback given to `InitFromParams` catches + # the converse, a retained address with no usable value. expected = particle.expected_reference_varnames strategy = if expected === nothing DynamicPPL.InitFromPrior() @@ -783,13 +785,7 @@ function AbstractMCMC.step( # The reference reproduces the retained trajectory by reusing its values (passed here and # consumed by `tilde_assume!!`), so it stays that trajectory even if the model was # re-conditioned since the last sweep. Its varinfo starts empty like any other particle. - reference = Particle( - model, - particle_varinfo(), - rewind!(deepcopy(state.rng)), - DynamicPPL.get_raw_values(state.varinfo), - copy(state.assumed_varnames), - ) + reference = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng)), state) particles = map(1:n) do i i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 7e5a3b97ba..6f7654111d 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -279,11 +279,7 @@ end while advance!(retained, false) !== nothing end reference = Particle( - model, - particle_varinfo(), - rewind!(deepcopy(retained.rng)), - get_raw_values(retained.varinfo), - copy(retained.assumed_varnames), + model, particle_varinfo(), rewind!(deepcopy(retained.rng)), retained ) particles = map( i -> @@ -331,11 +327,7 @@ end allok = true for _ in 1:nsteps ref = Particle( - model, - particle_varinfo(), - rewind!(deepcopy(state.rng)), - get_raw_values(state.varinfo), - copy(state.assumed_varnames), + model, particle_varinfo(), rewind!(deepcopy(state.rng)), state ) parts = map( i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, @@ -377,8 +369,7 @@ end reconditioned(2.0) | (@varname(a) => 5.0), # x's prior shifted far away particle_varinfo(), rewind!(deepcopy(retained.rng)), - retained_vals, - copy(retained.assumed_varnames), + retained, ) while advance!(reference, true) !== nothing end @@ -404,8 +395,7 @@ end branch_changes(false, 0.0), particle_varinfo(), rewind!(deepcopy(retained.rng)), - get_raw_values(retained.varinfo), - copy(retained.assumed_varnames), + retained, ) @test_throws "reference execution trace changed" advance!(reference, true) @@ -423,8 +413,7 @@ end branch_drops(false, 0.0), particle_varinfo(), rewind!(deepcopy(retained.rng)), - get_raw_values(retained.varinfo), - copy(retained.assumed_varnames), + retained, ) @test_throws "reference execution trace changed" begin while advance!(reference, true) !== nothing From bee0e0c78a5d1708a0b1a114b826b5d1235eaea2 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 01:18:43 +0100 Subject: [PATCH 41/58] Record the measured CSMC+ESS margin, which is thinner than the comment claimed Multinomial conditional sweeps took |err| on `s` from 0.024 to 0.061 against `atol = 0.1`. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- test/mcmc/ess.jl | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test/mcmc/ess.jl b/test/mcmc/ess.jl index 44dd75fe79..a5a62a224f 100644 --- a/test/mcmc/ess.jl +++ b/test/mcmc/ess.jl @@ -60,9 +60,13 @@ using Turing @testset "gdemo with CSMC + ESS" begin alg = Gibbs(:s => CSMC(15), :m => ESS()) - # CSMC mixes the variance `s` slowly, so the Monte Carlo error at 3_000 draws sits - # close to `atol` and an unlucky seed tips it over; the estimator is unbiased, so - # 10_000 draws simply give comfortable headroom. + # CSMC mixes the variance `s` slowly, so the Monte Carlo error at 3_000 draws + # exceeds `atol` on this seed (measured |err| 0.212); the estimator is unbiased, so + # 10_000 draws buy the headroom back. That headroom is not large: conditional + # sweeps draw their ancestors multinomially, which is noisier than the stratified + # scheme used before, and `s` moved from |err| 0.024 to 0.061 against `atol = 0.1` + # as a result. If this test starts failing, suspect the draw count before the + # sampler. chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 10_000) check_gdemo(chain; atol=0.1) end From e8ef7a0feb8160d0b0bde1e10ac1055e66a1af6b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 01:45:36 +0100 Subject: [PATCH 42/58] Stop AGENTS.md claiming the particle samplers wrap AdvancedPS This branch removed AdvancedPS and also creates AGENTS.md, so it contradicted itself on landing. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- AGENTS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 4922775aff..bcb41efba7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,7 +38,7 @@ CI matrix: Julia stable + min, Ubuntu/Windows/macOS, 1 and 2 threads. Most complexity is in DynamicPPL. Turing.jl contains: - - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (wrapping AdvancedPS), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. + - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (implemented natively in `particle_mcmc.jl` on top of Libtask coroutines and Random123 — not a wrapper), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. - **External sampler interface** (`src/mcmc/external_sampler.jl`): The `externalsampler()` wrapper lets any `AbstractMCMC.AbstractSampler` that implements `step` for `LogDensityModel` work with Turing models. This is the easier path for new samplers — it only requires a dependency on AbstractMCMC and the LogDensityProblems.jl interface, with no Turing internals. The tradeoff is less power: you can only interact with the model as a black-box log-density function, just like using `LogDensityFunction` directly. - **Variational inference** (`src/variational/`): Wraps AdvancedVI algorithms. - **Mode estimation** (`src/optimisation/`): MAP and MLE via Optimization.jl. From f7af068d782de57763fc52ed1783778ae712e3ab Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 01:45:46 +0100 Subject: [PATCH 43/58] Delegate the accloglikelihood!! overload instead of reimplementing it It shadows DynamicPPL's method for every `OnlyAccsVarInfo`; duplicating the body would let it drift from upstream unnoticed. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 35 +++++++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index c46c55d8c0..623c9b3bd1 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -343,21 +343,36 @@ end # `@addlogprob!` bypasses `tilde_observe!!`, so its produce is emitted here instead -- again # only once the accumulator has been updated. Gated on the producing accumulator, so outside # particle evaluation this reduces to the default (non-producing) method (issue #1996). +# +# This is type piracy -- both `accloglikelihood!!` and `OnlyAccsVarInfo` belong to DynamicPPL, +# and this method shadows DynamicPPL's own for every `OnlyAccsVarInfo`, which is what all +# samplers use. It is deliberate for want of an alternative: the produce has to happen once the +# accumulator has been updated, and no Turing-owned type appears anywhere in the signature. A +# DynamicPPL-side post-accumulate hook would let this go away. +# +# Everything except the produce is delegated to the general method via `@invoke`, so that this +# cannot silently drift from upstream: re-implementing that body here would leave every +# `OnlyAccsVarInfo` in the ecosystem running Turing's copy, while only particle sampling is +# covered by these tests. function DynamicPPL.accloglikelihood!!( vi::DynamicPPL.OnlyAccsVarInfo, logp; ignore_missing_accumulator=false ) acc_name = Val(:LogLikelihood) - if ignore_missing_accumulator && !DynamicPPL.hasacc(vi, acc_name) - return vi - end - is_particle = DynamicPPL.getacc(vi, acc_name) isa ProduceLogLikelihoodAccumulator - before = is_particle ? DynamicPPL.getloglikelihood(vi) : zero(DynamicPPL.LogProbType) - vi = DynamicPPL.map_accumulator!!(acc -> DynamicPPL.acclogp(acc, logp), vi, acc_name) - if is_particle - particle = Libtask.get_taped_globals(Particle) - particle.varinfo = vi - Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) + is_particle = + DynamicPPL.hasacc(vi, acc_name) && + DynamicPPL.getacc(vi, acc_name) isa ProduceLogLikelihoodAccumulator + if !is_particle + return @invoke DynamicPPL.accloglikelihood!!( + vi::DynamicPPL.AbstractVarInfo, logp; ignore_missing_accumulator + ) end + before = DynamicPPL.getloglikelihood(vi) + vi = @invoke DynamicPPL.accloglikelihood!!( + vi::DynamicPPL.AbstractVarInfo, logp; ignore_missing_accumulator + ) + particle = Libtask.get_taped_globals(Particle) + particle.varinfo = vi + Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) return vi end From 542b85c53b8f70803884509724023abe0ff1a3d8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 02:10:23 +0100 Subject: [PATCH 44/58] Assert MoGtest posterior means, not idealised cluster labels `[1,1,2,2,1,4]` are the labels, not the means; the 0.072 offset spent half the tolerance and made gibbs.jl flaky. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- test/test_utils/numerical_tests.jl | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/test/test_utils/numerical_tests.jl b/test/test_utils/numerical_tests.jl index 1322c5362f..a020bf6bca 100644 --- a/test/test_utils/numerical_tests.jl +++ b/test/test_utils/numerical_tests.jl @@ -24,6 +24,26 @@ function check_gdemo(chain; atol=0.2, rtol=0.0) ) end +# Exact posterior means of `MoGtest_default`, i.e. `MoGtest([1.0 1.0 4.0 4.0])`, for +# `[z1, z2, z3, z4, mu1, mu2]`. The four `z` are discrete, so enumerating all 2^4 cluster +# assignments and integrating `mu1`, `mu2` out analytically (given `z`, each cluster's +# observations are Gaussian with a Gaussian prior, so the marginal likelihood and the +# conditional posterior mean are both closed-form) gives these exactly. +# +# These are deliberately *not* the idealised labels `[1, 1, 2, 2, 1, 4]`. The data do not +# identify the clusters perfectly, so the truth sits 0.046 off each label and 0.072 off each +# prior mean. Asserting the idealised values spent that 0.072 of the tolerance before any Monte +# Carlo error -- half of a 0.15 `atol`, and 72% of the 0.1 one used in mcmc/ess.jl -- which made +# accurate samplers fail intermittently and made the rest look worse than they are. +const MOGTEST_DEFAULT_MEANS = [ + 1.0455583636758707, + 1.0455583636758707, + 1.9544416363241297, + 1.9544416363241297, + 1.071541120128846, + 3.928458879871154, +] + # Wrapper function to check MoGtest. function check_MoGtest_default(chain; atol=0.2, rtol=0.0) return check_numerical( @@ -36,7 +56,7 @@ function check_MoGtest_default(chain; atol=0.2, rtol=0.0) @varname(mu1), @varname(mu2) ], - [1.0, 1.0, 2.0, 2.0, 1.0, 4.0]; + MOGTEST_DEFAULT_MEANS; atol=atol, rtol=rtol, ) @@ -53,7 +73,7 @@ function check_MoGtest_default_z_vector(chain; atol=0.2, rtol=0.0) @varname(mu1), @varname(mu2) ], - [1.0, 1.0, 2.0, 2.0, 1.0, 4.0]; + MOGTEST_DEFAULT_MEANS; atol=atol, rtol=rtol, ) From 861d1ce8a1dd03032f542a1b5fc8fcd6382f94cc Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 02:15:58 +0100 Subject: [PATCH 45/58] Mention the reference-replay branch in the SMCContext docstring Drawing from the prior is only half of what `tilde_assume!!` does under it. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 623c9b3bd1..c53694738e 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -122,8 +122,9 @@ end """ SMCContext -Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws -from the prior using the particle's [`TracedRNG`](@ref), and `tilde_observe!!` scores the +Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws from +the prior using the particle's [`TracedRNG`](@ref) -- or, for a conditional-SMC reference, +replays the retained trajectory's value at that address -- and `tilde_observe!!` scores the observation and `Libtask.produce`s the increment as the particle's weight. """ struct SMCContext <: DynamicPPL.AbstractContext end From 7b12f5f08f9e7afa8bc6e1ce5855ff3a563bcaa8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 10:40:10 +0100 Subject: [PATCH 46/58] Skip Aqua's persistent_tasks check on Windows, where it is flaky The same tree both passed and failed in CI; a lingering Task is OS-independent. Co-Authored-By: Charles Knipp <32943413+charlesknipp@users.noreply.github.com> Co-Authored-By: Claude Code --- test/Aqua.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/Aqua.jl b/test/Aqua.jl index e5b655c6e0..e7a5ab6d1b 100644 --- a/test/Aqua.jl +++ b/test/Aqua.jl @@ -24,6 +24,8 @@ using Turing # meaningful in practice (in particular, to trigger this we would need to call `g(..., f)`, # which is incredibly unlikely). Aqua.test_ambiguities([Turing]; exclude=[Libtask.might_produce]) -Aqua.test_all(Turing; ambiguities=false) + +# `persistent_tasks` is flaky on Windows; a lingering Task is a code property, not an OS one. +Aqua.test_all(Turing; ambiguities=false, persistent_tasks=!Sys.iswindows()) end From 6c8247d295691d8141dafa27acfafc6fa5e85d06 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 18:09:45 +0100 Subject: [PATCH 47/58] Drop the reference's seed-replay machinery, which value-pinning made dead Pinning the CSMC reference to the retained trajectory by value made the seed-recording half of `TracedRNG` unreachable: every reference draw is now supplied by `InitFromParams`, so the reference never consults its own generator. Two checks confirm it. Scrambling the reference's seed before every step leaves the trajectory it regenerates untouched; and scrambling it throughout the particle suite breaks exactly one test, `"rng replay"`, which builds a reference with no retained particle -- a shape no sampler constructs. So `count`, `keys`, `save_state!`, `load_state!`, `inc_step!`, `rewind!` and `inner_key` go. That leaves `TracedRNG` with no state of its own, and since Random123's `seed!` already zeroes the counter the wrapper has nothing left to do either: particles now hold a plain `Random123.Philox2x` from `particle_rng`. Three `Random` forwarding methods go with it, including the one that needed an ambiguity fix. Being the reference then stops being a fact about a particle's slot and becomes one about the particle -- `expected_reference_varnames !== nothing`, exposed as `isreference`. `advance!` loses its `isref` argument, which lets `reweight!` lose `conditional`, which drops the index-versus-`n` arithmetic from the reweighting path entirely. The reference carries the retained generator forward instead of taking a fresh one, so the sweep draws nothing from the sampler rng on its behalf. Draws are therefore bit-identical to before across SMC, PG under two resampling schemes, and Gibbs(CSMC + HMC). Comments and section headers get a pass here too, since most of the long ones described the machinery being removed. The file banner becomes a level-1 header like every other section, the two longest sections gain subsection rules, and the hand-maintained "Sections below" index is gone -- the headers make it redundant, and it had already gone stale naming "traced RNG". Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 253 ++++++++++++++++--------------------- test/mcmc/particle_mcmc.jl | 137 ++++++++++---------- 2 files changed, 181 insertions(+), 209 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index c53694738e..0e444fb2d6 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -1,91 +1,39 @@ -### -### Particle filtering and particle MCMC samplers: SMC, PG / conditional SMC. -### -### Key design. -### A probabilistic model becomes a particle filter by reading each `observe` statement as one -### filtering step. Evaluated under `SMCContext`, every likelihood term calls -### `Libtask.produce`, so a *particle* is a suspended model execution: we `advance!` it to its -### next `observe`, take the produced log-likelihood as its weight, then resample. SMC is one -### such sweep; particle Gibbs (PG/CSMC) runs a *conditional* sweep -- one particle is a fixed -### reference trajectory -- inside an MCMC loop. -### -### The reference reproduces the retained trajectory by *reusing its values*: the sampler state -### carries those values, and the reference re-runs the model with `InitFromParams`, so it stays -### the retained trajectory even when the model is re-conditioned between Gibbs sweeps (e.g. a -### state-space transition prior that depends on a parameter owned by another Gibbs component). -### A particle forked from the reference forgets the remaining values (in `reseed!`), so -### branching samples fresh with no per-particle flag. Each particle's counter-based `TracedRNG` -### supplies splittable, version-stable seeds that decorrelate the fresh draws across particles -### and Julia versions. -### -### Sections below: traced RNG; model evaluation via Libtask; resampling schemes; the particle -### sweep; the SMC sampler; the PG/CSMC sampler; the Gibbs-component interface. -### -### Reference: Andrieu, Doucet & Holenstein, "Particle Markov chain Monte Carlo methods", -### Journal of the Royal Statistical Society: Series B 72(3), 269-342 (2010). -### +# +# Particle filtering and particle MCMC samplers: SMC, PG / conditional SMC +# + +# A probabilistic model becomes a particle filter by reading each `observe` statement as one +# filtering step. Evaluated under `SMCContext`, every likelihood term calls `Libtask.produce`, so +# a *particle* is a suspended model execution: we `advance!` it to its next `observe`, take the +# produced log-likelihood as its weight, then resample. SMC is one such sweep; particle Gibbs +# (PG/CSMC) runs a *conditional* sweep -- one particle is a fixed reference trajectory -- inside +# an MCMC loop. +# +# The reference reproduces the retained trajectory by *reusing its values*: the sampler state +# carries those values, and the reference re-runs the model with `InitFromParams`, so it stays the +# retained trajectory even when the model is re-conditioned between Gibbs sweeps (e.g. a +# state-space transition prior that depends on a parameter owned by another Gibbs component). A +# particle forked from the reference forgets the remaining values (in `reseed!`), so branching +# samples fresh with no per-particle flag. +# +# Reference: Andrieu, Doucet & Holenstein, "Particle Markov chain Monte Carlo methods", Journal of +# the Royal Statistical Society: Series B 72(3), 269-342 (2010). using StatsFuns: softmax, logsumexp import Random123 # -# Traced RNG +# Particle random number generation # -# A counter-based RNG that records the seed used at each model step, so that a particle's -# trajectory can be replayed exactly: the conditional-SMC reference regenerates itself by -# replaying its recorded seeds. This section comes first because `Particle` names `TracedRNG` -# in its type signature. - -""" - TracedRNG([rng = Random.default_rng()]) - -A `Random123.Philox2x` generator that remembers the seed (`key`) it used at each model step -in `keys`, indexed by the step counter `count`. - - - [`save_state!`](@ref) records the current seed (ordinary particles); - - [`load_state!`](@ref) restores `keys[count]`, replaying that step's randomness (the - reference trajectory). -""" -mutable struct TracedRNG{K<:Unsigned,T<:Random123.AbstractR123} <: Random.AbstractRNG - count::Int - rng::T - keys::Vector{K} -end - -function TracedRNG(inner::Random123.AbstractR123{T}) where {T<:Unsigned} - Random123.set_counter!(inner, 0) - return TracedRNG(1, inner, T[]) -end -function TracedRNG(rng::AbstractRNG=Random.default_rng()) - inner = Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) - return TracedRNG(inner) -end -Random.rng_native_52(trng::TracedRNG) = Random.rng_native_52(trng.rng) -Random.rand(trng::TracedRNG, ::Type{T}) where {T<:Unsigned} = Random.rand(trng.rng, T) +# Each particle owns a counter-based `Random123.Philox2x`. This section comes first because +# `Particle` names the generator type in its signature. -"The current seed of the inner generator." -inner_key(rng::Random123.Philox2x) = rng.key - -"Reseed and rewind the inner generator. The model-step counter is left untouched." -function Random.seed!(trng::TracedRNG, key::Integer) - Random.seed!(trng.rng, key) - Random123.set_counter!(trng.rng, 0) - return trng +"A fresh counter-based generator for one particle, seeded from `rng`." +function particle_rng(rng::AbstractRNG=Random.default_rng()) + return Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) end -"Record the seed used at the current step." -save_state!(trng::TracedRNG) = push!(trng.keys, inner_key(trng.rng)) - -"Replay the seed recorded at the current step." -load_state!(trng::TracedRNG) = Random.seed!(trng, trng.keys[trng.count]) - -"Advance the model-step counter by one." -inc_step!(trng::TracedRNG) = (trng.count += 1; trng) - -"Rewind the model-step counter to the first step, so a trajectory replays from the start." -rewind!(trng::TracedRNG) = (trng.count = 1; trng) - # Derive a fresh seed from `key`. Splitting one generator into many by re-seeding is fragile # in two ways: the derived seeds can yield *correlated* streams (Steele et al., "Fast # Splittable Pseudorandom Number Generators", OOPSLA 2014), and a stdlib `MersenneTwister` @@ -97,11 +45,12 @@ rewind!(trng::TracedRNG) = (trng.count = 1; trng) split_key(key::Integer) = rand(Random.seed!(Random123.Philox2x(), key), typeof(key)) "Reseed from the generator's own current state (used between steps when not resampling)." -refresh!(trng::TracedRNG) = Random.seed!(trng, split_key(inner_key(trng.rng))) +refresh!(rng::Random123.Philox2x) = Random.seed!(rng, split_key(rng.key)) # # Model evaluation via Libtask # + # A `Particle` is the only mutable state. It is stored as its `TapedTask`'s "taped globals", # so the tilde overloads reach it from *inside* a running model via `get_taped_globals`. # This keeps all state explicit on the particle -- no `task_local_storage`. @@ -123,9 +72,9 @@ end SMCContext Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws from -the prior using the particle's [`TracedRNG`](@ref) -- or, for a conditional-SMC reference, -replays the retained trajectory's value at that address -- and `tilde_observe!!` scores the -observation and `Libtask.produce`s the increment as the particle's weight. +the prior using the particle's own generator -- or, for a conditional-SMC reference, reuses the +retained trajectory's value at that address -- and `tilde_observe!!` scores the observation and +`Libtask.produce`s the increment as the particle's weight. """ struct SMCContext <: DynamicPPL.AbstractContext end @@ -134,22 +83,22 @@ struct SMCContext <: DynamicPPL.AbstractContext end DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any """ - Particle(model, varinfo, rng::TracedRNG) - Particle(model, varinfo, rng::TracedRNG, retained::Particle) + Particle(model, varinfo, rng) + Particle(model, varinfo, rng, retained::Particle) -A single particle: a suspended `model` execution together with its `varinfo`, its own -replayable `rng`, and an accumulated `logweight`. It also serves directly as the particle -Gibbs sampler state (there is no separate state struct). +A single particle: a suspended `model` execution together with its `varinfo`, its own `rng`, and +an accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there +is no separate state struct). Without `retained` the particle draws from the prior. Given the previous sweep's retained -particle, it becomes a conditional-SMC reference pinned to that trajectory: it replays the +particle, it becomes a conditional-SMC reference pinned to that trajectory: it reuses the retained values, and errors if its execution reaches an address the retained trajectory does not have, or finishes without reaching one that it does. Taking the whole particle -- rather than its values and addresses separately -- is what makes a half-specified reference unrepresentable; only the two pieces are kept, so the retained particle itself is not held alive across sweeps. """ -mutable struct Particle{RT<:TracedRNG,WT<:Real} +mutable struct Particle{RT<:AbstractRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go # through Libtask's (already type-unstable) taped globals, so this costs nothing extra. varinfo::DynamicPPL.AbstractVarInfo @@ -158,25 +107,22 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # is ever changed. logweight::WT # The retained trajectory's values, which the CSMC reference reproduces by reusing them - # (`InitFromParams` in `tilde_assume!!`); empty for ordinary particles. `reseed!` clears it - # so a particle forked off the reference samples fresh beyond the fork point. Reproducing - # by *value* (not by replayed RNG seeds) is what keeps the reference the retained trajectory - # when the model is re-conditioned between Gibbs sweeps. A draw is a deterministic function - # x = g(u; θ) of the RNG output `u` and the distribution parameters θ (canonically the - # inverse-CDF, x = F⁻¹(u; θ)). Seed-replay fixes `u` and recomputes x' = g(u; θ'); value-replay - # reuses x' = x. These agree only when θ' = θ. Re-conditioning updates a value θ depends on - # (owned by another block), so θ' ≠ θ and the replayed draw moves with the changed distribution - # -- e.g. x ~ Normal(μ, 1) draws x = μ + Φ⁻¹(u), so after μ → μ' the same `u` gives x + (μ' − μ), - # not x. + # (`InitFromParams` in `tilde_assume!!`); empty for ordinary particles. `reseed!` clears it so + # a particle forked off the reference samples fresh beyond the fork point. Reusing the *value* + # rather than replaying the RNG draw is what keeps the reference on the retained trajectory + # when Gibbs re-conditions the model: a draw is x = g(u; θ) in the RNG output u and the + # distribution parameters θ, so replaying u after θ → θ' yields g(u; θ') ≠ x -- e.g. + # x ~ Normal(μ, 1) is μ + Φ⁻¹(u), which shifts by μ' − μ. reference_values::DynamicPPL.VarNamedTuple # `nothing` for an ordinary particle; for a CSMC reference, the addresses the retained - # trajectory assumed. Two reasons this cannot be read off `reference_values`: a slice - # assume such as `x[1:2] ~ MvNormal(...)` is stored there under the keys `x[1]`, `x[2]` - # but assumed under the single address `x[1:2]`, so comparing against those keys would - # report a spurious trace change; and an empty `reference_values` cannot distinguish a - # reference with no latents from an ordinary particle. Without it, an address the retained - # trajectory never had would silently draw from the prior, corrupting the reference. It is - # set only alongside `reference_values`, by the reference constructor below. + # trajectory assumed. This cannot be read off `reference_values`: a slice assume such as + # `x[1:2] ~ MvNormal(...)` is stored there under the keys `x[1]`, `x[2]` but assumed under the + # single address `x[1:2]`, so comparing against those keys would report a spurious trace + # change; and an empty `reference_values` cannot distinguish a reference with no latents from + # an ordinary particle. Without it, an address the retained trajectory never had would + # silently draw from the prior, corrupting the reference. Being non-`nothing` is also what + # marks a particle as the reference. Set only alongside `reference_values`, by the reference + # constructor below. expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}} # Addresses assumed by this execution, in the same form as the set above. Survives forking, # so a particle that becomes the retained state hands the complete set to the next @@ -188,7 +134,7 @@ mutable struct Particle{RT<:TracedRNG,WT<:Real} # for that reason: `task` is left undefined here and set immediately after. function Particle( vi::DynamicPPL.AbstractVarInfo, rng::RT, retained::Union{Nothing,Particle}=nothing - ) where {RT<:TracedRNG} + ) where {RT<:AbstractRNG} w = zero(DynamicPPL.LogProbType) values, varnames = if retained === nothing DynamicPPL.VarNamedTuple(), nothing @@ -202,7 +148,7 @@ end function Particle( model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo, - rng::TracedRNG, + rng::AbstractRNG, retained::Union{Nothing,Particle}=nothing, ) model = DynamicPPL.setleafcontext(model, SMCContext()) @@ -215,13 +161,11 @@ end """ reseed!(particle, rng) -Restart `particle` as a fresh continuation seeded from `rng`: it switches from replaying to -sampling afresh, and `keys` is truncated to the steps already taken so a particle descended -from the reference forgets the reference's future. Mutates and returns `particle`. +Restart `particle` as a fresh continuation seeded from `rng`, so that a particle descended from +the reference stops reusing retained values and samples afresh. Mutates and returns `particle`. """ function reseed!(particle::Particle, rng::AbstractRNG) Random.seed!(particle.rng, rand(rng, UInt64)) - resize!(particle.rng.keys, particle.rng.count - 1) # A fork samples fresh from here on, so it must forget the reference's remaining values. particle.reference_values = DynamicPPL.VarNamedTuple() particle.expected_reference_varnames = nothing @@ -238,21 +182,24 @@ back-reference; [`reseed!`](@ref) then gives it its own random stream. fork(particle::Particle, rng::AbstractRNG) = reseed!(deepcopy(particle), rng) """ - advance!(particle, isref) -> Union{Real,Nothing} +Whether `particle` is a conditional-SMC reference, i.e. pinned to a retained trajectory. Carried +on the particle rather than inferred from its slot, so forking and resampling cannot get it wrong. +""" +isreference(particle::Particle) = particle.expected_reference_varnames !== nothing + +""" + advance!(particle) -> Union{Real,Nothing} Run the particle to its next `observe`, returning the incremental log-likelihood, or -`nothing` once the model finishes. An ordinary particle records the step's seed; the -reference (`isref = true`) replays its recorded seed instead. +`nothing` once the model finishes. """ -function advance!(particle::Particle, isref::Bool) - isref ? load_state!(particle.rng) : save_state!(particle.rng) - inc_step!(particle.rng) +function advance!(particle::Particle) score = Libtask.consume(particle.task) # `tilde_assume!!` already rejects any address outside the retained set, so once the # reference has run to completion the only discrepancy left to catch is a retained address # it never visited (e.g. a branch that stopped being taken after re-conditioning). expected = particle.expected_reference_varnames - if isref && score === nothing && expected !== nothing + if score === nothing && expected !== nothing dropped = setdiff(expected, particle.assumed_varnames) isempty(dropped) || error( "the reference execution trace changed while replaying retained values " * @@ -410,6 +357,7 @@ end # # Resampling schemes # + # For unconditional SMC, multinomial, stratified, and systematic resampling all have offspring # counts satisfying `E[Oᵏ] = N·Wᵏ`. Multinomial resampling is broadly consistent, and # stratified resampling is consistent under standard regularity conditions; systematic @@ -427,6 +375,8 @@ end # systematic resampling mixes noticeably better than multinomial in particle Gibbs -- so # implementing the conditional schemes properly would be a genuine improvement. +# ── Resampler interface ─────────────────────────────────────────────────────── + abstract type AbstractResampler end """Whether to resample given the normalized `weights`. Bare schemes always resample.""" @@ -435,6 +385,8 @@ should_resample(::AbstractResampler, weights) = true """Draw `n` ancestor indices from `1:length(weights)` with probabilities `weights`.""" function resample_indices end +# ── Schemes ─────────────────────────────────────────────────────────────────── + "Multinomial resampling: `n` independent draws from the categorical over `weights`." struct MultinomialResampler <: AbstractResampler end function resample_indices(rng::AbstractRNG, ::MultinomialResampler, weights, n::Integer) @@ -481,6 +433,8 @@ function resample_indices(rng::AbstractRNG, ::SystematicResampler, weights, n::I return indices end +# ── Effective-sample-size gating ────────────────────────────────────────────── + """ ESSResampler(threshold, scheme = StratifiedResampler()) @@ -503,9 +457,12 @@ end # # Particle sweep # -# In a conditional sweep the last particle is the reference: it is always retained and -# replays its recorded randomness, while the other `n-1` slots are resampled from all `n` -# particles (so they may descend from the reference). + +# In a conditional sweep the last particle is the reference: it is always retained and reuses the +# retained trajectory's values, while the other `n-1` slots are resampled from all `n` particles +# (so they may descend from the reference). + +# ── Weights and diagnostics ─────────────────────────────────────────────────── logweights(particles) = [p.logweight for p in particles] normalized_weights(particles) = softmax(logweights(particles)) @@ -513,11 +470,13 @@ log_normalizing_constant(particles) = logsumexp(logweights(particles)) "Effective sample size of a normalised weight vector, `1 / Σ wᵢ²`." ess(weights) = inv(sum(abs2, weights)) +# ── Reweighting ─────────────────────────────────────────────────────────────── + # Advance one particle by one observation, folding its incremental weight in; return `true` # once it has finished (produced nothing). Factored out so the serial and multithreaded loops in # `reweight!` share one body. -function advance_particle!(p::Particle, isref::Bool) - score = advance!(p, isref) +function advance_particle!(p::Particle) + score = advance!(p) score === nothing && return true p.logweight += score return false @@ -537,19 +496,19 @@ end # seeded serially in `resample_propagate!`, so the multithreaded loop is race-free and gives # results identical to the serial one. Only the model evaluations parallelise; the shared # sampler rng is untouched here. -function reweight!(particles, conditional::Bool, multithreaded::Bool) +function reweight!(particles, multithreaded::Bool) n = length(particles) if multithreaded # A shared counter would race, so collect per-particle results and tally afterwards. finished = Vector{Bool}(undef, n) Threads.@threads for i in 1:n - finished[i] = advance_particle!(particles[i], conditional && i == n) + finished[i] = advance_particle!(particles[i]) end n_done = count(finished) else n_done = 0 - for i in 1:n - n_done += advance_particle!(particles[i], conditional && i == n) + for p in particles + n_done += advance_particle!(p) end end n_done == 0 && return false @@ -559,6 +518,8 @@ function reweight!(particles, conditional::Bool, multithreaded::Bool) ) end +# ── Resample and propagate ──────────────────────────────────────────────────── + # Resample (if the scheme calls for it) and propagate the survivors, or -- when not # resampling -- refresh each ordinary particle's seed so the next step draws fresh randomness. function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional::Bool) @@ -579,7 +540,7 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional # Reuse each surviving parent's object for its first offspring; only extra # offspring -- and any offspring of the retained reference -- need the costly # `deepcopy`. Either way the child is reseeded to continue independently. - reuse = !seen[a] && !(conditional && a == n) + reuse = !seen[a] && !isreference(old[a]) seen[a] = true child = reuse ? reseed!(old[a], rng) : fork(old[a], rng) child.logweight = zero(DynamicPPL.LogProbType) @@ -588,16 +549,17 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional # reference retained, weight reset conditional && (particles[n].logweight = zero(DynamicPPL.LogProbType)) else - for (i, p) in enumerate(particles) - # Refresh every particle's seed except the reference, which keeps replaying. - if !(conditional && i == n) - refresh!(p.rng) - end + # The reference draws nothing (it reuses retained values), so only the others need a + # fresh seed for the next step. + for p in particles + isreference(p) || refresh!(p.rng) end end return nothing end +# ── One sweep ───────────────────────────────────────────────────────────────── + # Run a full particle sweep in place, returning the log-evidence estimate and the # per-observation effective sample sizes. function sweep!( @@ -610,7 +572,7 @@ function sweep!( while true resample_propagate!(rng, particles, resampler, conditional) logZ0 = log_normalizing_constant(particles) - done = reweight!(particles, conditional, multithreaded) + done = reweight!(particles, multithreaded) # Each observation contributes the log-ratio of total weight it adds; summed over the # sweep these telescope into an estimate of the model's log-evidence log p(y). logZ += log_normalizing_constant(particles) - logZ0 @@ -698,7 +660,9 @@ function AbstractMCMC.sample( if callback !== nothing @warn "SMC runs one sweep rather than an MCMC loop, so there are no per-iteration callbacks; `callback` is ignored." end - particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:nparticles] + particles = [ + Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:nparticles + ] logZ, ess_per_step = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) weights = normalized_weights(particles) # One final resampling step, so the returned particles are an equal-weight sample. The @@ -780,7 +744,8 @@ function AbstractMCMC.step( ) error_if_threadsafe_eval(model) particles = [ - Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:(sampler.nparticles) + Particle(model, particle_varinfo(), particle_rng(rng)) for + _ in 1:(sampler.nparticles) ] logZ, _ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) @@ -801,10 +766,16 @@ function AbstractMCMC.step( # The reference reproduces the retained trajectory by reusing its values (passed here and # consumed by `tilde_assume!!`), so it stays that trajectory even if the model was # re-conditioned since the last sweep. Its varinfo starts empty like any other particle. - reference = Particle(model, particle_varinfo(), rewind!(deepcopy(state.rng)), state) - particles = map(1:n) do i - i < n ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference - end + # + # Every reference draw is supplied by value, so its own generator is never read -- only its + # forks' are, and `reseed!` gives those fresh seeds. It therefore carries the retained + # generator forward rather than taking a fresh one: that way the sweep draws nothing from + # `rng` on the reference's behalf, keeping the sampler's stream independent of how the + # reference happens to be built. The copy just avoids aliasing `state`. + reference = Particle(model, particle_varinfo(), deepcopy(state.rng), state) + # `n - 1` fresh particles, with the reference last -- the slot `resample_propagate!` retains. + particles = [Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:(n - 1)] + push!(particles, reference) logZ, _ = sweep!( rng, particles, sampler.resampler, sampler.multithreaded; conditional=true ) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 6f7654111d..8567610b54 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -10,12 +10,10 @@ using Turing.Inference: MultinomialResampler, ESSResampler, Particle, - TracedRNG, + particle_rng, particle_varinfo, advance!, fork, - rewind!, - refresh!, sweep!, resample_indices, normalized_weights @@ -275,17 +273,14 @@ end model = drifting([0.3, -0.7, 1.1]) function conditional_sweep(scheme) rng = StableRNG(77) - retained = Particle(model, particle_varinfo(), TracedRNG(rng)) - while advance!(retained, false) !== nothing + retained = Particle(model, particle_varinfo(), particle_rng(rng)) + while advance!(retained) !== nothing end - reference = Particle( - model, particle_varinfo(), rewind!(deepcopy(retained.rng)), retained - ) - particles = map( - i -> - i < 5 ? Particle(model, particle_varinfo(), TracedRNG(rng)) : reference, - 1:5, - ) + reference = Particle(model, particle_varinfo(), particle_rng(rng), retained) + particles = [ + Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:4 + ] + push!(particles, reference) sweep!(StableRNG(78), particles, scheme, false; conditional=true) return map(p -> get_raw_values(p.varinfo), particles) end @@ -303,9 +298,9 @@ end @testset "ensuring reference consistency" begin # In conditional SMC the retained trajectory must be regenerated *exactly* by the - # reference particle on the next iteration -- this is what makes CSMC valid. Value - # replay must reach every latent address, and the traced-RNG step counter and recorded - # seeds must stay aligned with the observation boundaries. + # reference particle on the next iteration -- this is what makes CSMC valid. Reusing the + # retained values must reach every latent address, over many sweeps and after the + # reference has itself been resampled from. @model function state_space_model(y) ρ ~ Uniform(0, 1) x = Vector{Float64}(undef, length(y) + 1) @@ -321,30 +316,35 @@ end # the mutating loop out of test soft scope. function run_csmc(model, N, nsteps, rng) draw(ps) = ps[rand(rng, Categorical(normalized_weights(ps)))] - particles = [Particle(model, particle_varinfo(), TracedRNG(rng)) for _ in 1:N] + particles = [ + Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:N + ] sweep!(rng, particles, ESSResampler(0.5), false) state = draw(particles) allok = true + nlatents = 0 for _ in 1:nsteps - ref = Particle( - model, particle_varinfo(), rewind!(deepcopy(state.rng)), state - ) - parts = map( - i -> i < N ? Particle(model, particle_varinfo(), TracedRNG(rng)) : ref, - 1:N, - ) + ref = Particle(model, particle_varinfo(), particle_rng(rng), state) + parts = [ + Particle(model, particle_varinfo(), particle_rng(rng)) for + _ in 1:(N - 1) + ] + push!(parts, ref) sweep!(rng, parts, ESSResampler(0.5), false; conditional=true) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) state = draw(parts) + nlatents = length(state.assumed_varnames) end - return allok, length(state.rng.keys) + return allok, nlatents end rng = StableRNG(1234) y = randn(rng, 10) - allok, nkeys = run_csmc(state_space_model(y), 3, 30, rng) - @test allok # reference regenerated exactly every step - @test nkeys == length(y) + 1 # keys stay aligned with the trajectory length + # ρ plus x[1:length(y)+1]: the retained trajectory must span every latent, so that the + # next reference is pinned on all of them rather than silently redrawing the rest. + allok, nlatents = run_csmc(state_space_model(y), 3, 30, rng) + @test allok # reference regenerated exactly every step + @test nlatents == length(y) + 2 end @testset "reference is pinned to retained values under re-conditioning" begin @@ -360,18 +360,18 @@ end end rng = StableRNG(42) retained = Particle( - reconditioned(2.0) | (@varname(a) => 0.0), particle_varinfo(), TracedRNG(rng) + reconditioned(2.0) | (@varname(a) => 0.0), particle_varinfo(), particle_rng(rng) ) - while advance!(retained, false) !== nothing + while advance!(retained) !== nothing end retained_vals = get_raw_values(retained.varinfo) reference = Particle( reconditioned(2.0) | (@varname(a) => 5.0), # x's prior shifted far away particle_varinfo(), - rewind!(deepcopy(retained.rng)), + particle_rng(rng), retained, ) - while advance!(reference, true) !== nothing + while advance!(reference) !== nothing end @test get_raw_values(reference.varinfo) == retained_vals end @@ -388,16 +388,15 @@ end return y ~ Normal(μ, 1) end rng = StableRNG(91) - retained = Particle(branch_changes(true, 0.0), particle_varinfo(), TracedRNG(rng)) - while advance!(retained, false) !== nothing + retained = Particle( + branch_changes(true, 0.0), particle_varinfo(), particle_rng(rng) + ) + while advance!(retained) !== nothing end reference = Particle( - branch_changes(false, 0.0), - particle_varinfo(), - rewind!(deepcopy(retained.rng)), - retained, + branch_changes(false, 0.0), particle_varinfo(), particle_rng(rng), retained ) - @test_throws "reference execution trace changed" advance!(reference, true) + @test_throws "reference execution trace changed" advance!(reference) @model function branch_drops(flag, y) x ~ Normal() @@ -406,17 +405,14 @@ end end return y ~ Normal(x, 1) end - retained = Particle(branch_drops(true, 0.0), particle_varinfo(), TracedRNG(rng)) - while advance!(retained, false) !== nothing + retained = Particle(branch_drops(true, 0.0), particle_varinfo(), particle_rng(rng)) + while advance!(retained) !== nothing end reference = Particle( - branch_drops(false, 0.0), - particle_varinfo(), - rewind!(deepcopy(retained.rng)), - retained, + branch_drops(false, 0.0), particle_varinfo(), particle_rng(rng), retained ) @test_throws "reference execution trace changed" begin - while advance!(reference, true) !== nothing + while advance!(reference) !== nothing end end end @@ -555,25 +551,25 @@ end @testset "advance!" begin # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. - particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) - @test advance!(particle, false) ≈ -log(2) - @test advance!(particle, false) ≈ -log(2) # `0 ~ Bernoulli(0.5)` - @test advance!(particle, false) === nothing # model finished + particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) + @test advance!(particle) ≈ -log(2) + @test advance!(particle) ≈ -log(2) # `0 ~ Bernoulli(0.5)` + @test advance!(particle) === nothing # model finished end @testset "matches a direct evaluation" begin # A particle advanced without resampling draws from its RNG continuously, so it must # produce exactly the same values and log-likelihood as a plain DynamicPPL evaluation # seeded identically. - particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) - while advance!(particle, false) !== nothing + particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) + while advance!(particle) !== nothing end accs = DynamicPPL.OnlyAccsVarInfo() accs = DynamicPPL.setacc!!(accs, DynamicPPL.LogLikelihoodAccumulator()) accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) _, accs = DynamicPPL.init!!( - TracedRNG(Xoshiro(23)), + particle_rng(Xoshiro(23)), test(), accs, DynamicPPL.InitFromPrior(), @@ -585,16 +581,20 @@ end end @testset "fork" begin - particle = Particle(test(), particle_varinfo(), TracedRNG(Xoshiro(23))) - advance!(particle, false) + particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) + advance!(particle) child = fork(particle, Xoshiro(1)) # Independent continuations: advancing one does not touch the other. - @test advance!(child, false) ≈ -log(2) + @test advance!(child) ≈ -log(2) @test particle.varinfo !== child.varinfo - @test advance!(particle, false) ≈ -log(2) + @test advance!(particle) ≈ -log(2) end - @testset "rng replay" begin + @testset "reference consumes no randomness" begin + # The reference reproduces the retained trajectory purely from its values, so its own + # generator is never consulted. Pinning that down is what lets the reference be handed an + # ordinary generator instead of a replayable one: scrambling its seeds before every step + # must not perturb the trajectory it regenerates. @model function normal() a ~ Normal(0, 1) 3 ~ Normal(a, 2) @@ -603,19 +603,20 @@ end return a, b end - # Run a particle to completion, then replay it from its recorded seeds (as the - # reference trajectory of a conditional sweep does) and check it regenerates exactly. - # Replay relies on each step using a distinct seed, so we refresh before every step - # exactly as the sweep's no-resample path does. - particle = Particle(normal(), particle_varinfo(), TracedRNG(Xoshiro(23))) - while (refresh!(particle.rng); advance!(particle, false)) !== nothing + retained = Particle(normal(), particle_varinfo(), particle_rng(Xoshiro(23))) + while advance!(retained) !== nothing end - values = DynamicPPL.get_raw_values(particle.varinfo) + values = get_raw_values(retained.varinfo) - reference = Particle(normal(), particle_varinfo(), rewind!(deepcopy(particle.rng))) - while advance!(reference, true) !== nothing + scrambler = Xoshiro(99) + reference = Particle( + normal(), particle_varinfo(), particle_rng(Xoshiro(7)), retained + ) + while ( + Random.seed!(reference.rng, rand(scrambler, UInt64)); advance!(reference) + ) !== nothing end - @test DynamicPPL.get_raw_values(reference.varinfo) == values + @test get_raw_values(reference.varinfo) == values end end From 9efc1841694b227b236dc97f19d2552978c08c55 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 18:33:28 +0100 Subject: [PATCH 48/58] Rename ESSResampler and the weight-ESS helper away from the ESS sampler's name `Turing.Inference` had accumulated `ESS`, `ESSLikelihood`, `ESSPrior` and `TuringESSState`, all elliptical slice sampling, and then `ESSResampler`, which is about effective sample size. Three of the four `ESS*` names mean one thing and the fourth means another, so `SMC(ESSResampler(0.5))` reads ambiguously next to `Gibbs(:s => CSMC(15), :m => ESS())`. `ESSThresholdResampler` says what triggers the resampling and matches the `AdvancedPS.ResampleWithESSThreshold` it replaces, which also eases migration. Separately, `Turing.Inference.ess` and `Turing.ess` were two *different* functions: the former the weight degeneracy of one particle population, the latter `MCMCDiagnosticTools.ess` measuring a chain's autocorrelation. Renaming ours to `weight_ess` removes the collision, and its docstring now names the distinction so the next reader does not have to rediscover it. Both names are unexported, and 0.47.0 is unreleased, so nothing downstream depends on them yet. Co-Authored-By: Claude Code --- HISTORY.md | 2 +- src/mcmc/particle_mcmc.jl | 44 ++++++++++++++++++++++---------------- test/mcmc/particle_mcmc.jl | 18 ++++++++-------- 3 files changed, 36 insertions(+), 28 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 6d9bd005b8..bfc411a813 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -6,7 +6,7 @@ `SMC` and `PG` / `CSMC` have been reimplemented natively and no longer depend on AdvancedPS. -Resampling schemes are now types rather than functions — `StratifiedResampler()`, `SystematicResampler()`, and `MultinomialResampler()` (in `Turing.Inference`), optionally wrapped in `ESSResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.SystematicResampler())`, `SMC(0.5)`, or `PG(10, Turing.Inference.MultinomialResampler(), 0.5)`. +Resampling schemes are now types rather than functions — `StratifiedResampler()`, `SystematicResampler()`, and `MultinomialResampler()` (in `Turing.Inference`), optionally wrapped in `ESSThresholdResampler(threshold, scheme)` to resample only when the effective sample size falls below `threshold * nparticles`; for example `SMC(Turing.Inference.SystematicResampler())`, `SMC(0.5)`, or `PG(10, Turing.Inference.MultinomialResampler(), 0.5)`. The old function-based API (`resample_systematic`, `AdvancedPS.ResampleWithESSThreshold`, …) is gone. The default scheme is now **stratified** rather than systematic: it stays consistent as the number of particles grows, which systematic does not. diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 0e444fb2d6..c2f8d839bd 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -436,21 +436,25 @@ end # ── Effective-sample-size gating ────────────────────────────────────────────── """ - ESSResampler(threshold, scheme = StratifiedResampler()) + ESSThresholdResampler(threshold, scheme = StratifiedResampler()) Resample with `scheme`, but only when the effective sample size drops below `threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). """ -struct ESSResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler +struct ESSThresholdResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler threshold::T scheme::R end -ESSResampler(threshold::Real) = ESSResampler(threshold, StratifiedResampler()) +function ESSThresholdResampler(threshold::Real) + return ESSThresholdResampler(threshold, StratifiedResampler()) +end -function should_resample(resampler::ESSResampler, weights) - return ess(weights) ≤ resampler.threshold * length(weights) +function should_resample(resampler::ESSThresholdResampler, weights) + return weight_ess(weights) ≤ resampler.threshold * length(weights) end -function resample_indices(rng::AbstractRNG, resampler::ESSResampler, weights, n::Integer) +function resample_indices( + rng::AbstractRNG, resampler::ESSThresholdResampler, weights, n::Integer +) return resample_indices(rng, resampler.scheme, weights, n) end @@ -467,8 +471,12 @@ end logweights(particles) = [p.logweight for p in particles] normalized_weights(particles) = softmax(logweights(particles)) log_normalizing_constant(particles) = logsumexp(logweights(particles)) -"Effective sample size of a normalised weight vector, `1 / Σ wᵢ²`." -ess(weights) = inv(sum(abs2, weights)) +""" +Effective sample size of a normalised weight vector, `1 / Σ wᵢ²`. Named for the weights to keep it +distinct from `MCMCDiagnosticTools.ess`, which Turing re-exports and which measures a *chain's* +autocorrelation rather than a population's weight degeneracy. +""" +weight_ess(weights) = inv(sum(abs2, weights)) # ── Reweighting ─────────────────────────────────────────────────────────────── @@ -580,7 +588,7 @@ function sweep!( # Post-reweight ESS for this observation: a degeneracy diagnostic (low ESS means few # particles carry the weight). After the break, so the finishing pass -- which adds no # observation and leaves the weights unchanged -- contributes no spurious entry. - push!(ess_per_step, ess(normalized_weights(particles))) + push!(ess_per_step, weight_ess(normalized_weights(particles))) end return logZ, ess_per_step end @@ -611,7 +619,7 @@ struct SMC{R<:AbstractResampler} <: ParticleInference end """ - SMC([resampler = ESSResampler(0.5)]; multithreaded = false) + SMC([resampler = ESSThresholdResampler(0.5)]; multithreaded = false) SMC([scheme = StratifiedResampler(), ]threshold; multithreaded = false) Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the @@ -620,13 +628,13 @@ effective sample size drops below half the number of particles. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). -The resampling scheme types (`StratifiedResampler`, `SystematicResampler`, `MultinomialResampler`, `ESSResampler`) are +The resampling scheme types (`StratifiedResampler`, `SystematicResampler`, `MultinomialResampler`, `ESSThresholdResampler`) are not exported; refer to them as e.g. `Turing.Inference.SystematicResampler`. """ -SMC(; kwargs...) = SMC(ESSResampler(0.5); kwargs...) -SMC(threshold::Real; kwargs...) = SMC(ESSResampler(threshold); kwargs...) +SMC(; kwargs...) = SMC(ESSThresholdResampler(0.5); kwargs...) +SMC(threshold::Real; kwargs...) = SMC(ESSThresholdResampler(threshold); kwargs...) function SMC(scheme::AbstractResampler, threshold::Real; kwargs...) - return SMC(ESSResampler(threshold, scheme); kwargs...) + return SMC(ESSThresholdResampler(threshold, scheme); kwargs...) end # SMC is a single weighted sweep, not a Markov chain: rather than fake an iteration through @@ -713,7 +721,7 @@ struct PG{R<:AbstractResampler} <: ParticleInference end """ - PG(n, [resampler = ESSResampler(0.5)]; multithreaded = false) + PG(n, [resampler = ESSThresholdResampler(0.5)]; multithreaded = false) PG(n, [scheme = StratifiedResampler(), ]threshold; multithreaded = false) Particle Gibbs sampler with `n` particles. By default resampling is triggered whenever the @@ -725,10 +733,10 @@ different algorithm rather than the same draw with one output pinned. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). """ -PG(n::Int; kwargs...) = PG(n, ESSResampler(0.5); kwargs...) -PG(n::Int, threshold::Real; kwargs...) = PG(n, ESSResampler(threshold); kwargs...) +PG(n::Int; kwargs...) = PG(n, ESSThresholdResampler(0.5); kwargs...) +PG(n::Int, threshold::Real; kwargs...) = PG(n, ESSThresholdResampler(threshold); kwargs...) function PG(n::Int, scheme::AbstractResampler, threshold::Real; kwargs...) - return PG(n, ESSResampler(threshold, scheme); kwargs...) + return PG(n, ESSThresholdResampler(threshold, scheme); kwargs...) end "Conditional SMC, an alias for [`PG`](@ref)." diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 8567610b54..bc8a686ee2 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -8,7 +8,7 @@ using Turing.Inference: StratifiedResampler, SystematicResampler, MultinomialResampler, - ESSResampler, + ESSThresholdResampler, Particle, particle_rng, particle_varinfo, @@ -27,11 +27,11 @@ using Turing @testset "SMC" begin @testset "constructor" begin - @test SMC().resampler == ESSResampler(0.5) + @test SMC().resampler == ESSThresholdResampler(0.5) @test SMC().resampler.scheme isa StratifiedResampler # stratified is the default scheme - @test SMC(0.6).resampler == ESSResampler(0.6) + @test SMC(0.6).resampler == ESSThresholdResampler(0.6) @test SMC(MultinomialResampler(), 0.6).resampler == - ESSResampler(0.6, MultinomialResampler()) + ESSThresholdResampler(0.6, MultinomialResampler()) @test SMC(SystematicResampler()).resampler == SystematicResampler() @test SMC().multithreaded == false @test SMC(; multithreaded=true).multithreaded == true @@ -206,10 +206,10 @@ end @testset "PG" begin @testset "constructor" begin @test PG(10).nparticles == 10 - @test PG(10).resampler == ESSResampler(0.5) - @test PG(60, 0.6).resampler == ESSResampler(0.6) + @test PG(10).resampler == ESSThresholdResampler(0.5) + @test PG(60, 0.6).resampler == ESSThresholdResampler(0.6) @test PG(80, MultinomialResampler(), 0.6).resampler == - ESSResampler(0.6, MultinomialResampler()) + ESSThresholdResampler(0.6, MultinomialResampler()) @test PG(100, SystematicResampler()).resampler == SystematicResampler() @test PG(10).multithreaded == false @test PG(10; multithreaded=true).multithreaded == true @@ -319,7 +319,7 @@ end particles = [ Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:N ] - sweep!(rng, particles, ESSResampler(0.5), false) + sweep!(rng, particles, ESSThresholdResampler(0.5), false) state = draw(particles) allok = true nlatents = 0 @@ -330,7 +330,7 @@ end _ in 1:(N - 1) ] push!(parts, ref) - sweep!(rng, parts, ESSResampler(0.5), false; conditional=true) + sweep!(rng, parts, ESSThresholdResampler(0.5), false; conditional=true) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) state = draw(parts) nlatents = length(state.assumed_varnames) From 2ff6b2222a83469e22a37c6a00d982f4a2f483a8 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 18:45:28 +0100 Subject: [PATCH 49/58] Say that PG's log_normalizing_constant is biased, and test that it is MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SMC` and `PG` both put `log_normalizing_constant` in the chain, under one name and one per-iteration shape, and HISTORY.md called it the marginal-likelihood estimate for both. Only `SMC`'s deserves that description. A conditional sweep retains the reference whatever its weight, and the reference is a draw from the posterior rather than from the proposal, so it carries far more likelihood than a fresh particle and inflates the mean weight at every step. Measured against the exact p(y) of a linear Gaussian SSM: SMC's `E[Ẑ/p(y)]` is 0.945 ± 0.040 at 16 particles and 1.029 ± 0.020 at 64, both consistent with 1, while PG's is 1.80 and 1.16 -- an 80% overestimate of the marginal likelihood at 16 particles. It decays like 1/n but stays large at any practical n. `mean(exp.(chain[:log_normalizing_constant]))` was therefore silently wrong under PG, in the direction that favours whichever model was fitted with fewer particles. The existing test could not have caught this. Its model pins `x = 1` through `x ~ Bernoulli(1)`, so both observes contribute exactly `log(1/2)` for every particle; with zero weight variance the estimate is exact for conditional sweeps too, which is also why the test could assert that all iterations agree. That degeneracy is now spelled out where it could otherwise read as evidence of unbiasedness, and a second testset pins the bias against a closed-form Beta-Bernoulli p(y): SMC's ratio averages to 1, PG's exceeds 1.05. Co-Authored-By: Claude Code --- HISTORY.md | 4 +++- src/mcmc/particle_mcmc.jl | 9 +++++++++ test/mcmc/particle_mcmc.jl | 36 +++++++++++++++++++++++++++++++++++- 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index bfc411a813..eb0b1ec2e5 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -17,7 +17,9 @@ The rewrite also brings: - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. - - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry the log-normalizing-constant (marginal-likelihood) estimate `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). + - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). + + For `SMC` this is an unbiased estimate of the marginal likelihood `p(y)`. For `PG` / `CSMC` it is **not**, and must not be used for model comparison; see the `PG` docstring. # 0.46.0 diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index c2f8d839bd..d0e2abfc7f 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -732,6 +732,15 @@ different algorithm rather than the same draw with one output pinned. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). + +!!! warning "`log_normalizing_constant` is biased for PG" + PG chains carry `log_normalizing_constant`, but unlike [`SMC`](@ref)'s it does **not** estimate + `log p(y)` without bias, so it must not be used for model comparison. A conditional sweep + retains the reference whatever its weight, and the reference is a draw from the posterior rather + than from the proposal, so it usually carries far more likelihood than a fresh particle and + inflates the mean weight at every step. Measured against the exact `p(y)` of a linear Gaussian + SSM, `E[Ẑ]` overshoots by 80% at `n = 16` and 16% at `n = 64`; the bias decays like `1/n` but + stays large at practical `n`. Use `SMC` for an unbiased estimate. """ PG(n::Int; kwargs...) = PG(n, ESSThresholdResampler(0.5); kwargs...) PG(n::Int, threshold::Real; kwargs...) = PG(n, ESSThresholdResampler(threshold); kwargs...) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index bc8a686ee2..87f7045363 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -21,6 +21,7 @@ using Distributions: Bernoulli, Beta, Gamma, MvNormal, Normal, Uniform, Categori using FlexiChains: VNChain, has_same_data using LinearAlgebra: I using Random: Random, Xoshiro +using SpecialFunctions: logbeta using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset using Turing @@ -240,10 +241,43 @@ end @test all(isone, chains_pg[:x]) pg_log_normalizing_constant = mean(chains_pg[:log_normalizing_constant]) @test pg_log_normalizing_constant ≈ -2 * log(2) atol = 0.01 - # Should be the same for all iterations. + # Every particle scores the same here -- `x ~ Bernoulli(1)` pins `x = 1`, so both observes + # contribute exactly `log(1/2)` regardless of the trajectory. Zero weight variance is why + # the estimate is exact for PG too, and why all iterations agree. It is *not* evidence that + # PG's estimator is unbiased in general; the testset below covers that. @test chains_pg[:log_normalizing_constant] ≈ fill(pg_log_normalizing_constant, 100) end + @testset "log_normalizing_constant is biased upward for conditional sweeps" begin + # Unlike SMC's, PG's `log_normalizing_constant` is not an unbiased estimate of log p(y): + # a conditional sweep keeps the reference whatever its weight, and the reference is a + # posterior draw rather than a proposal draw, so it inflates the mean weight at each step. + # Pin the direction and rough size so a future change to the sweep cannot quietly alter it. + # + # Beta-Bernoulli, so p(y) is exact: with a Beta(1,1) prior, p(y) = B(1+s, 1+n-s)/B(1,1). + @model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end + end + obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1] + s, n = sum(obs), length(obs) + exact_logp = logbeta(1 + s, 1 + n - s) - logbeta(1, 1) + + # SMC's estimate is unbiased, so averaging Ẑ over independent sweeps lands on p(y). + smc_ratios = map(1:200) do i + chn = sample(StableRNG(900 + i), coinflip(obs), SMC(), 32) + exp(first(chn[:log_normalizing_constant]) - exact_logp) + end + @test mean(smc_ratios) ≈ 1 atol = 0.1 + + # PG's overshoots. Drop iteration 1, which is an unconditional sweep. + chn = sample(StableRNG(468), coinflip(obs), PG(8), 2_000) + pg_ratios = exp.(vec(collect(chn[:log_normalizing_constant]))[2:end] .- exact_logp) + @test mean(pg_ratios) > 1.05 + end + @testset "multithreaded execution matches serial" begin # Threading the reweighting must not perturb the reference-replay bookkeeping, so the # conditional sweeps have to reproduce the serial draws exactly, whatever the thread From 92b36387ff4ee4a7af141c475326c9fc2274a4ca Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 20:06:37 +0100 Subject: [PATCH 50/58] Check the particle samplers against two models with tractable posteriors Until now nothing pinned SMC/PG to a known answer. These add a scalar linear Gaussian SSM and a discrete HMM, both with closed-form posteriors, and check each twice: `PG` alone against the exact smoothing marginals, then `Gibbs(theta => NUTS/HMC, states => CSMC)` with a static parameter unknown. The second case is the one worth having -- the states' distribution depends on the theta the *other* Gibbs component owns, so it only passes if the CSMC reference stays pinned to its retained trajectory as the model is re-conditioned between sweeps. The exact parameter posterior comes from quadrature against the closed-form likelihood rather than from a second sampler, and the theta-mixed state marginals follow from the laws of total expectation and variance over the same grid. Tolerances are batch-means standard errors, not hard-coded `atol`, so they stay meaningful as mixing changes instead of going vacuous or flaky; HMM states rarer than 0.02 are skipped, since a bursty 0/1 indicator makes that estimate unreliable. `ExactSSM` validates itself rather than asking to be trusted: the closed-form Gaussian smoother is cross-checked against an independent Kalman/RTS recursion, and forward-backward against enumeration of all K^T state paths, both to 1e-12, as part of the suite. The models are defined at module scope deliberately. Inside a testset they would share a local scope with the simulated `x`/`z`, which makes those captured locals that the model body then rebinds -- so every particle mutates one shared array and the posterior is silently wrong. That cost an afternoon; the simulated truth is now named `xtrue`/`ztrue` and both sites say why. GeneralisedFilters would have been the natural source of ground truth, but it cannot go in the test environment: v0.4.2 pins CUDA to 5.0-5.11 while Mooncake, which the AD tests need, requires CUDA 6.x. Two things found while trying, worth knowing if it is ever revisited: its prior sits at t = 0 with one transition applied before the first observation, so comparing it against a naturally-written Turing model is silently wrong unless the prior is stationary (1.5e-3 in log p(y) here, 4.2e-2 for the HMM); and `smooth` throws a convert MethodError on v0.4.2. Co-Authored-By: Claude Code --- test/mcmc/particle_mcmc.jl | 206 ++++++++++++++++++++++++++++++++++- test/runtests.jl | 1 + test/test_utils/exact_ssm.jl | 191 ++++++++++++++++++++++++++++++++ 3 files changed, 395 insertions(+), 3 deletions(-) create mode 100644 test/test_utils/exact_ssm.jl diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 87f7045363..8925f2a288 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -3,6 +3,7 @@ module ParticleMCMCTests using ..Models: gdemo_default using ..SamplerTestUtils: test_chain_logp_metadata, test_rng_respected using ..NumericalTests: check_numerical +using ..ExactSSM: ExactSSM using DynamicPPL: DynamicPPL, extract_priors, get_raw_values, getloglikelihood using Turing.Inference: StratifiedResampler, @@ -17,7 +18,18 @@ using Turing.Inference: sweep!, resample_indices, normalized_weights -using Distributions: Bernoulli, Beta, Gamma, MvNormal, Normal, Uniform, Categorical, sample +using Distributions: + Bernoulli, + Beta, + Categorical, + Gamma, + InverseGamma, + LogNormal, + MvNormal, + Normal, + Uniform, + logpdf, + sample using FlexiChains: VNChain, has_same_data using LinearAlgebra: I using Random: Random, Xoshiro @@ -315,7 +327,7 @@ end Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:4 ] push!(particles, reference) - sweep!(StableRNG(78), particles, scheme, false; conditional=true) + sweep!(StableRNG(78), particles, scheme, false) return map(p -> get_raw_values(p.varinfo), particles) end multinomial = conditional_sweep(MultinomialResampler()) @@ -364,7 +376,7 @@ end _ in 1:(N - 1) ] push!(parts, ref) - sweep!(rng, parts, ESSThresholdResampler(0.5), false; conditional=true) + sweep!(rng, parts, ESSThresholdResampler(0.5), false) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) state = draw(parts) nlatents = length(state.assumed_varnames) @@ -654,4 +666,192 @@ end end end +# +# State space models with tractable posteriors +# +# These are the checks that pin the particle samplers to a known answer rather than to each other. A +# scalar linear Gaussian SSM and a discrete HMM both have closed-form posteriors, supplied by +# `ExactSSM` and validated there against brute force, so a disagreement beyond Monte Carlo error is a +# bug. +# +# Each model is exercised twice: `PG` alone against the exact smoothing marginals, then +# `Gibbs(θ => NUTS/HMC, states => CSMC)` with a static parameter unknown. The second is the case that +# matters, because the states' distribution depends on the θ owned by the *other* Gibbs component, so +# the CSMC reference has to stay pinned to its retained trajectory as the model is re-conditioned +# between sweeps. The exact θ posterior comes from quadrature against the closed-form likelihood, and +# the θ-mixed state marginals from the laws of total expectation and variance over the same grid. + +"Draws for one variable as a plain vector; chain indexing yields an iteration×chain matrix." +particle_draws(chn, vn) = vec(collect(chn[vn])) + +"Batch-means standard error of the mean of a correlated chain." +function batch_means_se(v; nbatches::Int=40) + n = length(v) ÷ nbatches + b = [mean(@view v[((i - 1) * n + 1):(i * n)]) for i in 1:nbatches] + return std(b) / sqrt(nbatches) +end + +""" +Assert that `samples` estimates `exact` to within `nsigma` batch-means standard errors. Using the +chain's own error estimate keeps the tolerance honest as mixing changes, rather than hard-coding an +`atol` that silently becomes either vacuous or flaky. +""" +function test_within_mc_error(exact, samples; nsigma=4) + @test abs(mean(samples) - exact) <= nsigma * batch_means_se(samples) + return nothing +end + +@model function lgssm(y, a, q, r) + # `typeof(q)` stays generic when HMC differentiates through `q`, without boxing every + # element the way `Vector{Real}` would. + x = Vector{typeof(q)}(undef, length(y)) + x[1] ~ Normal(0, sqrt(q / (1 - a^2))) + y[1] ~ Normal(x[1], sqrt(r)) + for t in 2:length(y) + x[t] ~ Normal(a * x[t - 1], sqrt(q)) + y[t] ~ Normal(x[t], sqrt(r)) + end +end + +@model function lgssm_unknown_q(y, a, r) + q ~ InverseGamma(3, 2) + x = Vector{typeof(q)}(undef, length(y)) + x[1] ~ Normal(0, sqrt(q / (1 - a^2))) + y[1] ~ Normal(x[1], sqrt(r)) + for t in 2:length(y) + x[t] ~ Normal(a * x[t - 1], sqrt(q)) + y[t] ~ Normal(x[t], sqrt(r)) + end +end + +const HMM_P = [0.80 0.15 0.05; 0.10 0.80 0.10; 0.05 0.15 0.80] +const HMM_MEANS = [-1.5, 0.0, 1.5] +const HMM_PI0 = ExactSSM.stationary_distribution(HMM_P) + +@model function hmm(y, sd) + z = Vector{Int}(undef, length(y)) + z[1] ~ Categorical(HMM_PI0) + y[1] ~ Normal(HMM_MEANS[z[1]], sd) + for t in 2:length(y) + z[t] ~ Categorical(HMM_P[z[t - 1], :]) + y[t] ~ Normal(HMM_MEANS[z[t]], sd) + end +end + +@model function hmm_unknown_sd(y) + sd ~ LogNormal(log(0.7), 0.4) + z = Vector{Int}(undef, length(y)) + z[1] ~ Categorical(HMM_PI0) + y[1] ~ Normal(HMM_MEANS[z[1]], sd) + for t in 2:length(y) + z[t] ~ Categorical(HMM_P[z[t - 1], :]) + y[t] ~ Normal(HMM_MEANS[z[t]], sd) + end +end + +@testset "linear Gaussian SSM" begin + ExactSSM.test_exact_ssm_reference() + + a, r, true_q, T = 0.8, 0.3, 0.5, 6 + s0(q) = q / (1 - a^2) # stationary, so the chain has no burn-in transient + + # Named `xtrue`, not `x`: a `@model` sharing a local scope with an `x` assignment captures it, + # so every particle would mutate one shared array -- silently, and catastrophically. + rng = StableRNG(1234) + xtrue = zeros(T) + xtrue[1] = sqrt(s0(true_q)) * randn(rng) + for t in 2:T + xtrue[t] = a * xtrue[t - 1] + sqrt(true_q) * randn(rng) + end + y = xtrue .+ sqrt(r) .* randn(rng, T) + + @testset "PG recovers the exact smoothing marginals" begin + means, vars = ExactSSM.lgssm_smoother(y, a, true_q, r, s0(true_q)) + chn = sample(StableRNG(24), lgssm(y, a, true_q, r), PG(32), 4_000) + for t in 1:T + xs = particle_draws(chn, @varname(x[t])) + test_within_mc_error(means[t], xs) + test_within_mc_error(vars[t], (xs .- mean(xs)) .^ 2) + end + end + + @testset "Gibbs(q => NUTS, x => CSMC) recovers the exact posterior" begin + prior = InverseGamma(3, 2) + qs = range(0.05, 4.0; length=400) + w = ExactSSM.grid_posterior( + prior, qs, q -> ExactSSM.lgssm_loglik(y, a, q, r, s0(q)) + ) + q_mean, q_sd = ExactSSM.grid_moments(w, qs) + smoothed = [ExactSSM.lgssm_smoother(y, a, q, r, s0(q)) for q in qs] + x_mean = sum(w[i] * first(smoothed[i]) for i in eachindex(w)) + x_second = sum( + w[i] * (last(smoothed[i]) .+ first(smoothed[i]) .^ 2) for i in eachindex(w) + ) + + alg = Gibbs(@varname(q) => NUTS(), @varname(x) => CSMC(32)) + chn = sample(StableRNG(31), lgssm_unknown_q(y, a, r), alg, 4_000) + + qd = particle_draws(chn, @varname(q)) + test_within_mc_error(q_mean, qd) + @test std(qd) ≈ q_sd rtol = 0.2 + for t in 1:T + xs = particle_draws(chn, @varname(x[t])) + test_within_mc_error(x_mean[t], xs) + @test var(xs) ≈ x_second[t] - x_mean[t]^2 rtol = 0.25 + end + end +end + +@testset "discrete HMM" begin + P, π0, means = HMM_P, HMM_PI0, HMM_MEANS + true_sd, K, T = 0.7, 3, 6 + obs_loglik(y, sd) = [logpdf(Normal(means[k], sd), y[t]) for t in 1:length(y), k in 1:K] + + # `ztrue`, not `z`, for the same reason as `xtrue` above. + rng = StableRNG(99) + ztrue = Vector{Int}(undef, T) + ztrue[1] = rand(rng, Categorical(π0)) + for t in 2:T + ztrue[t] = rand(rng, Categorical(P[ztrue[t - 1], :])) + end + y = [means[ztrue[t]] + true_sd * randn(rng) for t in 1:T] + + @testset "PG recovers the exact state marginals" begin + # Discrete states make this sharp: the reference is a probability vector, so any bias shows + # up directly instead of being absorbed into a mean. + post, _ = ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, true_sd)) + chn = sample(StableRNG(25), hmm(y, true_sd), PG(32), 4_000) + for t in 1:T, k in 1:K + post[t, k] < 0.02 && continue # too rare to resolve at this chain length + test_within_mc_error( + post[t, k], Float64.(particle_draws(chn, @varname(z[t])) .== k) + ) + end + end + + @testset "Gibbs(sd => HMC, z => CSMC) recovers the exact posterior" begin + prior = LogNormal(log(0.7), 0.4) + sds = range(0.2, 2.5; length=400) + w = ExactSSM.grid_posterior( + prior, sds, s -> last(ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, s))) + ) + sd_mean, sd_sd = ExactSSM.grid_moments(w, sds) + posts = [first(ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, s))) for s in sds] + mixed = sum(w[i] * posts[i] for i in eachindex(w)) + + alg = Gibbs(@varname(sd) => HMC(0.1, 12), @varname(z) => CSMC(32)) + chn = sample(StableRNG(32), hmm_unknown_sd(y), alg, 4_000) + + sdd = particle_draws(chn, @varname(sd)) + test_within_mc_error(sd_mean, sdd) + @test std(sdd) ≈ sd_sd rtol = 0.2 + for t in 1:T, k in 1:K + mixed[t, k] < 0.02 && continue + test_within_mc_error( + mixed[t, k], Float64.(particle_draws(chn, @varname(z[t])) .== k) + ) + end + end +end + end diff --git a/test/runtests.jl b/test/runtests.jl index b6065d74d0..bb140194b4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,6 +12,7 @@ seed!(23) include("test_utils/models.jl") include("test_utils/numerical_tests.jl") include("test_utils/sampler.jl") +include("test_utils/exact_ssm.jl") Turing.setprogress!(false) included_paths, excluded_paths = parse_args(ARGS) diff --git a/test/test_utils/exact_ssm.jl b/test/test_utils/exact_ssm.jl new file mode 100644 index 0000000000..33899ee9f8 --- /dev/null +++ b/test/test_utils/exact_ssm.jl @@ -0,0 +1,191 @@ +# +# Exact inference for the two state space models with tractable posteriors +# +# These give samplers something unarguable to be compared against: a Kalman filter and smoother for +# a scalar linear Gaussian model, and forward-backward for a discrete HMM. `test_exact_ssm_reference` +# validates both against brute force, so the reference is checked in CI rather than trusted. +# + +module ExactSSM + +using Distributions: Categorical, MvNormal, Normal, logpdf +using LinearAlgebra: I, Symmetric, diag, eigen +using Test: @test, @testset + +export lgssm_smoother, + lgssm_loglik, + hmm_forward_backward, + stationary_distribution, + grid_posterior, + grid_moments, + test_exact_ssm_reference + +## +## Scalar linear Gaussian SSM: x₁ ~ N(0, s0), xₜ = a·xₜ₋₁ + N(0, q), yₜ = xₜ + N(0, r) +## + +# Stacking x = x₁:T gives x = Lε with ε ~ N(0, D), L unit-lower-triangular from the AR(1) recursion +# and D = diag(s0, q, …, q). So Σx = L D Lᵀ and Σy = Σx + rI, and everything follows by conditioning +# a jointly Gaussian vector -- closed form, with no recursion to get subtly wrong. + +"Joint prior covariance of `x₁:T` for the AR(1) latent chain." +function lgssm_prior_cov(T::Integer, a::Real, q::Real, s0::Real) + F = typeof(one(a) * one(q) * one(s0)) + L = [j <= i ? F(a)^(i - j) : zero(F) for i in 1:T, j in 1:T] + D = fill(F(q), T) + D[1] = s0 + return Symmetric(L * (D .* L')) +end + +""" + lgssm_smoother(y, a, q, r, s0) -> (means, variances) + +Exact smoothing marginals `E[xₜ | y₁:T]` and `Var[xₜ | y₁:T]`. +""" +function lgssm_smoother(y::AbstractVector, a::Real, q::Real, r::Real, s0::Real) + Σx = lgssm_prior_cov(length(y), a, q, s0) + G = Σx / (Σx + r * I) # Σx Σy⁻¹; both means are zero a priori + return G * y, diag(Symmetric(Σx - G * Σx)) +end + +"Exact marginal log-likelihood `log p(y₁:T)`." +function lgssm_loglik(y::AbstractVector, a::Real, q::Real, r::Real, s0::Real) + Σy = lgssm_prior_cov(length(y), a, q, s0) + r * I + return logpdf(MvNormal(zeros(eltype(Σy), length(y)), Σy), y) +end + +"Kalman filter plus RTS smoother, kept only to cross-check the closed forms above." +function lgssm_kalman(y::AbstractVector, a::Real, q::Real, r::Real, s0::Real) + T = length(y) + F = typeof(one(a) * one(q) * one(r) * one(s0) * one(eltype(y))) + mp, Pp, mf, Pf = (zeros(F, T) for _ in 1:4) + ll = zero(F) + for t in 1:T + mp[t] = t == 1 ? zero(F) : a * mf[t - 1] + Pp[t] = t == 1 ? F(s0) : a^2 * Pf[t - 1] + q + S = Pp[t] + r + ll += -(log(2pi * S) + (y[t] - mp[t])^2 / S) / 2 + K = Pp[t] / S + mf[t] = mp[t] + K * (y[t] - mp[t]) + Pf[t] = (1 - K) * Pp[t] + end + ms, Ps = copy(mf), copy(Pf) + for t in (T - 1):-1:1 + C = Pf[t] * a / Pp[t + 1] + ms[t] = mf[t] + C * (ms[t + 1] - mp[t + 1]) + Ps[t] = Pf[t] + C^2 * (Ps[t + 1] - Pp[t + 1]) + end + return ms, Ps, ll +end + +## +## Discrete HMM: z₁ ~ Categorical(π0), zₜ | zₜ₋₁ ~ Categorical(P[zₜ₋₁, :]), yₜ | zₜ +## + +""" + hmm_forward_backward(π0, P, loglik_obs) -> (posterior, loglik) + +Forward-backward, where `loglik_obs[t, k] = log p(yₜ | zₜ = k)`. `posterior[t, k]` is +`p(zₜ = k | y₁:T)`. +""" +function hmm_forward_backward( + π0::AbstractVector, P::AbstractMatrix, loglik_obs::AbstractMatrix +) + T, K = size(loglik_obs) + F = promote_type(eltype(π0), eltype(P), eltype(loglik_obs)) + α = zeros(F, T, K) + c = zeros(F, T) # per-step normalisers, which give the log-likelihood + α[1, :] = π0 .* exp.(loglik_obs[1, :]) + c[1] = sum(α[1, :]) + α[1, :] ./= c[1] + for t in 2:T + α[t, :] = (P' * α[t - 1, :]) .* exp.(loglik_obs[t, :]) + c[t] = sum(α[t, :]) + α[t, :] ./= c[t] + end + β = ones(F, T, K) + for t in (T - 1):-1:1 + β[t, :] = P * (exp.(loglik_obs[t + 1, :]) .* β[t + 1, :]) ./ c[t + 1] + end + post = α .* β + return post ./ sum(post; dims=2), sum(log, c) +end + +"Brute-force HMM posterior and log-likelihood by enumerating all `K^T` state paths." +function hmm_brute_force(π0::AbstractVector, P::AbstractMatrix, loglik_obs::AbstractMatrix) + T, K = size(loglik_obs) + post = zeros(T, K) + total = 0.0 + for idx in 0:(K^T - 1) + z = [(idx ÷ K^(t - 1)) % K + 1 for t in 1:T] + lp = log(π0[z[1]]) + loglik_obs[1, z[1]] + for t in 2:T + lp += log(P[z[t - 1], z[t]]) + loglik_obs[t, z[t]] + end + w = exp(lp) + total += w + for t in 1:T + post[t, z[t]] += w + end + end + return post ./ total, log(total) +end + +## +## Shared helpers +## + +"Stationary distribution of a row-stochastic transition matrix." +function stationary_distribution(P::AbstractMatrix) + ev = eigen(collect(transpose(P))) + π0 = real.(ev.vectors[:, argmin(abs.(ev.values .- 1))]) + return π0 ./ sum(π0) +end + +""" + grid_posterior(prior, θs, loglik) -> weights + +Normalised posterior weights over a parameter grid, from `p(θ | y) ∝ p(θ)·p(y | θ)`. Given an exact +`loglik`, this makes the θ posterior exact rather than another Monte Carlo estimate. +""" +function grid_posterior(prior, θs, loglik) + logw = [logpdf(prior, θ) + loglik(θ) for θ in θs] + w = exp.(logw .- maximum(logw)) + return w ./ sum(w) +end + +"Mean and standard deviation of a grid posterior." +function grid_moments(w::AbstractVector, θs) + m = sum(w .* θs) + return m, sqrt(sum(w .* (θs .- m) .^ 2)) +end + +""" +Check the exact implementations against brute force, so that anything comparing a sampler to them is +comparing against something independently verified. The Gaussian closed form is checked against a +Kalman recursion, and forward-backward against enumeration of every state path. +""" +function test_exact_ssm_reference() + @testset "exact SSM reference" begin + a, q, r, s0, T = 0.8, 0.5, 0.3, 1.7, 7 + y = [0.4, -0.7, 1.1, 0.2, -0.5, 0.9, 0.1] + m_joint, v_joint = lgssm_smoother(y, a, q, r, s0) + m_kf, v_kf, ll_kf = lgssm_kalman(y, a, q, r, s0) + @test m_joint ≈ m_kf atol = 1e-12 + @test v_joint ≈ v_kf atol = 1e-12 + @test lgssm_loglik(y, a, q, r, s0) ≈ ll_kf atol = 1e-12 + + P = [0.7 0.2 0.1; 0.15 0.7 0.15; 0.1 0.3 0.6] + π0 = [0.5, 0.3, 0.2] + loglik_obs = log.([0.3 0.5 0.2; 0.6 0.1 0.3; 0.2 0.2 0.6; 0.4 0.4 0.2; 0.1 0.8 0.1]) + post_fb, ll_fb = hmm_forward_backward(π0, P, loglik_obs) + post_bf, ll_bf = hmm_brute_force(π0, P, loglik_obs) + @test post_fb ≈ post_bf atol = 1e-12 + @test ll_fb ≈ ll_bf atol = 1e-12 + + # A stationary π0 is a fixed point of the transition, which several tests rely on. + @test transpose(P) * stationary_distribution(P) ≈ stationary_distribution(P) + end +end + +end From 8390a7c7a47d26c0ad7b2f18dcf3314e478ad509 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 20:41:04 +0100 Subject: [PATCH 51/58] Produce the particle weight from the accumulator, dropping the type piracy `@addlogprob!` bypasses `tilde_observe!!`, so its weight was emitted from a method on `DynamicPPL.accloglikelihood!!(::DynamicPPL.OnlyAccsVarInfo, ...)` -- piracy, since both the function and the type belong to DynamicPPL, and it shadowed upstream's method for *every* `OnlyAccsVarInfo` in the ecosystem while only particle sampling exercised it. Every other sampler used Turing's copy and none of them tested it. There was a Turing-owned dispatch point available all along, and it is where `main` had this before 75d3c0b1f moved it out: `acclogp` on the accumulator. Both routes to the likelihood -- an `observe` via `accumulate_observe!!`, and `@addlogprob!` via `accloglikelihood!!` -> `map_accumulator!!` -- pass through exactly one `acclogp` call, so producing there emits exactly one weight per term and reaches `@addlogprob!` (issue #1996) without touching anything DynamicPPL owns. It is also what 75d3c0b1f was after. That commit moved the produce to the call sites so the emitted weight would agree with the accumulated log-likelihood, and enforced it by diffing `getloglikelihood` either side of the update. Inside `acclogp` the increment *is* the argument, so the agreement is structural: `tilde_observe!!` loses the before/after diff, the second produce site goes, and `ProduceLogLikelihoodAccumulator` stops being a marker that other code consults and becomes the mechanism. Three things worth recording, since they are what make this safe rather than merely shorter. Accumulator merging cannot fire a spurious produce: `combine` sums the two `logp`s directly instead of going through `acclogp`, which is also why `main` needed a `might_produce` method on `Base.:+` and this does not. `produce` suspends before the caller stores the varinfo back on the particle, so a suspended particle's total lags one term -- nothing reads it in that state, as the sweep reweights from `logweight` and every varinfo read happens after the model finishes. And draws are bit-identical to before across SMC, PG under two resampling schemes, and Gibbs(CSMC + HMC): the produce fires from a different place, in the same order, with the same values. Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 208 ++++++++++++++++---------------------- 1 file changed, 87 insertions(+), 121 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index d0e2abfc7f..9173b2cbef 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -90,13 +90,11 @@ A single particle: a suspended `model` execution together with its `varinfo`, it an accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there is no separate state struct). -Without `retained` the particle draws from the prior. Given the previous sweep's retained -particle, it becomes a conditional-SMC reference pinned to that trajectory: it reuses the -retained values, and errors if its execution reaches an address the retained trajectory does -not have, or finishes without reaching one that it does. Taking the whole particle -- rather -than its values and addresses separately -- is what makes a half-specified reference -unrepresentable; only the two pieces are kept, so the retained particle itself is not held -alive across sweeps. +Without `retained` the particle draws from the prior. Given the previous sweep's retained particle it +becomes a conditional-SMC reference pinned to that trajectory, erroring if its execution reaches an +address the retained trajectory lacks or finishes without reaching one it has. Taking the whole +particle, rather than its values and addresses separately, is what makes a half-specified reference +unrepresentable; only those two pieces are kept, so the retained particle is not held alive. """ mutable struct Particle{RT<:AbstractRNG,WT<:Real} # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go @@ -213,12 +211,10 @@ function DynamicPPL.tilde_assume!!( ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) - # A CSMC reference reuses the retained value (`InitFromParams`) at every address it visits. - # `expected_reference_varnames` is `nothing` for ordinary particles and is cleared by - # `reseed!` on a fork, so both draw from the prior. An address outside the retained set means - # the execution trace changed, which must error rather than silently drawing part of the - # nominally fixed reference afresh; the `nothing` fallback given to `InitFromParams` catches - # the converse, a retained address with no usable value. + # A reference reuses the retained value at every address it visits (see `reference_values`); + # ordinary particles and forks have no expected set, so they draw from the prior. The two error + # paths cover the two ways the trace can have moved: an address outside the retained set here, + # and -- via the `nothing` fallback -- a retained address with no usable value. expected = particle.expected_reference_varnames strategy = if expected === nothing DynamicPPL.InitFromPrior() @@ -236,11 +232,8 @@ function DynamicPPL.tilde_assume!!( return x, vi end -# Reweighting invariant: a particle's per-step score is `produce`d from here and from the -# `accloglikelihood!!` overload (for `@addlogprob!`), *after* the likelihood accumulator is -# updated, and equals the accumulator's increment. Producing after the update -- rather than -# inside `acclogp` -- keeps the produced weight in step with the accumulated log-likelihood -# (no one-step lag) and lets `@addlogprob!` terms reach the accumulator, not just the weight. +# Routes the observe through the particle's varinfo and stores the result back. The weight is not +# emitted here -- `ProduceLogLikelihoodAccumulator` produces it as it accumulates, below. function DynamicPPL.tilde_observe!!( ::SMCContext, dist::Distribution, @@ -250,23 +243,20 @@ function DynamicPPL.tilde_observe!!( ::DynamicPPL.AbstractVarInfo, ) particle = Libtask.get_taped_globals(Particle) - before = DynamicPPL.getloglikelihood(particle.varinfo) left, vi = DynamicPPL.tilde_observe!!( DynamicPPL.DefaultContext(), dist, left, vn, template, particle.varinfo ) particle.varinfo = vi - Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) return left, vi end """ ProduceLogLikelihoodAccumulator{T} <: LogProbAccumulator{T} -A marker likelihood accumulator: it accumulates exactly like `LogLikelihoodAccumulator`, but -its distinct type flags a varinfo as belonging to a particle, so the produce sites know to -emit. The produce happens in [`tilde_observe!!`](@ref) (observations) and the -`accloglikelihood!!` overload below (`@addlogprob!`, issue #1996) -- in each case from the -increase in accumulated log-likelihood, keeping the accumulator the single source of truth. +A likelihood accumulator that `Libtask.produce`s each increment as it accumulates it, which is what +turns a model evaluation into a particle filter: one produce per likelihood term, so the sweep sees +one filtering step per `observe`. Substituting it for `LogLikelihoodAccumulator` is the only thing +that distinguishes a particle's varinfo. """ struct ProduceLogLikelihoodAccumulator{T<:Real} <: DynamicPPL.LogProbAccumulator{T} logp::T @@ -274,8 +264,25 @@ end DynamicPPL.accumulator_name(::Type{<:ProduceLogLikelihoodAccumulator}) = :LogLikelihood DynamicPPL.logp(acc::ProduceLogLikelihoodAccumulator) = acc.logp -# `acclogp` is inherited from the generic `LogProbAccumulator` method (plain addition); the -# produce is handled by the produce sites, not here. + +# The produce lives here, in the single place the likelihood is accumulated, so `val` *is* the +# increment: "the produced weight equals the accumulator's increment" becomes structural rather than +# an invariant to keep two call sites in step with. Both routes reach the accumulator through exactly +# one `acclogp` -- an `observe` via `accumulate_observe!!`, and `@addlogprob!` via +# `accloglikelihood!!` -> `map_accumulator!!` -- so each emits exactly one weight, which is what +# gets `@addlogprob!` terms into the weight as well as the accumulator (issue #1996). +# +# Accumulator merging cannot fire a spurious produce: `combine` adds the two `logp`s directly rather +# than going through `acclogp`, so a submodel's varinfo folding into its parent's stays silent. +# +# `produce` suspends the task before the caller assigns the updated varinfo back onto the particle, so +# a *suspended* particle's accumulated total lags this term. Nothing reads it in that state -- the +# sweep reweights from `logweight`, and every varinfo read (`pg_transition_and_state`, +# `gibbs_get_raw_values`, SMC's bundling) happens once the model has run to completion. +function DynamicPPL.acclogp(acc::ProduceLogLikelihoodAccumulator, val) + Libtask.produce(val) + return ProduceLogLikelihoodAccumulator(DynamicPPL.logp(acc) + val) +end function DynamicPPL.accumulate_assume!!( acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, dist, template @@ -288,49 +295,12 @@ function DynamicPPL.accumulate_observe!!( return DynamicPPL.acclogp(acc, Distributions.loglikelihood(dist, left)) end -# `@addlogprob!` bypasses `tilde_observe!!`, so its produce is emitted here instead -- again -# only once the accumulator has been updated. Gated on the producing accumulator, so outside -# particle evaluation this reduces to the default (non-producing) method (issue #1996). -# -# This is type piracy -- both `accloglikelihood!!` and `OnlyAccsVarInfo` belong to DynamicPPL, -# and this method shadows DynamicPPL's own for every `OnlyAccsVarInfo`, which is what all -# samplers use. It is deliberate for want of an alternative: the produce has to happen once the -# accumulator has been updated, and no Turing-owned type appears anywhere in the signature. A -# DynamicPPL-side post-accumulate hook would let this go away. +# Tell Libtask which calls may contain a `produce`, so it instruments them. The produce itself is in +# `acclogp`; everything else here is marked because it sits on a path that reaches it. Over- +# approximating is safe (a wrongly-marked call is merely instrumented); missing a real one is not. # -# Everything except the produce is delegated to the general method via `@invoke`, so that this -# cannot silently drift from upstream: re-implementing that body here would leave every -# `OnlyAccsVarInfo` in the ecosystem running Turing's copy, while only particle sampling is -# covered by these tests. -function DynamicPPL.accloglikelihood!!( - vi::DynamicPPL.OnlyAccsVarInfo, logp; ignore_missing_accumulator=false -) - acc_name = Val(:LogLikelihood) - is_particle = - DynamicPPL.hasacc(vi, acc_name) && - DynamicPPL.getacc(vi, acc_name) isa ProduceLogLikelihoodAccumulator - if !is_particle - return @invoke DynamicPPL.accloglikelihood!!( - vi::DynamicPPL.AbstractVarInfo, logp; ignore_missing_accumulator - ) - end - before = DynamicPPL.getloglikelihood(vi) - vi = @invoke DynamicPPL.accloglikelihood!!( - vi::DynamicPPL.AbstractVarInfo, logp; ignore_missing_accumulator - ) - particle = Libtask.get_taped_globals(Particle) - particle.varinfo = vi - Libtask.produce(DynamicPPL.getloglikelihood(vi) - before) - return vi -end - -# Tell Libtask which calls may contain a `produce`, so it instruments them. The produce lives -# in `tilde_observe!!` and `accloglikelihood!!`; the rest of each chain is marked so Libtask -# tapes through to reach it. Over-approximating is safe (a wrongly-marked call just gets -# instrumented); missing a real one is not, so we err towards marking. -# -# observe: tilde_observe!! accumulates (accumulate_observe!! -> acclogp), then produces -# @addlogprob!: accloglikelihood!! accumulates (map_accumulator!! -> acclogp), then produces +# observe: tilde_observe!! -> accumulate_observe!! -> acclogp +# @addlogprob!: accloglikelihood!! -> map_accumulator!! -> acclogp # (the `@addlogprob! (; ...)` NamedTuple form routes through acclogp!! first) # Gibbs: GibbsContext turns a tilde_assume!! into a tilde_observe!! Libtask.@might_produce(DynamicPPL.tilde_observe!!) @@ -393,17 +363,19 @@ function resample_indices(rng::AbstractRNG, ::MultinomialResampler, weights, n:: return rand(rng, Distributions.Categorical(weights), n) end -"Stratified resampling: one independent uniform per stratum of width `1/n`." -struct StratifiedResampler <: AbstractResampler end -function resample_indices(rng::AbstractRNG, ::StratifiedResampler, weights, n::Integer) +# Stratified and systematic resampling are the same walk up the cumulative weights, differing only +# in where each stratum's offset comes from, so `offset(k)` supplies it. Both schemes draw their +# uniforms in the same order as a hand-written loop would -- note `rand(rng, n)` would *not* be +# equivalent, since Julia fills arrays through a SIMD path that yields a different stream. +function inverse_cdf_indices(weights, n::Integer, offset) v = n * weights[1] indices = Vector{Int}(undef, n) s = 1 for k in 1:n - u = oftype(v, (k - 1) + rand(rng)) - # `s < length(weights)` guards the last particle: if `weights` sums to slightly under - # one (softmax rounding), `v` can fall a hair short of `u` at the final stratum and the - # unguarded loop would index past the end. + u = oftype(v, offset(k)) + # `s < length(weights)` guards the last particle: if `weights` sums to slightly under one + # (softmax rounding), `v` can fall a hair short of `u` at the final stratum and the + # unguarded walk would index past the end. while s < length(weights) && v < u s += 1 v += n * weights[s] @@ -413,24 +385,17 @@ function resample_indices(rng::AbstractRNG, ::StratifiedResampler, weights, n::I return indices end +"Stratified resampling: one independent uniform per stratum of width `1/n`." +struct StratifiedResampler <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::StratifiedResampler, weights, n::Integer) + return inverse_cdf_indices(weights, n, k -> (k - 1) + rand(rng)) +end + "Systematic resampling: one shared uniform placed on a regular grid of `n` points." struct SystematicResampler <: AbstractResampler end function resample_indices(rng::AbstractRNG, ::SystematicResampler, weights, n::Integer) - v = n * weights[1] - u = oftype(v, rand(rng)) - indices = Vector{Int}(undef, n) - s = 1 - for k in 1:n - # See `StratifiedResampler`: `s < length(weights)` keeps the final stratum from - # indexing past the end when `weights` sums to slightly under one. - while s < length(weights) && v < u - s += 1 - v += n * weights[s] - end - indices[k] = s - u += one(u) - end - return indices + u = rand(rng) + return inverse_cdf_indices(weights, n, k -> (k - 1) + u) end # ── Effective-sample-size gating ────────────────────────────────────────────── @@ -514,10 +479,7 @@ function reweight!(particles, multithreaded::Bool) end n_done = count(finished) else - n_done = 0 - for p in particles - n_done += advance_particle!(p) - end + n_done = count(advance_particle!, particles) end n_done == 0 && return false n_done == n && return true @@ -528,10 +490,16 @@ end # ── Resample and propagate ──────────────────────────────────────────────────── -# Resample (if the scheme calls for it) and propagate the survivors, or -- when not -# resampling -- refresh each ordinary particle's seed so the next step draws fresh randomness. -function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional::Bool) +# Resample (if the scheme calls for it) and propagate the survivors, or -- when not resampling -- +# refresh each ordinary particle's seed so the next step draws fresh randomness. Returns whether it +# resampled, which tells `sweep!` what the total weight now is without recomputing it. +# +# Whether this is a conditional sweep is read off the particles rather than passed in: the reference +# always occupies the last slot, so `isreference` is the single source of truth and resampling cannot +# disagree with the rest of the sweep about which particle is pinned. +function resample_propagate!(rng::AbstractRNG, particles, resampler) n = length(particles) + conditional = isreference(last(particles)) weights = normalized_weights(particles) if should_resample(resampler, weights) # A conditional sweep draws the `n-1` free ancestors independently from the categorical @@ -556,34 +524,39 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler, conditional end # reference retained, weight reset conditional && (particles[n].logweight = zero(DynamicPPL.LogProbType)) + return true else # The reference draws nothing (it reuses retained values), so only the others need a # fresh seed for the next step. for p in particles isreference(p) || refresh!(p.rng) end + return false end - return nothing end # ── One sweep ───────────────────────────────────────────────────────────────── # Run a full particle sweep in place, returning the log-evidence estimate and the # per-observation effective sample sizes. -function sweep!( - rng::AbstractRNG, particles, resampler, multithreaded::Bool; conditional::Bool=false -) +function sweep!(rng::AbstractRNG, particles, resampler, multithreaded::Bool) logZ = zero(DynamicPPL.LogProbType) # The ESS values are computed from the particle weights, so they follow whatever # `DynamicPPL.LogProbType` is rather than being pinned to `Float64`. ess_per_step = DynamicPPL.LogProbType[] + # Total log weight entering the step. Resampling zeroes every weight, so it is then exactly + # `log(n)`; otherwise the weights are untouched and it is still last step's total. Either way + # there is nothing to recompute -- particles start at weight zero, hence `log(n)` initially. + logZ0 = log(oftype(logZ, length(particles))) while true - resample_propagate!(rng, particles, resampler, conditional) - logZ0 = log_normalizing_constant(particles) + resampled = resample_propagate!(rng, particles, resampler) + resampled && (logZ0 = log(oftype(logZ, length(particles)))) done = reweight!(particles, multithreaded) # Each observation contributes the log-ratio of total weight it adds; summed over the # sweep these telescope into an estimate of the model's log-evidence log p(y). - logZ += log_normalizing_constant(particles) - logZ0 + total = log_normalizing_constant(particles) + logZ += total - logZ0 + logZ0 = total done && break # Post-reweight ESS for this observation: a degeneracy diagnostic (low ESS means few # particles carry the weight). After the break, so the finishing pass -- which adds no @@ -725,10 +698,9 @@ end PG(n, [scheme = StratifiedResampler(), ]threshold; multithreaded = false) Particle Gibbs sampler with `n` particles. By default resampling is triggered whenever the -effective sample size drops below half the number of particles. The selected scheme is used -for the unconditional first sweep; conditional sweeps draw their ancestors from the categorical -over the weights, because the conditional version of stratified or systematic resampling is a -different algorithm rather than the same draw with one output pinned. +effective sample size drops below half the number of particles. The selected scheme applies to the +unconditional first sweep only; conditional sweeps draw their ancestors from the categorical over +the weights, for the reason given in the resampling-schemes section of this file. Set `multithreaded = true` to evaluate the particles across threads within each sweep; results are unchanged (start Julia with multiple threads, e.g. `julia -t auto`, for this to have effect). @@ -780,22 +752,16 @@ function AbstractMCMC.step( ) error_if_threadsafe_eval(model) n = sampler.nparticles - # The reference reproduces the retained trajectory by reusing its values (passed here and - # consumed by `tilde_assume!!`), so it stays that trajectory even if the model was - # re-conditioned since the last sweep. Its varinfo starts empty like any other particle. - # - # Every reference draw is supplied by value, so its own generator is never read -- only its - # forks' are, and `reseed!` gives those fresh seeds. It therefore carries the retained - # generator forward rather than taking a fresh one: that way the sweep draws nothing from - # `rng` on the reference's behalf, keeping the sampler's stream independent of how the - # reference happens to be built. The copy just avoids aliasing `state`. + # Passing `state` makes this the reference, pinned to the retained trajectory by value (see + # `reference_values`). Its own generator is never read, since every draw is supplied by value -- + # only its forks' are, and `reseed!` gives those fresh seeds. So it carries the retained + # generator forward rather than taking a fresh one, which keeps the sweep from drawing anything + # from `rng` on its behalf; the copy just avoids aliasing `state`. reference = Particle(model, particle_varinfo(), deepcopy(state.rng), state) # `n - 1` fresh particles, with the reference last -- the slot `resample_propagate!` retains. particles = [Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:(n - 1)] push!(particles, reference) - logZ, _ = sweep!( - rng, particles, sampler.resampler, sampler.multithreaded; conditional=true - ) + logZ, _ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) end From bbfc692c2edc6345e3b915259f586260c2d5670d Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 21:26:03 +0100 Subject: [PATCH 52/58] Tidy the particle samplers and their tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A pass over the branch for duplication and derivable state. None of it changes behaviour: draws are bit-identical to the previous commit across SMC, PG under two resampling schemes, and Gibbs(CSMC + HMC), checked before and after each source edit. In the sampler: - Stratified and systematic resampling were the same walk up the cumulative weights, differing only in where each stratum's offset comes from. They now share `inverse_cdf_indices`, and the guard against `weights` summing to slightly under one is stated once instead of twice. The offsets stay scalar `rand` draws inside the loop -- `rand(rng, n)` would read better but Julia fills arrays through a SIMD path that yields a different stream, silently changing every result. - `resample_propagate!` no longer takes `conditional`, nor `sweep!` a keyword to thread it: the reference occupies the last slot, so `isreference(last(particles))` is the single source of truth, which is what `isreference` exists for. - `reference_values` and `expected_reference_varnames` had to be set and cleared together, with only a comment enforcing it. They are now one `reference::Union{Nothing,@NamedTuple{…}}` field, so a half-specified reference is unrepresentable and `reseed!` is a single assignment. - `Particle` built its own varinfo at all 23 call sites, which could take only one value, via a `deepcopy` of a freshly allocated object. It now calls `particle_varinfo()` itself. - `sweep!` recomputed the entering total weight every step. Resampling zeroes every weight, so it is then exactly `log(n)`; otherwise the weights are untouched and it is still last step's total. - `ess_per_step` is now behind `ess=true`, which only `SMC` passes. `PG` discarded it on every sweep -- a gather, a softmax, a `sum(abs2)` and a `push!` per observation, thrown away thousands of times per chain. In the tests, `coinflip` was defined five times, `test()` three, `normal()` twice and the threadsafe model twice; those move to module scope with a shared observation vector, and eight copies of `while advance!(p) !== nothing end` become `run_to_end!`. Two of these were not mechanical: the reference-replay testset's `normal()` is centred at zero rather than four, so it became `centred_normal()` rather than being folded into the others; and the CSMC-consistency test was reimplementing `pg_transition_and_state`'s ancestor draw, so it would have kept passing if the real selection rule drifted -- it now calls the sampler's own function. Deliberately not done: hoisting the per-resample scratch buffers and the weight vector into `sweep!`. Measured, `fork`'s `deepcopy` is about 2.3 MB per sweep step against 2.7 kB for all the weight machinery combined, so that is a 0.1% change bought with persistent mutable state threaded through two functions. Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 120 ++++++++-------- test/mcmc/particle_mcmc.jl | 268 ++++++++++++++--------------------- test/test_utils/exact_ssm.jl | 25 ++-- 3 files changed, 173 insertions(+), 240 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 9173b2cbef..d88b5ce564 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -83,12 +83,12 @@ struct SMCContext <: DynamicPPL.AbstractContext end DynamicPPL.get_param_eltype(::DynamicPPL.AbstractVarInfo, ::SMCContext) = Any """ - Particle(model, varinfo, rng) - Particle(model, varinfo, rng, retained::Particle) + Particle(model, rng) + Particle(model, rng, retained::Particle) -A single particle: a suspended `model` execution together with its `varinfo`, its own `rng`, and -an accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there -is no separate state struct). +A single particle: a suspended `model` execution together with its `varinfo`, its own `rng`, and an +accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there is no +separate state struct). Without `retained` the particle draws from the prior. Given the previous sweep's retained particle it becomes a conditional-SMC reference pinned to that trajectory, erroring if its execution reaches an @@ -104,26 +104,26 @@ mutable struct Particle{RT<:AbstractRNG,WT<:Real} # `logweight` tracks whatever `DynamicPPL.LogProbType` is, so weights follow suit if it # is ever changed. logweight::WT - # The retained trajectory's values, which the CSMC reference reproduces by reusing them - # (`InitFromParams` in `tilde_assume!!`); empty for ordinary particles. `reseed!` clears it so - # a particle forked off the reference samples fresh beyond the fork point. Reusing the *value* - # rather than replaying the RNG draw is what keeps the reference on the retained trajectory - # when Gibbs re-conditions the model: a draw is x = g(u; θ) in the RNG output u and the - # distribution parameters θ, so replaying u after θ → θ' yields g(u; θ') ≠ x -- e.g. - # x ~ Normal(μ, 1) is μ + Φ⁻¹(u), which shifts by μ' − μ. - reference_values::DynamicPPL.VarNamedTuple - # `nothing` for an ordinary particle; for a CSMC reference, the addresses the retained - # trajectory assumed. This cannot be read off `reference_values`: a slice assume such as - # `x[1:2] ~ MvNormal(...)` is stored there under the keys `x[1]`, `x[2]` but assumed under the - # single address `x[1:2]`, so comparing against those keys would report a spurious trace - # change; and an empty `reference_values` cannot distinguish a reference with no latents from - # an ordinary particle. Without it, an address the retained trajectory never had would - # silently draw from the prior, corrupting the reference. Being non-`nothing` is also what - # marks a particle as the reference. Set only alongside `reference_values`, by the reference - # constructor below. - expected_reference_varnames::Union{Nothing,Set{DynamicPPL.VarName}} - # Addresses assumed by this execution, in the same form as the set above. Survives forking, - # so a particle that becomes the retained state hands the complete set to the next + # `nothing` unless this particle is a CSMC reference; one field rather than two so the pair + # cannot get out of step, and so `isreference` has a single thing to test. + # + # `values` is the retained trajectory, which the reference reproduces by reusing it + # (`InitFromParams` in `tilde_assume!!`). Reusing the *value* rather than replaying the RNG draw + # is what keeps the reference on that trajectory when Gibbs re-conditions the model: a draw is + # x = g(u; θ) in the RNG output u and the distribution parameters θ, so replaying u after θ → θ' + # yields g(u; θ') ≠ x -- e.g. x ~ Normal(μ, 1) is μ + Φ⁻¹(u), which shifts by μ' − μ. + # + # `varnames` is the set of addresses the retained trajectory assumed, and cannot be recovered + # from `values`: a slice assume such as `x[1:2] ~ MvNormal(...)` is stored under the keys `x[1]`, + # `x[2]` but assumed under the single address `x[1:2]`, so comparing against those keys would + # report a spurious trace change. Without it an address the retained trajectory never had would + # silently draw from the prior, corrupting the reference. + reference::Union{ + Nothing, + @NamedTuple{values::DynamicPPL.VarNamedTuple, varnames::Set{DynamicPPL.VarName}} + } + # Addresses assumed by this execution, in the same form as `reference.varnames`. Survives + # forking, so a particle that becomes the retained state hands the complete set to the next # reference; a reference must finish having assumed exactly that set. assumed_varnames::Set{DynamicPPL.VarName} task::Libtask.TapedTask @@ -134,24 +134,25 @@ mutable struct Particle{RT<:AbstractRNG,WT<:Real} vi::DynamicPPL.AbstractVarInfo, rng::RT, retained::Union{Nothing,Particle}=nothing ) where {RT<:AbstractRNG} w = zero(DynamicPPL.LogProbType) - values, varnames = if retained === nothing - DynamicPPL.VarNamedTuple(), nothing + reference = if retained === nothing + nothing else - DynamicPPL.get_raw_values(retained.varinfo), copy(retained.assumed_varnames) + (; + values=DynamicPPL.get_raw_values(retained.varinfo), + varnames=copy(retained.assumed_varnames), + ) end - return new{RT,typeof(w)}(vi, rng, w, values, varnames, Set{DynamicPPL.VarName}()) + return new{RT,typeof(w)}(vi, rng, w, reference, Set{DynamicPPL.VarName}()) end end function Particle( - model::DynamicPPL.Model, - varinfo::DynamicPPL.AbstractVarInfo, - rng::AbstractRNG, - retained::Union{Nothing,Particle}=nothing, + model::DynamicPPL.Model, rng::AbstractRNG, retained::Union{Nothing,Particle}=nothing ) model = DynamicPPL.setleafcontext(model, SMCContext()) + varinfo = particle_varinfo() args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - particle = Particle(deepcopy(varinfo), rng, retained) + particle = Particle(varinfo, rng, retained) particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) return particle end @@ -165,8 +166,7 @@ the reference stops reusing retained values and samples afresh. Mutates and retu function reseed!(particle::Particle, rng::AbstractRNG) Random.seed!(particle.rng, rand(rng, UInt64)) # A fork samples fresh from here on, so it must forget the reference's remaining values. - particle.reference_values = DynamicPPL.VarNamedTuple() - particle.expected_reference_varnames = nothing + particle.reference = nothing return particle end @@ -183,7 +183,7 @@ fork(particle::Particle, rng::AbstractRNG) = reseed!(deepcopy(particle), rng) Whether `particle` is a conditional-SMC reference, i.e. pinned to a retained trajectory. Carried on the particle rather than inferred from its slot, so forking and resampling cannot get it wrong. """ -isreference(particle::Particle) = particle.expected_reference_varnames !== nothing +isreference(particle::Particle) = particle.reference !== nothing """ advance!(particle) -> Union{Real,Nothing} @@ -196,9 +196,9 @@ function advance!(particle::Particle) # `tilde_assume!!` already rejects any address outside the retained set, so once the # reference has run to completion the only discrepancy left to catch is a retained address # it never visited (e.g. a branch that stopped being taken after re-conditioning). - expected = particle.expected_reference_varnames - if score === nothing && expected !== nothing - dropped = setdiff(expected, particle.assumed_varnames) + reference = particle.reference + if score === nothing && reference !== nothing + dropped = setdiff(reference.varnames, particle.assumed_varnames) isempty(dropped) || error( "the reference execution trace changed while replaying retained values " * "(retained addresses never reached: $(collect(dropped)))", @@ -211,19 +211,19 @@ function DynamicPPL.tilde_assume!!( ::SMCContext, dist::Distribution, vn::VarName, template, ::DynamicPPL.AbstractVarInfo ) particle = Libtask.get_taped_globals(Particle) - # A reference reuses the retained value at every address it visits (see `reference_values`); + # A reference reuses the retained value at every address it visits (see the `reference` field); # ordinary particles and forks have no expected set, so they draw from the prior. The two error # paths cover the two ways the trace can have moved: an address outside the retained set here, # and -- via the `nothing` fallback -- a retained address with no usable value. - expected = particle.expected_reference_varnames - strategy = if expected === nothing + reference = particle.reference + strategy = if reference === nothing DynamicPPL.InitFromPrior() else - vn in expected || error( + vn in reference.varnames || error( "the reference execution trace changed while replaying retained values " * "(new address: $vn)", ) - DynamicPPL.InitFromParams(particle.reference_values, nothing) + DynamicPPL.InitFromParams(reference.values, nothing) end ctx = DynamicPPL.InitContext(particle.rng, strategy, DynamicPPL.UnlinkAll()) x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, particle.varinfo) @@ -537,9 +537,12 @@ end # ── One sweep ───────────────────────────────────────────────────────────────── -# Run a full particle sweep in place, returning the log-evidence estimate and the -# per-observation effective sample sizes. -function sweep!(rng::AbstractRNG, particles, resampler, multithreaded::Bool) +# Run a full particle sweep in place, returning the log-evidence estimate and -- when `ess` is set -- +# the per-observation effective sample sizes. Only `SMC` reports those, and `PG` runs thousands of +# sweeps, so computing them unconditionally would be pure waste on the sampler that sweeps most. +function sweep!( + rng::AbstractRNG, particles, resampler, multithreaded::Bool; ess::Bool=false +) logZ = zero(DynamicPPL.LogProbType) # The ESS values are computed from the particle weights, so they follow whatever # `DynamicPPL.LogProbType` is rather than being pinned to `Float64`. @@ -561,7 +564,7 @@ function sweep!(rng::AbstractRNG, particles, resampler, multithreaded::Bool) # Post-reweight ESS for this observation: a degeneracy diagnostic (low ESS means few # particles carry the weight). After the break, so the finishing pass -- which adds no # observation and leaves the weights unchanged -- contributes no spurious entry. - push!(ess_per_step, weight_ess(normalized_weights(particles))) + ess && push!(ess_per_step, weight_ess(normalized_weights(particles))) end return logZ, ess_per_step end @@ -641,10 +644,10 @@ function AbstractMCMC.sample( if callback !== nothing @warn "SMC runs one sweep rather than an MCMC loop, so there are no per-iteration callbacks; `callback` is ignored." end - particles = [ - Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:nparticles - ] - logZ, ess_per_step = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) + particles = [Particle(model, particle_rng(rng)) for _ in 1:nparticles] + logZ, ess_per_step = sweep!( + rng, particles, sampler.resampler, sampler.multithreaded; ess=true + ) weights = normalized_weights(particles) # One final resampling step, so the returned particles are an equal-weight sample. The # sweep ends on a reweight, leaving the population weighted; resampling once here makes the @@ -732,10 +735,7 @@ function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, sampler::PG; discard_sample=false, kwargs... ) error_if_threadsafe_eval(model) - particles = [ - Particle(model, particle_varinfo(), particle_rng(rng)) for - _ in 1:(sampler.nparticles) - ] + particles = [Particle(model, particle_rng(rng)) for _ in 1:(sampler.nparticles)] logZ, _ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) end @@ -753,13 +753,13 @@ function AbstractMCMC.step( error_if_threadsafe_eval(model) n = sampler.nparticles # Passing `state` makes this the reference, pinned to the retained trajectory by value (see - # `reference_values`). Its own generator is never read, since every draw is supplied by value -- + # the `reference` field). Its own generator is never read, since every draw is supplied by value -- # only its forks' are, and `reseed!` gives those fresh seeds. So it carries the retained # generator forward rather than taking a fresh one, which keeps the sweep from drawing anything # from `rng` on its behalf; the copy just avoids aliasing `state`. - reference = Particle(model, particle_varinfo(), deepcopy(state.rng), state) + reference = Particle(model, deepcopy(state.rng), state) # `n - 1` fresh particles, with the reference last -- the slot `resample_propagate!` retains. - particles = [Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:(n - 1)] + particles = [Particle(model, particle_rng(rng)) for _ in 1:(n - 1)] push!(particles, reference) logZ, _ = sweep!(rng, particles, sampler.resampler, sampler.multithreaded) return pg_transition_and_state(rng, particles, logZ, discard_sample) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 8925f2a288..15afe6fe8d 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -12,12 +12,12 @@ using Turing.Inference: ESSThresholdResampler, Particle, particle_rng, - particle_varinfo, advance!, fork, sweep!, resample_indices, - normalized_weights + normalized_weights, + pg_transition_and_state using Distributions: Bernoulli, Beta, @@ -38,6 +38,61 @@ using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset using Turing +# Models shared across the testsets below. Defined at module scope, not inside a testset: a `@model` +# sharing a local scope with a same-named variable captures it, and every particle then mutates one +# shared array (see the `xtrue`/`ztrue` note further down). + +@model function coinflip(y) + p ~ Beta(1, 1) + for t in eachindex(y) + y[t] ~ Bernoulli(p) + end +end +const COIN_OBS = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1] + +# `x ~ Bernoulli(1)` pins x = 1, so both observes contribute exactly log(1/2) whatever the +# trajectory -- zero weight variance, which several log_normalizing_constant tests rely on. +@model function test() + a ~ Normal(0, 1) + x ~ Bernoulli(1) + b ~ Gamma(2, 3) + 1 ~ Bernoulli(x / 2) + c ~ Beta() + 0 ~ Bernoulli(x / 2) + return x +end + +@model function normal() + a ~ Normal(4, 5) + 3 ~ Normal(a, 2) + b ~ Normal(a, 1) + 1.5 ~ Normal(b, 2) + return a, b +end + +# Nondeterministic evaluation order, which the particle samplers must refuse. +# As `normal()` but centred at zero; used where the replay test wants a different trajectory. +@model function centred_normal() + a ~ Normal(0, 1) + 3 ~ Normal(a, 2) + b ~ Normal(a, 1) + 1.5 ~ Normal(b, 2) + return a, b +end + +# Nondeterministic evaluation order, which the particle samplers must refuse. +@model function threadsafe_model(y) + x ~ Normal() + Threads.@threads for i in eachindex(y) + y[i] ~ Normal(x) + end +end + +"Run a particle to completion." +run_to_end!(p) = (while advance!(p) !== nothing +end; +p) + @testset "SMC" begin @testset "constructor" begin @test SMC().resampler == ESSThresholdResampler(0.5) @@ -52,24 +107,11 @@ using Turing end @testset "basic model" begin - @model function normal() - a ~ Normal(4, 5) - 3 ~ Normal(a, 2) - b ~ Normal(a, 1) - 1.5 ~ Normal(b, 2) - return a, b - end tested = sample(normal(), SMC(), 100) end @testset "resampling schemes" begin - @model function coinflip(y) - p ~ Beta(1, 1) - for t in eachindex(y) - y[t] ~ Bernoulli(p) - end - end - obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1] + obs = COIN_OBS coin_model = coinflip(obs) prior = extract_priors(coin_model)[@varname(p)] exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs)) @@ -119,16 +161,6 @@ using Turing end @testset "log_normalizing_constant" begin - @model function test() - a ~ Normal(0, 1) - x ~ Bernoulli(1) - b ~ Gamma(2, 3) - 1 ~ Bernoulli(x / 2) - c ~ Beta() - 0 ~ Bernoulli(x / 2) - return x - end - chains_smc = sample(StableRNG(100), test(), SMC(), 100) @test all(isone, chains_smc[:x]) @@ -145,13 +177,7 @@ using Turing @testset "multithreaded execution matches serial" begin # Particles are seeded serially before the parallel reweighting, so `multithreaded=true` # must reproduce the serial draws exactly (bit for bit), whatever the thread count. - @model function coinflip(y) - p ~ Beta(1, 1) - for t in eachindex(y) - y[t] ~ Bernoulli(p) - end - end - model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + model = coinflip(COIN_OBS) serial = sample(StableRNG(23), model, SMC(), 200) multithreaded = sample(StableRNG(23), model, SMC(; multithreaded=true), 200) @test serial[@varname(p)] == multithreaded[@varname(p)] @@ -160,25 +186,11 @@ using Turing @testset "refuses to run threadsafe eval" begin # SMC can't run models that have nondeterministic evaluation order, # so it should refuse to run models marked as threadsafe. - @model function f(y) - x ~ Normal() - Threads.@threads for i in eachindex(y) - y[i] ~ Normal(x) - end - end - model = setthreadsafe(f(randn(10)), true) + model = setthreadsafe(threadsafe_model(randn(10)), true) @test_throws ArgumentError sample(model, SMC(), 100) end @testset "discard_initial, thinning, initial_params and callback are ignored" begin - @model function normal() - a ~ Normal(4, 5) - 3 ~ Normal(a, 2) - b ~ Normal(a, 1) - 1.5 ~ Normal(b, 2) - return a, b - end - @test_logs (:warn, r"initial_params.*ignored") match_mode = :any sample( normal(), SMC(), 10; initial_params=(; a=1.0) ) @@ -238,16 +250,6 @@ end end @testset "log_normalizing_constant" begin - @model function test() - a ~ Normal(0, 1) - x ~ Bernoulli(1) - b ~ Gamma(2, 3) - 1 ~ Bernoulli(x / 2) - c ~ Beta() - 0 ~ Bernoulli(x / 2) - return x - end - chains_pg = sample(StableRNG(468), test(), PG(10), 100) @test all(isone, chains_pg[:x]) @@ -267,13 +269,7 @@ end # Pin the direction and rough size so a future change to the sweep cannot quietly alter it. # # Beta-Bernoulli, so p(y) is exact: with a Beta(1,1) prior, p(y) = B(1+s, 1+n-s)/B(1,1). - @model function coinflip(y) - p ~ Beta(1, 1) - for t in eachindex(y) - y[t] ~ Bernoulli(p) - end - end - obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1] + obs = COIN_OBS s, n = sum(obs), length(obs) exact_logp = logbeta(1 + s, 1 + n - s) - logbeta(1, 1) @@ -294,13 +290,7 @@ end # Threading the reweighting must not perturb the reference-replay bookkeeping, so the # conditional sweeps have to reproduce the serial draws exactly, whatever the thread # count. - @model function coinflip(y) - p ~ Beta(1, 1) - for t in eachindex(y) - y[t] ~ Bernoulli(p) - end - end - model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + model = coinflip(COIN_OBS) serial = sample(StableRNG(23), model, PG(10), 200) multithreaded = sample(StableRNG(23), model, PG(10; multithreaded=true), 200) @test serial[@varname(p)] == multithreaded[@varname(p)] @@ -319,13 +309,10 @@ end model = drifting([0.3, -0.7, 1.1]) function conditional_sweep(scheme) rng = StableRNG(77) - retained = Particle(model, particle_varinfo(), particle_rng(rng)) - while advance!(retained) !== nothing - end - reference = Particle(model, particle_varinfo(), particle_rng(rng), retained) - particles = [ - Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:4 - ] + retained = Particle(model, particle_rng(rng)) + run_to_end!(retained) + reference = Particle(model, particle_rng(rng), retained) + particles = [Particle(model, particle_rng(rng)) for _ in 1:4] push!(particles, reference) sweep!(StableRNG(78), particles, scheme, false) return map(p -> get_raw_values(p.varinfo), particles) @@ -361,20 +348,17 @@ end # N) and check it reproduces the trajectory we retained. Wrapped in a function to keep # the mutating loop out of test soft scope. function run_csmc(model, N, nsteps, rng) - draw(ps) = ps[rand(rng, Categorical(normalized_weights(ps)))] - particles = [ - Particle(model, particle_varinfo(), particle_rng(rng)) for _ in 1:N - ] + # The sampler's own selection rule, not a reimplementation of it: if `pg_transition_and` + # `_state` ever changes how the retained particle is chosen, this test follows. + draw(ps) = last(pg_transition_and_state(rng, ps, 0.0, true)) + particles = [Particle(model, particle_rng(rng)) for _ in 1:N] sweep!(rng, particles, ESSThresholdResampler(0.5), false) state = draw(particles) allok = true nlatents = 0 for _ in 1:nsteps - ref = Particle(model, particle_varinfo(), particle_rng(rng), state) - parts = [ - Particle(model, particle_varinfo(), particle_rng(rng)) for - _ in 1:(N - 1) - ] + ref = Particle(model, particle_rng(rng), state) + parts = [Particle(model, particle_rng(rng)) for _ in 1:(N - 1)] push!(parts, ref) sweep!(rng, parts, ESSThresholdResampler(0.5), false) allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) @@ -405,20 +389,15 @@ end return y ~ Normal(x, 1) end rng = StableRNG(42) - retained = Particle( - reconditioned(2.0) | (@varname(a) => 0.0), particle_varinfo(), particle_rng(rng) - ) - while advance!(retained) !== nothing - end + retained = Particle(reconditioned(2.0) | (@varname(a) => 0.0), particle_rng(rng)) + run_to_end!(retained) retained_vals = get_raw_values(retained.varinfo) reference = Particle( reconditioned(2.0) | (@varname(a) => 5.0), # x's prior shifted far away - particle_varinfo(), particle_rng(rng), retained, ) - while advance!(reference) !== nothing - end + run_to_end!(reference) @test get_raw_values(reference.varinfo) == retained_vals end @@ -434,14 +413,9 @@ end return y ~ Normal(μ, 1) end rng = StableRNG(91) - retained = Particle( - branch_changes(true, 0.0), particle_varinfo(), particle_rng(rng) - ) - while advance!(retained) !== nothing - end - reference = Particle( - branch_changes(false, 0.0), particle_varinfo(), particle_rng(rng), retained - ) + retained = Particle(branch_changes(true, 0.0), particle_rng(rng)) + run_to_end!(retained) + reference = Particle(branch_changes(false, 0.0), particle_rng(rng), retained) @test_throws "reference execution trace changed" advance!(reference) @model function branch_drops(flag, y) @@ -451,15 +425,11 @@ end end return y ~ Normal(x, 1) end - retained = Particle(branch_drops(true, 0.0), particle_varinfo(), particle_rng(rng)) - while advance!(retained) !== nothing - end - reference = Particle( - branch_drops(false, 0.0), particle_varinfo(), particle_rng(rng), retained - ) + retained = Particle(branch_drops(true, 0.0), particle_rng(rng)) + run_to_end!(retained) + reference = Particle(branch_drops(false, 0.0), particle_rng(rng), retained) @test_throws "reference execution trace changed" begin - while advance!(reference) !== nothing - end + run_to_end!(reference) end end @@ -556,25 +526,13 @@ end @testset "refuses to run threadsafe eval" begin # PG can't run models that have nondeterministic evaluation order, # so it should refuse to run models marked as threadsafe. - @model function f(y) - x ~ Normal() - Threads.@threads for i in eachindex(y) - y[i] ~ Normal(x) - end - end - model = setthreadsafe(f(randn(10)), true) + model = setthreadsafe(threadsafe_model(randn(10)), true) @test_throws ArgumentError sample(model, PG(10), 100) end end @testset "parallel chains (MCMCThreads)" begin - @model function coinflip(y) - p ~ Beta(1, 1) - for t in eachindex(y) - y[t] ~ Bernoulli(p) - end - end - model = coinflip([0, 1, 0, 1, 1, 1, 1, 1, 1, 1]) + model = coinflip(COIN_OBS) # Multiple chains through AbstractMCMC's thread-based ensemble stay reproducible under a # fixed seed (genuinely parallel only when Julia is started with more than one thread). for sampler in (SMC(), PG(10)) @@ -585,19 +543,9 @@ end end @testset "particle container" begin - @model function test() - a ~ Normal(0, 1) - x ~ Bernoulli(1) - b ~ Gamma(2, 3) - 1 ~ Bernoulli(x / 2) - c ~ Beta() - 0 ~ Bernoulli(x / 2) - return x - end - @testset "advance!" begin # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. - particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) + particle = Particle(test(), particle_rng(Xoshiro(23))) @test advance!(particle) ≈ -log(2) @test advance!(particle) ≈ -log(2) # `0 ~ Bernoulli(0.5)` @test advance!(particle) === nothing # model finished @@ -607,9 +555,8 @@ end # A particle advanced without resampling draws from its RNG continuously, so it must # produce exactly the same values and log-likelihood as a plain DynamicPPL evaluation # seeded identically. - particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) - while advance!(particle) !== nothing - end + particle = Particle(test(), particle_rng(Xoshiro(23))) + run_to_end!(particle) accs = DynamicPPL.OnlyAccsVarInfo() accs = DynamicPPL.setacc!!(accs, DynamicPPL.LogLikelihoodAccumulator()) @@ -627,7 +574,7 @@ end end @testset "fork" begin - particle = Particle(test(), particle_varinfo(), particle_rng(Xoshiro(23))) + particle = Particle(test(), particle_rng(Xoshiro(23))) advance!(particle) child = fork(particle, Xoshiro(1)) # Independent continuations: advancing one does not touch the other. @@ -641,23 +588,12 @@ end # generator is never consulted. Pinning that down is what lets the reference be handed an # ordinary generator instead of a replayable one: scrambling its seeds before every step # must not perturb the trajectory it regenerates. - @model function normal() - a ~ Normal(0, 1) - 3 ~ Normal(a, 2) - b ~ Normal(a, 1) - 1.5 ~ Normal(b, 2) - return a, b - end - - retained = Particle(normal(), particle_varinfo(), particle_rng(Xoshiro(23))) - while advance!(retained) !== nothing - end + retained = Particle(centred_normal(), particle_rng(Xoshiro(23))) + run_to_end!(retained) values = get_raw_values(retained.varinfo) scrambler = Xoshiro(99) - reference = Particle( - normal(), particle_varinfo(), particle_rng(Xoshiro(7)), retained - ) + reference = Particle(normal(), particle_rng(Xoshiro(7)), retained) while ( Random.seed!(reference.rng, rand(scrambler, UInt64)); advance!(reference) ) !== nothing @@ -821,11 +757,12 @@ end # up directly instead of being absorbed into a mean. post, _ = ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, true_sd)) chn = sample(StableRNG(25), hmm(y, true_sd), PG(32), 4_000) - for t in 1:T, k in 1:K - post[t, k] < 0.02 && continue # too rare to resolve at this chain length - test_within_mc_error( - post[t, k], Float64.(particle_draws(chn, @varname(z[t])) .== k) - ) + for t in 1:T + zs = particle_draws(chn, @varname(z[t])) + for k in 1:K + post[t, k] < 0.02 && continue # too rare to resolve at this chain length + test_within_mc_error(post[t, k], Float64.(zs .== k)) + end end end @@ -845,11 +782,12 @@ end sdd = particle_draws(chn, @varname(sd)) test_within_mc_error(sd_mean, sdd) @test std(sdd) ≈ sd_sd rtol = 0.2 - for t in 1:T, k in 1:K - mixed[t, k] < 0.02 && continue - test_within_mc_error( - mixed[t, k], Float64.(particle_draws(chn, @varname(z[t])) .== k) - ) + for t in 1:T + zs = particle_draws(chn, @varname(z[t])) + for k in 1:K + mixed[t, k] < 0.02 && continue + test_within_mc_error(mixed[t, k], Float64.(zs .== k)) + end end end end diff --git a/test/test_utils/exact_ssm.jl b/test/test_utils/exact_ssm.jl index 33899ee9f8..6154f3a1cf 100644 --- a/test/test_utils/exact_ssm.jl +++ b/test/test_utils/exact_ssm.jl @@ -12,13 +12,7 @@ using Distributions: Categorical, MvNormal, Normal, logpdf using LinearAlgebra: I, Symmetric, diag, eigen using Test: @test, @testset -export lgssm_smoother, - lgssm_loglik, - hmm_forward_backward, - stationary_distribution, - grid_posterior, - grid_moments, - test_exact_ssm_reference +# Everything here is reached qualified (`ExactSSM.foo`), so there is no export list. ## ## Scalar linear Gaussian SSM: x₁ ~ N(0, s0), xₜ = a·xₜ₋₁ + N(0, q), yₜ = xₜ + N(0, r) @@ -93,19 +87,20 @@ function hmm_forward_backward( ) T, K = size(loglik_obs) F = promote_type(eltype(π0), eltype(P), eltype(loglik_obs)) + lik = exp.(loglik_obs) # both passes need it; exponentiate once α = zeros(F, T, K) c = zeros(F, T) # per-step normalisers, which give the log-likelihood - α[1, :] = π0 .* exp.(loglik_obs[1, :]) - c[1] = sum(α[1, :]) + α[1, :] = π0 .* @view lik[1, :] + c[1] = sum(@view α[1, :]) α[1, :] ./= c[1] for t in 2:T - α[t, :] = (P' * α[t - 1, :]) .* exp.(loglik_obs[t, :]) - c[t] = sum(α[t, :]) + α[t, :] = (P' * @view(α[t - 1, :])) .* @view lik[t, :] + c[t] = sum(@view α[t, :]) α[t, :] ./= c[t] end β = ones(F, T, K) for t in (T - 1):-1:1 - β[t, :] = P * (exp.(loglik_obs[t + 1, :]) .* β[t + 1, :]) ./ c[t + 1] + β[t, :] = P * (@view(lik[t + 1, :]) .* @view(β[t + 1, :])) ./ c[t + 1] end post = α .* β return post ./ sum(post; dims=2), sum(log, c) @@ -116,8 +111,7 @@ function hmm_brute_force(π0::AbstractVector, P::AbstractMatrix, loglik_obs::Abs T, K = size(loglik_obs) post = zeros(T, K) total = 0.0 - for idx in 0:(K^T - 1) - z = [(idx ÷ K^(t - 1)) % K + 1 for t in 1:T] + for z in CartesianIndices(ntuple(_ -> K, T)) lp = log(π0[z[1]]) + loglik_obs[1, z[1]] for t in 2:T lp += log(P[z[t - 1], z[t]]) + loglik_obs[t, z[t]] @@ -184,7 +178,8 @@ function test_exact_ssm_reference() @test ll_fb ≈ ll_bf atol = 1e-12 # A stationary π0 is a fixed point of the transition, which several tests rely on. - @test transpose(P) * stationary_distribution(P) ≈ stationary_distribution(P) + π0_stat = stationary_distribution(P) + @test transpose(P) * π0_stat ≈ π0_stat end end From 3f822ffd0a70d767f01d8be7a4917189559c466e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 21:44:44 +0100 Subject: [PATCH 53/58] Derive the fixed-parameter test models, and give the files a header hierarchy Each state space model was written twice, once with its parameter sampled and once with it passed in, so the PG test and the Gibbs test could drift into testing different models. The fixed-parameter tests now `fix` the parameter on the single sampled model instead. `fix` and not `condition`, and the difference is not cosmetic: `fix` substitutes the value without adding a log-density term, whereas `condition` turns the assume into an observe and so adds a produce -- another filtering step in the sweep. Measured on the linear Gaussian model with a fixed seed, the explicitly-fixed model and `fix` give bit-identical draws while `condition` gives different ones (x[1] mean 0.5851 against 0.4925). The comment at the model records this. Separately, the section headers now follow the repo conventions properly. `particle_mcmc.jl` has two depths of structure -- major regions, and sections within them -- but rendered the inner depth as level-3 one-line rules directly beneath level-1 frames, skipping level 2; those seven are now `##` frames. The test file had one header in 780 lines, so its real region boundaries (shared models, SMC, PG, chain-level parallelism, particle mechanics, the SSM checks) now have level-1 headers. Two headers had prose butted against the closing frame, against the blank-line rule. The source change here is comments only, and draws stay bit-identical. Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 28 ++++++++++++---- test/mcmc/particle_mcmc.jl | 63 ++++++++++++++++++++---------------- test/test_utils/exact_ssm.jl | 1 + 3 files changed, 57 insertions(+), 35 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index d88b5ce564..1336c2d605 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -345,7 +345,9 @@ end # systematic resampling mixes noticeably better than multinomial in particle Gibbs -- so # implementing the conditional schemes properly would be a genuine improvement. -# ── Resampler interface ─────────────────────────────────────────────────────── +## +## Resampler interface +## abstract type AbstractResampler end @@ -355,7 +357,9 @@ should_resample(::AbstractResampler, weights) = true """Draw `n` ancestor indices from `1:length(weights)` with probabilities `weights`.""" function resample_indices end -# ── Schemes ─────────────────────────────────────────────────────────────────── +## +## Schemes +## "Multinomial resampling: `n` independent draws from the categorical over `weights`." struct MultinomialResampler <: AbstractResampler end @@ -398,7 +402,9 @@ function resample_indices(rng::AbstractRNG, ::SystematicResampler, weights, n::I return inverse_cdf_indices(weights, n, k -> (k - 1) + u) end -# ── Effective-sample-size gating ────────────────────────────────────────────── +## +## Effective-sample-size gating +## """ ESSThresholdResampler(threshold, scheme = StratifiedResampler()) @@ -431,7 +437,9 @@ end # retained trajectory's values, while the other `n-1` slots are resampled from all `n` particles # (so they may descend from the reference). -# ── Weights and diagnostics ─────────────────────────────────────────────────── +## +## Weights and diagnostics +## logweights(particles) = [p.logweight for p in particles] normalized_weights(particles) = softmax(logweights(particles)) @@ -443,7 +451,9 @@ autocorrelation rather than a population's weight degeneracy. """ weight_ess(weights) = inv(sum(abs2, weights)) -# ── Reweighting ─────────────────────────────────────────────────────────────── +## +## Reweighting +## # Advance one particle by one observation, folding its incremental weight in; return `true` # once it has finished (produced nothing). Factored out so the serial and multithreaded loops in @@ -488,7 +498,9 @@ function reweight!(particles, multithreaded::Bool) ) end -# ── Resample and propagate ──────────────────────────────────────────────────── +## +## Resample and propagate +## # Resample (if the scheme calls for it) and propagate the survivors, or -- when not resampling -- # refresh each ordinary particle's seed so the next step draws fresh randomness. Returns whether it @@ -535,7 +547,9 @@ function resample_propagate!(rng::AbstractRNG, particles, resampler) end end -# ── One sweep ───────────────────────────────────────────────────────────────── +## +## One sweep +## # Run a full particle sweep in place, returning the log-evidence estimate and -- when `ess` is set -- # the per-observation effective sample sizes. Only `SMC` reports those, and `PG` runs thousands of diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 15afe6fe8d..b13c33c5ea 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -38,6 +38,10 @@ using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset using Turing +# +# Shared models +# + # Models shared across the testsets below. Defined at module scope, not inside a testset: a `@model` # sharing a local scope with a same-named variable captures it, and every particle then mutates one # shared array (see the `xtrue`/`ztrue` note further down). @@ -93,6 +97,10 @@ run_to_end!(p) = (while advance!(p) !== nothing end; p) +# +# SMC +# + @testset "SMC" begin @testset "constructor" begin @test SMC().resampler == ESSThresholdResampler(0.5) @@ -228,6 +236,10 @@ p) end end +# +# PG / conditional SMC +# + @testset "PG" begin @testset "constructor" begin @test PG(10).nparticles == 10 @@ -531,6 +543,10 @@ end end end +# +# Chain-level parallelism +# + @testset "parallel chains (MCMCThreads)" begin model = coinflip(COIN_OBS) # Multiple chains through AbstractMCMC's thread-based ensemble stay reproducible under a @@ -542,6 +558,10 @@ end end end +# +# Particle mechanics +# + @testset "particle container" begin @testset "advance!" begin # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. @@ -605,6 +625,7 @@ end # # State space models with tractable posteriors # + # These are the checks that pin the particle samplers to a known answer rather than to each other. A # scalar linear Gaussian SSM and a discrete HMM both have closed-form posteriors, supplied by # `ExactSSM` and validated there against brute force, so a disagreement beyond Monte Carlo error is a @@ -637,20 +658,14 @@ function test_within_mc_error(exact, samples; nsigma=4) return nothing end -@model function lgssm(y, a, q, r) - # `typeof(q)` stays generic when HMC differentiates through `q`, without boxing every - # element the way `Vector{Real}` would. - x = Vector{typeof(q)}(undef, length(y)) - x[1] ~ Normal(0, sqrt(q / (1 - a^2))) - y[1] ~ Normal(x[1], sqrt(r)) - for t in 2:length(y) - x[t] ~ Normal(a * x[t - 1], sqrt(q)) - y[t] ~ Normal(x[t], sqrt(r)) - end -end - -@model function lgssm_unknown_q(y, a, r) +# One model per SSM, with the parameter always sampled; the fixed-parameter tests `fix` it instead of +# duplicating the body. `fix` substitutes the value without adding a log-density term, so the sweep +# still sees one filtering step per observation -- `condition` would turn the assume into an observe +# and add a produce, which measurably changes the draws. +@model function lgssm(y, a, r) q ~ InverseGamma(3, 2) + # `typeof(q)` stays generic when HMC differentiates through `q`, without boxing every element the + # way `Vector{Real}` would. x = Vector{typeof(q)}(undef, length(y)) x[1] ~ Normal(0, sqrt(q / (1 - a^2))) y[1] ~ Normal(x[1], sqrt(r)) @@ -664,17 +679,7 @@ const HMM_P = [0.80 0.15 0.05; 0.10 0.80 0.10; 0.05 0.15 0.80] const HMM_MEANS = [-1.5, 0.0, 1.5] const HMM_PI0 = ExactSSM.stationary_distribution(HMM_P) -@model function hmm(y, sd) - z = Vector{Int}(undef, length(y)) - z[1] ~ Categorical(HMM_PI0) - y[1] ~ Normal(HMM_MEANS[z[1]], sd) - for t in 2:length(y) - z[t] ~ Categorical(HMM_P[z[t - 1], :]) - y[t] ~ Normal(HMM_MEANS[z[t]], sd) - end -end - -@model function hmm_unknown_sd(y) +@model function hmm(y) sd ~ LogNormal(log(0.7), 0.4) z = Vector{Int}(undef, length(y)) z[1] ~ Categorical(HMM_PI0) @@ -703,7 +708,9 @@ end @testset "PG recovers the exact smoothing marginals" begin means, vars = ExactSSM.lgssm_smoother(y, a, true_q, r, s0(true_q)) - chn = sample(StableRNG(24), lgssm(y, a, true_q, r), PG(32), 4_000) + chn = sample( + StableRNG(24), fix(lgssm(y, a, r), @varname(q) => true_q), PG(32), 4_000 + ) for t in 1:T xs = particle_draws(chn, @varname(x[t])) test_within_mc_error(means[t], xs) @@ -725,7 +732,7 @@ end ) alg = Gibbs(@varname(q) => NUTS(), @varname(x) => CSMC(32)) - chn = sample(StableRNG(31), lgssm_unknown_q(y, a, r), alg, 4_000) + chn = sample(StableRNG(31), lgssm(y, a, r), alg, 4_000) qd = particle_draws(chn, @varname(q)) test_within_mc_error(q_mean, qd) @@ -756,7 +763,7 @@ end # Discrete states make this sharp: the reference is a probability vector, so any bias shows # up directly instead of being absorbed into a mean. post, _ = ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, true_sd)) - chn = sample(StableRNG(25), hmm(y, true_sd), PG(32), 4_000) + chn = sample(StableRNG(25), fix(hmm(y), @varname(sd) => true_sd), PG(32), 4_000) for t in 1:T zs = particle_draws(chn, @varname(z[t])) for k in 1:K @@ -777,7 +784,7 @@ end mixed = sum(w[i] * posts[i] for i in eachindex(w)) alg = Gibbs(@varname(sd) => HMC(0.1, 12), @varname(z) => CSMC(32)) - chn = sample(StableRNG(32), hmm_unknown_sd(y), alg, 4_000) + chn = sample(StableRNG(32), hmm(y), alg, 4_000) sdd = particle_draws(chn, @varname(sd)) test_within_mc_error(sd_mean, sdd) diff --git a/test/test_utils/exact_ssm.jl b/test/test_utils/exact_ssm.jl index 6154f3a1cf..138a117dfb 100644 --- a/test/test_utils/exact_ssm.jl +++ b/test/test_utils/exact_ssm.jl @@ -1,6 +1,7 @@ # # Exact inference for the two state space models with tractable posteriors # + # These give samplers something unarguable to be compared against: a Kalman filter and smoother for # a scalar linear Gaussian model, and forward-backward for a discrete HMM. `test_exact_ssm_reference` # validates both against brute force, so the reference is checked in CI rather than trusted. From 5cd20ea6959be680565b6276a3052d0cc0b0d8df Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 22:40:31 +0100 Subject: [PATCH 54/58] Drop the dead normalized_weights import from the particle tests Orphaned when the test switched to calling pg_transition_and_state instead of reimplementing the transition itself; nothing in the file references it now. Co-Authored-By: Claude Code --- test/mcmc/particle_mcmc.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index b13c33c5ea..5b5bd87bbd 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -16,7 +16,6 @@ using Turing.Inference: fork, sweep!, resample_indices, - normalized_weights, pg_transition_and_state using Distributions: Bernoulli, From 9dce0035cdac1c8c7031dff852461f890d267f74 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 22:40:31 +0100 Subject: [PATCH 55/58] Correct two comments that misdescribe how particles carry state particle_varinfo's comment claimed the prior and Jacobian terms are recomputed downstream from the raw values, so accumulating them per particle would be wasted work. Both halves are wrong: OnlyAccsVarInfo ships LogPrior and LogJacobian by default, so every particle does accumulate them, and ParamsWithStats reads them straight off this varinfo to fill a chain's logprior and logjoint columns. Because that read is guarded by hasacc, acting on the old comment and dropping the accumulators would not error -- it would silently omit those columns, and fail test_chain_logp_metadata. The PG step comment still described the reference as regenerating its trajectory by replaying state.rng from the first step. That is the seed-replay scheme removed in 6c8247d29; the reference reuses the retained values, which the comment a dozen lines below it already explains. Co-Authored-By: Claude Code --- src/mcmc/particle_mcmc.jl | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index 1336c2d605..ea46a6c0c4 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -314,9 +314,11 @@ Libtask.@might_produce(DynamicPPL.acclogp!!) # See https://github.com/TuringLang/Libtask.jl/issues/217. Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true -# A particle needs only the produce-aware likelihood accumulator (which drives reweighting) -# and the raw sampled values. The prior/Jacobian terms shown in chain metadata are recomputed -# downstream from the raw values, so accumulating them per particle would be wasted work. +# Swap the default likelihood accumulator for the produce-aware one that drives reweighting, and +# add the raw sampled values. `OnlyAccsVarInfo`'s defaults also bring `LogPrior` and `LogJacobian` +# along, and they are kept on purpose: `ParamsWithStats` reads them straight off this varinfo to +# fill a chain's `logprior` and `logjoint` columns. Dropping them to save the per-particle logpdf +# work would not error -- the read is guarded by `hasacc` -- it would silently omit those columns. function particle_varinfo() vi = DynamicPPL.OnlyAccsVarInfo() vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator()) @@ -755,7 +757,7 @@ function AbstractMCMC.step( end # Subsequent iterations: conditional SMC given the retained trajectory, which the reference -# particle regenerates by replaying `state.rng` from the first step. +# particle reproduces by reusing the retained values (see the `reference` field). function AbstractMCMC.step( rng::AbstractRNG, model::DynamicPPL.Model, From e4628df9499bcf8a70c17dfbc0c818dcc17a935a Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 23:17:49 +0100 Subject: [PATCH 56/58] Cover latents whose dimension varies between executions The existing value-replay tests cover a trace that changed and must be rejected. Nothing covered one that is legitimately different on every execution and must simply work, which is the property particle samplers uniquely support: k[t] decides how many jumps step t has, so the reference must reuse k[t] before it can reuse a jump vector of the matching length. The target is exact and needs no reference implementation, because tilting k ~ Poisson(1) by c^k is exactly Poisson(c). The tilt is the only informative term and is a function of the trace's shape, so a reference that replayed or reweighted the varying-length part wrongly would move E[k[t]] off c. Measured: with the tilt, E[k[t]] is 2.00-2.07 across the four steps; dropping it lands on the prior mean 1.0, roughly eight standard errors outside the tolerance, so the assertion separates the two rather than passing either way. Costs about 28s in the PG block. Co-Authored-By: Claude Code --- test/mcmc/particle_mcmc.jl | 47 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 5b5bd87bbd..b9205f6d2f 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -21,13 +21,16 @@ using Distributions: Bernoulli, Beta, Categorical, + Exponential, Gamma, InverseGamma, LogNormal, MvNormal, Normal, + Poisson, Uniform, logpdf, + product_distribution, sample using FlexiChains: VNChain, has_same_data using LinearAlgebra: I @@ -459,6 +462,50 @@ end @test size(chn, 1) == 20 end + @testset "latents whose dimension varies between executions" begin + # The two testsets above cover traces that *changed* and must be rejected. This covers one + # that is legitimately different on every execution and must simply work: `k[t]` decides how + # many jumps step `t` has, so the reference must reuse `k[t]` before it can reuse a jump + # vector of the matching length. + # + # The target is exact without a reference implementation, and still informative about the + # varying dimension. Tilting `k[t] ~ Poisson(1)` by `c^k[t]` gives exactly `Poisson(c)`, + # since `e⁻¹c^k/k!` normalises to `e⁻ᶜc^k/k!`. The tilt is the only term that carries + # information, and it is a function of the trace's shape, so a reference that replayed or + # reweighted the varying-length part wrongly would move `E[k[t]]` off `c`. The observation + # itself ignores the latents; it is there to give the sweep its produce points. + tilt = 2.0 + @model function random_dimension(y, c) + k = Vector{Int}(undef, length(y)) + jumps = Vector{Vector{Float64}}(undef, length(y)) + for t in eachindex(y) + k[t] ~ Poisson(1.0) + if k[t] > 0 + jumps[t] ~ product_distribution(fill(Exponential(1.0), k[t])) + else + # A zero-length `product_distribution` is not usable, and `k[t] = 0` has + # probability e⁻¹, so this branch is taken constantly. + jumps[t] = Float64[] + end + @addlogprob! k[t] * log(c) + y[t] ~ Normal(0.0, 1.0) + end + end + + ndraws = 2_000 + chn = sample(StableRNG(106), random_dimension(zeros(4), tilt), PG(16), ndraws) + @test size(chn, 1) == ndraws + ks = reduce(vcat, (reshape(collect(k), 1, :) for k in collect(chn[@varname(k)]))) + @test size(ks) == (ndraws, 4) + # Four standard errors of a Poisson(`tilt`) mean over `ndraws` draws, which also leaves room + # for the autocorrelation a Markov chain carries. + tol = 4 * sqrt(tilt / ndraws) + for t in 1:4 + @test mean(@view ks[:, t]) ≈ tilt atol = tol + @test mean(==(0), @view ks[:, t]) ≈ exp(-tilt) atol = tol + end + end + @testset "addlogprob leads to reweighting" begin # Make sure that PG takes @addlogprob! into account. It didn't use to: # https://github.com/TuringLang/Turing.jl/issues/1996 From 65350945a69d67b3dab5f7bcaab51934aa037e1e Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Wed, 29 Jul 2026 23:23:04 +0100 Subject: [PATCH 57/58] Widen the varying-dimension tolerance to allow for chain autocorrelation The tolerance was eight iid standard errors' worth of slack expressed as four, because sqrt(tilt/ndraws) assumes independent draws and a PG chain is not independent. Measured across eight seeds, the batch-means standard error of E[k[t]] is about 2.3x the iid figure, which left the old bound only ~1.75 real standard errors wide; two of the eight seeds came within 13% of failing. Widening costs nothing in sensitivity. The failure mode being guarded against -- dropping the weight that depends on the trace's shape -- puts E[k[t]] at the prior mean 1.0, still four times the widened tolerance away from the target. Co-Authored-By: Claude Code --- test/mcmc/particle_mcmc.jl | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index b9205f6d2f..19a932e85e 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -497,9 +497,13 @@ end @test size(chn, 1) == ndraws ks = reduce(vcat, (reshape(collect(k), 1, :) for k in collect(chn[@varname(k)]))) @test size(ks) == (ndraws, 4) - # Four standard errors of a Poisson(`tilt`) mean over `ndraws` draws, which also leaves room - # for the autocorrelation a Markov chain carries. - tol = 4 * sqrt(tilt / ndraws) + # `sqrt(tilt / ndraws)` is the standard error this mean would have from independent draws. + # A PG chain is autocorrelated, and measured across eight seeds its batch-means standard + # error runs about 2.3x that, so eight iid errors is roughly 3.5 real ones. The margin is + # there to keep the test from flaking, and it costs nothing here: the failure mode this + # guards against -- dropping the weight that depends on the trace's shape -- lands on the + # prior mean 1.0, four times the tolerance away. + tol = 8 * sqrt(tilt / ndraws) for t in 1:4 @test mean(@view ks[:, t]) ≈ tilt atol = tol @test mean(==(0), @view ks[:, t]) ≈ exp(-tilt) atol = tol From 8663adff308a40818dae733b23759fe51c6eb471 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Thu, 30 Jul 2026 00:01:29 +0100 Subject: [PATCH 58/58] Fold the normalizing-constant caveat into its bullet, for the formatter HISTORY.md was failing the format check. With format_markdown = true, a list item containing a second paragraph makes JuliaFormatter rewrite the list: it indents the separating blank line to four spaces, adding trailing whitespace, and spaces out the sibling items to match. `format(".")` therefore reported a diff on a clean tree, which the Format workflow would flag. Merging the caveat into the bullet it belongs to removes the multi-paragraph item and leaves `format(".")` clean and idempotent. It also removes an ambiguity worth fixing on its own: "For SMC this is an unbiased estimate" had ess_per_step as its nearest antecedent, when what is unbiased is the normalizing constant, which the sentence now names. Co-Authored-By: Claude Code --- HISTORY.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index eb0b1ec2e5..d5b63c0968 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -17,9 +17,7 @@ The rewrite also brings: - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. - - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). - - For `SMC` this is an unbiased estimate of the marginal likelihood `p(y)`. For `PG` / `CSMC` it is **not**, and must not be used for model comparison; see the `PG` docstring. + - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). For `SMC` the normalizing constant is an unbiased estimate of the marginal likelihood `p(y)`; for `PG` / `CSMC` it is **not**, and must not be used for model comparison — see the `PG` docstring. # 0.46.0