implement AbstractMCMC.getstats (#119)

* Implement AbstractMCMC.getstats

* Add tests

* Bump compat etc

* Disable JuliaPre workflow

* Add testing on 1.10 (lts)

Author: penelopeysm

.github/workflows/CI.yml

@@ -17,27 +17,25 @@ jobs:
           # Current stable version
           - version: '1'
             os: ubuntu-latest
-            arch: x64
+          # LTS
+          - version: 'lts'
+            os: ubuntu-latest
           # Minimum supported version
           - version: 'min'
             os: ubuntu-latest
-            arch: x64
           # Windows
           - version: '1'
             os: windows-latest
-            arch: x64
           # macOS
           - version: '1'
             os: macos-latest
-            arch: aarch64
 
     steps:
       - uses: actions/checkout@v4
 
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.runner.version }}
-          arch: ${{ matrix.runner.arch }}
 
       - uses: julia-actions/cache@v2
 

.github/workflows/JuliaPre.yml

@@ -1,10 +1,12 @@
 name: JuliaPre
 
 on:
-  push:
-    branches:
-      - main
-  pull_request:
+  workflow_dispatch:
+  # Disabled for now because there is no available prerelease of Julia.
+  # push:
+  #   branches:
+  #     - main
+  # pull_request:
 
 # needed to allow julia-actions/cache to delete old caches that it has created
 permissions:

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.8.8"
+version = "0.8.9"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -24,7 +24,7 @@ AdvancedMHMCMCChainsExt = "MCMCChains"
 AdvancedMHStructArraysExt = "StructArrays"
 
 [compat]
-AbstractMCMC = "5.6"
+AbstractMCMC = "5.9"
 DiffResults = "1"
 Distributions = "0.25"
 DocStringExtensions = "0.9"

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,

test/runtests.jl

@@ -53,6 +53,15 @@ include("util.jl")
         end
     end
 
+    @testset "AbstractMCMC.getstats" begin
+        t1, _ = AbstractMCMC.step(Random.default_rng(), model, StaticMH([Normal(0, 1), Normal(0, 1)]))
+        t2, _ = AbstractMCMC.step(Random.default_rng(), model, MALA(x -> MvNormal(x, I)); initial_params=ones(2))
+        for t in [t1, t2]
+            stats = AbstractMCMC.getstats(t)
+            @test stats == (accepted = t.accepted,)
+        end
+    end
+
     @testset "StaticMH" begin
         # Set up our sampler with initial parameters.
         spl1 = StaticMH([Normal(0,1), Normal(0, 1)])