diff --git a/Project.toml b/Project.toml index f35b1df4..4f1bdd84 100644 --- a/Project.toml +++ b/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" diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl index fcaab095..1895f8cd 100644 --- a/src/AdvancedHMC.jl +++ b/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 diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl index 40585909..c3a73778 100644 --- a/src/abstractmcmc.jl +++ b/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, diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index b5546051..f8f2ab32 100644 --- a/src/hamiltonian.jl +++ b/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. diff --git a/src/metric.jl b/src/metric.jl index 2afbd629..017f5ac5 100644 --- a/src/metric.jl +++ b/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) diff --git a/src/sampler.jl b/src/sampler.jl index 7d1b7eb5..9d46d276 100644 --- a/src/sampler.jl +++ b/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 diff --git a/src/trajectory.jl b/src/trajectory.jl index 4b17422f..56aedd1a 100644 --- a/src/trajectory.jl +++ b/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." diff --git a/test/adaptation.jl b/test/adaptation.jl index 1a644c71..0be29505 100644 --- a/test/adaptation.jl +++ b/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 diff --git a/test/sampler.jl b/test/sampler.jl index 109f01a4..dd439008 100644 --- a/test/sampler.jl +++ b/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