Add getparams and setparams!! following AbstractMCMC v5.5 and v5.6 (#378)

* add `getparams` and `setparams!!`

* undo formatting

* Update test/abstractmcmc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* add some new tests

* update `setparams!!`

* Update src/abstractmcmc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* add comment

* Update abstractmcmc.jl

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

* format

* update implementation

* bump AbstractMCMC

* update test

* Update src/abstractmcmc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix method ambiguity

* Update src/abstractmcmc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/abstractmcmc.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix test error

* fix more test error

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -30,7 +30,7 @@ AdvancedHMCMCMCChainsExt = "MCMCChains"
 AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
 
 [compat]
-AbstractMCMC = "5"
+AbstractMCMC = "5.6"
 ArgCheck = "1, 2"
 CUDA = "3, 4, 5"
 DocStringExtensions = "0.8, 0.9"

research/tests/runtests.jl

@@ -11,6 +11,6 @@ include("../src/riemannian_hmc.jl")
 include("relativistic_hmc.jl")
 include("riemannian_hmc.jl")
 
-@main function runtests(patterns...; dry::Bool = false)
+Comonicon.@main function runtests(patterns...; dry::Bool = false)
     retest(patterns...; dry = dry, verbose = Inf)
 end

src/abstractmcmc.jl

@@ -30,6 +30,24 @@ getadaptor(state::HMCState) = state.adaptor
 getmetric(state::HMCState) = state.metric
 getintegrator(state::HMCState) = state.κ.τ.integrator
 
+function AbstractMCMC.getparams(state::HMCState)
+    return state.transition.z.θ
+end
+
+function AbstractMCMC.setparams!!(
+    model::AbstractMCMC.LogDensityModel,
+    state::HMCState,
+    params,
+)
+    hamiltonian = AdvancedHMC.Hamiltonian(state.metric, model)
+    return Setfield.@set state.transition.z = AdvancedHMC.phasepoint(
+        hamiltonian,
+        params,
+        state.transition.z.r;
+        ℓκ = state.transition.z.ℓκ,
+    )
+end
+
 """
     $(TYPEDSIGNATURES)
 

test/abstractmcmc.jl

@@ -21,6 +21,26 @@ using Statistics: mean
         LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), ℓπ_gdemo),
     )
 
+    @testset "getparams and setparams!!" begin
+        t, s = AbstractMCMC.step(rng, model, nuts;)
+
+        θ = AbstractMCMC.getparams(s)
+        @test θ == t.z.θ
+        new_state = AbstractMCMC.setparams!!(model, s, θ)
+        @test new_state.transition.z.θ == θ
+        new_state_logπ = new_state.transition.z.ℓπ
+        @test new_state_logπ.value == s.transition.z.ℓπ.value
+        @test new_state_logπ.gradient == s.transition.z.ℓπ.gradient
+        new_state_logκ = new_state.transition.z.ℓκ
+        @test new_state_logκ.value == s.transition.z.ℓκ.value
+        @test new_state_logκ.gradient == s.transition.z.ℓκ.gradient
+        @test new_state.transition.z.r == s.transition.z.r
+
+        new_θ = randn(rng, 2)
+        new_state = AbstractMCMC.setparams!!(model, s, new_θ)
+        @test AbstractMCMC.getparams(new_state) == new_θ
+    end
+
     samples_nuts = AbstractMCMC.sample(
         rng,
         model,