diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 9d9bf0804..7b344b46e 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -84,13 +84,14 @@ abstract type ZygoteProtocol <: AutodiffProtocol end abstract type EnzymeProtocol <: AutodiffProtocol end =# -function _get_proposal(encoded_prior::ParameterDistribution, encoder_schedule::VV) where {VV <: AbstractVector} - # We use the prior covariance to shape the proposal (in the encoded space), - # as proposals are based on increments we do not need to shift the mean too +function _get_cholesky_factor(encoded_prior::ParameterDistribution, encoder_schedule::VV) where {VV <: AbstractVector} + # We use the prior covariance to shape the proposal (in the encoded space), + # as proposals are based on increments we do not need to shift the mean too. + # The Cholesky factor L (C = LL') is the single source of truth every sampler below + # uses both to draw its noise and to evaluate its transition (log-)density, so the + # two can never drift apart. C = cov(encoded_prior) - Σ = cholesky(Symmetric(C)) - - return AdvancedMH.RandomWalkProposal(Σ.L * MvNormal(zeros(size(Σ, 1)), I)) + return cholesky(Symmetric(C)).L end @@ -112,15 +113,10 @@ struct RWMHSampling{T <: AutodiffProtocol} <: MCMCProtocol end RWMHSampling() = RWMHSampling{GradFreeProtocol}() -struct RWMetropolisHastings{PT, ADT <: AutodiffProtocol} <: AdvancedMH.MHSampler - proposal::PT +struct RWMetropolisHastings{LT, ADT <: AutodiffProtocol} <: AdvancedMH.MHSampler + "Lower Cholesky factor `L` of the (encoded) prior covariance `C = LL'`, shaping the proposal noise." + cholesky_L::LT end -# Define method needed by AdvancedMH for new Sampler -AdvancedMH.logratio_proposal_density( - sampler::RWMetropolisHastings, - transition_prev::AdvancedMH.AbstractTransition, - candidate, -) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate) """ $(TYPEDSIGNATURES) @@ -142,8 +138,8 @@ function MetropolisHastingsSampler( encoded_prior::ParameterDistribution, encoder_schedule::VV, ) where {T <: AutodiffProtocol, VV <: AbstractVector} - proposal = _get_proposal(encoded_prior, encoder_schedule) - return RWMetropolisHastings{typeof(proposal), T}(proposal) + L = _get_cholesky_factor(encoded_prior, encoder_schedule) + return RWMetropolisHastings{typeof(L), T}(L) end """ @@ -157,22 +153,17 @@ on the covariance of `prior`. struct pCNMHSampling{T <: AutodiffProtocol} <: MCMCProtocol end pCNMHSampling() = pCNMHSampling{GradFreeProtocol}() -struct pCNMetropolisHastings{D, T <: AutodiffProtocol} <: AdvancedMH.MHSampler - proposal::D +struct pCNMetropolisHastings{LT, T <: AutodiffProtocol} <: AdvancedMH.MHSampler + "Lower Cholesky factor `L` of the (encoded) prior covariance `C = LL'`, shaping the proposal noise." + cholesky_L::LT end -# Define method needed by AdvancedMH for new Sampler -AdvancedMH.logratio_proposal_density( - sampler::pCNMetropolisHastings, - transition_prev::AdvancedMH.AbstractTransition, - candidate, -) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate) function MetropolisHastingsSampler( ::pCNMHSampling{T}, encoded_prior::ParameterDistribution, encoder_schedule::VV, ) where {T <: AutodiffProtocol, VV <: AbstractVector} - proposal = _get_proposal(encoded_prior, encoder_schedule) - return pCNMetropolisHastings{typeof(proposal), T}(proposal) + L = _get_cholesky_factor(encoded_prior, encoder_schedule) + return pCNMetropolisHastings{typeof(L), T}(L) end #------ The following are gradient-based samplers @@ -186,23 +177,54 @@ new parameters according to the Barker proposal. struct BarkerSampling{T <: AutodiffProtocol} <: MCMCProtocol end BarkerSampling() = BarkerSampling{ForwardDiffProtocol}() -struct BarkerMetropolisHastings{D, T <: AutodiffProtocol} <: AdvancedMH.MHSampler - proposal::D +struct BarkerMetropolisHastings{LT, T <: AutodiffProtocol} <: AdvancedMH.MHSampler + "Lower Cholesky factor `L` of the (encoded) prior covariance `C = LL'`, shaping the proposal noise." + cholesky_L::LT end -# Define method needed by AdvancedMH for new Sampler -AdvancedMH.logratio_proposal_density( - sampler::BarkerMetropolisHastings, - transition_prev::AdvancedMH.AbstractTransition, - candidate, -) = AdvancedMH.logratio_proposal_density(sampler.proposal, transition_prev.params, candidate) function MetropolisHastingsSampler( ::BarkerSampling{T}, encoded_prior::ParameterDistribution, encoder_schedule::VV, ) where {T <: AutodiffProtocol, VV <: AbstractVector} - proposal = _get_proposal(encoded_prior, encoder_schedule) - return BarkerMetropolisHastings{typeof(proposal), T}(proposal) + L = _get_cholesky_factor(encoded_prior, encoder_schedule) + return BarkerMetropolisHastings{typeof(L), T}(L) +end + +# ------------------------------------------------------------------------------------------ +# Shared Metropolis-Hastings transition-density interface. +# +# Every sampler above is required to implement a function, +# log_transition_density(sampler, model, θ_from, θ_to; stepsize) = log q(θ_to | θ_from) +# the forward transition log-density that sampler's `propose` draws from. +# there is no fallback + +# AdvancedMH.logratio_proposal_density, fed into the MH acceptance ratio, is +# a shared logic among all AdvancedMH.MHSamplers. Applying the Hastings correction +# log q(θ_from | θ_to) − log q(θ_to | θ_from) (reverse − forward), by +# evaluating log_transition_density in both directions. +function log_transition_density( + sampler::MHS, + model, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {MHS <: AdvancedMH.MHSampler, FT <: AbstractFloat} + _throw_log_transition_density_not_implemented(MHS) +end + +function AdvancedMH.logratio_proposal_density( + sampler::MHS, + model::AdvancedMH.DensityModel, + transition_prev::AdvancedMH.AbstractTransition, + candidate; + stepsize::FT = 1.0, +) where {MHS <: AdvancedMH.MHSampler, FT <: AbstractFloat} + θ_from = transition_prev.params + θ_to = candidate + log_q_forward = log_transition_density(sampler, model, θ_from, θ_to; stepsize = stepsize) + log_q_reverse = log_transition_density(sampler, model, θ_to, θ_from; stepsize = stepsize) + return log_q_reverse - log_q_forward end @@ -328,29 +350,128 @@ function AdvancedMH.transition( return MCMCState(params, log_density, true) end -# method extending AdvancedMH.propose() to vanilla random walk with explicitly given stepsize +# ------------------------------------------------------------------------------------------ +# Markov transition kernels as Distributions.jl objects. +# +# For samplers whose transition kernel q(·|θ_from) is an ordinary Distributions.jl +# distribution, we go one step further than log_transition_density: rather than hand-writing +# `propose` (a sample) and `log_transition_density` (a logpdf) as two separate expressions that +# merely happen to describe the same kernel, we build the kernel ONCE as a Distributions.jl +# object and derive both `rand` and `logpdf` from it generically. The mean/covariance/ρ formulas +# then exist in exactly one place (`transition_kernel`) instead of two, so they cannot drift +# apart even in principle — this is strictly stronger than the log_transition_density interface +# above, which only guarantees forward/reverse agree with *each other*, not that `propose` +# agrees with either. +const DistributionKernelSampler = Union{RWMetropolisHastings, pCNMetropolisHastings} + function AdvancedMH.propose( rng::Random.AbstractRNG, - sampler::RWMetropolisHastings, + sampler::MHS, model::AdvancedMH.DensityModel, current_state::MCMCState; stepsize::FT = 1.0, +) where {MHS <: DistributionKernelSampler, FT <: AbstractFloat} + return rand(rng, transition_kernel(sampler, model, current_state.params; stepsize = stepsize)) +end + +function log_transition_density( + sampler::MHS, + model::AdvancedMH.DensityModel, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {MHS <: DistributionKernelSampler, FT <: AbstractFloat} + return logpdf(transition_kernel(sampler, model, θ_from; stepsize = stepsize), θ_to) +end + +# θ_to | θ_from ~ N(θ_from, stepsize² C) — the symmetric random-walk kernel. +function transition_kernel( + sampler::RWMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from; + stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - return current_state.params + stepsize * rand(rng, sampler.proposal) + L = sampler.cholesky_L + return MvNormal(θ_from, Symmetric((stepsize^2) .* (L * L'))) end -# method extending AdvancedMH.propose() for preconditioned Crank-Nicholson -function AdvancedMH.propose( - rng::Random.AbstractRNG, +# θ_to | θ_from ~ N(ρθ_from, (1-ρ²)C) — the asymmetric pCN kernel (Beskos et al. 2017 for the +# ρ-vs-Euler-stepsize relation). Evaluating this honestly in both directions (via the generic +# logratio_proposal_density) is what recovers the pCN cancellation log q(θ_from|θ_to) - +# log q(θ_to|θ_from) = logprior(θ_from) - logprior(θ_to), instead of the silently-symmetric 0 +# the previous implementation returned. +function transition_kernel( sampler::pCNMetropolisHastings, model::AdvancedMH.DensityModel, - current_state::MCMCState; + θ_from; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - # Use prescription in Beskos et al (2017) "Geometric MCMC for infinite-dimensional - # inverse problems." for relating ρ to Euler stepsize: ρ = (1 - stepsize / 4) / (1 + stepsize / 4) - return ρ * current_state.params .+ sqrt(1 - ρ^2) * rand(rng, sampler.proposal) + L = sampler.cholesky_L + return MvNormal(ρ .* θ_from, Symmetric((1 - ρ^2) .* (L * L'))) +end + +# The Barker transition kernel q(·|θ_from), as a lightweight (non-Distributions.jl) struct: the +# whitened-coordinate machinery (Livingstone & Zanella, 2022) needs a state- and model-dependent +# gradient computed once, so we bundle exactly what _barker_rand/_barker_logpdf need to stay +# consistent, and no more — unlike RW/pCN this isn't a Distributions.jl distribution (its density +# isn't a standard family), so we write the two functions directly instead of reimplementing +# Distributions.jl's dispatch machinery (_rand!/_logpdf/insupport/...) for a one-off type. +struct BarkerKernel{VT, MT, FT <: AbstractFloat} + "Current point θ_from that the kernel proposes away from." + θ_from::VT + "Gradient of the target log-density at θ_from, rotated into whitened coordinates: L'∇log π(θ_from)." + grad_white::VT + "Cholesky factor L of the prior covariance (C = LL'), used to whiten/unwhiten." + L::MT + stepsize::FT +end + +function transition_kernel( + sampler::BarkerMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + L = sampler.cholesky_L + grad_white = L' * autodiff_gradient(model, θ_from, sampler) + return BarkerKernel(θ_from, grad_white, L, stepsize) +end + +# Livingstone and Zanella (2022). The elementwise sigmoid selection that defines the Barker +# proposal is only reversible for *independent* coordinates, so we apply it in the whitened +# coordinate system u = L⁻¹θ (where the prior-covariance preconditioning L is decorrelated to the +# identity) and map the increment back with L. We also flip the sign of the whitened noise z +# rather than zeroing it out — the standard Barker kernel moves by ±z, never by exactly 0, so its +# transition density (_barker_logpdf below) is an ordinary, atom-free density; zeroing out +# (keep-or-drop) instead creates a mixed discrete/continuous kernel with no closed form. +# +# The flip decision must use the *actual* proposed whitened increment z = stepsize * η, not the +# unscaled η: the Barker flip probability is P(keep sign | z) = σ(grad_white · z). Using η instead +# of z here would leave the flip probability insensitive to `stepsize`, which defeats the +# stepsize-robustness that is this proposal's whole point (Barker's acceptance-vs-stepsize +# behaviour should be governed by the same z that appears in the increment). +function _barker_rand(rng::Random.AbstractRNG, kernel::BarkerKernel) + n = length(kernel.θ_from) + η = randn(rng, n) + z = kernel.stepsize .* η + flip_prob = 1 ./ (1 .+ exp.(-kernel.grad_white .* z)) + sign = 2 .* (rand(rng, n) .< flip_prob) .- 1 + ζ = sign .* z + return kernel.θ_from .+ kernel.L * ζ +end + +# For a symmetric noise density g (here N(0, stepsize²)) and whitened increment Δ = θ_to_white - +# θ_from_white, the marginal density of the ± flip-sign move is q(Δ | θ_from) = 2 g(Δ) σ(grad_white·Δ) +# (Barker, 1965; Livingstone & Zanella, 2022): summing the probability that Δ itself was kept (w.p. +# σ(grad_white·Δ)) with the probability that -Δ was drawn and flipped (w.p. σ(grad_white·Δ) too, +# since g(-Δ) = g(Δ)). See _barker_rand above for how θ_to is generated from θ_from — Δ there is +# `ζ`, the same quantity the flip decision is based on. +function _barker_logpdf(kernel::BarkerKernel, θ_to) + Δ = kernel.L \ (θ_to .- kernel.θ_from) + e = Δ ./ kernel.stepsize + return sum(log(2) .+ logpdf.(Normal(), e) .+ log.(1 ./ (1 .+ exp.(-kernel.grad_white .* Δ)))) - + length(e) * log(kernel.stepsize) end # method extending AdvancedMH.propose() for the Barker proposal @@ -361,12 +482,17 @@ function AdvancedMH.propose( current_state::MCMCState; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - # Livingstone and Zanella (2022) - # Compute the gradient of the log-density at the current state - n = length(current_state.params) - log_gradient = autodiff_gradient(model, current_state.params, sampler) - xi = rand(rng, sampler.proposal) - return current_state.params .+ (stepsize .* ((rand(rng, n) .< 1 ./ (1 .+ exp.(-log_gradient .* xi))) .* xi)) + return _barker_rand(rng, transition_kernel(sampler, model, current_state.params; stepsize = stepsize)) +end + +function log_transition_density( + sampler::BarkerMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + return _barker_logpdf(transition_kernel(sampler, model, θ_from; stepsize = stepsize), θ_to) end # Copy a MCMCState and set accepted = false @@ -393,7 +519,8 @@ function AbstractMCMC.step( AdvancedMH.logdensity(model, current_state) log_α = - new_log_density - current_log_density + AdvancedMH.logratio_proposal_density(sampler, current_state, new_params) + new_log_density - current_log_density + + AdvancedMH.logratio_proposal_density(sampler, model, current_state, new_params; stepsize = stepsize) # Decide whether to return the previous params or the new one. new_state = if -Random.randexp(rng) < log_α @@ -694,11 +821,8 @@ $(TYPEDSIGNATURES) Fraction of MC proposals in `chain` which were accepted (according to Metropolis-Hastings.) """ function accept_ratio(chain::MCMCChains.Chains) - if :accepted in names(chain, :internals) - return mean(chain, :accepted) - else - error("MH `:accepted` not recorded in MCMC chain — available internals: $(names(chain, :internals)).") - end + :accepted in names(chain, :internals) || _throw_accepted_not_recorded(chain) + return mean(chain, :accepted) end function _find_mcmc_step_log(mcmc::MCMCWrapper) @@ -722,11 +846,28 @@ end """ $(TYPEDSIGNATURES) -Uses a heuristic to return a stepsize for the `mh_proposal_sampler` element of +Uses a bracket-and-bisect search to return a stepsize for the `mh_proposal_sampler` element of [`MCMCWrapper`](@ref) which yields fast convergence of the Markov chain. -The criterion used is that Metropolis-Hastings proposals should be accepted between 15% and -35% of the time. +The criterion used is that Metropolis-Hastings proposals should be accepted between +`target_acc - tol` and `target_acc + tol` of the time (`tol = 0.1` by default). Acceptance rate +is assumed to (eventually) decrease as `stepsize` increases, but is not assumed to do so at any +particular rate: some proposals (e.g. gradient-informed ones such as the Barker proposal) can +have a much steeper acceptance-vs-stepsize transition than a vanilla random walk. The search +therefore first exponentially expands a bracket `[lo, hi]` around `init_stepsize` until `target_acc` +is known to lie strictly between the two trial acceptance rates, then bisects (in log-stepsize +space, i.e. by geometric mean) within that bracket. Bisection halves the bracket every iteration +regardless of how steep the transition is, so this converges reliably in cases where a fixed-factor +doubling/halving search can overshoot the target window and fail to converge within `max_iter`. + +The exponential expansion phase is capped separately at `max_expansions` doublings/halvings +(default 10, i.e. a ~1000x range around `init_stepsize`). Acceptance rate computed from a noisy, +numerically-evaluated log-density (e.g. from an emulator) need not actually decrease monotonically +all the way to stepsize → 0 — very small proposals can be dominated by the log-density's own +floating-point noise floor rather than its true (smooth) small-stepsize limit, degrading rather +than improving acceptance. If expansion hasn't bracketed `target_acc` within `max_expansions` +steps, that is reported directly (`target_acc` may not be achievable in this range) rather than +continuing to shrink/grow the stepsize toward a pathological value. """ function optimize_stepsize( rng::Random.AbstractRNG, @@ -734,51 +875,77 @@ function optimize_stepsize( init_stepsize = 1.0, N = 2000, max_iter = 20, + max_expansions = 10, target_acc = 0.25, + tol = 0.1, sample_kwargs..., ) - increase = false - decrease = false - stepsize = init_stepsize - factor = [1.0] - step_history = [true, true] _find_mcmc_step_log(mcmc) - for it in 1:max_iter + + n_evals = 0 + function acc_at(stepsize) + n_evals += 1 + n_evals <= max_iter || _throw_max_iter_exceeded(n_evals, max_iter, target_acc, tol, stepsize) trial_chain = sample(rng, mcmc, N; stepsize = stepsize, sample_kwargs...) acc_ratio = accept_ratio(trial_chain) - _find_mcmc_step_log(it, stepsize, acc_ratio, trial_chain) + _find_mcmc_step_log(n_evals, stepsize, acc_ratio, trial_chain) + return acc_ratio + end - change_step = true - if acc_ratio < target_acc - 0.1 - decrease = true - elseif acc_ratio > target_acc + 0.1 - increase = true - else - change_step = false - end + function done(stepsize) + @printf "Returning optimized stepsize: %.3g\n" stepsize + return stepsize + end + within_tol(acc) = abs(acc - target_acc) <= tol - if increase && decrease - factor[1] /= 2 - increase = false - decrease = false + stepsize = init_stepsize + acc = acc_at(stepsize) + within_tol(acc) && return done(stepsize) + + # Phase 1: exponentially expand a bracket [lo, hi] around init_stepsize, keeping the tightest + # known too-high-acceptance point (lo) and too-low-acceptance point (hi), until both are known. + # Capped at max_expansions: if acceptance hasn't crossed target_acc by then, it likely isn't + # monotonic in this range (rather than simply needing one more doubling/halving), so we stop + # and report that clearly instead of continuing toward a pathological stepsize. + if acc > target_acc + lo, hi = stepsize, stepsize + acc_hi = acc + n_expand = 0 + while acc_hi > target_acc + n_expand += 1 + n_expand <= max_expansions || + _throw_max_expansions_exceeded(:above, n_expand, max_expansions, target_acc, tol, hi) + lo = hi + hi *= 2 + acc_hi = acc_at(hi) + within_tol(acc_hi) && return done(hi) end - - if acc_ratio < target_acc - 0.1 - stepsize *= 2^(-factor[1]) - elseif acc_ratio > target_acc + 0.1 - stepsize *= 2^(factor[1]) + else + lo, hi = stepsize, stepsize + acc_lo = acc + n_expand = 0 + while acc_lo < target_acc + n_expand += 1 + n_expand <= max_expansions || + _throw_max_expansions_exceeded(:below, n_expand, max_expansions, target_acc, tol, lo) + hi = lo + lo /= 2 + acc_lo = acc_at(lo) + within_tol(acc_lo) && return done(lo) end + end - if change_step - @printf "Set sampler to new stepsize: %.3g\n" stepsize + # Phase 2: bisect within [lo, hi] (geometric mean, since stepsize tuning is multiplicative). + while true + mid = sqrt(lo * hi) + acc_mid = acc_at(mid) + within_tol(acc_mid) && return done(mid) + if acc_mid > target_acc + lo = mid else - @printf "Returning optimized stepsize: %.3g\n" stepsize - return stepsize + hi = mid end end - error( - "optimize_stepsize: acceptance rate $(round(acc_ratio; sigdigits = 3)) did not reach target $(target_acc) ± 0.1 within $(max_iter) iterations — last stepsize: $(round(stepsize; sigdigits = 3)).", - ) end # use default rng if none given optimize_stepsize(mcmc::MCMCWrapper; kwargs...) = optimize_stepsize(Random.GLOBAL_RNG, mcmc; kwargs...) @@ -872,6 +1039,94 @@ Suggestion: """)) end +@noinline function _throw_log_transition_density_not_implemented(sampler_type) + throw(ArgumentError(""" +`log_transition_density` is not implemented for sampler type $(sampler_type). + +Expected: + Every `AdvancedMH.MHSampler` used with this module must implement its own + `log_transition_density(sampler, model, θ_from, θ_to; stepsize)`, returning + `log q(θ_to | θ_from)`. There is no generic/symmetric fallback, since silently + assuming one previously caused a prior-double-counting bug. + +Got: + typeof(sampler) = $(sampler_type) + +Suggestion: + Define `log_transition_density(sampler::$(sampler_type), model, θ_from, θ_to; stepsize = 1.0)` + for your sampler type. +""")) +end + +@noinline function _throw_accepted_not_recorded(chain::MCMCChains.Chains) + throw(ArgumentError(""" +Chain is missing the `:accepted` internal required by `accept_ratio`. + +Expected: + A `Chains` object produced by this module's `sample`/`optimize_stepsize`, which always + records `:accepted` as an internal. + +Got: + available internals = $(names(chain, :internals)) + +Suggestion: + Only pass `Chains` objects returned by this module's `sample`/`optimize_stepsize`; do + not construct or filter the internals section by hand. +""")) +end + +@noinline function _throw_max_iter_exceeded(n_evals, max_iter, target_acc, tol, stepsize) + throw(ArgumentError(""" +optimize_stepsize did not reach the target acceptance rate within the iteration budget. + +Expected: + Acceptance rate within target_acc ± tol = $(target_acc) ± $(tol), reached within + max_iter = $(max_iter) `sample()` calls. + +Loop context: + n_evals = $(n_evals) (max_iter = $(max_iter)) + last stepsize tried = $(round(stepsize; sigdigits = 3)) + +Suggestion: + Increase `max_iter`, or widen `tol` if a looser acceptance-rate window is acceptable. +""")) +end + +@noinline function _throw_max_expansions_exceeded( + direction::Symbol, + n_expand, + max_expansions, + target_acc, + tol, + bound_stepsize, +) + above = direction == :above + action = above ? "doublings" : "halvings" + stayed = above ? "stayed above" : "stayed below" + monotonic_note = + above ? "rather than simply needing a larger stepsize" : + "(e.g. a numerically noisy log-density at very small stepsizes) rather than simply needing a smaller stepsize" + throw( + ArgumentError( + """ +optimize_stepsize: acceptance rate $(stayed) target $(target_acc) ± $(tol) after $(max_expansions) $(action) without crossing it. + +Expected: + Acceptance rate to cross target_acc ± tol = $(target_acc) ± $(tol) within + max_expansions = $(max_expansions) $(action). + +Loop context: + n_expand = $(n_expand) (max_expansions = $(max_expansions)) + stepsize reached = $(round(bound_stepsize; sigdigits = 3)) + +Suggestion: + Increase `max_expansions` to search farther. Otherwise, the acceptance-vs-stepsize + relationship may not be monotonic in this range, $(monotonic_note). +""", + ), + ) +end + diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index 17e9e5fb2..9a63646e7 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -3,6 +3,9 @@ using LinearAlgebra using Distributions using GaussianProcesses using Test +using AdvancedMH +using AbstractMCMC +using MCMCChains using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.MarkovChainMonteCarlo @@ -16,6 +19,10 @@ using CalibrateEmulateSample.Utilities # range 0->2, with lengthscale of transition 0.5, and y=1 at x=2 G(x) = 5 * (tanh.((x .- 2) ./ 0.5) .+ 1) +# A minimal MHSampler that deliberately does not implement log_transition_density, used to check +# that the shared interface errors loudly instead of silently falling back to a wrong value. +struct _DummyMHSampler <: AdvancedMH.MHSampler end + function test_data(prior; rng_seed = 41, n = 80, var_y = 0.05, rest...) # Seed for pseudo-random number generator rng = Random.MersenneTwister(rng_seed) @@ -670,7 +677,11 @@ end for alg in mcmc_algs mcmc_params_ad = deepcopy(mcmc_params) mcmc_params_ad[:mcmc_alg] = alg - mcmc_params_ad[:target_acc] = 0.6 # should be > 0.5 + # 0.4 sits comfortably inside the tight, low-variance region of Barker's + # acceptance-vs-stepsize curve for this problem; 0.6 sits right at its edge (measured + # ceiling ~0.48-0.50 here), which occasionally made optimize_stepsize unable to find a + # valid stepsize at all. + mcmc_params_ad[:target_acc] = 0.4 @info "testing algorithm: $(typeof(alg))" new_step, posterior_mean, chain = mcmc_test_template(prior, σ2_y, em_1b; mcmc_alg = alg, mcmc_params_ad...) @@ -682,4 +693,300 @@ end end end end + + @testset "optimize_stepsize: branch coverage" begin + # optimize_stepsize's bracket-and-bisect search has several distinct code paths + mcmc = MCMCWrapper( + RWMHSampling(), + mcmc_params[:obs_sample], + prior, + em_1; + init_params = vec(collect(mcmc_params[:init_params])), + ) + + @testset "returns immediately if the initial stepsize is already within tolerance" begin + step = optimize_stepsize(Random.MersenneTwister(1), mcmc; init_stepsize = 0.125, N = 800, target_acc = 0.25) + @test isapprox(step, 0.125; atol = 1e-8) # no expansion/bisection needed at all + end + + @testset "phase 1, 'acceptance too high' branch (expands stepsize upward)" begin + step = optimize_stepsize(Random.MersenneTwister(2), mcmc; init_stepsize = 0.001, N = 800, target_acc = 0.25) + @test 0.15 <= accept_ratio(MCMC.sample(Random.MersenneTwister(2), mcmc, 800; stepsize = step)) <= 0.35 + end + + @testset "phase 1, 'acceptance too low' branch (expands stepsize downward)" begin + step = optimize_stepsize(Random.MersenneTwister(3), mcmc; init_stepsize = 4.0, N = 800, target_acc = 0.25) + @test 0.15 <= accept_ratio(MCMC.sample(Random.MersenneTwister(3), mcmc, 800; stepsize = step)) <= 0.35 + end + + @testset "phase 2: bisection is actually reached" begin + + step = optimize_stepsize( + Random.MersenneTwister(4), + mcmc; + init_stepsize = 0.0625, + N = 1500, + target_acc = 0.33, + tol = 0.05, + ) + @test 0.0625 < step < 0.125 # strictly inside the bracket -> bisection ran, not just phase 1 + end + + @testset "max_expansions exceeded: 'stayed above target' (expand-up) branch" begin + let thrown = @test_throws ArgumentError optimize_stepsize( + Random.MersenneTwister(5), + mcmc; + init_stepsize = 0.25, + N = 300, + target_acc = -0.5, + tol = 0.05, + max_expansions = 3, + ) + @test contains(thrown.value.msg, "stayed above") + @test contains(thrown.value.msg, "doublings") + @test contains(thrown.value.msg, "max_expansions") + @test contains(thrown.value.msg, "-0.5") + end + end + + @testset "max_expansions exceeded: 'stayed below target' (expand-down) branch" begin + # target_acc set above the achievable range (>1), symmetric to the case above. + let thrown = @test_throws ArgumentError optimize_stepsize( + Random.MersenneTwister(6), + mcmc; + init_stepsize = 0.25, + N = 300, + target_acc = 1.5, + tol = 0.05, + max_expansions = 3, + ) + @test contains(thrown.value.msg, "stayed below") + @test contains(thrown.value.msg, "halvings") + @test contains(thrown.value.msg, "max_expansions") + @test contains(thrown.value.msg, "1.5") + end + end + + @testset "overall max_iter budget exhausted" begin + # A reachable target, but with only 1 sample() call allowed in total: the very first + # (non-converging) phase-1 expansion step must hit the n_evals > max_iter guard. + let thrown = @test_throws ArgumentError optimize_stepsize( + Random.MersenneTwister(7), + mcmc; + init_stepsize = 4.0, + N = 300, + target_acc = 0.25, + max_iter = 1, + ) + @test contains(thrown.value.msg, "iteration budget") + @test contains(thrown.value.msg, "max_iter") + end + end + end + + @testset "Sampler transition-density interface (pCN Hastings-term fix)" begin + # Regression tests for: pCN's (and Barker's) log_transition_density must reflect the + # *actual* asymmetric proposal kernel, not the symmetric additive random-walk kernel. + # Constructed directly at the sampler level (bypassing Emulator/MCMCWrapper) so the + # underlying transition-density math can be checked against closed-form references. + rng = Random.MersenneTwister(2026) + n = 3 + # A genuinely correlated (non-diagonal) covariance, so the tests exercise the + # whitening/off-diagonal structure, not just a diagonal special case. + C = [2.0 0.5 0.1; 0.5 1.5 0.3; 0.1 0.3 1.0] + L = cholesky(Symmetric(C)).L + prior_dist = MvNormal(zeros(n), Symmetric(C)) + + a = randn(rng, n) + b = randn(rng, n) + dummy_model = AdvancedMH.DensityModel(θ -> 0.0) + + @testset "RW: transition density is symmetric" begin + rw = MCMC.RWMetropolisHastings{typeof(L), MCMC.GradFreeProtocol}(L) + for s in (0.1, 1.0, 2.5) + @test isapprox( + MCMC.log_transition_density(rw, dummy_model, a, b; stepsize = s), + MCMC.log_transition_density(rw, dummy_model, b, a; stepsize = s), + ) + end + end + + @testset "pCN: Hastings term exactly cancels the prior ratio" begin + # This is the core bug: reverse - forward must equal logprior(a) - logprior(b) + # for every stepsize (every ρ), not silently 0. + pcn = MCMC.pCNMetropolisHastings{typeof(L), MCMC.GradFreeProtocol}(L) + prev_state = MCMC.MCMCState(a, 0.0, true) + for s in (0.05, 0.5, 1.5, 3.9) + hastings = AdvancedMH.logratio_proposal_density(pcn, dummy_model, prev_state, b; stepsize = s) + @test isapprox(hastings, logpdf(prior_dist, a) - logpdf(prior_dist, b); atol = 1e-8) + end + end + + @testset "pCN: flat-likelihood chain always accepts" begin + # If the "likelihood" is constant, the full posterior log-density is just the prior, + # and the correct pCN acceptance log-ratio is then EXACTLY 0 for every proposal + # (loglik cancels trivially, and the Hastings term cancels the prior ratio). Under + # the pre-fix code (Hastings term silently 0), this would instead fluctuate with the + # prior ratio and cause spurious rejections. + pcn = MCMC.pCNMetropolisHastings{typeof(L), MCMC.GradFreeProtocol}(L) + flat_model = AdvancedMH.DensityModel(θ -> logpdf(prior_dist, θ)) + θ0 = zeros(n) + state = MCMC.MCMCState(θ0, AdvancedMH.logdensity(flat_model, θ0), true) + for i in 1:200 + state, _ = AbstractMCMC.step(rng, flat_model, pcn, state; stepsize = 0.6) + @test state.accepted + end + end + + @testset "RW/pCN: propose draws agree with transition_kernel's analytic mean" begin + # propose() and log_transition_density() are now both derived from the single + # transition_kernel(...) object (a Distributions.jl MvNormal), so this is really a + # sanity check that rand(transition_kernel(...)) behaves as documented, rather than a + # check that two independent formulas happen to agree. + # + # n_draws and atol are chosen together for a ~5σ-safe check (not flaky, but tight + # enough to catch a real scaling bug), via the Monte Carlo mean's norm-error RMS + # sqrt(tr(Var)/n_draws): RW's Var = stepsize²C gives 5·RMS ≈ 0.053; pCN's Var = + # (1-ρ²)C gives 5·RMS ≈ 0.067 (both at stepsize=0.5, n_draws=10_000). + rw = MCMC.RWMetropolisHastings{typeof(L), MCMC.GradFreeProtocol}(L) + pcn = MCMC.pCNMetropolisHastings{typeof(L), MCMC.GradFreeProtocol}(L) + state = MCMC.MCMCState(a, 0.0, true) + n_draws = 10_000 + rw_draws = [AdvancedMH.propose(rng, rw, dummy_model, state; stepsize = 0.5) for _ in 1:n_draws] + pcn_draws = [AdvancedMH.propose(rng, pcn, dummy_model, state; stepsize = 0.5) for _ in 1:n_draws] + @test isapprox(mean(rw_draws), a; atol = 0.06) + ρ = (1 - 0.5 / 4) / (1 + 0.5 / 4) + @test isapprox(mean(pcn_draws), ρ .* a; atol = 0.07) + end + + @testset "Barker: transition density matches closed-form gradient formula" begin + # Target chosen as the same Gaussian, so ∇log π(θ) = -C⁻¹θ is known in closed form, + # letting us check the whitened, flip-sign transition density against hand-derived + # algebra rather than trusting the autodiff call alone. + target_logdensity(θ) = -0.5 * dot(θ, C \ θ) + barker_model = AdvancedMH.DensityModel(target_logdensity) + barker = MCMC.BarkerMetropolisHastings{typeof(L), MCMC.ForwardDiffProtocol}(L) + prev_state = MCMC.MCMCState(a, 0.0, true) + for s in (0.3, 1.0, 2.0) + hastings = AdvancedMH.logratio_proposal_density(barker, barker_model, prev_state, b; stepsize = s) + + gw_a = L' * (-(C \ a)) + gw_b = L' * (-(C \ b)) + Δ = L \ (b .- a) # whitened increment (unscaled by stepsize) — the flip decision is a + # function of the actual proposed increment, not a stepsize-independent standardization. + sigmoid(x) = 1 / (1 + exp(-x)) + manual = sum(log.(sigmoid.(-gw_b .* Δ)) .- log.(sigmoid.(gw_a .* Δ))) + @test isapprox(hastings, manual; atol = 1e-8) + end + end + + @testset "Barker: _barker_rand's marginal matches the 2φ(e)σ(d·e) density" begin + # Unlike the Hastings-term check above (which only tests reversibility), this checks + # that _barker_rand's actual empirical distribution matches the closed-form q(e) = + # 2φ(e)σ(d·e) claimed in its docstring, via a fine-grid quadrature reference — an + # independent numerical check, not just internal self-consistency. + d = 1.3 + L1 = reshape([1.0], 1, 1) + kernel = MCMC.BarkerKernel([0.0], [d], L1, 1.0) + + grid = -12:0.001:12 + dx = step(grid) + q(e) = 2 * pdf(Normal(), e) * (1 / (1 + exp(-d * e))) + total_mass = sum(q(e) for e in grid) * dx + mean_theory = sum(e * q(e) for e in grid) * dx / total_mass + @test isapprox(total_mass, 1.0; atol = 1e-6) + + # n_draws/atol chosen for a ~5σ-safe check: Var[e] under q (via the same quadrature) + # is ≈0.76, so SE(n_draws=10_000) ≈ 0.0087 and 5·SE ≈ 0.044. + n_draws = 10_000 + draws = [MCMC._barker_rand(rng, kernel)[1] for _ in 1:n_draws] + @test isapprox(mean_theory, sum(draws) / n_draws; atol = 0.05) + + for etest in (-2.0, -0.5, 0.0, 0.7, 3.0) + @test isapprox(exp(MCMC._barker_logpdf(kernel, [etest])), q(etest); atol = 1e-8) + end + end + + @testset "Barker: flip probability responds to stepsize (regression for stepsize-blind sigmoid bug)" begin + # Regression test for a bug where the flip-probability sigmoid used the unscaled + # whitened noise η instead of the actual proposed increment Δ = stepsize*η, making the + # flip decision (and hence Barker's stepsize-robustness) blind to `stepsize`. The + # correct marginal density of the whitened increment Δ is q(Δ) = 2 φ(Δ; 0, stepsize²) + # σ(d·Δ) — checked here at stepsize ≠ 1, where the old (buggy) and correct formulas + # diverge (they coincide at stepsize = 1, which is why the test above doesn't catch this). + d = 1.3 + s = 2.5 + L1 = reshape([1.0], 1, 1) + kernel = MCMC.BarkerKernel([0.0], [d], L1, s) + + q(Δ) = 2 * pdf(Normal(0, s), Δ) * (1 / (1 + exp(-d * Δ))) + for Δtest in (-3.0, -1.0, 0.0, 1.5, 4.0) + @test isapprox(exp(MCMC._barker_logpdf(kernel, [Δtest])), q(Δtest); atol = 1e-8) + end + + grid = -30:0.001:30 + dx = step(grid) + mean_theory = sum(Δ * q(Δ) for Δ in grid) * dx + n_draws = 10_000 + draws = [MCMC._barker_rand(Random.MersenneTwister(99), kernel)[1] for _ in 1:n_draws] + @test isapprox(mean_theory, sum(draws) / n_draws; atol = 0.1) + end + + @testset "Barker: recovers the exact 1D conjugate-Gaussian posterior (regression for wrong-variance bug)" begin + # Regression test at the full-sampler level, matching a bug report that measured the + # posterior variance ~2x too small on an earlier version of this sampler. Prior + # θ ~ N(0, 1), single "observation" y ~ N(θ, 0.5) with a zero-error (exact) forward map, + # so the posterior is available in closed form: mean = y τ²/(τ²+σ²), var = τ²σ²/(τ²+σ²). + # Chosen so mean = 5/3, var = 1/3 exactly. + τ2 = 1.0 + σ2 = 0.5 + y_obs = 2.5 + post_mean = 5.0 / 3.0 + post_var = 1.0 / 3.0 + + target_logdensity(θ) = logpdf(Normal(0, sqrt(τ2)), θ[1]) + logpdf(Normal(θ[1], sqrt(σ2)), y_obs) + model = AdvancedMH.DensityModel(target_logdensity) + L1 = reshape([1.0], 1, 1) + barker = MCMC.BarkerMetropolisHastings{typeof(L1), MCMC.ForwardDiffProtocol}(L1) + + rng_local = Random.MersenneTwister(2024) + stepsize = 0.7 + state = MCMC.MCMCState([0.0], AdvancedMH.logdensity(model, [0.0]), true) + for _ in 1:1000 + state, _ = AbstractMCMC.step(rng_local, model, barker, state; stepsize = stepsize) + end + n_samples = 20_000 + draws = zeros(n_samples) + for i in 1:n_samples + state, _ = AbstractMCMC.step(rng_local, model, barker, state; stepsize = stepsize) + draws[i] = state.params[1] + end + @test isapprox(mean(draws), post_mean; rtol = 0.05) + @test isapprox(var(draws), post_var; rtol = 0.05) + end + + @testset "log_transition_density errors loudly for an unimplemented sampler" begin + # A hypothetical new sampler that forgets to implement log_transition_density must + # fail loudly (ArgumentError) rather than silently falling back to a symmetric, + # possibly-wrong correction — this is what the shared interface guards against. + let thrown = @test_throws ArgumentError MCMC.log_transition_density(_DummyMHSampler(), dummy_model, a, b) + @test contains(thrown.value.msg, "_DummyMHSampler") + @test contains(thrown.value.msg, "log_transition_density") + end + end + end + + @testset "accept_ratio errors on a chain missing the :accepted internal" begin + # A Chains object not produced by this module's sample()/optimize_stepsize() (e.g. hand-built, + # or with internals filtered out) must fail loudly rather than silently mis-reporting. + bad_chain = MCMCChains.Chains( + rand(5, 3, 1), + [:a, :b, :other_internal], + (parameters = [:a, :b], internals = [:other_internal]), + ) + let thrown = @test_throws ArgumentError accept_ratio(bad_chain) + @test contains(thrown.value.msg, "accept_ratio") + @test contains(thrown.value.msg, "other_internal") + end + end end