From 3fb8cb1263c38baf4ce7ed78754c59de809b0db1 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 16:33:54 -0700 Subject: [PATCH 1/8] proposals for individual MCMC methods, shared asymmetric transition logic --- src/MarkovChainMonteCarlo.jl | 185 +++++++++++++++++++------ test/MarkovChainMonteCarlo/runtests.jl | 88 ++++++++++++ 2 files changed, 227 insertions(+), 46 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 9d9bf0804..7e1acc284 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,65 @@ 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 ONE function, log_transition_density(sampler, +# model, θ_from, θ_to; stepsize) = log q(θ_to | θ_from) — the honest forward transition +# log-density of whatever Markov kernel that sampler's `propose` actually draws from. There is +# no default/fallback implementation: a sampler that doesn't implement it errors loudly instead +# of silently inheriting a (possibly wrong) symmetric random-walk correction. +# +# AdvancedMH.logratio_proposal_density, which feeds directly into the MH acceptance ratio, is +# then defined ONCE, generically over every AdvancedMH.MHSampler (so it applies automatically to +# any sampler added to this module in future, not just the three below), as the textbook +# Hastings correction log q(θ_from | θ_to) − log q(θ_to | θ_from) (reverse − forward), by +# evaluating log_transition_density in both directions. A sampler can never update its `propose` +# without also updating the quantity that determines its acceptance ratio, because there is only +# one function (log_transition_density) doing double duty for both. +function log_transition_density( + sampler::MHS, + model, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {MHS <: AdvancedMH.MHSampler, FT <: AbstractFloat} + throw( + ArgumentError( + "log_transition_density not implemented for $(MHS). Every MHSampler in this module " * + "must implement its own log_transition_density(sampler, model, θ_from, θ_to; stepsize) " * + "= log q(θ_to | θ_from); there is no generic/symmetric fallback, since silently assuming " * + "one caused the prior-double-counting bug this interface replaces.", + ), + ) +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 @@ -336,7 +369,20 @@ function AdvancedMH.propose( current_state::MCMCState; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - return current_state.params + stepsize * rand(rng, sampler.proposal) + L = sampler.cholesky_L + return current_state.params .+ stepsize .* (L * randn(rng, size(L, 1))) +end + +# θ_to | θ_from ~ N(θ_from, stepsize² C), the symmetric kernel that propose() above draws from. +function log_transition_density( + sampler::RWMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + L = sampler.cholesky_L + return logpdf(MvNormal(θ_from, Symmetric((stepsize^2) .* (L * L'))), θ_to) end # method extending AdvancedMH.propose() for preconditioned Crank-Nicholson @@ -347,10 +393,27 @@ function AdvancedMH.propose( current_state::MCMCState; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - # Use prescription in Beskos et al (2017) "Geometric MCMC for infinite-dimensional + # 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 ρ .* current_state.params .+ sqrt(1 - ρ^2) .* (L * randn(rng, size(L, 1))) +end + +# θ_to | θ_from ~ N(ρθ_from, (1-ρ²)C), the asymmetric kernel that propose() above draws from. +# 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 log_transition_density( + sampler::pCNMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + ρ = (1 - stepsize / 4) / (1 + stepsize / 4) + L = sampler.cholesky_L + return logpdf(MvNormal(ρ .* θ_from, Symmetric((1 - ρ^2) .* (L * L'))), θ_to) end # method extending AdvancedMH.propose() for the Barker proposal @@ -361,12 +424,41 @@ 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)) + # 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: the gradient is rotated into whitened + # coordinates via L'∇log π(θ) (chain rule for θ = θ₀ + Lu). We also flip the sign of the + # whitened noise η rather than zeroing it out — the standard Barker kernel moves by ±η, never + # by exactly 0, so its transition density (below) is an ordinary, atom-free density; zeroing + # out (keep-or-drop) instead creates a mixed discrete/continuous kernel with no closed form. + θ = current_state.params + n = length(θ) + L = sampler.cholesky_L + grad_white = L' * autodiff_gradient(model, θ, sampler) + η = randn(rng, n) + flip_prob = 1 ./ (1 .+ exp.(-grad_white .* η)) + sign = 2 .* (rand(rng, n) .< flip_prob) .- 1 + ζ = stepsize .* sign .* η + return θ .+ L * ζ +end + +# For a symmetric noise density g (here standard normal) and whitened increment e = ζ/stepsize, +# the marginal density of the ± flip-sign move is q(e | θ_from) = 2 g(e) σ(grad_white·e) (Barker, +# 1965; Livingstone & Zanella, 2022): summing the probability that e itself was kept (w.p. +# σ(grad_white·e)) with the probability that -e was drawn and flipped (w.p. σ(grad_white·e) too, +# since g(-e) = g(e)). See propose() above for how θ_to is generated from θ_from. +function log_transition_density( + sampler::BarkerMetropolisHastings, + model::AdvancedMH.DensityModel, + θ_from, + θ_to; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + L = sampler.cholesky_L + grad_white = L' * autodiff_gradient(model, θ_from, sampler) + e = (L \ (θ_to .- θ_from)) ./ stepsize + return sum(log(2) .+ logpdf.(Normal(), e) .+ log.(1 ./ (1 .+ exp.(-grad_white .* e)))) - length(e) * log(stepsize) end # Copy a MCMCState and set accepted = false @@ -393,7 +485,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_α diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index 17e9e5fb2..2cbf04afe 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -3,6 +3,8 @@ using LinearAlgebra using Distributions using GaussianProcesses using Test +using AdvancedMH +using AbstractMCMC using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.MarkovChainMonteCarlo @@ -16,6 +18,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) @@ -682,4 +688,86 @@ end 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 "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)) + e = (L \ (b .- a)) ./ s + sigmoid(x) = 1 / (1 + exp(-x)) + manual = sum(log.(sigmoid.(-gw_b .* e)) .- log.(sigmoid.(gw_a .* e))) + @test isapprox(hastings, manual; atol = 1e-8) + end + 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. + @test_throws ArgumentError MCMC.log_transition_density(_DummyMHSampler(), dummy_model, a, b) + end + end end From 926a50f66d3c10f1700d0ce08caf01add7f80202 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 16:50:39 -0700 Subject: [PATCH 2/8] MCMC now creates the markov kernel, not the individual proposal and transition density --- src/MarkovChainMonteCarlo.jl | 82 ++++++++++++++++++------------------ 1 file changed, 42 insertions(+), 40 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 7e1acc284..43a114800 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -194,19 +194,15 @@ end # ------------------------------------------------------------------------------------------ # Shared Metropolis-Hastings transition-density interface. # -# Every sampler above is required to implement ONE function, log_transition_density(sampler, -# model, θ_from, θ_to; stepsize) = log q(θ_to | θ_from) — the honest forward transition -# log-density of whatever Markov kernel that sampler's `propose` actually draws from. There is -# no default/fallback implementation: a sampler that doesn't implement it errors loudly instead -# of silently inheriting a (possibly wrong) symmetric random-walk correction. -# -# AdvancedMH.logratio_proposal_density, which feeds directly into the MH acceptance ratio, is -# then defined ONCE, generically over every AdvancedMH.MHSampler (so it applies automatically to -# any sampler added to this module in future, not just the three below), as the textbook -# Hastings correction log q(θ_from | θ_to) − log q(θ_to | θ_from) (reverse − forward), by -# evaluating log_transition_density in both directions. A sampler can never update its `propose` -# without also updating the quantity that determines its acceptance ratio, because there is only -# one function (log_transition_density) doing double duty for both. +# 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, @@ -361,59 +357,65 @@ 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 {FT <: AbstractFloat} - L = sampler.cholesky_L - return current_state.params .+ stepsize .* (L * randn(rng, size(L, 1))) +) where {MHS <: DistributionKernelSampler, FT <: AbstractFloat} + return rand(rng, transition_kernel(sampler, model, current_state.params; stepsize = stepsize)) end -# θ_to | θ_from ~ N(θ_from, stepsize² C), the symmetric kernel that propose() above draws from. function log_transition_density( - sampler::RWMetropolisHastings, + sampler::MHS, model::AdvancedMH.DensityModel, θ_from, θ_to; stepsize::FT = 1.0, -) where {FT <: AbstractFloat} - L = sampler.cholesky_L - return logpdf(MvNormal(θ_from, Symmetric((stepsize^2) .* (L * L'))), θ_to) +) where {MHS <: DistributionKernelSampler, FT <: AbstractFloat} + return logpdf(transition_kernel(sampler, model, θ_from; stepsize = stepsize), θ_to) end -# method extending AdvancedMH.propose() for preconditioned Crank-Nicholson -function AdvancedMH.propose( - rng::Random.AbstractRNG, - sampler::pCNMetropolisHastings, +# θ_to | θ_from ~ N(θ_from, stepsize² C) — the symmetric random-walk kernel. +function transition_kernel( + sampler::RWMetropolisHastings, 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) L = sampler.cholesky_L - return ρ .* current_state.params .+ sqrt(1 - ρ^2) .* (L * randn(rng, size(L, 1))) + return MvNormal(θ_from, Symmetric((stepsize^2) .* (L * L'))) end -# θ_to | θ_from ~ N(ρθ_from, (1-ρ²)C), the asymmetric kernel that propose() above draws from. -# 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 log_transition_density( +# θ_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, - θ_from, - θ_to; + θ_from; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} ρ = (1 - stepsize / 4) / (1 + stepsize / 4) L = sampler.cholesky_L - return logpdf(MvNormal(ρ .* θ_from, Symmetric((1 - ρ^2) .* (L * L'))), θ_to) + return MvNormal(ρ .* θ_from, Symmetric((1 - ρ^2) .* (L * L'))) end # method extending AdvancedMH.propose() for the Barker proposal From d352c1ba61ff9708171a401e2d9ae62e697ead4a Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 17:16:31 -0700 Subject: [PATCH 3/8] add kernel struct for Barker that is a little more flexible than using Distribtuions.jl --- src/MarkovChainMonteCarlo.jl | 76 +++++++++++++++++++++++++----------- 1 file changed, 53 insertions(+), 23 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 43a114800..73a62c756 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -418,38 +418,71 @@ function transition_kernel( return MvNormal(ρ .* θ_from, Symmetric((1 - ρ^2) .* (L * L'))) end -# method extending AdvancedMH.propose() for the Barker proposal -function AdvancedMH.propose( - rng::Random.AbstractRNG, +# 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, - current_state::MCMCState; + θ_from; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - # 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: the gradient is rotated into whitened - # coordinates via L'∇log π(θ) (chain rule for θ = θ₀ + Lu). We also flip the sign of the - # whitened noise η rather than zeroing it out — the standard Barker kernel moves by ±η, never - # by exactly 0, so its transition density (below) is an ordinary, atom-free density; zeroing - # out (keep-or-drop) instead creates a mixed discrete/continuous kernel with no closed form. - θ = current_state.params - n = length(θ) L = sampler.cholesky_L - grad_white = L' * autodiff_gradient(model, θ, sampler) + 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 η +# rather than zeroing it out — the standard Barker kernel moves by ±η, 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. +function _barker_rand(rng::Random.AbstractRNG, kernel::BarkerKernel) + n = length(kernel.θ_from) η = randn(rng, n) - flip_prob = 1 ./ (1 .+ exp.(-grad_white .* η)) + flip_prob = 1 ./ (1 .+ exp.(-kernel.grad_white .* η)) sign = 2 .* (rand(rng, n) .< flip_prob) .- 1 - ζ = stepsize .* sign .* η - return θ .+ L * ζ + ζ = kernel.stepsize .* sign .* η + return kernel.θ_from .+ kernel.L * ζ end # For a symmetric noise density g (here standard normal) and whitened increment e = ζ/stepsize, # the marginal density of the ± flip-sign move is q(e | θ_from) = 2 g(e) σ(grad_white·e) (Barker, # 1965; Livingstone & Zanella, 2022): summing the probability that e itself was kept (w.p. # σ(grad_white·e)) with the probability that -e was drawn and flipped (w.p. σ(grad_white·e) too, -# since g(-e) = g(e)). See propose() above for how θ_to is generated from θ_from. +# since g(-e) = g(e)). See _barker_rand above for how θ_to is generated from θ_from. +function _barker_logpdf(kernel::BarkerKernel, θ_to) + e = (kernel.L \ (θ_to .- kernel.θ_from)) ./ kernel.stepsize + return sum(log(2) .+ logpdf.(Normal(), e) .+ log.(1 ./ (1 .+ exp.(-kernel.grad_white .* e)))) - + length(e) * log(kernel.stepsize) +end + +# method extending AdvancedMH.propose() for the Barker proposal +function AdvancedMH.propose( + rng::Random.AbstractRNG, + sampler::BarkerMetropolisHastings, + model::AdvancedMH.DensityModel, + current_state::MCMCState; + stepsize::FT = 1.0, +) where {FT <: AbstractFloat} + return _barker_rand(rng, transition_kernel(sampler, model, current_state.params; stepsize = stepsize)) +end + function log_transition_density( sampler::BarkerMetropolisHastings, model::AdvancedMH.DensityModel, @@ -457,10 +490,7 @@ function log_transition_density( θ_to; stepsize::FT = 1.0, ) where {FT <: AbstractFloat} - L = sampler.cholesky_L - grad_white = L' * autodiff_gradient(model, θ_from, sampler) - e = (L \ (θ_to .- θ_from)) ./ stepsize - return sum(log(2) .+ logpdf.(Normal(), e) .+ log.(1 ./ (1 .+ exp.(-grad_white .* e)))) - length(e) * log(stepsize) + return _barker_logpdf(transition_kernel(sampler, model, θ_from; stepsize = stepsize), θ_to) end # Copy a MCMCState and set accepted = false From 30a886497e2cf34fa530b8ace9f57d431ec97eff Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 17:30:30 -0700 Subject: [PATCH 4/8] unit tests --- test/MarkovChainMonteCarlo/runtests.jl | 48 ++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index 2cbf04afe..3caab93a6 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -743,6 +743,27 @@ end 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 @@ -763,6 +784,33 @@ end 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 "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, From 85baa9000500e7e223de68fe204f7766927a7750 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 18:21:23 -0700 Subject: [PATCH 5/8] add better bisecting stepsize optimization --- src/MarkovChainMonteCarlo.jl | 95 +++++++++++++++++++++++------------- 1 file changed, 61 insertions(+), 34 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 73a62c756..e4d55aff0 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -847,11 +847,19 @@ 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`. """ function optimize_stepsize( rng::Random.AbstractRNG, @@ -860,50 +868,69 @@ function optimize_stepsize( N = 2000, max_iter = 20, 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 + if n_evals > max_iter + error( + "optimize_stepsize: acceptance rate did not reach target $(target_acc) ± $(tol) " * + "within $(max_iter) iterations — last stepsize tried: $(round(stepsize; sigdigits = 3)).", + ) + end 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. + if acc > target_acc + lo, hi = stepsize, stepsize + acc_hi = acc + while acc_hi > target_acc + 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 + while acc_lo < target_acc + 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...) From 07525bc6cc98a6a719443c171421562895cbf757 Mon Sep 17 00:00:00 2001 From: odunbar Date: Tue, 21 Jul 2026 18:27:23 -0700 Subject: [PATCH 6/8] add tests for stepsize, and reduce barker stepsize to 0.4 (0.6 for some reason is unreachable) --- src/MarkovChainMonteCarlo.jl | 34 ++++++++++++++++++++++++++ test/MarkovChainMonteCarlo/runtests.jl | 6 ++++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index e4d55aff0..70f512cf0 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -860,6 +860,15 @@ is known to lie strictly between the two trial acceptance rates, then bisects (i 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, @@ -867,6 +876,7 @@ function optimize_stepsize( init_stepsize = 1.0, N = 2000, max_iter = 20, + max_expansions = 10, target_acc = 0.25, tol = 0.1, sample_kwargs..., @@ -900,10 +910,23 @@ function optimize_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 + if n_expand > max_expansions + error( + "optimize_stepsize: acceptance rate stayed above target $(target_acc) ± $(tol) " * + "after $(max_expansions) doublings (up to stepsize $(round(hi; sigdigits = 3))) " * + "without crossing it — the acceptance-vs-stepsize relationship may not be monotonic " * + "in this range rather than simply needing a larger stepsize.", + ) + end lo = hi hi *= 2 acc_hi = acc_at(hi) @@ -912,7 +935,18 @@ function optimize_stepsize( else lo, hi = stepsize, stepsize acc_lo = acc + n_expand = 0 while acc_lo < target_acc + n_expand += 1 + if n_expand > max_expansions + error( + "optimize_stepsize: acceptance rate stayed below target $(target_acc) ± $(tol) " * + "after $(max_expansions) halvings (down to stepsize $(round(lo; sigdigits = 3))) " * + "without crossing it — the acceptance-vs-stepsize relationship may not be monotonic " * + "in this range (e.g. a numerically noisy log-density at very small stepsizes) rather " * + "than simply needing a smaller stepsize.", + ) + end hi = lo lo /= 2 acc_lo = acc_at(lo) diff --git a/test/MarkovChainMonteCarlo/runtests.jl b/test/MarkovChainMonteCarlo/runtests.jl index 3caab93a6..902e707ca 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -676,7 +676,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...) From b56663d125bd07c6a9c9581d1edffd90ce4f85a5 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 22 Jul 2026 12:33:56 -0700 Subject: [PATCH 7/8] improve coverage and error messages --- src/MarkovChainMonteCarlo.jl | 121 +++++++++++++++++-------- test/MarkovChainMonteCarlo/runtests.jl | 100 +++++++++++++++++++- 2 files changed, 184 insertions(+), 37 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 70f512cf0..485436d59 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -210,14 +210,7 @@ function log_transition_density( θ_to; stepsize::FT = 1.0, ) where {MHS <: AdvancedMH.MHSampler, FT <: AbstractFloat} - throw( - ArgumentError( - "log_transition_density not implemented for $(MHS). Every MHSampler in this module " * - "must implement its own log_transition_density(sampler, model, θ_from, θ_to; stepsize) " * - "= log q(θ_to | θ_from); there is no generic/symmetric fallback, since silently assuming " * - "one caused the prior-double-counting bug this interface replaces.", - ), - ) + _throw_log_transition_density_not_implemented(MHS) end function AdvancedMH.logratio_proposal_density( @@ -819,11 +812,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) @@ -886,12 +876,7 @@ function optimize_stepsize( n_evals = 0 function acc_at(stepsize) n_evals += 1 - if n_evals > max_iter - error( - "optimize_stepsize: acceptance rate did not reach target $(target_acc) ± $(tol) " * - "within $(max_iter) iterations — last stepsize tried: $(round(stepsize; sigdigits = 3)).", - ) - end + 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(n_evals, stepsize, acc_ratio, trial_chain) @@ -919,14 +904,8 @@ function optimize_stepsize( n_expand = 0 while acc_hi > target_acc n_expand += 1 - if n_expand > max_expansions - error( - "optimize_stepsize: acceptance rate stayed above target $(target_acc) ± $(tol) " * - "after $(max_expansions) doublings (up to stepsize $(round(hi; sigdigits = 3))) " * - "without crossing it — the acceptance-vs-stepsize relationship may not be monotonic " * - "in this range rather than simply needing a larger stepsize.", - ) - end + 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) @@ -938,15 +917,8 @@ function optimize_stepsize( n_expand = 0 while acc_lo < target_acc n_expand += 1 - if n_expand > max_expansions - error( - "optimize_stepsize: acceptance rate stayed below target $(target_acc) ± $(tol) " * - "after $(max_expansions) halvings (down to stepsize $(round(lo; sigdigits = 3))) " * - "without crossing it — the acceptance-vs-stepsize relationship may not be monotonic " * - "in this range (e.g. a numerically noisy log-density at very small stepsizes) rather " * - "than simply needing a smaller stepsize.", - ) - end + 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) @@ -1058,6 +1030,83 @@ 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 902e707ca..337fd7ad4 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -5,6 +5,7 @@ using GaussianProcesses using Test using AdvancedMH using AbstractMCMC +using MCMCChains using CalibrateEmulateSample.EnsembleKalmanProcesses using CalibrateEmulateSample.MarkovChainMonteCarlo @@ -693,6 +694,90 @@ 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. @@ -819,7 +904,20 @@ end # 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. - @test_throws ArgumentError MCMC.log_transition_density(_DummyMHSampler(), dummy_model, a, b) + 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 From ce1391c106e76c0b13fc99493ecfb8cf5c168324 Mon Sep 17 00:00:00 2001 From: odunbar Date: Wed, 22 Jul 2026 17:20:54 -0700 Subject: [PATCH 8/8] scale Barker kernel by step-sizecorrectly for robustness --- src/MarkovChainMonteCarlo.jl | 48 +++++++++++----- test/MarkovChainMonteCarlo/runtests.jl | 79 ++++++++++++++++++++++++-- 2 files changed, 108 insertions(+), 19 deletions(-) diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl index 485436d59..7b344b46e 100644 --- a/src/MarkovChainMonteCarlo.jl +++ b/src/MarkovChainMonteCarlo.jl @@ -441,27 +441,36 @@ 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 η -# rather than zeroing it out — the standard Barker kernel moves by ±η, never by exactly 0, so its +# 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) - flip_prob = 1 ./ (1 .+ exp.(-kernel.grad_white .* η)) + z = kernel.stepsize .* η + flip_prob = 1 ./ (1 .+ exp.(-kernel.grad_white .* z)) sign = 2 .* (rand(rng, n) .< flip_prob) .- 1 - ζ = kernel.stepsize .* sign .* η + ζ = sign .* z return kernel.θ_from .+ kernel.L * ζ end -# For a symmetric noise density g (here standard normal) and whitened increment e = ζ/stepsize, -# the marginal density of the ± flip-sign move is q(e | θ_from) = 2 g(e) σ(grad_white·e) (Barker, -# 1965; Livingstone & Zanella, 2022): summing the probability that e itself was kept (w.p. -# σ(grad_white·e)) with the probability that -e was drawn and flipped (w.p. σ(grad_white·e) too, -# since g(-e) = g(e)). See _barker_rand above for how θ_to is generated from θ_from. +# 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) - e = (kernel.L \ (θ_to .- kernel.θ_from)) ./ kernel.stepsize - return sum(log(2) .+ logpdf.(Normal(), e) .+ log.(1 ./ (1 .+ exp.(-kernel.grad_white .* e)))) - + Δ = 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 @@ -1083,14 +1092,23 @@ Suggestion: """)) end -@noinline function _throw_max_expansions_exceeded(direction::Symbol, n_expand, max_expansions, target_acc, tol, bound_stepsize) +@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(""" + throw( + ArgumentError( + """ optimize_stepsize: acceptance rate $(stayed) target $(target_acc) ± $(tol) after $(max_expansions) $(action) without crossing it. Expected: @@ -1104,7 +1122,9 @@ Loop context: 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 337fd7ad4..9a63646e7 100644 --- a/test/MarkovChainMonteCarlo/runtests.jl +++ b/test/MarkovChainMonteCarlo/runtests.jl @@ -696,7 +696,13 @@ 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]))) + 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) @@ -714,7 +720,7 @@ end end @testset "phase 2: bisection is actually reached" begin - + step = optimize_stepsize( Random.MersenneTwister(4), mcmc; @@ -866,9 +872,10 @@ end gw_a = L' * (-(C \ a)) gw_b = L' * (-(C \ b)) - e = (L \ (b .- a)) ./ s + Δ = 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 .* e)) .- log.(sigmoid.(gw_a .* e))) + manual = sum(log.(sigmoid.(-gw_b .* Δ)) .- log.(sigmoid.(gw_a .* Δ))) @test isapprox(hastings, manual; atol = 1e-8) end end @@ -900,6 +907,64 @@ end 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, @@ -914,7 +979,11 @@ 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])) + 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")