Implement AbstractMCMC.getstats

Author: penelopeysm

src/AdvancedMH.jl

@@ -146,6 +146,9 @@ end
 function AbstractMCMC.getparams(t::Transition)
     return t.params
 end
+function AbstractMCMC.getstats(t::Transition)
+    return (accepted=t.accepted,)
+end
 
 # TODO (sunxd): remove `DensityModel` in favor of `AbstractMCMC.LogDensityModel`
 function AbstractMCMC.setparams!!(model::DensityModelOrLogDensityModel, t::Transition, params)

src/MALA.jl

@@ -23,6 +23,9 @@ logdensity(model::DensityModelOrLogDensityModel, t::GradientTransition) = t.lp
 function AbstractMCMC.getparams(t::GradientTransition)
     return t.params
 end
+function AbstractMCMC.getstats(t::GradientTransition)
+    return (accepted=t.accepted,)
+end
 
 function AbstractMCMC.setparams!!(model::DensityModelOrLogDensityModel, t::GradientTransition, params)
     lp, gradient = logdensity_and_gradient(model, params)

src/RobustAdaptiveMetropolis.jl

@@ -116,9 +116,18 @@ end
 AbstractMCMC.getparams(state::RobustAdaptiveMetropolisState) = state.x
 function AbstractMCMC.setparams!!(state::RobustAdaptiveMetropolisState, x)
     return RobustAdaptiveMetropolisState(
-        x, state.logprob, state.S, state.logα, state.η, state.iteration, state.isaccept
+        x,
+        state.logprob,
+        state.S,
+        state.logα,
+        state.η,
+        state.iteration,
+        state.isaccept,
     )
 end
+function AbstractMCMC.getstats(state::RobustAdaptiveMetropolisState)
+    return (logα = state.logα, η = state.η, accepted = state.isaccept)
+end
 
 function ram_step_inner(
     rng::Random.AbstractRNG,