Deepcopy adaptor before starting sampling

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-version = "0.6.3"
+version = "0.6.5"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/AdvancedHMC.jl

@@ -6,8 +6,7 @@ using Statistics: mean, var, middle
 using LinearAlgebra:
     Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky, UniformScaling
 using StatsFuns: logaddexp, logsumexp
-import Random
-using Random: GLOBAL_RNG, AbstractRNG
+import Random: Random, AbstractRNG
 using ProgressMeter: ProgressMeter
 using SimpleUnPack: @unpack
 

src/abstractmcmc.jl

@@ -55,7 +55,7 @@ A convenient wrapper around `AbstractMCMC.sample` avoiding explicit construction
 """
 
 function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     model::AbstractMCMC.LogDensityModel,
     sampler::AbstractHMCSampler,
     N::Integer;
@@ -92,7 +92,7 @@ function AbstractMCMC.sample(
 end
 
 function AbstractMCMC.sample(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     model::AbstractMCMC.LogDensityModel,
     sampler::AbstractHMCSampler,
     parallel::AbstractMCMC.AbstractMCMCEnsemble,
@@ -318,7 +318,7 @@ end
 #########
 
 function make_step_size(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     spl::HMCSampler,
     hamiltonian::Hamiltonian,
     initial_params,
@@ -329,7 +329,7 @@ function make_step_size(
 end
 
 function make_step_size(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     spl::AbstractHMCSampler,
     hamiltonian::Hamiltonian,
     initial_params,
@@ -340,7 +340,7 @@ function make_step_size(
 end
 
 function make_step_size(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     integrator::AbstractIntegrator,
     T::Type,
     hamiltonian::Hamiltonian,
@@ -356,7 +356,7 @@ function make_step_size(
 end
 
 function make_step_size(
-    rng::Random.AbstractRNG,
+    rng::AbstractRNG,
     integrator::Symbol,
     T::Type,
     hamiltonian::Hamiltonian,

src/hamiltonian.jl

@@ -108,6 +108,10 @@ Base.isfinite(v::DualValue) = all(isfinite, v.value) && all(isfinite, v.gradient
 Base.isfinite(v::AbstractVecOrMat) = all(isfinite, v)
 Base.isfinite(z::PhasePoint) = isfinite(z.ℓπ) && isfinite(z.ℓκ)
 
+Base.isnan(v::DualValue) = any(isnan, v.value) || any(isnan, v.gradient)
+Base.isnan(v::AbstractVecOrMat) = any(isnan, v)
+Base.isnan(z::PhasePoint) = isnan(z.ℓπ) || isnan(z.ℓκ)
+
 ###
 ### Negative energy (or log probability) functions.
 ### NOTE: the general form (i.e. non-Euclidean) of K depends on both θ and r.

src/metric.jl

@@ -138,7 +138,7 @@ Base.rand(
     kinetic::AbstractKinetic,
 ) = _rand(rng, metric, kinetic)
 Base.rand(metric::AbstractMetric, kinetic::AbstractKinetic) =
-    rand(GLOBAL_RNG, metric, kinetic)
+    rand(Random.default_rng(), metric, kinetic)
 
 # ignore θ by default unless defined by the specific kinetic (i.e. not position-dependent)
 Base.rand(
@@ -154,4 +154,4 @@ Base.rand(
     θ::AbstractVecOrMat,
 ) = rand(rng, metric, kinetic)
 Base.rand(metric::AbstractMetric, kinetic::AbstractKinetic, θ::AbstractVecOrMat) =
-    rand(metric, kinetic)
+    rand(Random.default_rng(), metric, kinetic)

src/sampler.jl

@@ -118,7 +118,7 @@ sample(
     progress::Bool = false,
     (pm_next!)::Function = pm_next!,
 ) = sample(
-    GLOBAL_RNG,
+    Random.default_rng(),
     h,
     κ,
     θ,
@@ -146,7 +146,7 @@ sample(
     )
 Sample `n_samples` samples using the proposal `κ` under Hamiltonian `h`.
 - The randomness is controlled by `rng`. 
-    - If `rng` is not provided, `GLOBAL_RNG` will be used.
+    - If `rng` is not provided, `Random.default_rng()` will be used.
 - The initial point is given by `θ`.
 - The adaptor is set by `adaptor`, for which the default is no adaptation.
     - It will perform `n_adapts` steps of adaptation, for which the default is the minimum of `1_000` and 10% of `n_samples`
@@ -166,6 +166,42 @@ function sample(
     verbose::Bool = true,
     progress::Bool = false,
     (pm_next!)::Function = pm_next!,
+) where {T<:AbstractVecOrMat{<:AbstractFloat}}
+    # Prevent adaptor from being mutated
+    adaptor = deepcopy(adaptor)
+    # Then call sample_mutating_adaptor with the same arguments
+    return sample_mutating_adaptor(
+        rng,
+        h,
+        κ,
+        θ,
+        n_samples,
+        adaptor,
+        n_adapts;
+        drop_warmup = drop_warmup,
+        verbose = verbose,
+        progress = progress,
+        (pm_next!) = pm_next!,
+    )
+end
+
+"""
+    sample_mutating_adaptor(args...; kwargs...)
+
+The same as `sample`, but mutates the `adaptor` argument.
+"""
+function sample_mutating_adaptor(
+    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
+    h::Hamiltonian,
+    κ::HMCKernel,
+    θ::T,
+    n_samples::Int,
+    adaptor::AbstractAdaptor = NoAdaptation(),
+    n_adapts::Int = min(div(n_samples, 10), 1_000);
+    drop_warmup = false,
+    verbose::Bool = true,
+    progress::Bool = false,
+    (pm_next!)::Function = pm_next!,
 ) where {T<:AbstractVecOrMat{<:AbstractFloat}}
     @assert !(drop_warmup && (adaptor isa Adaptation.NoAdaptation)) "Cannot drop warmup samples if there is no adaptation phase."
     # Prepare containers to store sampling results

src/trajectory.jl

@@ -3,7 +3,7 @@
 ####
 #### Developers' Notes
 ####
-#### Not all functions that use `rng` require a fallback function with `GLOBAL_RNG`
+#### Not all functions that use `rng` require a fallback function with `Random.default_rng()`
 #### as default. In short, only those exported to other libries need such a fallback
 #### function. Internal uses shall always use the explict `rng` version. (Kai Xu 6/Jul/19)
 
@@ -241,10 +241,10 @@ $(SIGNATURES)
 
 Make a MCMC transition from phase point `z` using the trajectory `τ` under Hamiltonian `h`.
 
-NOTE: This is a RNG-implicit fallback function for `transition(GLOBAL_RNG, τ, h, z)`
+NOTE: This is a RNG-implicit fallback function for `transition(Random.default_rng(), τ, h, z)`
 """
 function transition(τ::Trajectory, h::Hamiltonian, z::PhasePoint)
-    return transition(GLOBAL_RNG, τ, h, z)
+    return transition(Random.default_rng(), τ, h, z)
 end
 
 ###
@@ -834,7 +834,7 @@ function find_good_stepsize(
     θ::AbstractVector{<:AbstractFloat};
     max_n_iters::Int = 100,
 )
-    return find_good_stepsize(GLOBAL_RNG, h, θ; max_n_iters = max_n_iters)
+    return find_good_stepsize(Random.default_rng(), h, θ; max_n_iters = max_n_iters)
 end
 
 "Perform MH acceptance based on energy, i.e. negative log probability."

test/adaptation.jl

@@ -14,7 +14,17 @@ function runnuts(ℓπ, metric; n_samples = 3_000)
     integrator = AdvancedHMC.make_integrator(nuts, step_size)
     κ = AdvancedHMC.make_kernel(nuts, integrator)
     adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator)
-    samples, stats = sample(h, κ, θ_init, n_samples, adaptor, n_adapts; verbose = false)
+    # Use mutating version of sample() here
+    samples, stats = AdvancedHMC.sample_mutating_adaptor(
+        rng,
+        h,
+        κ,
+        θ_init,
+        n_samples,
+        adaptor,
+        n_adapts;
+        verbose = false,
+    )
     return (samples = samples, stats = stats, adaptor = adaptor)
 end
 

test/sampler.jl

@@ -62,6 +62,7 @@ end
     n_steps = 10
     n_samples = 22_000
     n_adapts = 4_000
+
     @testset "$metricsym" for (metricsym, metric) in Dict(
         :UnitEuclideanMetric => UnitEuclideanMetric(D),
         :DiagEuclideanMetric => DiagEuclideanMetric(D),
@@ -157,6 +158,7 @@ end
             end
         end
     end
+
     @testset "drop_warmup" begin
         nuts = NUTS(0.8)
         metric = DiagEuclideanMetric(D)
@@ -191,4 +193,32 @@ end
         @test length(samples) == n_samples
         @test length(stats) == n_samples
     end
+
+    @testset "reproducibility" begin
+        # Multiple calls to sample() should yield the same results
+        nuts = NUTS(0.8)
+        metric = DiagEuclideanMetric(D)
+        h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
+        integrator = Leapfrog(ϵ)
+        κ = AdvancedHMC.make_kernel(nuts, integrator)
+        adaptor = AdvancedHMC.make_adaptor(nuts, metric, integrator)
+
+        all_samples = []
+        for i = 1:5
+            samples, stats = sample(
+                Random.MersenneTwister(42),
+                h,
+                κ,
+                θ_init,
+                100, # n_samples -- don't need so many
+                adaptor,
+                50, # n_adapts -- likewise
+                verbose = false,
+                progress = false,
+                drop_warmup = true,
+            )
+            push!(all_samples, samples)
+        end
+        @test all(map(s -> s ≈ all_samples[1], all_samples[2:end]))
+    end
 end