add stereographic slice sampling

Author: Red-Portal

Project.toml

@@ -10,6 +10,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
@@ -27,6 +28,7 @@ FillArrays = "1"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1"
+LogExpFunctions = "0.3"
 Random = "1"
 Requires = "1"
 Turing = "0.36"

src/SliceSampling.jl

@@ -7,6 +7,7 @@ using Distributions
 using FillArrays
 using LinearAlgebra
 using LogDensityProblems
+using LogExpFunctions
 using Random
 
 # The following is necessary because Turing wraps all models with
@@ -116,6 +117,10 @@ include("multivariate/latent.jl")
 export GibbsPolarSlice
 include("multivariate/gibbspolar.jl")
 
+# Stereographic Slice Sampling
+export StereographicSlice
+include("multivariate/stereographic.jl")
+
 # Turing Compatibility
 
 if !isdefined(Base, :get_extension)

src/multivariate/stereographic.jl

@@ -0,0 +1,112 @@
+
+"""
+    StereographicSlice(; max_proposals)
+
+Stereographic slice sampling algorithm by Bell, Latuszynski, and Roberts[^BLR].
+
+# Keyword Arguments
+- `max_proposals::Int`: Maximum number of proposals allowed until throwing an error (default: `$(DEFAULT_MAX_PROPOSALS)`).
+"""
+@kwdef struct StereographicSlice{RType <: Real} <: AbstractMultivariateSliceSampling
+    max_proposals :: Int   = DEFAULT_MAX_PROPOSALS
+end
+
+struct StereographicSliceState{T<:Transition}
+    "Current [`Transition`](@ref)."
+    transition::T
+end
+
+function rand_uniform_sphere_orthogonal_subspace(
+    rng::Random.AbstractRNG, subspace_vector::AbstractVector
+)
+    z      = subspace_vector
+    d      = length(subspace_vector)
+    v      = randn(rng, d)
+    v_proj = dot(z, v)/sum(abs2, z)*z
+    v_orth = v - v_proj
+    v_orth / norm(v_orth)
+end
+
+function stereographic_projection(z::AbstractVector)
+    d = length(z) - 1
+    return z[1:d] ./ (1 - z[d+1]) 
+end
+
+function stereographic_inverse_projection(x::AbstractVector)
+    d       = length(x)
+    z       = zeros(d + 1)
+    x_norm2 = sum(abs2, x)
+    z[1:d]  = 2*x / (x_norm2 + 1)
+    z[d+1]  = (x_norm2 - 1)/(x_norm2 + 1)
+    z
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    sampler::StereographicSlice;
+    initial_params=nothing,
+    kwargs...,
+)
+    logdensitymodel = model.logdensity
+    x = initial_params === nothing ? initial_sample(rng, logdensitymodel) : initial_params
+    lp = LogDensityProblems.logdensity(logdensitymodel, x)
+    t = Transition(x, lp, NamedTuple())
+    return t, t
+end
+
+function logdensity_sphere(ℓπ::Real, x::AbstractVector)
+    d  = length(x)
+    return ℓπ + d*log(1 + sum(abs2, x))
+end
+
+function AbstractMCMC.step(
+    rng::Random.AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    sampler::StereographicSlice,
+    state::Transition;
+    kwargs...,
+)
+    logdensitymodel = model.logdensity
+    max_proposals   = sampler.max_proposals
+
+    ℓp        = state.lp
+    x         = state.params
+    z         = stereographic_inverse_projection(x)
+    v         = rand_uniform_sphere_orthogonal_subspace(rng, z)
+    ℓp_sphere = logdensity_sphere(ℓp, x)
+    ℓw        = ℓp_sphere - Random.randexp(rng, eltype(x))
+
+    θ     = rand(rng, Uniform(0, 2π))
+    θ_max = θ
+    θ_min = θ - 2π
+
+    props = 0
+    while true
+        props += 1
+
+        x_prop         = stereographic_projection(z*cos(θ) + v*sin(θ))
+        ℓp_prop        = LogDensityProblems.logdensity(logdensitymodel, x_prop)
+        ℓp_sphere_prop = logdensity_sphere(ℓp_prop, x_prop)
+
+        if ℓw < ℓp_sphere_prop
+            ℓp = ℓp_prop
+            x  = x_prop
+            break
+        end
+
+        if props > max_proposals
+            exceeded_max_prop(max_proposals)
+        end
+
+        if θ < 0
+            θ_min = θ
+        else
+            θ_max = θ
+        end
+
+        θ = rand(rng, Uniform(θ_min, θ_max))
+    end
+    t = Transition(x, ℓp, (num_proposals=props,))
+    return t, t
+end