Update for newer AbstractMCMC/Turing interface

Author: penelopeysm

Project.toml

@@ -1,6 +1,6 @@
 name = "SliceSampling"
 uuid = "43f4d3e8-9711-4a8c-bd1b-03ac73a255cf"
-version = "0.7.10"
+version = "0.7.11"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -16,12 +16,12 @@ Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 SliceSamplingTuringExt = ["Turing"]
 
 [compat]
-AbstractMCMC = "4, 5"
+AbstractMCMC = "5.9"
 Distributions = "0.25"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 Random = "1"
-Turing = "0.41"
+Turing = "0.42"
 julia = "1.10"
 
 [extras]

docs/Project.toml

@@ -15,7 +15,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
-AbstractMCMC = "5"
+AbstractMCMC = "5.9"
 Distributions = "0.25"
 Documenter = "1"
 FillArrays = "1"
@@ -28,5 +28,5 @@ Random = "1"
 SliceSampling = "0.7.1"
 StableRNGs = "1"
 Statistics = "1"
-Turing = "0.41, 0.42"
+Turing = "0.42"
 julia = "1.10"

ext/SliceSamplingTuringExt.jl

@@ -6,35 +6,6 @@ using Random
 using SliceSampling
 using Turing
 
-# Required for using the slice samplers as `externalsampler`s in Turing
-# begin
-function Turing.Inference.getparams(
-    ::Turing.DynamicPPL.Model, sample::SliceSampling.Transition
-)
-    return sample.params
-end
-# end
-
-# Required for using the slice samplers as `Gibbs` samplers in Turing
-# begin
-Turing.Inference.isgibbscomponent(::SliceSampling.RandPermGibbs) = true
-Turing.Inference.isgibbscomponent(::SliceSampling.HitAndRun) = true
-Turing.Inference.isgibbscomponent(::SliceSampling.Slice) = true
-Turing.Inference.isgibbscomponent(::SliceSampling.SliceSteppingOut) = true
-Turing.Inference.isgibbscomponent(::SliceSampling.SliceDoublingOut) = true
-
-const SliceSamplingStates = Union{
-    SliceSampling.UnivariateSliceState,
-    SliceSampling.GibbsState,
-    SliceSampling.HitAndRunState,
-    SliceSampling.LatentSliceState,
-    SliceSampling.GibbsPolarSliceState,
-}
-function Turing.Inference.getparams(::Turing.DynamicPPL.Model, sample::SliceSamplingStates)
-    return sample.transition.params
-end
-# end
-
 function SliceSampling.initial_sample(rng::Random.AbstractRNG, ℓ::Turing.LogDensityFunction)
     n_max_attempts = 1000
 
@@ -50,7 +21,7 @@ function SliceSampling.initial_sample(rng::Random.AbstractRNG, ℓ::Turing.LogDe
             @warn "Failed to find valid initial parameters after $(init_attempt_count) attempts; consider providing explicit initial parameters using the `initial_params` keyword"
         end
 
-        # NOTE: This will sample in the unconstrained space if ℓ.varinfo is linked
+       # NOTE: This will sample in the unconstrained space if ℓ.varinfo is linked
         vi_spl = last(Turing.DynamicPPL.init!!(rng, model, vi, Turing.InitFromUniform()))
         ℓp = ℓ.getlogdensity(vi_spl)
         θ = vi_spl[:]

src/SliceSampling.jl

@@ -2,6 +2,7 @@
 module SliceSampling
 
 using AbstractMCMC
+using Accessors: Accessors
 using Distributions
 using LinearAlgebra
 using LogDensityProblems
@@ -37,6 +38,23 @@ struct Transition{P,L<:Real,I<:NamedTuple}
     info::I
 end
 
+"""
+    abstract type AbstractStateWithTransition
+
+Base type for MCMC states that contain a `Transition` stored in the `transition` field.
+"""
+abstract type AbstractStateWithTransition end
+get_transition(state::AbstractStateWithTransition) = state.transition
+AbstractMCMC.getparams(state::AbstractStateWithTransition) = get_transition(state).params
+AbstractMCMC.getstats(state::AbstractStateWithTransition) = get_transition(state).info
+function AbstractMCMC.setparams(
+    model::AbstractMCMC.LogDensityModel, state::AbstractStateWithTransition, params
+)
+    new_lp = LogDensityProblems.logdensity(model.logdensity, params)
+    new_transition = Transition(params, new_lp, NamedTuple())
+    return Accessors.@set state.transition = new_transition
+end
+
 """
     initial_sample(rng, model)
 

src/multivariate/gibbspolar.jl

@@ -29,7 +29,7 @@ function GibbsPolarSlice(w::Real; max_proposals::Int=DEFAULT_MAX_PROPOSALS)
     return GibbsPolarSlice(w, max_proposals)
 end
 
-struct GibbsPolarSliceState{T<:Transition,R<:Real,D<:AbstractVector}
+struct GibbsPolarSliceState{T<:Transition,R<:Real,D<:AbstractVector} <: AbstractStateWithTransition
     "Current [`Transition`](@ref)."
     transition::T
 

src/multivariate/hitandrun.jl

@@ -12,18 +12,11 @@ struct HitAndRun{S<:AbstractUnivariateSliceSampling} <: AbstractMultivariateSlic
     unislice::S
 end
 
-struct HitAndRunState{T<:Transition}
+struct HitAndRunState{T<:Transition} <: AbstractStateWithTransition
     "Current [`Transition`](@ref)."
     transition::T
 end
 
-function AbstractMCMC.setparams!!(
-    model::AbstractMCMC.LogDensityModel, state::HitAndRunState, params
-)
-    lp = LogDensityProblems.logdensity(model.logdensity, params)
-    return HitAndRunState(Transition(params, lp, NamedTuple()))
-end
-
 struct HitAndRunTarget{Model,Vec<:AbstractVector}
     model     :: Model
     direction :: Vec

src/multivariate/latent.jl

@@ -20,7 +20,7 @@ function LatentSlice(beta::Real; max_proposals::Int=DEFAULT_MAX_PROPOSALS)
     return LatentSlice(beta, max_proposals)
 end
 
-struct LatentSliceState{T<:Transition,S<:AbstractVector}
+struct LatentSliceState{T<:Transition,S<:AbstractVector} <: AbstractStateWithTransition
     "Current [`Transition`](@ref)."
     transition::T
 

src/multivariate/randpermgibbs.jl

@@ -20,16 +20,11 @@ struct RandPermGibbs{
     unislice::S
 end
 
-struct GibbsState{T<:Transition}
+struct GibbsState{T<:Transition} <: AbstractStateWithTransition
     "Current [`Transition`](@ref)."
     transition::T
 end
 
-function AbstractMCMC.setparams!!(model::AbstractMCMC.LogDensityModel, ::GibbsState, params)
-    lp = LogDensityProblems.logdensity(model.logdensity, params)
-    return GibbsState(Transition(params, lp, NamedTuple()))
-end
-
 struct GibbsTarget{Model,Idx<:Integer,Vec<:AbstractVector}
     model :: Model
     idx   :: Idx

src/univariate/univariate.jl

@@ -24,18 +24,11 @@ function slice_sampling_univariate(
     return exceeded_max_prop(max_prop)
 end
 
-struct UnivariateSliceState{T<:Transition}
+struct UnivariateSliceState{T<:Transition} <: AbstractStateWithTransition
     "Current [`Transition`](@ref)."
     transition::T
 end
 
-function AbstractMCMC.setparams!!(
-    model::AbstractMCMC.LogDensityModel, state::UnivariateSliceState, params
-)
-    lp = LogDensityProblems.logdensity(model.logdensity, params)
-    return UnivariateSliceState(Transition(params, lp, NamedTuple()))
-end
-
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::AbstractMCMC.LogDensityModel,

test/Project.toml

@@ -10,13 +10,13 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [compat]
-AbstractMCMC = "4, 5"
+AbstractMCMC = "5.9"
 Accessors = "0.1"
 Distributions = "0.25"
 LogDensityProblems = "2"
 MCMCTesting = "0.3"
 Random = "1"
 StableRNGs = "1"
 Test = "1"
-Turing = "0.41, 0.42"
+Turing = "0.42"
 julia = "1.10"