diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..bcb41efba7 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,100 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Project Overview + +Turing.jl is the user-facing entry point for the [TuringLang](https://github.com/TuringLang) probabilistic programming ecosystem. It is largely a translation layer between DynamicPPL models — which work with named, structured parameters — and inference algorithms that expect flat, vectorised samples (e.g. HMC/NUTS operate on `AbstractVector{<:Real}`). DynamicPPL's `LogDensityFunction` handles most of this translation; Turing provides the sampler wrappers that set it up and manage state across iterations. + +Model definition lives in [DynamicPPL.jl](https://github.com/TuringLang/DynamicPPL.jl), parameter transformations in [Bijectors.jl](https://github.com/TuringLang/Bijectors.jl), and sampling interfaces in [AbstractMCMC.jl](https://github.com/TuringLang/AbstractMCMC.jl). Turing re-exports their APIs and provides concrete sampler implementations that wire everything together. + +## Building and Testing + +Code formatting uses [JuliaFormatter.jl](https://github.com/domluna/JuliaFormatter.jl) v1 (not v2) with the **Blue style** (configured in `.JuliaFormatter.toml`). CI enforces formatting on all PRs. JuliaFormatter must be installed in the **global** Julia environment, not the project environment — do not use `--project`. See the [formatting guide](https://turinglang.org/docs/contributing/code-formatting/) for setup details. + +```bash +julia -e 'using JuliaFormatter; format(".")' +``` + +Tests use `SelectiveTests.jl` (in `test/test_utils/`) to filter by path. CI splits the suite into four shards: `mcmc/gibbs.jl`, `mcmc/Inference.jl`, `ad.jl`, and everything else. To run a subset locally: + +```bash +julia --project -e 'using Pkg; Pkg.test(; test_args=["mcmc/hmc.jl"])' +``` + +Use `--skip` to exclude files: + +```bash +julia --project -e 'using Pkg; Pkg.test(; test_args=["--skip", "mcmc/gibbs.jl", "ad.jl"])' +``` + +CI matrix: Julia stable + min, Ubuntu/Windows/macOS, 1 and 2 threads. + +`test/test_utils/sampler.jl` provides generic test helpers (`test_rng_respected`, `test_sampler_analytical`, `test_chain_logp_metadata`) that should work for any sampler. Beyond these, sampler-specific tests are needed to capture the properties you care about — there is no standardised test template yet. + +## Architecture + +### What lives here vs elsewhere + +Most complexity is in DynamicPPL. Turing.jl contains: + + - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (implemented natively in `particle_mcmc.jl` on top of Libtask coroutines and Random123 — not a wrapper), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. + - **External sampler interface** (`src/mcmc/external_sampler.jl`): The `externalsampler()` wrapper lets any `AbstractMCMC.AbstractSampler` that implements `step` for `LogDensityModel` work with Turing models. This is the easier path for new samplers — it only requires a dependency on AbstractMCMC and the LogDensityProblems.jl interface, with no Turing internals. The tradeoff is less power: you can only interact with the model as a black-box log-density function, just like using `LogDensityFunction` directly. + - **Variational inference** (`src/variational/`): Wraps AdvancedVI algorithms. + - **Mode estimation** (`src/optimisation/`): MAP and MLE via Optimization.jl. + - **Custom distributions** (`src/stdlib/`): `Flat`, `FlatPos`, `BinomialLogit`, `OrderedLogistic`, `LogPoisson`, and Dirichlet/Chinese Restaurant processes. + +For how the model and inference machinery works under the hood, see the [DynamicPPL docs](https://turinglang.org/DynamicPPL.jl/stable/) and the [developer guides](https://turinglang.org/docs/developers/). + +### Gibbs sampler + +The Gibbs sampler (`src/mcmc/gibbs.jl`) is the most complex piece in Turing.jl. It maintains a global `VarNamedTuple` of raw values for all variables. On each iteration, it conditions the model on the non-target variables via `GibbsContext`, runs the component sampler, and updates the global state. + +To plug a sampler into Gibbs, implement: + + - `gibbs_get_raw_values(state)` — return a `VarNamedTuple` of raw values for the variables this sampler is responsible for. + - `gibbs_update_state!!(sampler, state, model, global_vals)` — update the sampler's state to reflect new conditioned values. For samplers that use `LogDensityFunction`, the helper `gibbs_recompute_ldf_and_params` handles the common case. + - Optionally, `isgibbscomponent(sampler)` — return `false` to disallow use in Gibbs (the default is `true`). + +### Extension + +`ext/TuringDynamicHMCExt` provides the DynamicHMC.jl integration (loaded when DynamicHMC is imported). + +## Review Guidelines + +### Use `OnlyAccsVarInfo`, not `VarInfo` + +Sampler state should use `OnlyAccsVarInfo` (with appropriate accumulators), not `VarInfo`. `VarInfo` is being phased out across the ecosystem. + +Most gradient-based samplers (HMC, NUTS, external samplers) go through `LogDensityFunction`, which handles the model interaction. `LogDensityFunction` works well when the model structure is static (the set of variables is fixed across evaluations) and the sampler only needs a scalar log-density value. However, LDF is hard to use when the sampler needs extra accumulators beyond log-probability — for example, MH uses custom accumulators to capture proposal distributions and linked values, so it works directly with `OnlyAccsVarInfo` + `init!!` instead. Either approach is fine; the key constraint is no `VarInfo`. + +Note: "linked" and "unconstrained" are synonymous in this codebase. Linking transforms constrained parameters to unconstrained (Euclidean) space for gradient-based sampling. + +### `VarNamedTuple` for parameter collections + +Interfaces that accept or return named parameter collections should use `VarNamedTuple`, not `NamedTuple` or `Dict{VarName}`. `NamedTuple` and `Dict{VarName}` are accepted as user-facing input but should be converted to `VarNamedTuple` at the boundary (see `_to_varnamedtuple` in `src/common.jl`). Don't propagate them through internal code. + +### `getlogjoint_internal` vs `getlogjoint` + +Samplers operating in unconstrained space should use `getlogjoint_internal`, which includes the Jacobian correction from the linking transform. This is the default and what you almost always want. The exceptions are ESS (which needs the likelihood in constrained space, per the algorithm) and optimisation (where the Jacobian term should not influence the objective). + +### AD backend handling + +Gradient-based samplers accept an `adtype::ADTypes.AbstractADType` keyword (default: `AutoForwardDiff()`). When reviewing sampler code, check that `adtype` is threaded through to `LogDensityFunction` and not hardcoded. The AD backend is the user's choice, not the sampler's. + +### `initial_params` conversion + +User-facing functions accept `initial_params` as a convenience. `_convert_initial_params` in `src/common.jl` converts `NamedTuple`/`Dict{VarName}` to `InitFromParams`. Raw vectors are no longer supported and will error. Don't bypass this conversion or accept raw vectors in new code. + +### Discrete variables + +`allow_discrete_variables(sampler)` defaults to `true`. Gradient-based samplers (all `Hamiltonian` subtypes) override this to `false`. `_check_model` uses this to validate the model before sampling. If adding a new sampler that requires continuous variables, override `allow_discrete_variables` to return `false`. + +### GibbsContext is not ConditionContext + +`GibbsContext` is distinct from `condition`/`ConditionContext`. For non-target variables, `GibbsContext.tilde_assume!!` calls `tilde_observe!!` — this means particle samplers (PG/CSMC) will correctly resample on conditioned variables. The key difference from `condition` is that `GibbsContext` obtains the conditioned values from the global `VarNamedTuple` rather than from the model's conditioning, and it handles the bookkeeping needed for Gibbs (e.g. updating the global VNT when new variables appear). + +## Contributing + + - Non-breaking changes target `main`; breaking changes target the `breaking` branch. + - Julia ≥ 1.10.8 required (see `[compat]` in `Project.toml`). diff --git a/CLAUDE.md b/CLAUDE.md index 4922775aff..43c994c2d3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,100 +1 @@ -# CLAUDE.md - -This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. - -## Project Overview - -Turing.jl is the user-facing entry point for the [TuringLang](https://github.com/TuringLang) probabilistic programming ecosystem. It is largely a translation layer between DynamicPPL models — which work with named, structured parameters — and inference algorithms that expect flat, vectorised samples (e.g. HMC/NUTS operate on `AbstractVector{<:Real}`). DynamicPPL's `LogDensityFunction` handles most of this translation; Turing provides the sampler wrappers that set it up and manage state across iterations. - -Model definition lives in [DynamicPPL.jl](https://github.com/TuringLang/DynamicPPL.jl), parameter transformations in [Bijectors.jl](https://github.com/TuringLang/Bijectors.jl), and sampling interfaces in [AbstractMCMC.jl](https://github.com/TuringLang/AbstractMCMC.jl). Turing re-exports their APIs and provides concrete sampler implementations that wire everything together. - -## Building and Testing - -Code formatting uses [JuliaFormatter.jl](https://github.com/domluna/JuliaFormatter.jl) v1 (not v2) with the **Blue style** (configured in `.JuliaFormatter.toml`). CI enforces formatting on all PRs. JuliaFormatter must be installed in the **global** Julia environment, not the project environment — do not use `--project`. See the [formatting guide](https://turinglang.org/docs/contributing/code-formatting/) for setup details. - -```bash -julia -e 'using JuliaFormatter; format(".")' -``` - -Tests use `SelectiveTests.jl` (in `test/test_utils/`) to filter by path. CI splits the suite into four shards: `mcmc/gibbs.jl`, `mcmc/Inference.jl`, `ad.jl`, and everything else. To run a subset locally: - -```bash -julia --project -e 'using Pkg; Pkg.test(; test_args=["mcmc/hmc.jl"])' -``` - -Use `--skip` to exclude files: - -```bash -julia --project -e 'using Pkg; Pkg.test(; test_args=["--skip", "mcmc/gibbs.jl", "ad.jl"])' -``` - -CI matrix: Julia stable + min, Ubuntu/Windows/macOS, 1 and 2 threads. - -`test/test_utils/sampler.jl` provides generic test helpers (`test_rng_respected`, `test_sampler_analytical`, `test_chain_logp_metadata`) that should work for any sampler. Beyond these, sampler-specific tests are needed to capture the properties you care about — there is no standardised test template yet. - -## Architecture - -### What lives here vs elsewhere - -Most complexity is in DynamicPPL. Turing.jl contains: - - - **Sampler implementations** (`src/mcmc/`): HMC/NUTS/HMCDA (wrapping AdvancedHMC), MH (wrapping AdvancedMH), particle samplers SMC/PG/CSMC (wrapping AdvancedPS), ESS (wrapping EllipticalSliceSampling), SGLD/SGHMC, Emcee, and Gibbs. - - **External sampler interface** (`src/mcmc/external_sampler.jl`): The `externalsampler()` wrapper lets any `AbstractMCMC.AbstractSampler` that implements `step` for `LogDensityModel` work with Turing models. This is the easier path for new samplers — it only requires a dependency on AbstractMCMC and the LogDensityProblems.jl interface, with no Turing internals. The tradeoff is less power: you can only interact with the model as a black-box log-density function, just like using `LogDensityFunction` directly. - - **Variational inference** (`src/variational/`): Wraps AdvancedVI algorithms. - - **Mode estimation** (`src/optimisation/`): MAP and MLE via Optimization.jl. - - **Custom distributions** (`src/stdlib/`): `Flat`, `FlatPos`, `BinomialLogit`, `OrderedLogistic`, `LogPoisson`, and Dirichlet/Chinese Restaurant processes. - -For how the model and inference machinery works under the hood, see the [DynamicPPL docs](https://turinglang.org/DynamicPPL.jl/stable/) and the [developer guides](https://turinglang.org/docs/developers/). - -### Gibbs sampler - -The Gibbs sampler (`src/mcmc/gibbs.jl`) is the most complex piece in Turing.jl. It maintains a global `VarNamedTuple` of raw values for all variables. On each iteration, it conditions the model on the non-target variables via `GibbsContext`, runs the component sampler, and updates the global state. - -To plug a sampler into Gibbs, implement: - - - `gibbs_get_raw_values(state)` — return a `VarNamedTuple` of raw values for the variables this sampler is responsible for. - - `gibbs_update_state!!(sampler, state, model, global_vals)` — update the sampler's state to reflect new conditioned values. For samplers that use `LogDensityFunction`, the helper `gibbs_recompute_ldf_and_params` handles the common case. - - Optionally, `isgibbscomponent(sampler)` — return `false` to disallow use in Gibbs (the default is `true`). - -### Extension - -`ext/TuringDynamicHMCExt` provides the DynamicHMC.jl integration (loaded when DynamicHMC is imported). - -## Review Guidelines - -### Use `OnlyAccsVarInfo`, not `VarInfo` - -Sampler state should use `OnlyAccsVarInfo` (with appropriate accumulators), not `VarInfo`. `VarInfo` is being phased out across the ecosystem. - -Most gradient-based samplers (HMC, NUTS, external samplers) go through `LogDensityFunction`, which handles the model interaction. `LogDensityFunction` works well when the model structure is static (the set of variables is fixed across evaluations) and the sampler only needs a scalar log-density value. However, LDF is hard to use when the sampler needs extra accumulators beyond log-probability — for example, MH uses custom accumulators to capture proposal distributions and linked values, so it works directly with `OnlyAccsVarInfo` + `init!!` instead. Either approach is fine; the key constraint is no `VarInfo`. - -Note: "linked" and "unconstrained" are synonymous in this codebase. Linking transforms constrained parameters to unconstrained (Euclidean) space for gradient-based sampling. - -### `VarNamedTuple` for parameter collections - -Interfaces that accept or return named parameter collections should use `VarNamedTuple`, not `NamedTuple` or `Dict{VarName}`. `NamedTuple` and `Dict{VarName}` are accepted as user-facing input but should be converted to `VarNamedTuple` at the boundary (see `_to_varnamedtuple` in `src/common.jl`). Don't propagate them through internal code. - -### `getlogjoint_internal` vs `getlogjoint` - -Samplers operating in unconstrained space should use `getlogjoint_internal`, which includes the Jacobian correction from the linking transform. This is the default and what you almost always want. The exceptions are ESS (which needs the likelihood in constrained space, per the algorithm) and optimisation (where the Jacobian term should not influence the objective). - -### AD backend handling - -Gradient-based samplers accept an `adtype::ADTypes.AbstractADType` keyword (default: `AutoForwardDiff()`). When reviewing sampler code, check that `adtype` is threaded through to `LogDensityFunction` and not hardcoded. The AD backend is the user's choice, not the sampler's. - -### `initial_params` conversion - -User-facing functions accept `initial_params` as a convenience. `_convert_initial_params` in `src/common.jl` converts `NamedTuple`/`Dict{VarName}` to `InitFromParams`. Raw vectors are no longer supported and will error. Don't bypass this conversion or accept raw vectors in new code. - -### Discrete variables - -`allow_discrete_variables(sampler)` defaults to `true`. Gradient-based samplers (all `Hamiltonian` subtypes) override this to `false`. `_check_model` uses this to validate the model before sampling. If adding a new sampler that requires continuous variables, override `allow_discrete_variables` to return `false`. - -### GibbsContext is not ConditionContext - -`GibbsContext` is distinct from `condition`/`ConditionContext`. For non-target variables, `GibbsContext.tilde_assume!!` calls `tilde_observe!!` — this means particle samplers (PG/CSMC) will correctly resample on conditioned variables. The key difference from `condition` is that `GibbsContext` obtains the conditioned values from the global `VarNamedTuple` rather than from the model's conditioning, and it handles the bookkeeping needed for Gibbs (e.g. updating the global VNT when new variables appear). - -## Contributing - - - Non-breaking changes target `main`; breaking changes target the `breaking` branch. - - Julia ≥ 1.10.8 required (see `[compat]` in `Project.toml`). +@AGENTS.md diff --git a/HISTORY.md b/HISTORY.md index 580932cb71..d5b63c0968 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,24 @@ +# 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; `PG` / `CSMC` draw the ancestors of a conditional sweep from the categorical over the weights, since a correct conditional version of stratified or systematic resampling is scheme-specific rather than "pin one draw and keep the rest". +Exact draws may therefore differ from previous releases, but remain statistically consistent (the same target distribution). + +The rewrite also brings: + + - **Reproducibility.** Internal seeds are derived through a counter-based (Philox) generator, so a fixed user seed gives the same draws on every Julia version and platform, and splitting one stream into many is better decorrelated. Previously, results could drift between Julia versions even under a `StableRNG` (https://github.com/TuringLang/Turing.jl/issues/2781). + - **Parallelism** at two independent levels. *Across chains*, SMC/PG work with AbstractMCMC's `MCMCThreads()` / `MCMCDistributed()` like any other sampler — each chain is an independent run. *Within a single sweep*, `SMC(; multithreaded=true)` / `PG(n; multithreaded=true)` spread that sweep's particles across threads. These are separate knobs: the ensemble does not parallelise a sweep, `multithreaded` does not parallelise chains, and they compose. Neither changes the results; start Julia with multiple threads (e.g. `julia -t auto`) for the thread-based paths to take effect. + - **Equal-weight draws & diagnostics.** `SMC` resamples once at the end of the sweep so the returned particles are an equal-weight sample — `mean(chain[...])` and other summaries need no weighting. `SMC`, `PG`, and `CSMC` chains all carry `log_normalizing_constant`; `SMC` chains additionally carry `ess_per_step`, the per-observation effective sample size across the sweep (a degeneracy diagnostic). For `SMC` the normalizing constant is an unbiased estimate of the marginal likelihood `p(y)`; for `PG` / `CSMC` it is **not**, and must not be used for model comparison — see the `PG` docstring. + # 0.46.0 ## Breaking changes diff --git a/Project.toml b/Project.toml index e96c2dd443..c113c5bbb9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.46.0" +version = "0.47.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -9,7 +9,6 @@ AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -AdvancedPS = "576499cb-2369-40b2-a588-c64705576edc" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" @@ -29,6 +28,7 @@ OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Random123 = "74087812-796a-5b5d-8853-05524746bad3" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" @@ -52,7 +52,6 @@ AbstractPPL = "0.15" Accessors = "0.1" AdvancedHMC = "0.8.3" AdvancedMH = "0.8.9" -AdvancedPS = "0.7.2" AdvancedVI = "0.7" BangBang = "0.4.2" Bijectors = "0.15.17, 0.16" @@ -74,6 +73,7 @@ OptimizationOptimJL = "0.1 - 0.4" OrderedCollections = "1, 2" Printf = "1" Random = "1" +Random123 = "1.7.1" Reexport = "0.2, 1" SciMLBase = "2, 3" SpecialFunctions = "0.7.2, 0.8, 0.9, 0.10, 1, 2" diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl index b9aa08e8f6..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 2bef555edc..ea46a6c0c4 100644 --- a/src/mcmc/particle_mcmc.jl +++ b/src/mcmc/particle_mcmc.jl @@ -1,9 +1,62 @@ -### -### 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*: the sampler state +# carries those values, and the reference re-runs the model with `InitFromParams`, so it stays the +# retained trajectory even when the model is re-conditioned between Gibbs sweeps (e.g. a +# state-space transition prior that depends on a parameter owned by another Gibbs component). A +# particle forked from the reference forgets the remaining values (in `reseed!`), so branching +# samples fresh with no per-particle flag. +# +# Reference: Andrieu, Doucet & Holenstein, "Particle Markov chain Monte Carlo methods", Journal of +# the Royal Statistical Society: Series B 72(3), 269-342 (2010). + +using StatsFuns: softmax, logsumexp +import Random123 + +# +# Particle random number generation +# + +# Each particle owns a counter-based `Random123.Philox2x`. This section comes first because +# `Particle` names the generator type in its signature. + +"A fresh counter-based generator for one particle, seeded from `rng`." +function particle_rng(rng::AbstractRNG=Random.default_rng()) + return Random.seed!(Random123.Philox2x(), rand(rng, Random.Sampler(rng, UInt64))) +end + +# Derive a fresh seed from `key`. Splitting one generator into many by re-seeding is fragile +# in two ways: the derived seeds can yield *correlated* streams (Steele et al., "Fast +# Splittable Pseudorandom Number Generators", OOPSLA 2014), and a stdlib `MersenneTwister` +# derivation is not identical across Julia versions (Julia does not guarantee reproducible +# streams), which made SMC/PG drift between versions even under a StableRNG. Both bit the +# previous AdvancedPS implementation (#2781, AdvancedPS.jl#110). Philox is a counter-based +# generator with a fixed, portable algorithm and strong avalanche, so deriving the seed +# through it is both well-decorrelated from its parent and version-stable. +split_key(key::Integer) = rand(Random.seed!(Random123.Philox2x(), key), typeof(key)) + +"Reseed from the generator's own current state (used between steps when not resampling)." +refresh!(rng::Random123.Philox2x) = Random.seed!(rng, split_key(rng.key)) + +# +# Model evaluation via Libtask +# + +# A `Particle` is the only mutable state. It is stored as its `TapedTask`'s "taped globals", +# so the tilde overloads reach it from *inside* a running model via `get_taped_globals`. +# This keeps all state explicit on the particle -- no `task_local_storage`. + +# 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,543 +68,748 @@ function error_if_threadsafe_eval(model::DynamicPPL.Model) return nothing end -### AdvancedPS models and interface +""" + SMCContext -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 +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 + +# `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 -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. +""" + Particle(model, rng) + Particle(model, rng, retained::Particle) + +A single particle: a suspended `model` execution together with its `varinfo`, its own `rng`, and an +accumulated `logweight`. It also serves directly as the particle Gibbs sampler state (there is no +separate state struct). + +Without `retained` the particle draws from the prior. Given the previous sweep's retained particle it +becomes a conditional-SMC reference pinned to that trajectory, erroring if its execution reaches an +address the retained trajectory lacks or finishes without reaching one it has. Taking the whole +particle, rather than its values and addresses separately, is what makes a half-specified reference +unrepresentable; only those two pieces are kept, so the retained particle is not held alive. +""" +mutable struct Particle{RT<:AbstractRNG,WT<:Real} + # Abstract on purpose: the VarInfo type can change during PG-inside-Gibbs. Accesses go + # 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; one field rather than two so the pair + # cannot get out of step, and so `isreference` has a single thing to test. # - # 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) + # `values` is the retained trajectory, which the reference reproduces by reusing it + # (`InitFromParams` in `tilde_assume!!`). Reusing the *value* rather than replaying the RNG draw + # is what keeps the reference on that trajectory when Gibbs re-conditions the model: a draw is + # x = g(u; θ) in the RNG output u and the distribution parameters θ, so replaying u after θ → θ' + # yields g(u; θ') ≠ x -- e.g. x ~ Normal(μ, 1) is μ + Φ⁻¹(u), which shifts by μ' − μ. + # + # `varnames` is the set of addresses the retained trajectory assumed, and cannot be recovered + # from `values`: a slice assume such as `x[1:2] ~ MvNormal(...)` is stored under the keys `x[1]`, + # `x[2]` but assumed under the single address `x[1:2]`, so comparing against those keys would + # report a spurious trace change. Without it an address the retained trajectory never had would + # silently draw from the prior, corrupting the reference. + reference::Union{ + Nothing, + @NamedTuple{values::DynamicPPL.VarNamedTuple, varnames::Set{DynamicPPL.VarName}} + } + # Addresses assumed by this execution, in the same form as `reference.varnames`. Survives + # forking, so a particle that becomes the retained state hands the complete set to the next + # reference; a reference must finish having assumed exactly that set. + assumed_varnames::Set{DynamicPPL.VarName} + task::Libtask.TapedTask + # `task` is filled in once the particle exists, because the task must capture the + # particle as its taped globals (a back-reference). This has to be an inner constructor + # for that reason: `task` is left undefined here and set immediately after. + function Particle( + vi::DynamicPPL.AbstractVarInfo, rng::RT, retained::Union{Nothing,Particle}=nothing + ) where {RT<:AbstractRNG} + w = zero(DynamicPPL.LogProbType) + reference = if retained === nothing + nothing + else + (; + values=DynamicPPL.get_raw_values(retained.varinfo), + varnames=copy(retained.assumed_varnames), + ) + end + return new{RT,typeof(w)}(vi, rng, w, reference, Set{DynamicPPL.VarName}()) + end end -function AdvancedPS.advance!( - trace::AdvancedPS.Trace{<:AdvancedPS.LibtaskModel{<:TracedModel}}, isref::Bool=false +function Particle( + model::DynamicPPL.Model, rng::AbstractRNG, retained::Union{Nothing,Particle}=nothing ) - # 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. - return TracedModel(trace.model, trace.varinfo, true, trace.fargs, trace.kwargs) -end - -function AdvancedPS.reset_model(trace::TracedModel) - return trace -end - -function Libtask.TapedTask(taped_globals, model::TracedModel) - return Libtask.TapedTask( - taped_globals, model.fargs[1], model.fargs[2:end]...; model.kwargs... - ) + model = DynamicPPL.setleafcontext(model, SMCContext()) + varinfo = particle_varinfo() + args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo) + particle = Particle(varinfo, rng, retained) + particle.task = Libtask.TapedTask(particle, model.f, args...; kwargs...) + return particle end -abstract type ParticleInference <: AbstractSampler end - -#### -#### Generic Sequential Monte Carlo sampler. -#### - """ -$(TYPEDEF) - -Sequential Monte Carlo sampler. - -# 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 any address outside the retained set, so once the + # reference has run to completion the only discrepancy left to catch is a retained address + # it never visited (e.g. a branch that stopped being taken after re-conditioning). + reference = particle.reference + if score === nothing && reference !== nothing + dropped = setdiff(reference.varnames, particle.assumed_varnames) + isempty(dropped) || error( + "the reference execution trace changed while replaying retained values " * + "(retained addresses never reached: $(collect(dropped)))", + ) + 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, ::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 + vn in reference.varnames || error( + "the reference execution trace changed while replaying retained values " * + "(new address: $vn)", + ) + DynamicPPL.InitFromParams(reference.values, 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, particle.varinfo) + particle.varinfo = vi + push!(particle.assumed_varnames, vn) + return x, vi end -function AbstractMCMC.step( - rng::AbstractRNG, - model::DynamicPPL.Model, - spl::SMC; - nparticles::Int, - initial_params, - discard_sample=false, - kwargs..., +# Routes the observe through the particle's varinfo and stores the result back. The weight is not +# emitted here -- `ProduceLogLikelihoodAccumulator` produces it as it accumulates, below. +function DynamicPPL.tilde_observe!!( + ::SMCContext, + dist::Distribution, + left, + vn::Union{VarName,Nothing}, + template, + ::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, particle.varinfo ) + 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 here, in the single place the likelihood is accumulated, so `val` *is* the +# increment: "the produced weight equals the accumulator's increment" becomes structural rather than +# an invariant to keep two call sites in step with. Both routes reach the accumulator through exactly +# one `acclogp` -- an `observe` via `accumulate_observe!!`, and `@addlogprob!` via +# `accloglikelihood!!` -> `map_accumulator!!` -- so each emits exactly one weight, which is what +# gets `@addlogprob!` terms into the weight as well as the accumulator (issue #1996). +# +# Accumulator merging cannot fire a spurious produce: `combine` adds the two `logp`s directly rather +# than going through `acclogp`, so a submodel's varinfo folding into its parent's stays silent. +# +# `produce` suspends the task before the caller assigns the updated varinfo back onto the particle, so +# a *suspended* particle's accumulated total lags this term. Nothing reads it in that state -- the +# sweep reweights from `logweight`, and every varinfo read (`pg_transition_and_state`, +# `gibbs_get_raw_values`, SMC's bundling) happens once the model has run to completion. +function DynamicPPL.acclogp(acc::ProduceLogLikelihoodAccumulator, val) + Libtask.produce(val) + return ProduceLogLikelihoodAccumulator(DynamicPPL.logp(acc) + val) end -function 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 +) + return acc +end +function DynamicPPL.accumulate_observe!!( + acc::ProduceLogLikelihoodAccumulator, dist, left, vn, template ) - # Extract the index of the current particle. - index = state.particleindex + return DynamicPPL.acclogp(acc, Distributions.loglikelihood(dist, left)) +end - # Extract the current particle and its weight. - particles = state.particles - particle = particles.vals[index] - weight = AdvancedPS.getweight(particles, index) +# 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!! -> map_accumulator!! -> acclogp +# (the `@addlogprob! (; ...)` NamedTuple form routes through acclogp!! first) +# Gibbs: GibbsContext turns a tilde_assume!! into a tilde_observe!! +Libtask.@might_produce(DynamicPPL.tilde_observe!!) +Libtask.@might_produce(DynamicPPL.accumulate_observe!!) +Libtask.@might_produce(DynamicPPL.acclogp) +Libtask.@might_produce(DynamicPPL.tilde_assume!!) +Libtask.@might_produce(DynamicPPL.accloglikelihood!!) +Libtask.@might_produce(DynamicPPL.map_accumulator!!) +Libtask.@might_produce(DynamicPPL.acclogp!!) +# Every model / submodel evaluator takes a `DynamicPPL.Model`, so this covers them all. +# See https://github.com/TuringLang/Libtask.jl/issues/217. +Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true - # 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) +# Swap the default likelihood accumulator for the produce-aware one that drives reweighting, and +# add the raw sampled values. `OnlyAccsVarInfo`'s defaults also bring `LogPrior` and `LogJacobian` +# along, and they are kept on purpose: `ParamsWithStats` reads them straight off this varinfo to +# fill a chain's `logprior` and `logjoint` columns. Dropping them to save the per-particle logpdf +# work would not error -- the read is guarded by `hasacc` -- it would silently omit those columns. +function particle_varinfo() + vi = DynamicPPL.OnlyAccsVarInfo() + vi = DynamicPPL.setacc!!(vi, ProduceLogLikelihoodAccumulator()) + vi = DynamicPPL.setacc!!(vi, DynamicPPL.RawValueAccumulator(true)) + return vi +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 - nextstate = SMCState(state.particles, index + 1, state.average_logevidence) - - return transition, nextstate + return indices end -#### -#### Particle Gibbs sampler. -#### - -""" -$(TYPEDEF) - -Particle Gibbs sampler. - -# Fields - -$(TYPEDFIELDS) -""" -struct PG{R} <: ParticleInference - """Number of particles.""" - nparticles::Int - """Resampling algorithm.""" - resampler::R +"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 -""" -PG(n, [resampler = AdvancedPS.ResampleWithESSThreshold()]) -PG(n, [resampler = AdvancedPS.resample_systematic, ]threshold) - -Create a Particle Gibbs sampler of type [`PG`](@ref) with `n` particles. - -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()) +"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 -# Convenient constructors with ESS threshold -function PG(nparticles::Int, resampler, threshold::Real) - return PG(nparticles, AdvancedPS.ResampleWithESSThreshold(resampler, threshold)) -end -function PG(nparticles::Int, threshold::Real) - return PG(nparticles, AdvancedPS.resample_systematic, threshold) -end +## +## Effective-sample-size gating +## """ - CSMC(...) + ESSThresholdResampler(threshold, scheme = StratifiedResampler()) -Equivalent to [`PG`](@ref). +Resample with `scheme`, but only when the effective sample size drops below +`threshold * nparticles`. This is the default for [`SMC`](@ref) and [`PG`](@ref). """ -const CSMC = PG # type alias of PG as Conditional SMC - -struct PGState{V<:DynamicPPL.AbstractVarInfo,R<:Random.AbstractRNG} - vi::V - rng::R +struct ESSThresholdResampler{T<:Real,R<:AbstractResampler} <: AbstractResampler + threshold::T + scheme::R +end +function ESSThresholdResampler(threshold::Real) + return ESSThresholdResampler(threshold, StratifiedResampler()) 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 +## +## Weights and diagnostics +## + +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). Factored out so the serial and multithreaded loops in +# `reweight!` share one body. +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. +# +# `multithreaded` is *within-sweep* parallelism -- spreading this sweep's particle evaluations +# across threads. It is a separate axis from AbstractMCMC's chain-level ensemble +# (`MCMCThreads`/`MCMCDistributed`, which runs whole chains independently); the two compose. +# Only threading is offered here, not distribution: particles resample every step (all-to-all) +# and are live Libtask tasks, so spreading one sweep across processes would be communication- +# bound rather than a speed-up. +# +# Each particle advances only its own state (rng, varinfo, task), and its rng was already +# seeded serially in `resample_propagate!`, so the multithreaded loop is race-free and gives +# results identical to the serial one. Only the model evaluations parallelise; the shared +# sampler rng is untouched here. +function reweight!(particles, multithreaded::Bool) + n = length(particles) + if multithreaded + # A shared counter would race, so collect per-particle results and tally afterwards. + finished = Vector{Bool}(undef, n) + Threads.@threads for i in 1:n + finished[i] = advance_particle!(particles[i]) + end + n_done = count(finished) else - DynamicPPL.ParamsWithStats(deepcopy(_vi), (; logevidence=logevidence)) + n_done = count(advance_particle!, particles) end - - return transition, PGState(_vi, reference.rng) + 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 -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) +## +## Resample and propagate +## + +# Resample (if the scheme calls for it) and propagate the survivors, or -- when not resampling -- +# refresh each ordinary particle's seed so the next step draws fresh randomness. Returns whether it +# resampled, which tells `sweep!` what the total weight now is without recomputing it. +# +# Whether this is a conditional sweep is read off the particles rather than passed in: the reference +# always occupies the last slot, so `isreference` is the single source of truth and resampling cannot +# disagree with the rest of the sweep about which particle is pinned. +function resample_propagate!(rng::AbstractRNG, particles, resampler) + n = length(particles) + conditional = isreference(last(particles)) + weights = normalized_weights(particles) + if should_resample(resampler, weights) + # A conditional sweep draws the `n-1` free ancestors independently from the categorical + # 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 - return reference + resample_indices(rng, resampler, weights, n) end - 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] - - # Compute the transition. - _vi = newreference.model.f.varinfo - transition = if discard_sample - nothing + old = copy(particles) + seen = falses(n) + for (slot, a) in enumerate(ancestors) + # Reuse each surviving parent's object for its first offspring; only extra + # offspring -- and any offspring of the retained reference -- need the costly + # `deepcopy`. Either way the child is reseeded to continue independently. + reuse = !seen[a] && !isreference(old[a]) + seen[a] = true + child = reuse ? reseed!(old[a], rng) : fork(old[a], rng) + child.logweight = zero(DynamicPPL.LogProbType) + particles[slot] = child + end + # reference retained, weight reset + conditional && (particles[n].logweight = zero(DynamicPPL.LogProbType)) + return true else - DynamicPPL.ParamsWithStats(deepcopy(_vi), (; logevidence=logevidence)) + # The reference draws nothing (it reuses retained values), so only the others need a + # fresh seed for the next step. + for p in particles + isreference(p) || refresh!(p.rng) + end + return false end - - return transition, PGState(_vi, newreference.rng) 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, and `PG` runs thousands of +# sweeps, so computing them unconditionally would be pure waste on the sampler that sweeps most. +function sweep!( + rng::AbstractRNG, particles, resampler, multithreaded::Bool; ess::Bool=false +) + logZ = zero(DynamicPPL.LogProbType) + # The ESS values are computed from the particle weights, so they follow whatever + # `DynamicPPL.LogProbType` is rather than being pinned to `Float64`. + ess_per_step = DynamicPPL.LogProbType[] + # Total log weight entering the step. Resampling zeroes every weight, so it is then exactly + # `log(n)`; otherwise the weights are untouched and it is still last step's total. Either way + # there is nothing to recompute -- particles start at weight zero, hence `log(n)` initially. + logZ0 = log(oftype(logZ, length(particles))) + while true + resampled = resample_propagate!(rng, particles, resampler) + resampled && (logZ0 = log(oftype(logZ, length(particles)))) + done = reweight!(particles, multithreaded) + # Each observation contributes the log-ratio of total weight it adds; summed over the + # sweep these telescope into an estimate of the model's log-evidence log p(y). + total = log_normalizing_constant(particles) + logZ += total - logZ0 + logZ0 = total + done && break + # Post-reweight ESS for this observation: a degeneracy diagnostic (low ESS means few + # particles carry the weight). After the break, so the finishing pass -- which adds no + # observation and leaves the weights unchanged -- contributes no spurious entry. + ess && push!(ess_per_step, weight_ess(normalized_weights(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 -end +$(TYPEDEF) -""" - get_trace_local_rng() +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). -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 +# SMC is a single weighted sweep, not a Markov chain: rather than fake an iteration through +# AbstractMCMC's step loop (returning the population one particle at a time), we run the sweep +# and bundle the whole population into the chain in one shot. `discard_initial`/`thinning` +# therefore have nothing to apply to. +function AbstractMCMC.sample( + rng::AbstractRNG, + model::DynamicPPL.Model, + sampler::SMC, + nparticles::Integer; + check_model=true, + chain_type=DEFAULT_CHAIN_TYPE, + discard_initial=0, + thinning=1, + initial_params=nothing, + callback=nothing, + verbose=false, + kwargs..., ) - # 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) + 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 - ctx = DynamicPPL.InitContext(trng, init_strat, DynamicPPL.UnlinkAll()) - x, vi = DynamicPPL.tilde_assume!!(ctx, dist, vn, template, vi) - # Set the varinfo back in the trace. - set_trace_local_varinfo(vi) - return x, vi -end - -function DynamicPPL.tilde_observe!!( - ::ParticleMCMCContext, - right::Distribution, - left, - vn::Union{VarName,Nothing}, - template::Any, - vi::AbstractVarInfo, -) - vi = get_trace_local_varinfo() - left, vi = DynamicPPL.tilde_observe!!(DefaultContext(), right, left, vn, template, vi) - set_trace_local_varinfo(vi) - return left, vi + if initial_params !== nothing && !(initial_params isa DynamicPPL.InitFromPrior) + @warn "SMC draws its initial population from the prior; `initial_params` is ignored." + end + # Accepted only so it can be reported as ignored: AbstractMCMC's contract is one callback + # per step, and SMC is a single sweep, so there is no iteration to call back from. + if callback !== nothing + @warn "SMC runs one sweep rather than an MCMC loop, so there are no per-iteration callbacks; `callback` is ignored." + end + particles = [Particle(model, particle_rng(rng)) for _ in 1:nparticles] + logZ, ess_per_step = sweep!( + rng, particles, sampler.resampler, sampler.multithreaded; ess=true + ) + weights = normalized_weights(particles) + # One final resampling step, so the returned particles are an equal-weight sample. The + # sweep ends on a reweight, leaving the population weighted; resampling once here makes the + # result a standard unweighted chain (so `mean(chain[...])` and friends need no weighting), + # at the cost of a little resampling variance. Unconditional -- unlike the ESS-gated + # resampling inside the sweep. + ancestors = resample_indices(rng, sampler.resampler, weights, nparticles) + # `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 = AbstractMCMC.bundle_samples( + transitions, model, sampler, nothing, chain_type; kwargs... + ) + post_sample_hook(chain, sampler; verbose) + return chain end -# Convenient constructor -function AdvancedPS.Trace( - model::Model, varinfo::AbstractVarInfo, rng::AdvancedPS.TracedRNG, resample::Bool -) - newvarinfo = deepcopy(varinfo) - tmodel = TracedModel(model, newvarinfo, rng, resample) - newtrace = AdvancedPS.Trace(tmodel, rng) - return newtrace -end +# +# 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::Int, resampler::R; multithreaded::Bool=false + ) where {R<:AbstractResampler} + 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). + +!!! warning "`log_normalizing_constant` is biased for PG" + PG chains carry `log_normalizing_constant`, but unlike [`SMC`](@ref)'s it does **not** estimate + `log p(y)` without bias, so it must not be used for model comparison. A conditional sweep + retains the reference whatever its weight, and the reference is a draw from the posterior rather + than from the proposal, so it usually carries far more likelihood than a fresh particle and + inflates the mean weight at every step. Measured against the exact `p(y)` of a linear Gaussian + SSM, `E[Ẑ]` overshoots by 80% at `n = 16` and 16% at `n = 64`; the bias decays like `1/n` but + stays large at practical `n`. Use `SMC` for an unbiased estimate. +""" +PG(n::Int; kwargs...) = PG(n, ESSThresholdResampler(0.5); kwargs...) +PG(n::Int, threshold::Real; kwargs...) = PG(n, ESSThresholdResampler(threshold); kwargs...) +function PG(n::Int, 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 -end -function DynamicPPL.accumulate_observe!!( - acc::ProduceLogLikelihoodAccumulator, right, left, vn, template +"Conditional SMC, an alias for [`PG`](@ref)." +const CSMC = PG + +# PG's sampler state is just the retained `Particle`: it already carries the reference +# trajectory's `varinfo` and `rng` (its `task`/`logweight` are then unused), so there is no +# dedicated state struct. + +# First iteration: an ordinary (unconditional) particle sweep. +function AbstractMCMC.step( + rng::AbstractRNG, model::DynamicPPL.Model, sampler::PG; discard_sample=false, kwargs... ) - return DynamicPPL.acclogp(acc, Distributions.loglikelihood(right, left)) + error_if_threadsafe_eval(model) + 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::Particle; + discard_sample=false, + kwargs..., ) - return true + error_if_threadsafe_eval(model) + n = sampler.nparticles + # Passing `state` makes this the reference, pinned to the retained trajectory by value (see + # the `reference` field). Its own generator is never read, since every draw is supplied by value -- + # only its forks' are, and `reseed!` gives those fresh seeds. So it carries the retained + # generator forward rather than taking a fresh one, which keeps the sweep from drawing anything + # from `rng` on its behalf; the copy just avoids aliasing `state`. + reference = Particle(model, deepcopy(state.rng), state) + # `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 + DynamicPPL.ParamsWithStats( + deepcopy(retained.varinfo), (; log_normalizing_constant=logZ) + ) + end + return transition, retained 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::Particle) = DynamicPPL.get_raw_values(state.varinfo) function gibbs_update_state!!( - ::PG, state::PGState, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple + ::PG, state::Particle, model::DynamicPPL.Model, global_vals::DynamicPPL.VarNamedTuple ) - init_strat = DynamicPPL.InitFromParams(global_vals, nothing) - new_vi = last(DynamicPPL.init!!(model, state.vi, init_strat, DynamicPPL.UnlinkAll())) - return PGState(new_vi, state.rng) + init = DynamicPPL.InitFromParams(global_vals, nothing) + # Re-initialise the reference varinfo with the values conditioned by other Gibbs + # components. Mutating in place is safe: the caller replaces this state with the value we + # return and never reads the pre-update one again. + state.varinfo = last( + DynamicPPL.init!!(model, state.varinfo, init, DynamicPPL.UnlinkAll()) + ) + return state end diff --git a/test/Aqua.jl b/test/Aqua.jl index e5b655c6e0..e7a5ab6d1b 100644 --- a/test/Aqua.jl +++ b/test/Aqua.jl @@ -24,6 +24,8 @@ using Turing # meaningful in practice (in particular, to trigger this we would need to call `g(..., f)`, # which is incredibly unlikely). Aqua.test_ambiguities([Turing]; exclude=[Libtask.might_produce]) -Aqua.test_all(Turing; ambiguities=false) + +# `persistent_tasks` is flaky on Windows; a lingering Task is a code property, not an OS one. +Aqua.test_all(Turing; ambiguities=false, persistent_tasks=!Sys.iswindows()) end diff --git a/test/Project.toml b/test/Project.toml index 1110301a92..0e0bdcfb69 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,7 +3,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf" AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -AdvancedPS = "576499cb-2369-40b2-a588-c64705576edc" AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" @@ -45,7 +44,6 @@ ADTypes = "1" AbstractMCMC = "5.13" AbstractPPL = "0.14, 0.15" AdvancedMH = "0.8.9" -AdvancedPS = "0.7.2" AdvancedVI = "0.7" Aqua = "0.8" BangBang = "0.4" diff --git a/test/essential/container.jl b/test/essential/container.jl 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..a5a62a224f 100644 --- a/test/mcmc/ess.jl +++ b/test/mcmc/ess.jl @@ -60,7 +60,14 @@ using Turing @testset "gdemo with CSMC + ESS" begin alg = Gibbs(:s => CSMC(15), :m => ESS()) - chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 3_000) + # CSMC mixes the variance `s` slowly, so the Monte Carlo error at 3_000 draws + # exceeds `atol` on this seed (measured |err| 0.212); the estimator is unbiased, so + # 10_000 draws buy the headroom back. That headroom is not large: conditional + # sweeps draw their ancestors multinomially, which is noisier than the stratified + # scheme used before, and `s` moved from |err| 0.024 to 0.061 against `atol = 0.1` + # as a result. If this test starts failing, suspect the draw count before the + # sampler. + chain = sample(StableRNG(seed), gdemo(1.5, 2.0), alg, 10_000) check_gdemo(chain; atol=0.1) end diff --git a/test/mcmc/particle_mcmc.jl b/test/mcmc/particle_mcmc.jl index 2cbc4dbece..19a932e85e 100644 --- a/test/mcmc/particle_mcmc.jl +++ b/test/mcmc/particle_mcmc.jl @@ -1,41 +1,153 @@ module ParticleMCMCTests using ..Models: gdemo_default -using ..SamplerTestUtils: test_chain_logp_metadata -using AdvancedPS: ResampleWithESSThreshold, resample_systematic, resample_multinomial -using Distributions: Bernoulli, Beta, 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 Random: Random, Xoshiro +using SpecialFunctions: logbeta using StableRNGs: StableRNG using Test: @test, @test_logs, @test_throws, @testset using Turing -@testset "SMC" begin - @testset "constructor" begin - s = SMC() - @test s.resampler == ResampleWithESSThreshold() +# +# 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 - s = SMC(0.6) - @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6) +@model function normal() + a ~ Normal(4, 5) + 3 ~ Normal(a, 2) + b ~ Normal(a, 1) + 1.5 ~ Normal(b, 2) + return a, b +end - s = SMC(resample_multinomial, 0.6) - @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6) +# Nondeterministic evaluation order, which the particle samplers must refuse. +# As `normal()` but centred at zero; used where the replay test wants a different trajectory. +@model function centred_normal() + a ~ Normal(0, 1) + 3 ~ Normal(a, 2) + b ~ Normal(a, 1) + 1.5 ~ Normal(b, 2) + return a, b +end - s = SMC(resample_systematic) - @test s.resampler === resample_systematic +# Nondeterministic evaluation order, which the particle samplers must refuse. +@model function threadsafe_model(y) + x ~ Normal() + Threads.@threads for i in eachindex(y) + y[i] ~ Normal(x) + end +end + +"Run a particle to completion." +run_to_end!(p) = (while advance!(p) !== nothing +end; +p) + +# +# SMC +# + +@testset "SMC" begin + @testset "constructor" begin + @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 end @testset "basic model" begin - @model function normal() - a ~ Normal(4, 5) - 3 ~ Normal(a, 2) - b ~ Normal(a, 1) - 1.5 ~ Normal(b, 2) - return a, b - end tested = sample(normal(), SMC(), 100) end + @testset "resampling schemes" begin + 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 + + @testset "stratified/systematic resampling never index past the end" begin + # softmax can return weights summing to slightly under one; the cumulative walk must + # not run off the end of the last stratum. Exaggerate the undersum so the (otherwise + # astronomically rare) overrun is hit on every seed. + weights = fill(0.9 / 8, 8) + for scheme in (StratifiedResampler(), SystematicResampler()) + @test all( + all(in(1:8), resample_indices(Xoshiro(s), scheme, weights, 8)) for + s in 1:1000 + ) + end + end + @testset "errors when number of observations is not fixed" begin @model function fail_smc() a ~ Normal(4, 5) @@ -54,50 +166,58 @@ 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 "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 - end + @testset "discard_initial, thinning, initial_params and callback are 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() + ) + + # The ensemble wrapper injects the sampler's own default `InitFromPrior()` per chain. + # That is not a user-specified initialisation, so it must not warn. + @test_logs sample(Xoshiro(1), normal(), SMC(), MCMCSerial(), 10, 2; progress=false) + + # A callback is accepted only to be reported as ignored, and must not run. + called = false + @test_logs (:warn, r"callback.*ignored") sample( + normal(), SMC(), 10; callback=(args...; kwargs...) -> (called = true) + ) + @test !called @test_logs (:warn, r"ignored") sample(normal(), SMC(), 10; discard_initial=5) chn = sample(normal(), SMC(), 10; discard_initial=5) @@ -118,47 +238,102 @@ using Turing 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 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), retained) + particles = [Particle(model, particle_rng(rng)) for _ in 1:4] + push!(particles, reference) + sweep!(StableRNG(78), particles, scheme, false) + return map(p -> get_raw_values(p.varinfo), particles) + 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 +343,173 @@ end @test length(unique(c[:s])) == 1 end + @testset "ensuring reference consistency" begin + # In conditional SMC the retained trajectory must be regenerated *exactly* by the + # reference particle on the next iteration -- this is what makes CSMC valid. 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) + # The sampler's own selection rule, not a reimplementation of it: if `pg_transition_and` + # `_state` ever changes how the retained particle is chosen, this test follows. + draw(ps) = last(pg_transition_and_state(rng, ps, 0.0, true)) + particles = [Particle(model, particle_rng(rng)) for _ in 1:N] + sweep!(rng, particles, ESSThresholdResampler(0.5), false) + state = draw(particles) + allok = true + nlatents = 0 + for _ in 1:nsteps + ref = Particle(model, particle_rng(rng), state) + parts = [Particle(model, particle_rng(rng)) for _ in 1:(N - 1)] + push!(parts, ref) + sweep!(rng, parts, ESSThresholdResampler(0.5), false) + allok &= get_raw_values(parts[N].varinfo) == get_raw_values(state.varinfo) + state = draw(parts) + nlatents = length(state.assumed_varnames) + end + return allok, nlatents + 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, nlatents = run_csmc(state_space_model(y), 3, 30, rng) + @test allok # reference regenerated exactly every step + @test nlatents == length(y) + 2 + end + + @testset "reference is pinned to retained values under re-conditioning" begin + # Finding 1 regression. In Gibbs the model is re-conditioned between sweeps, so the + # CSMC reference must reproduce the *retained values* rather than re-draw them from the + # (now different) prior. Retain a trajectory under one conditioning, rebuild the + # reference under another; value-pinning keeps the trajectory, whereas re-drawing from + # the prior -- the pre-fix behaviour -- would follow the shifted prior instead. + @model function reconditioned(y) + a ~ Normal(0, 10) + x ~ Normal(a, 1) # x's prior depends on a, owned by another component + return y ~ Normal(x, 1) + end + rng = StableRNG(42) + retained = Particle(reconditioned(2.0) | (@varname(a) => 0.0), particle_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, + ) + 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), retained) + @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), retained) + @test_throws "reference execution trace changed" begin + run_to_end!(reference) + end + end + + @testset "value replay tracks slice assumes by their assumed address" begin + # `x[1:2] ~ MvNormal(...)` is assumed under the single address `x[1:2]` but stored in + # the retained values under the keys `x[1]`, `x[2]`. The expected-address set must + # therefore come from what the trajectory assumed, not from the retained values' keys, + # or every slice assume looks like a changed trace. + @model function slice_assume(y) + x = Vector{Float64}(undef, 2) + x[1:2] ~ MvNormal(zeros(2), I) + y[1] ~ Normal(x[1], 0.5) + return y[2] ~ Normal(x[2], 0.5) + end + chn = sample(StableRNG(105), slice_assume([0.4, -0.4]), PG(5), 20) + @test size(chn, 1) == 20 + end + + @testset "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 "addlogprob leads to reweighting" begin # Make sure that PG takes @addlogprob! into account. It didn't use to: # https://github.com/TuringLang/Turing.jl/issues/1996 @@ -184,6 +526,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 @@ -235,21 +581,271 @@ 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) - x ~ Normal() - Threads.@threads for i in eachindex(y) - y[i] ~ Normal(x) + 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 "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)), retained) + 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(32), 4_000 + ) + for t in 1:T + xs = particle_draws(chn, @varname(x[t])) + test_within_mc_error(means[t], xs) + test_within_mc_error(vars[t], (xs .- mean(xs)) .^ 2) + end + end + + @testset "Gibbs(q => NUTS, x => CSMC) recovers the exact posterior" begin + prior = InverseGamma(3, 2) + qs = range(0.05, 4.0; length=400) + w = ExactSSM.grid_posterior( + prior, qs, q -> ExactSSM.lgssm_loglik(y, a, q, r, s0(q)) + ) + q_mean, q_sd = ExactSSM.grid_moments(w, qs) + smoothed = [ExactSSM.lgssm_smoother(y, a, q, r, s0(q)) for q in qs] + x_mean = sum(w[i] * first(smoothed[i]) for i in eachindex(w)) + x_second = sum( + w[i] * (last(smoothed[i]) .+ first(smoothed[i]) .^ 2) for i in eachindex(w) + ) + + alg = Gibbs(@varname(q) => NUTS(), @varname(x) => CSMC(32)) + chn = sample(StableRNG(31), lgssm(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(32), 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.2, 2.5; length=400) + w = ExactSSM.grid_posterior( + prior, sds, s -> last(ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, s))) + ) + sd_mean, sd_sd = ExactSSM.grid_moments(w, sds) + posts = [first(ExactSSM.hmm_forward_backward(π0, P, obs_loglik(y, s))) for s in sds] + mixed = sum(w[i] * posts[i] for i in eachindex(w)) + + alg = Gibbs(@varname(sd) => HMC(0.1, 12), @varname(z) => CSMC(32)) + chn = sample(StableRNG(32), hmm(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..a020bf6bca 100644 --- a/test/test_utils/numerical_tests.jl +++ b/test/test_utils/numerical_tests.jl @@ -24,6 +24,26 @@ function check_gdemo(chain; atol=0.2, rtol=0.0) ) end +# Exact posterior means of `MoGtest_default`, i.e. `MoGtest([1.0 1.0 4.0 4.0])`, for +# `[z1, z2, z3, z4, mu1, mu2]`. The four `z` are discrete, so enumerating all 2^4 cluster +# assignments and integrating `mu1`, `mu2` out analytically (given `z`, each cluster's +# observations are Gaussian with a Gaussian prior, so the marginal likelihood and the +# conditional posterior mean are both closed-form) gives these exactly. +# +# These are deliberately *not* the idealised labels `[1, 1, 2, 2, 1, 4]`. The data do not +# identify the clusters perfectly, so the truth sits 0.046 off each label and 0.072 off each +# prior mean. Asserting the idealised values spent that 0.072 of the tolerance before any Monte +# Carlo error -- half of a 0.15 `atol`, and 72% of the 0.1 one used in mcmc/ess.jl -- which made +# accurate samplers fail intermittently and made the rest look worse than they are. +const MOGTEST_DEFAULT_MEANS = [ + 1.0455583636758707, + 1.0455583636758707, + 1.9544416363241297, + 1.9544416363241297, + 1.071541120128846, + 3.928458879871154, +] + # Wrapper function to check MoGtest. function check_MoGtest_default(chain; atol=0.2, rtol=0.0) return check_numerical( @@ -36,7 +56,7 @@ function check_MoGtest_default(chain; atol=0.2, rtol=0.0) @varname(mu1), @varname(mu2) ], - [1.0, 1.0, 2.0, 2.0, 1.0, 4.0]; + MOGTEST_DEFAULT_MEANS; atol=atol, rtol=rtol, ) @@ -53,7 +73,7 @@ function check_MoGtest_default_z_vector(chain; atol=0.2, rtol=0.0) @varname(mu1), @varname(mu2) ], - [1.0, 1.0, 2.0, 2.0, 1.0, 4.0]; + MOGTEST_DEFAULT_MEANS; atol=atol, rtol=rtol, )