diff --git a/AGENTS.md b/AGENTS.md index 646de380df..a78cccf4d1 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), 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. diff --git a/HISTORY.md b/HISTORY.md index 861dc78de1..d325001698 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,31 @@ +# 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 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. +The selected scheme applies to unconditional sweeps only; `PG` / `CSMC` draw a conditional sweep's ancestors from the categorical over the weights. +Exact draws may therefore differ from previous releases, but remain statistically consistent (the same target distribution). + +Chain statistics have changed: `chain[:logevidence]` is now `chain[:log_normalizing_constant]`, and SMC's per-particle `weight` is gone, since the returned particles are now equal-weight. + +`SMC` no longer has a sampler state, because it runs one sweep rather than an MCMC loop: `save_state` and `initial_state` are now warned about and ignored, so `loadstate` has nothing to return for an `SMC` chain. `PG` / `CSMC` are unaffected. + +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. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). + - Parallelism within a sweep. `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads without changing the results; start Julia with multiple threads (e.g. `julia -t auto`) for this to take effect. It is independent of `MCMCThreads()` / `MCMCDistributed()`, which parallelise whole chains and work with SMC/PG as with any other sampler. + - Equal-weight draws. `SMC` resamples once at the end of the sweep, so `mean(chain[...])` and other summaries need no weighting. + - A degeneracy diagnostic. `SMC` chains carry `ess_per_step`, the effective sample size after each filtering step; one entry per likelihood term, so an `@addlogprob!` adds one alongside the observations. MCMCChains exposes the entries as `ess_per_step[1]`, `ess_per_step[2]`, and so on. + - For `SMC`, `exp(log_normalizing_constant)` is an unbiased estimator of the marginal likelihood `p(y)` under the usual particle-filter assumptions. For `PG` / `CSMC` it is biased and must not be used for model comparison (see the `PG` docstring). + # 0.46.1 Fixed a bug, present since v0.41.0, that biased `PG` / `CSMC` posteriors, whether sampled on their own or as a Gibbs component. diff --git a/Project.toml b/Project.toml index 0e09c7ad6c..c113c5bbb9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.46.1" +version = "0.47.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -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/ext/TuringMCMCChainsExt.jl b/ext/TuringMCMCChainsExt.jl index 18416c433f..5b326d8868 100644 --- a/ext/TuringMCMCChainsExt.jl +++ b/ext/TuringMCMCChainsExt.jl @@ -2,10 +2,10 @@ module TuringMCMCChainsExt using Turing using Turing: AbstractMCMC, DynamicPPL -using Turing.Inference: HMC, NUTS, HMCDA, Emcee, EmceeState, _get_n_walkers +using Turing.Inference: HMC, NUTS, HMCDA, Emcee, EmceeState, SMC, _get_n_walkers using MCMCChains: MCMCChains -import Turing.Inference: post_sample_hook +import Turing.Inference: bundle_smc_samples, post_sample_hook """ loadstate(chain::MCMCChains.Chains) @@ -44,6 +44,28 @@ function AbstractMCMC.bundle_samples( return AbstractMCMC.chainscat(chains...) end +function bundle_smc_samples( + transitions::Vector{<:DynamicPPL.ParamsWithStats}, + model::DynamicPPL.Model, + sampler::SMC, + state, + ::Type{MCMCChains.Chains}; + kwargs..., +) + # MCMCChains drops non-scalar statistics, so give each filtering step its own internal. + ess_per_step = first(transitions).stats.ess_per_step + ess_names = ntuple(i -> Symbol("ess_per_step[$i]"), length(ess_per_step)) + ess_stats = NamedTuple{ess_names}(Tuple(ess_per_step)) + scalar_transitions = map(transitions) do transition + names = filter(!=(:ess_per_step), keys(transition.stats)) + stats = merge(NamedTuple{names}(transition.stats), ess_stats) + DynamicPPL.ParamsWithStats(transition.params, stats) + end + return AbstractMCMC.bundle_samples( + scalar_transitions, model, sampler, state, MCMCChains.Chains; kwargs... + ) +end + """ post_sample_hook(chain::MCMCChains.Chains, sampler::Union{HMC,NUTS,HMCDA}; kwargs...) diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl index b9aa08e8f6..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 AdvancedPS import EllipticalSliceSampling import LogDensityProblems import Random diff --git a/src/mcmc/particle_mcmc.jl b/src/mcmc/particle_mcmc.jl index ee4f86c325..926fac6432 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -1,9 +1,43 @@ -### -### Particle Filtering and Particle MCMC Samplers. -### - -using Accessors: Accessors - +# +# 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 rather than replaying its +# random draws (see the `reference` field for why); a fork forgets them in `reseed!`. +# +# 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 + +# Each particle owns a counter-based `Random123.Philox2x` seeded from the sampler's own +# generator. Splitting one generator into many by reseeding can yield correlated streams +# (Steele et al., "Fast Splittable Pseudorandom Number Generators", OOPSLA 2014), and a +# `MersenneTwister` derivation is not stable across Julia versions, which drifted SMC/PG +# results even under a StableRNG (#2781, AdvancedPS.jl#110). Philox is counter-based with a +# fixed algorithm, so the derived streams are both decorrelated and version-stable. +# `Philox2x()` would draw a throwaway seed from the OS, hence the explicit constructor. +"A fresh counter-based generator for one particle, seeded from `rng`." +particle_rng(rng::AbstractRNG) = Random123.Philox2x(UInt64, rand(rng, UInt64)) + +# +# 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( @@ -15,545 +49,849 @@ function error_if_threadsafe_eval(model::DynamicPPL.Model) return nothing end -### AdvancedPS models and interface - -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. -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 -) - model = DynamicPPL.setleafcontext(model, ParticleMCMCContext(rng)) - args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) - fargs = (model.f, args...) - return TracedModel(model, varinfo, resample, fargs, kwargs) -end - -function AdvancedPS.advance!( - trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false -) - # 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 -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. This has to mutate: - # `AdvancedPS.fork` calls it for its side effect and discards the return value. - trace.resample = true - return trace -end +""" + SMCContext -function AdvancedPS.reset_model(trace::TracedModel) - return trace -end +Leaf context marking a model evaluation as a particle-filter step: `tilde_assume!!` draws from +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 -function Libtask.TapedTask(taped_globals, model::TracedModel) - return Libtask.TapedTask( - taped_globals, model.fargs[1], model.fargs[2:end]...; model.kwargs... - ) -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, ::SMCContext) = Any -abstract type ParticleInference <: AbstractSampler end +""" + Particle(model, rng) + Particle(model, rng, reference::DynamicPPL.VarNamedTuple) -#### -#### Generic Sequential Monte Carlo sampler. -#### +A single particle: a suspended `model` execution together with its `varinfo`, its own `rng`, and an +accumulated `logweight`. +Without `reference` the particle draws from the prior. Given a retained trajectory's raw values it +becomes a conditional-SMC reference pinned to that trajectory, erroring if its execution reaches an +address that trajectory lacks or finishes without reaching one it has. """ -$(TYPEDEF) +mutable struct Particle{RT<:AbstractRNG,WT<:Real} + # Abstract on purpose: the varinfo type changes as the execution proceeds, since each new + # address widens the raw-value accumulator's `VarNamedTuple`. Accesses go through Libtask's + # (already type-unstable) taped globals, so this costs nothing extra. + varinfo::DynamicPPL.AbstractVarInfo + rng::RT + # `logweight` tracks whatever `DynamicPPL.LogProbType` is, so weights follow suit if it + # is ever changed. + logweight::WT + # `nothing` unless this particle is a CSMC reference; otherwise the retained trajectory's raw + # values, which it reuses (`InitFromParams` in `tilde_assume!!`). Reusing the *value* rather than + # replaying the RNG draw is what survives Gibbs re-conditioning: a draw is x = g(u; θ), so + # replaying u after θ → θ' gives g(u; θ') ≠ x -- μ + Φ⁻¹(u) shifts by μ' − μ. + reference::Union{Nothing,DynamicPPL.VarNamedTuple} + # Left undefined here and set immediately after construction: the task has to capture the + # particle as its taped globals, a back-reference the particle cannot supply before it exists. + task::Libtask.TapedTask + function Particle( + vi::DynamicPPL.AbstractVarInfo, + rng::RT, + reference::Union{Nothing,DynamicPPL.VarNamedTuple}=nothing, + ) where {RT<:AbstractRNG} + w = zero(DynamicPPL.LogProbType) + return new{RT,typeof(w)}(vi, rng, w, reference) + end +end -Sequential Monte Carlo sampler. +function Particle( + model::DynamicPPL.Model, + rng::AbstractRNG, + reference::Union{Nothing,DynamicPPL.VarNamedTuple}=nothing, +) + model = DynamicPPL.setleafcontext(model, SMCContext()) + varinfo = particle_varinfo() + args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) + particle = Particle(varinfo, rng, reference) + particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) + return particle +end -# Fields +""" + reseed!(particle, rng) -$(TYPEDFIELDS) +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`. """ -struct SMC{R} <: ParticleInference - resampler::R +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 = nothing + return particle end """ - SMC([resampler = AdvancedPS.ResampleWithESSThreshold()]) - SMC([resampler = AdvancedPS.resample_systematic, ]threshold) + fork(particle, rng) -Create a sequential Monte Carlo sampler of type [`SMC`](@ref). +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) -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. """ -SMC() = SMC(AdvancedPS.ResampleWithESSThreshold()) +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.reference !== nothing -# 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) -end +""" + advance!(particle) -> Union{Real,Nothing} -struct SMCState{P,F<:AbstractFloat} - particles::P - particleindex::Int - # The logevidence after aggregating all samples together. - average_logevidence::F +Run the particle to its next `observe`, returning the incremental log-likelihood, or +`nothing` once the model finishes. +""" +function advance!(particle::Particle) + score = Libtask.consume(particle.task) + # `tilde_assume!!` already rejects an address the retained values lack, so the only + # discrepancy left to catch here is a retained address the reference never visited (e.g. a + # branch that stopped being taken after re-conditioning). + reference = particle.reference + if score === nothing && reference !== nothing + dropped = setdiff( + keys(reference), keys(DynamicPPL.get_raw_values(particle.varinfo)) + ) + isempty(dropped) || error( + "the reference execution trace changed while replaying retained values " * + "(retained addresses never reached: $(collect(dropped)))", + ) + end + return score end -function AbstractMCMC.sample( - rng::AbstractRNG, - model::DynamicPPL.Model, - sampler::SMC, - N::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, - kwargs..., +function DynamicPPL.tilde_assume!!( + ::SMCContext, dist::Distribution, vn::VarName, template, vi::DynamicPPL.AbstractVarInfo ) - 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). - if discard_initial > 0 || thinning > 1 - @warn "SMC samplers do not support `discard_initial` or `thinning`. These keyword arguments will be ignored." + particle = Libtask.get_taped_globals(Particle) + # 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. + reference = particle.reference + strategy = if reference === nothing + DynamicPPL.InitFromPrior() + else + haskey(reference, vn) || error( + "the reference execution trace changed while replaying retained values " * + "(new address: $vn)", + ) + DynamicPPL.InitFromParams(reference, nothing) end - chn = AbstractMCMC.mcmcsample( - rng, - model, - sampler, - N; - chain_type=chain_type, - initial_params=initial_params, - progress=progress, - nparticles=N, - kwargs..., - ) - post_sample_hook(chn, sampler; verbose) - return chn + ctx = DynamicPPL.InitContext(particle.rng, strategy, DynamicPPL.UnlinkAll()) + x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, vi) + particle.varinfo = vi + return x, vi end -function AbstractMCMC.step( - rng::AbstractRNG, - model::DynamicPPL.Model, - spl::SMC; - nparticles::Int, - initial_params, - discard_sample=false, - kwargs..., +# Both tilde handlers thread the varinfo the model passed in, and mirror the result onto the +# particle for the sweep to read between produces. Reading `particle.varinfo` instead would discard +# whatever the model body accumulated since the previous tilde: `@addlogprob!` reaches the varinfo +# directly rather than through a handler, so its term would never reach a chain's `loglikelihood`. +function DynamicPPL.tilde_observe!!( + ::SMCContext, + dist::Distribution, + left, + vn::Union{VarName,Nothing}, + template, + vi::DynamicPPL.AbstractVarInfo, ) - # 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, + particle = Libtask.get_taped_globals(Particle) + left, vi = DynamicPPL.tilde_observe!!( + DynamicPPL.DefaultContext(), dist, left, vn, template, vi ) + particle.varinfo = vi + return left, vi +end - # Perform particle sweep. - logevidence = AdvancedPS.sweep!(rng, particles, spl.resampler, spl) +""" + ProduceLogLikelihoodAccumulator{T} <: LogProbAccumulator{T} - # Extract the first particle and its weight. - particle = particles.vals[1] - weight = AdvancedPS.getweight(particles, 1) +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 +end - # 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) +DynamicPPL.accumulator_name(::Type{<:ProduceLogLikelihoodAccumulator}) = :LogLikelihood +DynamicPPL.logp(acc::ProduceLogLikelihoodAccumulator) = acc.logp - return transition, state +# The produce lives in the one place the likelihood is accumulated, so `val` *is* the increment: +# "produced weight equals accumulated increment" is structural rather than an invariant two call +# sites must keep in step. `observe` and `@addlogprob!` both arrive through exactly one `acclogp`, +# which is what gets `@addlogprob!` into the weight (issue #1996). `combine` adds the two `logp`s +# directly rather than through `acclogp`, so folding a submodel's varinfo into its parent's stays +# silent. `produce` suspends before the caller stores the updated varinfo, so a *suspended* +# particle's total lags this term; every varinfo read happens after the model has run to completion. +function DynamicPPL.acclogp(acc::ProduceLogLikelihoodAccumulator, val) + Libtask.produce(val) + return ProduceLogLikelihoodAccumulator(DynamicPPL.logp(acc) + val) end -function AbstractMCMC.step( - ::AbstractRNG, - model::DynamicPPL.Model, - spl::SMC, - state::SMCState; - discard_sample=false, - kwargs..., +function DynamicPPL.accumulate_assume!!( + acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, dist, template ) - # 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) - transition = if discard_sample - nothing - else - DynamicPPL.ParamsWithStats(deepcopy(particle.model.f.varinfo), stats) + return acc +end +function DynamicPPL.accumulate_observe!!( + acc::ProduceLogLikelihoodAccumulator, dist, left, vn, template +) + return DynamicPPL.acclogp(acc, Distributions.loglikelihood(dist, left)) +end + +# `@addlogprob!` and `:=` update the varinfo without passing through a tilde handler, so a trailing +# one has nothing after it to mirror the result onto the particle, and the final varinfo is +# unreachable: `Libtask.consume` discards the evaluator's return value once the model is done. The +# methods below therefore mirror as the update happens, keyed on the produce-aware accumulator, +# which only a particle carries (`trajectory_varinfo` is the set without it) and so identifies a +# particle varinfo. +# +# `:=` is dispatched on the context, so its method below takes `SMCContext`. `@addlogprob!` has no +# such hook, so the accumulator methods extend DynamicPPL's own. They dispatch on +# `OnlyAccsVarInfo`, which every sampler uses, so their non-particle behaviour must stay identical +# to DynamicPPL's; they belong there once it offers a context-dispatched entry point. +function is_particle_varinfo(vi::DynamicPPL.OnlyAccsVarInfo) + acc_name = Val(:LogLikelihood) + return DynamicPPL.hasacc(vi, acc_name) && + DynamicPPL.getacc(vi, acc_name) isa ProduceLogLikelihoodAccumulator +end + +function acclogp_and_mirror!!(vi, acc_name, logp, ignore_missing_accumulator) + if ignore_missing_accumulator && !DynamicPPL.hasacc(vi, acc_name) + return vi end - nextstate = SMCState(state.particles, index + 1, state.average_logevidence) - - return transition, nextstate + vi = DynamicPPL.map_accumulator!!(acc -> DynamicPPL.acclogp(acc, logp), vi, acc_name) + is_particle_varinfo(vi) && (Libtask.get_taped_globals(Particle).varinfo = vi) + return vi end -#### -#### Particle Gibbs sampler. -#### - -""" -$(TYPEDEF) - -Particle Gibbs sampler. - -# Fields - -$(TYPEDFIELDS) -""" -struct PG{R} <: ParticleInference - """Number of particles.""" - nparticles::Int - """Resampling algorithm.""" - resampler::R +function DynamicPPL.accloglikelihood!!( + vi::DynamicPPL.OnlyAccsVarInfo, logp; ignore_missing_accumulator=false +) + return acclogp_and_mirror!!(vi, Val(:LogLikelihood), logp, ignore_missing_accumulator) +end + +# Particles are drawn from the prior, so an ordinary `x ~ Dist` term needs no weight of its own -- +# the draw already accounts for it. A `logprior` term added by hand never influenced a draw, so it +# has to enter the importance weight explicitly. The produce is spelled out here only because +# `LogPrior` is DynamicPPL's ordinary accumulator; a `loglikelihood` field needs none, since +# `ProduceLogLikelihoodAccumulator` produces as it accumulates. A two-field `@addlogprob!` therefore +# produces twice: the same weight, and one more point at which the sweep may resample. +# +# The `hasacc` guard is repeated rather than delegated so that a produced weight cannot outlive the +# increment it stands for. Keeping the produce in this method also keeps it in a function Libtask +# instruments: a keyword argument on `acclogp_and_mirror!!` would route the call through +# `Core.kwcall`, which `@might_produce` does not name. +function DynamicPPL.acclogprior!!( + vi::DynamicPPL.OnlyAccsVarInfo, logp; ignore_missing_accumulator=false +) + acc_name = Val(:LogPrior) + if ignore_missing_accumulator && !DynamicPPL.hasacc(vi, acc_name) + return vi + end + is_particle_varinfo(vi) && Libtask.produce(logp) + return acclogp_and_mirror!!(vi, acc_name, logp, false) end -""" -PG(n, [resampler = AdvancedPS.ResampleWithESSThreshold()]) -PG(n, [resampler = AdvancedPS.resample_systematic, ]threshold) +function DynamicPPL.store_coloneq_value!!( + ::SMCContext, vn::VarName, right, template, vi::DynamicPPL.AbstractVarInfo +) + vi = DynamicPPL.store_coloneq_value!!( + DynamicPPL.DefaultContext(), vn, right, template, vi + ) + Libtask.get_taped_globals(Particle).varinfo = vi + return vi +end + +# 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. +# +# observe: tilde_observe!! -> accumulate_observe!! -> acclogp +# @addlogprob!: accloglikelihood!! -> acclogp_and_mirror!! -> map_accumulator!! -> acclogp +# (the `@addlogprob! (; ...)` NamedTuple form arrives via acclogp!!, and its +# `logprior` field produces in acclogprior!! itself) +# Gibbs: GibbsContext turns a tilde_assume!! into a tilde_observe!! +Libtask.@might_produce(DynamicPPL.tilde_observe!!) +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.acclogprior!!) +Libtask.@might_produce(acclogp_and_mirror!!) +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 -Create a Particle Gibbs sampler of type [`PG`](@ref) with `n` particles. +# Add the raw sampled values (`:=` included, so a chain reports them as under any other sampler; +# they take no part in a sweep). `OnlyAccsVarInfo`'s default `LogPrior` and `LogJacobian` are kept +# on purpose: `ParamsWithStats` reads them off this varinfo to fill a chain's `logprior` and +# `logjoint` columns, and dropping them would silently omit those columns rather than error. +function trajectory_varinfo() + return DynamicPPL.setacc!!( + DynamicPPL.OnlyAccsVarInfo(), DynamicPPL.RawValueAccumulator(true) + ) +end -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. -""" -function PG(nparticles::Int) - return PG(nparticles, AdvancedPS.ResampleWithESSThreshold()) +# A particle's varinfo is the same set with the produce-aware likelihood accumulator in place of the +# default one. Outside a sweep use `trajectory_varinfo`: there a `produce` has no task to suspend, +# and the mirroring keyed on that accumulator would write onto an unrelated particle. +function particle_varinfo() + return DynamicPPL.setacc!!(trajectory_varinfo(), ProduceLogLikelihoodAccumulator()) +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 +# 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. + +## +## Resampler interface +## + +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 + +## +## Schemes +## + +"Multinomial resampling: `n` independent draws from the categorical over `weights`." +struct MultinomialResampler <: AbstractResampler end +function resample_indices(rng::AbstractRNG, ::MultinomialResampler, weights, n::Integer) + return rand(rng, Distributions.Categorical(weights), n) +end + +# 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, 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] + end + indices[k] = s + end + return indices end -# Convenient constructors with ESS threshold -function PG(nparticles::Int, resampler, threshold::Real) - return PG(nparticles, AdvancedPS.ResampleWithESSThreshold(resampler, threshold)) +"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 -function PG(nparticles::Int, threshold::Real) - return PG(nparticles, AdvancedPS.resample_systematic, threshold) + +"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) + u = rand(rng) + return inverse_cdf_indices(weights, n, k -> (k - 1) + u) end -""" - CSMC(...) +## +## Effective-sample-size gating +## -Equivalent to [`PG`](@ref). """ -const CSMC = PG # type alias of PG as Conditional SMC + ESSThresholdResampler(threshold, scheme = StratifiedResampler()) -struct PGState{V<:DynamicPPL.AbstractVarInfo,R<:Random.AbstractRNG} - vi::V - rng::R +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 ESSThresholdResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler + threshold::T + scheme::R + function ESSThresholdResampler( + threshold::T, scheme::R=StratifiedResampler() + ) where {T<:Real,R<:AbstractResampler} + # The effective sample size lies in `[1, nparticles]`, so anything outside `[0, 1]` is + # silently "never resample" or "always resample". + 0 <= threshold <= 1 || + throw(ArgumentError("ESS threshold must lie in [0, 1]; got $threshold")) + return new{T,R}(threshold, scheme) + end end -function AbstractMCMC.step( - rng::AbstractRNG, model::DynamicPPL.Model, spl::PG; discard_sample=false, kwargs... +function should_resample(resampler::ESSThresholdResampler, weights) + return weight_ess(weights) ≤ resampler.threshold * length(weights) +end +function resample_indices( + rng::AbstractRNG, resampler::ESSThresholdResampler, weights, n::Integer ) - 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, - ) + return resample_indices(rng, resampler.scheme, weights, n) +end - # Perform a particle sweep. - logevidence = AdvancedPS.sweep!(rng, particles, spl.resampler, spl) +# +# Particle sweep +# - # Pick a particle to be retained. - Ws = AdvancedPS.getweights(particles) - index = AdvancedPS.randcat(rng, Ws) - reference = particles.vals[index] +# 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). - # 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) -end +## +## Weights and diagnostics +## -function AbstractMCMC.step( - rng::AbstractRNG, - model::DynamicPPL.Model, - spl::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 +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ᵢ²`. 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 +## + +# Advance one particle by one observation, folding its incremental weight in; return `true` +# once it has finished (produced nothing). +function advance_particle!(p::Particle) + score = advance!(p) + 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. +# +# Each particle advances only its own state (rng, varinfo, task), and its rng was already seeded +# serially in `resample_propagate!`, so the multithreaded loop is race-free and gives results +# identical to the serial one; the shared sampler rng is untouched here. +function reweight!(particles, multithreaded::Bool) + n = length(particles) + n_done = if multithreaded + # A shared counter would race, so record per particle and tally afterwards. + finished = Vector{Bool}(undef, n) + Threads.@threads for i in 1:n + finished[i] = advance_particle!(particles[i]) end + count(finished) + else + count(advance_particle!, particles) 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] + 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 - # Compute the transition. - _vi = newreference.model.f.varinfo - transition = if discard_sample - nothing +## +## Resample and propagate +## + +# Resample, if the scheme calls for it, and propagate the survivors. Returns whether it resampled, +# which tells `sweep!` what the total weight now is without recomputing it. +# +# A conditional sweep is recognised by its reference particle, which always occupies the last slot, +# rather than by a flag passed in, so resampling cannot disagree with the rest of the sweep about +# which particle is pinned. +function resample_propagate!(rng::AbstractRNG, particles, resampler) + weights = normalized_weights(particles) + # Not resampling leaves every particle to continue on its own generator, untouched. + should_resample(resampler, weights) || return false + + n = length(particles) + conditional = isreference(last(particles)) + # 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 - DynamicPPL.ParamsWithStats(deepcopy(_vi), (; logevidence=logevidence)) + resample_indices(rng, resampler, weights, n) end - - return transition, PGState(_vi, newreference.rng) + parents = copy(particles) + taken = falses(n) + for (slot, a) in enumerate(ancestors) + # A parent's first offspring continues in the parent's own object; extra offspring need + # the costly `deepcopy`, and so does any offspring of the reference, which has to survive + # the sweep intact. Either way the child is reseeded to continue independently. + first_offspring = !taken[a] && !isreference(parents[a]) + taken[a] = true + child = first_offspring ? reseed!(parents[a], rng) : fork(parents[a], rng) + child.logweight = zero(DynamicPPL.LogProbType) + particles[slot] = child + end + conditional && (particles[n].logweight = zero(DynamicPPL.LogProbType)) + return true end -""" - get_trace_local_varinfo() +## +## One sweep +## -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 +# 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; `PG` runs thousands of +# sweeps and would pay for them every time. +function sweep!( + rng::AbstractRNG, particles, resampler, multithreaded::Bool; ess::Bool=false +) + logZ = zero(DynamicPPL.LogProbType) + 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))) + nobs = 0 + while true + done = reweight!(particles, multithreaded) + nobs += 1 + # 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). + total = log_normalizing_constant(particles) + # A non-finite total means every particle is dead, which would otherwise surface as a + # `Categorical` domain error over all-`NaN` weights. `nobs` names a real observation: the + # finishing pass leaves the weights alone, so it can never be the first to fail here. + isfinite(total) || error( + "all $(length(particles)) particles have zero probability at observation $nobs " * + "(total log weight $total), so the sweep cannot continue.", + ) + logZ += total - logZ0 + logZ0 = total + done && break + # Post-reweight ESS for this filtering step: a degeneracy diagnostic, one entry per produce, + # so an `@addlogprob!` contributes one alongside the observations. After the break, so the + # finishing pass -- which produces nothing and leaves the weights alone -- adds no entry. + ess && push!(ess_per_step, weight_ess(normalized_weights(particles))) + # Resample between observations, never before the first: there the weights are still all + # equal, so a multinomial draw -- what every conditional sweep uses -- would duplicate + # particles before any data. One resample is still spent after the last observation, since + # a particle reveals that it has finished only by producing nothing, one pass later. + resampled = resample_propagate!(rng, particles, resampler) + resampled && (logZ0 = log(oftype(logZ, length(particles)))) + end + return logZ, ess_per_step end -""" - get_trace_local_resampled() - -Get the `resample` flag stored in the 'taped globals' of a `Libtask.TapedTask`. +# +# Sequential Monte Carlo +# -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. +abstract type ParticleInference <: AbstractSampler end -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 +function require_positive_particle_count(nparticles::Integer) + nparticles > 0 || + throw(ArgumentError("number of particles must be positive; got $nparticles")) + return nothing end """ - get_trace_local_rng() +$(TYPEDEF) + +Sequential Monte Carlo sampler. -Get the RNG stored in the 'taped globals' of a `Libtask.TapedTask`, if one exists. +# Fields -This function is meant to be called from *inside* the TapedTask itself. +$(TYPEDFIELDS) """ -function get_trace_local_rng() - return Libtask.get_taped_globals(Any).rng +struct SMC{R<:AbstractResampler} <: ParticleInference + "resampling scheme" + resampler::R + "reweight the particles across threads within each sweep" + multithreaded::Bool + function SMC(resampler::R; multithreaded::Bool=false) where {R<:AbstractResampler} + return new{R}(resampler, multithreaded) + end end """ - set_trace_local_varinfo(vi::AbstractVarInfo) + SMC([resampler = ESSThresholdResampler(0.5)]; multithreaded = false) + SMC([scheme = StratifiedResampler(), ]threshold; multithreaded = false) -Set the `varinfo` stored in Libtask's taped globals. The 'other' taped global in Libtask -is expected to be an `AdvancedPS.Trace`. +Sequential Monte Carlo sampler. By default stratified resampling is triggered whenever the +effective sample size drops below half the number of particles. -Returns `nothing`. +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). +Threads are the only option here: a suspended particle is a live Libtask task and cannot be +serialised, so a single sweep cannot be spread across processes. Passing `MCMCThreads()` or +`MCMCDistributed()` to [`sample`](@ref) parallelises whole chains instead, which is a separate axis +and composes with this one. -This function is meant to be called from *inside* the TapedTask itself. +The resampling scheme types (`StratifiedResampler`, `SystematicResampler`, `MultinomialResampler`, `ESSThresholdResampler`) are +not exported; refer to them as e.g. `Turing.Inference.SystematicResampler`. """ -function set_trace_local_varinfo(vi::AbstractVarInfo) - trace = Libtask.get_taped_globals(Any).other - trace.model.f.varinfo = vi - return nothing +SMC(; kwargs...) = SMC(ESSThresholdResampler(0.5); kwargs...) +SMC(threshold::Real; kwargs...) = SMC(ESSThresholdResampler(threshold); kwargs...) +function SMC(scheme::AbstractResampler, threshold::Real; kwargs...) + return SMC(ESSThresholdResampler(threshold, scheme); kwargs...) 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) +# Neither sampler has anywhere to put a user-supplied starting point: both draw their particles from +# the prior. `InitFromPrior` is what the ensemble wrapper injects per chain, not a user request. +function warn_initial_params_ignored(name, initial_params) + if initial_params !== nothing && !(initial_params isa DynamicPPL.InitFromPrior) + @warn "$name draws its particles from the prior; `initial_params` is ignored." 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 + return nothing 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 +# An extension hook: `TuringMCMCChainsExt` overrides this to flatten `ess_per_step` into scalars. +# It cannot specialise `AbstractMCMC.bundle_samples` on `SMC` instead, because the flattened +# transitions have to be handed back to the generic method and that call would recurse. +function bundle_smc_samples(transitions, model, sampler, state, chain_type; kwargs...) + return AbstractMCMC.bundle_samples( + transitions, model, sampler, state, chain_type; kwargs... + ) end -# Convenient constructor -function AdvancedPS.Trace( - model::Model, varinfo::AbstractVarInfo, rng::AdvancedPS.TracedRNG, resample::Bool +# 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` and the +# state keywords therefore have nothing to apply to. +function AbstractMCMC.sample( + rng::AbstractRNG, + model::DynamicPPL.Model, + sampler::SMC, + nparticles::Integer; + check_model=true, + chain_type=DEFAULT_CHAIN_TYPE, + discard_initial=0, + thinning=1, + initial_params=nothing, + initial_state=nothing, + save_state::Bool=false, + callback=nothing, + verbose::Bool=true, + kwargs..., ) - newvarinfo = deepcopy(varinfo) - tmodel = TracedModel(model, newvarinfo, rng, resample) - newtrace = AdvancedPS.Trace(tmodel, rng) - return newtrace + require_positive_particle_count(nparticles) + check_model && Turing._check_model(model, sampler) + error_if_threadsafe_eval(model) + if discard_initial > 0 || thinning > 1 + @warn "SMC does not support `discard_initial` or `thinning`; they are ignored." + end + # Caught here rather than forwarded: `save_state` would otherwise store `nothing` as the + # chain's sampler state, and `initial_state` would be dropped without a word. + if save_state || initial_state !== nothing + @warn "SMC runs one sweep and keeps no sampler state; `save_state` and `initial_state` are ignored." + end + warn_initial_params_ignored("SMC", initial_params) + # 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_rng(rng)) for _ in 1:nparticles] + logZ, ess_per_step = sweep!( + rng, particles, sampler.resampler, sampler.multithreaded; ess=true + ) + # The sweep ends on a reweight, so resample once more -- unconditionally, unlike the ESS-gated + # resampling inside it -- to return an equal-weight sample that needs no weighting downstream. + # Unless it is one already, which it is whenever the sweep's last act was a resample: drawing + # again over equal weights would only add duplicates. + ancestors = if allequal(logweights(particles)) + eachindex(particles) + else + resample_indices(rng, sampler.resampler, normalized_weights(particles), nparticles) + end + # `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, (; log_normalizing_constant=logZ, ess_per_step) + ) + end + chain = bundle_smc_samples(transitions, model, sampler, nothing, chain_type; kwargs...) + post_sample_hook(chain, sampler; verbose) + return chain end +# +# Particle Gibbs / conditional SMC +# + """ -ProduceLogLikelihoodAccumulator{T<:Real} <: AbstractAccumulator +$(TYPEDEF) -Exactly like `LogLikelihoodAccumulator`, but calls `Libtask.produce` on change of value. +Particle Gibbs (conditional SMC) sampler. # Fields + $(TYPEDFIELDS) """ -struct ProduceLogLikelihoodAccumulator{T<:Real} <: DynamicPPL.LogProbAccumulator{T} - "the scalar log likelihood value" - logp::T +struct PG{R<:AbstractResampler} <: ParticleInference + "number of particles" + nparticles::Int + "resampling scheme" + resampler::R + "reweight the particles across threads within each sweep" + multithreaded::Bool + function PG( + nparticles::Integer, resampler::R; multithreaded::Bool=false + ) where {R<:AbstractResampler} + require_positive_particle_count(nparticles) + return new{R}(nparticles, resampler, multithreaded) + end 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) +""" + 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 +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). +Threads are the only option here: a suspended particle is a live Libtask task and cannot be +serialised, so a single sweep cannot be spread across processes. Passing `MCMCThreads()` or +`MCMCDistributed()` to [`sample`](@ref) parallelises whole chains instead, which is a separate axis +and composes with this one. + +!!! warning "`log_normalizing_constant` is biased for PG" + PG chains carry `log_normalizing_constant`, but its exponential is not an unbiased estimator of + `p(y)` and must not be used for model comparison. A conditional sweep retains the reference + whatever its weight. Because the reference is a posterior draw rather than a proposal draw, it + usually has much higher likelihood than a fresh particle and inflates the mean weight at each + step. In a linear Gaussian SSM with known `p(y)`, `E[Ẑ]` exceeded `p(y)` by 80% at `n = 16` + and 16% at `n = 64`; the bias decreased approximately as `1/n` in that experiment but remained + substantial at these particle counts. Use [`SMC`](@ref) when an unbiased estimator of `p(y)` is + required. Its likelihood-scale estimate `exp(log_normalizing_constant)` is unbiased under the + usual particle-filter assumptions. +""" +PG(n::Integer; kwargs...) = PG(n, ESSThresholdResampler(0.5); kwargs...) +function PG(n::Integer, threshold::Real; kwargs...) + return PG(n, ESSThresholdResampler(threshold); kwargs...) +end +function PG(n::Integer, scheme::AbstractResampler, threshold::Real; kwargs...) + return PG(n, ESSThresholdResampler(threshold, scheme); kwargs...) end -function DynamicPPL.accumulate_assume!!( - acc::ProduceLogLikelihoodAccumulator, val, tval, logjac, vn, right, template -) - return acc +"Conditional SMC, an alias for [`PG`](@ref)." +const CSMC = PG + +""" + PGState(trajectory) + +Particle Gibbs sampler state: the retained trajectory's raw values, which the next sweep's +reference particle reuses. Plain data, because sampler state has to survive `save_state=true` and +`MCMCDistributed()`, whereas the [`Particle`](@ref) it is read off owns a live `Libtask.TapedTask` +that cannot be serialised. Nothing else needs carrying over: the reference consumes no randomness of +its own, and every other particle is seeded from the sampler's `rng`. +""" +struct PGState{V<:DynamicPPL.VarNamedTuple} + trajectory::V end -function DynamicPPL.accumulate_observe!!( - acc::ProduceLogLikelihoodAccumulator, right, left, vn, template + +# First iteration: an ordinary (unconditional) particle sweep. +function AbstractMCMC.step( + rng::AbstractRNG, + model::DynamicPPL.Model, + sampler::PG; + initial_params=nothing, + discard_sample=false, + kwargs..., ) - return DynamicPPL.acclogp(acc, Distributions.loglikelihood(right, left)) + error_if_threadsafe_eval(model) + warn_initial_params_ignored("PG", initial_params) + 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 -# 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, - }, - }, +# Subsequent iterations: conditional SMC given the retained trajectory, which the reference +# particle reproduces by reusing the retained values (see the `reference` field). +function AbstractMCMC.step( + rng::AbstractRNG, + model::DynamicPPL.Model, + sampler::PG, + state::PGState; + discard_sample=false, + kwargs..., ) - return true + error_if_threadsafe_eval(model) + n = sampler.nparticles + # Passing the retained trajectory makes this the reference (see the `reference` field). It is + # handed a generator like any other particle, but never reads it: every draw comes from the + # retained values, and its forks are reseeded from `rng` in `reseed!`. + reference = Particle(model, particle_rng(rng), state.trajectory) + # `n - 1` fresh particles, with the reference last -- the slot `resample_propagate!` retains. + 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) +end + +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 + # Copied because `ParamsWithStats` densifies the raw values, and `!!` lets it do so in + # place; the state below has to keep the trajectory the next reference will replay. + DynamicPPL.ParamsWithStats( + deepcopy(retained.varinfo), (; log_normalizing_constant=logZ) + ) + end + return transition, PGState(DynamicPPL.get_raw_values(retained.varinfo)) 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 -#### +# +# Gibbs interface +# -function gibbs_get_raw_values(state::PGState) - return DynamicPPL.get_raw_values(state.vi) -end +gibbs_get_raw_values(state::PGState) = state.trajectory 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())) - return PGState(new_vi, state.rng) + # Re-derive the retained trajectory under the values the other Gibbs components have since + # updated, keeping only the addresses this model still visits. The `nothing` fallback errors on + # an address `global_vals` lacks rather than inventing one; `GibbsContext` puts every address + # the component owns into the global values before we get here. + init = DynamicPPL.InitFromParams(global_vals, nothing) + vi = last(DynamicPPL.init!!(model, trajectory_varinfo(), init, DynamicPPL.UnlinkAll())) + return PGState(DynamicPPL.get_raw_values(vi)) end diff --git a/test/Project.toml b/test/Project.toml index 1110301a92..d2a33ca604 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" @@ -32,6 +31,7 @@ PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -45,7 +45,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" @@ -73,6 +72,7 @@ PDMats = "0.10, 0.11" Pkg = "1" Random = "1" ReverseDiff = "1.4.2" +Serialization = "1" SpecialFunctions = "0.10.3, 1, 2" StableRNGs = "1" Statistics = "1" diff --git a/test/essential/container.jl b/test/essential/container.jl deleted file mode 100644 index 3554526c97..0000000000 --- a/test/essential/container.jl +++ /dev/null @@ -1,61 +0,0 @@ -module ContainerTests - -using AdvancedPS: AdvancedPS -using Distributions: Bernoulli, Beta, Gamma, Normal -using DynamicPPL: DynamicPPL, @model -using Test: @test, @testset -using Turing - -@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 "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) - end - - @testset "fork" 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 - 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) - - newtrace = AdvancedPS.forkr(trace) - # Catch broken replay mechanism - @test AdvancedPS.advance!(trace) ≈ AdvancedPS.advance!(newtrace) - end -end - -end 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)]) diff --git a/test/mcmc/ess.jl b/test/mcmc/ess.jl index ab81287868..78b5624299 100644 --- a/test/mcmc/ess.jl +++ b/test/mcmc/ess.jl @@ -60,7 +60,9 @@ 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 `s` slowly, and multinomial ancestor draws in a conditional sweep are + # noisier still: 3_000 draws leave |err| 0.212 against `atol = 0.1`, 10_000 give 0.061. + chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 10_000) check_gdemo(chain; atol=0.1) end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 2cecc39d75..8853ec6c03 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -1,39 +1,178 @@ module ParticleMCMCTests using ..Models: gdemo_default -using ..SamplerTestUtils: test_chain_logp_metadata -using AdvancedPS: ResampleWithESSThreshold, resample_systematic, resample_multinomial -using Distributions: Bernoulli, Beta, Categorical, Gamma, Normal, sample -using FlexiChains: VNChain -using Random: Random +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, + SystematicResampler, + MultinomialResampler, + ESSThresholdResampler, + Particle, + particle_rng, + advance!, + fork, + sweep!, + resample_indices, + pg_transition_and_state +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 +using MCMCChains: MCMCChains +using Random: Random, Xoshiro +using Serialization: deserialize, serialize +using SpecialFunctions: logbeta 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). + +@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 + +# 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." +function run_to_end!(p) + while advance!(p) !== nothing + end + return p +end + +# +# SMC +# + @testset "SMC" begin @testset "constructor" begin - s = SMC() - @test s.resampler == ResampleWithESSThreshold() - - s = SMC(0.6) - @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6) + @test SMC().resampler == ESSThresholdResampler(0.5) + @test SMC().resampler.scheme isa StratifiedResampler # stratified is the default scheme + @test SMC(0.6).resampler == ESSThresholdResampler(0.6) + @test SMC(MultinomialResampler(), 0.6).resampler == + ESSThresholdResampler(0.6, MultinomialResampler()) + @test SMC(SystematicResampler()).resampler == SystematicResampler() + @test SMC().multithreaded == false + @test SMC(; multithreaded=true).multithreaded == true + @test SMC(SystematicResampler(); multithreaded=true).multithreaded == true + for nparticles in (0, -1) + @test_throws "number of particles must be positive" sample( + normal(), SMC(), nparticles + ) + end + for threshold in (-0.1, 1.5) + @test_throws "ESS threshold must lie in [0, 1]" SMC(threshold) + end + end - s = SMC(resample_multinomial, 0.6) - @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6) + @testset "resampling schemes" begin + obs = COIN_OBS + 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(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. + @test chn_strat[@varname(p)] != chn_multi[@varname(p)] + end - s = SMC(resample_systematic) - @test s.resampler === resample_systematic + @testset "offspring counts follow the weights" begin + # `E[Oᵏ] = n·Wᵏ` is what makes a scheme unbiased; for systematic resampling it holds + # within one offspring on every single draw, which is a far sharper check than the mean. + weights = [0.05, 0.35, 0.1, 0.4, 0.1] + n = 20 + expected = n .* weights + counts(scheme, seed) = [ + count(==(k), resample_indices(Xoshiro(seed), scheme, weights, n)) for + k in eachindex(weights) + ] + @test all( + all(floor.(expected) .<= counts(SystematicResampler(), s) .<= ceil.(expected)) + for s in 1:200 + ) + for scheme in (StratifiedResampler(), SystematicResampler(), MultinomialResampler()) + average = mean(counts(scheme, s) for s in 1:4_000) + @test average ≈ expected atol = 0.1 + end 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 + @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 - tested = sample(normal(), SMC(), 100) end @testset "errors when number of observations is not fixed" begin @@ -46,7 +185,6 @@ using Turing end return a, b end - @test_throws ErrorException sample(fail_smc(), SMC(), 100) @test_throws "number of observations" sample(fail_smc(), SMC(), 100) end @@ -54,111 +192,213 @@ using Turing test_chain_logp_metadata(SMC()) end - @testset "logevidence" 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 "rng is respected" begin + test_rng_respected(SMC()) + end + @testset "log_normalizing_constant" begin 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 "MCMCChains preserves ess_per_step" begin + flexi = sample(StableRNG(101), test(), SMC(), 32) + mcmcchains = sample(StableRNG(101), test(), SMC(), 32; chain_type=MCMCChains.Chains) + ess_names = MCMCChains.namesingroup(mcmcchains, :ess_per_step) + ess = [first(mcmcchains[name]) for name in ess_names] + @test ess == first(flexi[:ess_per_step]) + @test length(ess_names) == 2 + @test :ess_per_step ∉ MCMCChains.names(mcmcchains, :internals) + end + + @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 = 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)] 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. - @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 and thinning 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 + @testset "does not resample an already equal-weight population" begin + # A bare scheme resamples after the last observation, so the sweep ends equal-weight and + # the closing resample has nothing to correct: a second multinomial draw would only + # duplicate particles. Replay the sweep on the same seed -- the chain must be its + # population, particle for particle. + n, seed = 8, 6 + rng = StableRNG(seed) + particles = [Particle(normal(), particle_rng(rng)) for _ in 1:n] + sweep!(rng, particles, MultinomialResampler(), false) + expected = sort([get_raw_values(p.varinfo)[@varname(a)] for p in particles]) + chn = sample(StableRNG(seed), normal(), SMC(MultinomialResampler()), n) + @test sort(vec(chn[@varname(a)])) == expected + end + + @testset "reports a population that has died out" begin + # Every particle scoring `-Inf` leaves nothing to resample from. The sweep has to say so, + # naming the observation, rather than let `Categorical` reject the all-`NaN` weights. + @model function impossible() + x ~ Normal() + 1.5 ~ Normal(x) + return 1 ~ Bernoulli(0) end + @test_throws "zero probability at observation 2" sample(impossible(), SMC(), 5) + end + + @testset "unsupported keywords are reported as ignored" begin + @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() + ) - @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 - @test chn isa VNChain + # 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) - @test_logs (:warn, r"ignored") sample(normal(), SMC(), 10; thinning=3) - chn2 = sample(normal(), SMC(), 10; thinning=3) - @test size(chn2, 1) == 10 - @test chn2 isa VNChain + # 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=2, thinning=2 + # SMC keeps no sampler state, so neither state keyword has anything to act on; without + # this, `save_state=true` would silently store `nothing` in the chain. + @test_logs (:warn, r"save_state.*ignored") sample( + normal(), SMC(), 10; save_state=true ) - chn3 = sample(normal(), SMC(), 10; discard_initial=2, thinning=2) - @test size(chn3, 1) == 10 - @test chn3 isa VNChain + @test_logs (:warn, r"initial_state.*ignored") sample( + normal(), SMC(), 10; initial_state=1 + ) + + for kw in ((; discard_initial=5), (; thinning=3), (; discard_initial=2, thinning=2)) + chn = @test_logs (:warn, r"ignored") sample(normal(), SMC(), 10; kw...) + @test size(chn, 1) == 10 + @test chn isa VNChain + end end end +# +# PG / conditional SMC +# + @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 == ESSThresholdResampler(0.5) + @test PG(60, 0.6).resampler == ESSThresholdResampler(0.6) + @test PG(80, MultinomialResampler(), 0.6).resampler == + ESSThresholdResampler(0.6, MultinomialResampler()) + @test PG(100, SystematicResampler()).resampler == SystematicResampler() + @test PG(10).multithreaded == false + @test PG(10; multithreaded=true).multithreaded == true + @test PG(80, MultinomialResampler(), 0.6; multithreaded=true).multithreaded == true + for nparticles in (0, -1) + @test_throws "number of particles must be positive" PG(nparticles) + end + for threshold in (-0.1, 1.5) + @test_throws "ESS threshold must lie in [0, 1]" PG(10, threshold) + end end @testset "chain log-density metadata" begin test_chain_logp_metadata(PG(10)) end - @testset "logevidence" 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 "rng is respected" begin + test_rng_respected(PG(10)) + end + @testset "log_normalizing_constant" begin 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 - # Should be the same for all iterations. - @test chains_pg[:logevidence] ≈ fill(pg_logevidence, 100) + pg_log_normalizing_constant = mean(chains_pg[:log_normalizing_constant]) + @test pg_log_normalizing_constant ≈ -2 * log(2) atol = 0.01 + # 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). + obs = COIN_OBS + 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 + # count. + 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)] + 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_rng(rng)) + run_to_end!(retained) + reference = Particle(model, particle_rng(rng), get_raw_values(retained.varinfo)) + 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) + end + multinomial = conditional_sweep(MultinomialResampler()) + @test conditional_sweep(StratifiedResampler()) == multinomial + @test conditional_sweep(SystematicResampler()) == multinomial end # https://github.com/TuringLang/Turing.jl/issues/1598 @@ -168,6 +408,196 @@ end @test length(unique(c[:s])) == 1 end + @testset "initial_params is ignored" begin + # PG's first sweep draws from the prior, so there is nowhere to put a starting point. + @test_logs (:warn, r"initial_params.*ignored") match_mode = :any sample( + normal(), PG(5), 10; initial_params=(; a=1.0) + ) + # `InitFromPrior()` is what the ensemble wrapper injects per chain, not a user request. + @test_logs sample(Xoshiro(1), normal(), PG(5), MCMCSerial(), 10, 2; progress=false) + end + + @testset "the saved state survives serialisation" begin + chn = sample(StableRNG(24), gdemo_default, PG(5), 10; save_state=true) + io = IOBuffer() + serialize(io, chn) + seekstart(io) + state = only(loadstate(deserialize(io))) + # With one particle the reference is the whole population, so every draw is the retained + # trajectory: the round-tripped state has to still carry it. + resumed = sample(StableRNG(25), gdemo_default, PG(1), 5; initial_state=state) + @test length(unique(resumed[:m])) == 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. 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) + 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) + # Call the sampler's own selection rule rather than reimplement it, so this test + # follows any change to how the retained particle is chosen. + 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 + nretained = 0 + for _ in 1:nsteps + ref = Particle(model, particle_rng(rng), state.trajectory) + 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) == state.trajectory + state = draw(parts) + nretained = length(keys(state.trajectory)) + end + return allok, nretained + end + + rng = StableRNG(1234) + y = randn(rng, 10) + # ρ 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, nretained = run_csmc(state_space_model(y), 3, 30, rng) + @test allok # reference regenerated exactly every step + @test nretained == length(y) + 2 + end + + @testset "reference is pinned to retained values under re-conditioning" begin + # 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, then rebuild the reference under another: + # value-pinning keeps the trajectory, whereas re-drawing would follow the shifted prior. + @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_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_rng(rng), + retained_vals, + ) + run_to_end!(reference) + @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_rng(rng)) + run_to_end!(retained) + reference = Particle( + branch_changes(false, 0.0), particle_rng(rng), get_raw_values(retained.varinfo) + ) + @test_throws "reference execution trace changed" advance!(reference) + + @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_rng(rng)) + run_to_end!(retained) + reference = Particle( + branch_drops(false, 0.0), particle_rng(rng), get_raw_values(retained.varinfo) + ) + @test_throws "reference execution trace changed" begin + run_to_end!(reference) + end + end + + @testset "value replay handles a slice assume" 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 reference has to reuse it anyway, which + # it does because `haskey` resolves the slice against those keys. + @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 "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) + # `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 + end + end + @testset "conditional sweeps target the exact posterior" begin # Conditional SMC is invariant only if the reference is exactly the retained path, and # this model is shaped so that either way of getting that wrong biases the marginals. @@ -227,6 +657,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 @@ -278,21 +712,292 @@ 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 @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) + model = setthreadsafe(threadsafe_model(randn(10)), true) + @test_throws ArgumentError sample(model, PG(10), 100) + 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 + # 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 + +# +# Particle mechanics +# + +@testset "particle container" begin + @testset "advance!" begin + # `x ~ Bernoulli(1)` forces `x = 1`, so the first observe is `1 ~ Bernoulli(0.5)`. + 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 + 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_rng(Xoshiro(23))) + run_to_end!(particle) + + accs = DynamicPPL.OnlyAccsVarInfo() + accs = DynamicPPL.setacc!!(accs, DynamicPPL.LogLikelihoodAccumulator()) + accs = DynamicPPL.setacc!!(accs, DynamicPPL.RawValueAccumulator(true)) + _, accs = DynamicPPL.init!!( + particle_rng(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_rng(Xoshiro(23))) + advance!(particle) + child = fork(particle, Xoshiro(1)) + # Independent continuations: advancing one does not touch the other. + @test advance!(child) ≈ -log(2) + @test particle.varinfo !== child.varinfo + @test advance!(particle) ≈ -log(2) + end + + @testset "both fields of a named addlogprob are weighted" begin + # DynamicPPL's `acclogp!!` splits a two-field `@addlogprob!` across `acclogprior!!` and + # `accloglikelihood!!`; both have to produce, or the sweep silently drops one term from the + # weight. Sorted because which arrives first is DynamicPPL's to decide. + @model function both_terms() x ~ Normal() - Threads.@threads for i in eachindex(y) - y[i] ~ Normal(x) + @addlogprob! (; loglikelihood=-2.0, logprior=-1.0) + end + particle = Particle(both_terms(), particle_rng(Xoshiro(23))) + scores = Float64[] + while (score = advance!(particle)) !== nothing + push!(scores, score) + end + @test sort(scores) == [-2.0, -1.0] + end + + @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. + 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_rng(Xoshiro(7)), values) + while ( + Random.seed!(reference.rng, rand(scrambler, UInt64)); advance!(reference) + ) !== nothing + end + @test get_raw_values(reference.varinfo) == values + 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 + +# 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)) + 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 ~ 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), fix(lgssm(y, a, r), @varname(q) => true_q), PG(16), 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.005, 12.0; length=1_200) + w = ExactSSM.grid_posterior( + prior, qs, q -> ExactSSM.lgssm_loglik(y, a, q, r, s0(q)) + ) + # Quadrature is exact only if the grid covers the posterior. At these bounds the mean + # and sd agree to seven digits with a far wider grid; a grid ending at 4.0 leaves + # weight 2e-5 on its last point and understates the sd by 6%. + @test w[1] + w[end] < 1e-6 + 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(16)) + chn = sample(StableRNG(31), lgssm(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), fix(hmm(y), @varname(sd) => true_sd), PG(16), 4_000) + 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 + + @testset "Gibbs(sd => HMC, z => CSMC) recovers the exact posterior" begin + prior = LogNormal(log(0.7), 0.4) + sds = range(0.05, 5.0; length=1_000) + w = ExactSSM.grid_posterior( + prior, sds, s -> last(ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, s))) + ) + @test w[1] + w[end] < 1e-6 # the grid has to cover the posterior; see above + 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(16)) + chn = sample(StableRNG(32), hmm(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 + 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 - model = setthreadsafe(f(randn(10)), true) - @test_throws ArgumentError sample(model, PG(10), 100) end end diff --git a/test/runtests.jl b/test/runtests.jl index c4faea5a37..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) @@ -38,10 +39,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") diff --git a/test/test_utils/exact_ssm.jl b/test/test_utils/exact_ssm.jl new file mode 100644 index 0000000000..138a117dfb --- /dev/null +++ b/test/test_utils/exact_ssm.jl @@ -0,0 +1,187 @@ +# +# 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 + +# 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) +## + +# 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)) + 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 .* @view lik[1, :] + c[1] = sum(@view α[1, :]) + α[1, :] ./= c[1] + for t in 2: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 * (@view(lik[t + 1, :]) .* @view(β[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 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]] + 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. + π0_stat = stationary_distribution(P) + @test transpose(P) * π0_stat ≈ π0_stat + end +end + +end diff --git a/test/test_utils/numerical_tests.jl b/test/test_utils/numerical_tests.jl index 1322c5362f..a08ce30c4d 100644 --- a/test/test_utils/numerical_tests.jl +++ b/test/test_utils/numerical_tests.jl @@ -25,6 +25,9 @@ function check_gdemo(chain; atol=0.2, rtol=0.0) end # Wrapper function to check MoGtest. +# The reference values are the exact posterior means, not the idealised labels +# [1, 1, 2, 2, 1, 4]: the data do not identify the clusters perfectly, so asserting the +# labels spends 0.07 of the tolerance before any Monte Carlo error. function check_MoGtest_default(chain; atol=0.2, rtol=0.0) return check_numerical( chain, @@ -36,7 +39,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]; + [1.05, 1.05, 1.95, 1.95, 1.07, 3.93]; atol=atol, rtol=rtol, ) @@ -53,7 +56,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]; + [1.05, 1.05, 1.95, 1.95, 1.07, 3.93]; atol=atol, rtol=rtol, )