diff --git a/Project.toml b/Project.toml index cd60ab6974..12776485d7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Turing" uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" -version = "0.42.7" +version = "0.42.8" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/mcmc/Inference.jl b/src/mcmc/Inference.jl index b02d0887f7..031974f0f3 100644 --- a/src/mcmc/Inference.jl +++ b/src/mcmc/Inference.jl @@ -115,6 +115,26 @@ function mh_accept(logp_current::Real, logp_proposal::Real, log_proposal_ratio:: return log(rand()) + logp_current ≤ logp_proposal + log_proposal_ratio end +# Directly overload the constructor of `AbstractMCMC.ParamsWithStats` so that we don't +# hit the default method, which uses `getparams(state)` and `getstats(state)`. For Turing's +# MCMC samplers, the state might contain results that are in linked space. Using the +# outputs of the transition here ensures that parameters and logprobs are provided in +# user space (similar to chains output). +function AbstractMCMC.ParamsWithStats( + model, + sampler, + transition::DynamicPPL.ParamsWithStats, + state; + params::Bool=true, + stats::Bool=false, + extras::Bool=false, +) + p = params ? [string(k) => v for (k, v) in transition.params] : nothing + s = stats ? transition.stats : NamedTuple() + e = extras ? NamedTuple() : NamedTuple() + return AbstractMCMC.ParamsWithStats(p, s, e) +end + ####################################### # Concrete algorithm implementations. # ####################################### diff --git a/test/Project.toml b/test/Project.toml index f3cb375110..b2dfb189fa 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -40,7 +40,7 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [compat] ADTypes = "1" -AbstractMCMC = "5.9" +AbstractMCMC = "5.13" AbstractPPL = "0.11, 0.12, 0.13" AdvancedMH = "0.8.9" AdvancedPS = "0.7.2" diff --git a/test/mcmc/callbacks.jl b/test/mcmc/callbacks.jl new file mode 100644 index 0000000000..c62e9d133a --- /dev/null +++ b/test/mcmc/callbacks.jl @@ -0,0 +1,65 @@ +module CallbacksTests + +using Test, Turing, AbstractMCMC, Random, Distributions, LinearAlgebra + +@model function test_normals() + x ~ Normal() + return y ~ MvNormal(zeros(3), I) +end + +@testset "AbstractMCMC Callbacks Interface" begin + rng = Random.default_rng() + model = test_normals() + + samplers = [ + ("NUTS", NUTS(10, 0.65)), + ("HMC", HMC(0.1, 5)), + ("MH", MH()), + ("ESS", ESS()), + ("Gibbs", Gibbs(:x => HMC(0.1, 5), :y => MH())), + ("SGHMC", SGHMC(; learning_rate=0.01, momentum_decay=1e-2)), + ("PG", PG(10)), + ] + + for (name, sampler) in samplers + @testset "$name" begin + t1, s1 = AbstractMCMC.step( + rng, model, sampler; initial_params=Turing.Inference.init_strategy(sampler) + ) + + # ParamsWithStats returns named params (not θ[i]) + pws = AbstractMCMC.ParamsWithStats( + model, sampler, t1, s1; params=true, stats=true + ) + pairs_dict = Dict(k => v for (k, v) in Base.pairs(pws)) + # Keys are Symbols since ParamsWithStats stores NamedTuple internally + @test haskey(pairs_dict, Symbol("x")) + @test haskey(pairs_dict, Symbol("y")) + @test pairs_dict[Symbol("y")] isa AbstractVector + @test length(pairs_dict[Symbol("y")]) == 3 + + # Check stats contain lp + @test haskey(pairs_dict, :lp) || haskey(pairs_dict, :logjoint) + end + end + + # NUTS second step has full AHMC transition metrics + @testset "NUTS Transition Metrics" begin + sampler = NUTS(10, 0.65) + t1, s1 = AbstractMCMC.step( + rng, model, sampler; initial_params=Turing.Inference.init_strategy(sampler) + ) + t2, s2 = AbstractMCMC.step(rng, model, sampler, s1) + + pws = AbstractMCMC.ParamsWithStats(model, sampler, t2, s2; params=true, stats=true) + pairs_dict = Dict(k => v for (k, v) in Base.pairs(pws)) + + # Keys are Symbols from NamedTuple + @test haskey(pairs_dict, :tree_depth) + @test haskey(pairs_dict, :n_steps) + @test haskey(pairs_dict, :acceptance_rate) + @test haskey(pairs_dict, :hamiltonian_energy) + end +end + +end diff --git a/test/runtests.jl b/test/runtests.jl index a4a60d8e29..2a6c00c312 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -44,6 +44,7 @@ end @testset "samplers (without AD)" verbose = true begin @timeit_include("mcmc/abstractmcmc.jl") + @timeit_include("mcmc/callbacks.jl") @timeit_include("mcmc/particle_mcmc.jl") @timeit_include("mcmc/emcee.jl") @timeit_include("mcmc/ess.jl")