improve coverage and error messages

odunbar · odunbar · commit b56663d125bd · 2026-07-22T12:33:56.000-07:00
diff --git 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
@@ -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
