Add support for custom RNG

Author: devmotion

src/AdvancedMH.jl

@@ -5,7 +5,7 @@ import AbstractMCMC
 using Distributions
 using Requires
 
-using Random
+import Random
 
 # Exports
 export MetropolisHastings, DensityModel, RWMH, StaticMH, StaticProposal, RandomWalkProposal
@@ -51,7 +51,7 @@ logdensity(model::DensityModel, t::Transition) = t.lp
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::AbstractRNG, 
+    rng::Random.AbstractRNG, 
     model::DensityModel, 
     s::Metropolis, 
     N::Integer,
@@ -64,7 +64,7 @@ function AbstractMCMC.bundle_samples(
 end
 
 function AbstractMCMC.bundle_samples(
-    rng::AbstractRNG, 
+    rng::Random.AbstractRNG, 
     model::DensityModel, 
     s::Metropolis, 
     N::Integer, 

src/mcmcchains-connect.jl

@@ -2,7 +2,7 @@ import .MCMCChains: Chains
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::AbstractRNG, 
+    rng::Random.AbstractRNG, 
     model::DensityModel, 
     s::Metropolis, 
     N::Integer, 

src/mh-core.jl

@@ -48,62 +48,71 @@ end
 StaticMH(d) = MetropolisHastings(StaticProposal(d))
 RWMH(d) = MetropolisHastings(RandomWalkProposal(d))
 
+# default function without RNG
+propose(spl::MetropolisHastings, args...) = propose(Random.GLOBAL_RNG, spl, args...)
+
 # Propose from a vector of proposals
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:AbstractArray},
     model::DensityModel
 )
-    proposal = map(p -> propose(p, model), spl.proposal)
+    proposal = map(p -> propose(rng, p, model), spl.proposal)
     return Transition(model, proposal)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:AbstractArray},
     model::DensityModel,
     params_prev::Transition
 )
     proposal = map(spl.proposal, params_prev.params) do p, params
-        propose(p, model, params)
+        propose(rng, p, model, params)
     end
     return Transition(model, proposal)
 end
 
 # Make a proposal from one Proposal struct.
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:Proposal},
     model::DensityModel
 )
-    proposal = propose(spl.proposal, model)
+    proposal = propose(rng, spl.proposal, model)
     return Transition(model, proposal)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:Proposal},
     model::DensityModel,
     params_prev::Transition
 )
-    proposal = propose(spl.proposal, model, params_prev.params)
+    proposal = propose(rng, spl.proposal, model, params_prev.params)
     return Transition(model, proposal)
 end
 
 # Make a proposal from a NamedTuple of Proposal.
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:NamedTuple},
     model::DensityModel
 )
     proposal = map(spl.proposal) do p
-        propose(p, model)
+        propose(rng, p, model)
     end
     return Transition(model, proposal)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     spl::MetropolisHastings{<:NamedTuple},
     model::DensityModel,
     params_prev::Transition
 )
     proposal = map(spl.proposal, params_prev.params) do p, params
-        propose(p, model, params)
+        propose(rng, p, model, params)
     end
     return Transition(model, proposal)
 end
@@ -143,7 +152,7 @@ end
 # beginning of sampling. Return the initial parameter used
 # to define the sampler.
 function AbstractMCMC.step!(
-    rng::AbstractRNG,
+    rng::Random.AbstractRNG,
     model::DensityModel,
     spl::MetropolisHastings,
     N::Integer,
@@ -152,7 +161,7 @@ function AbstractMCMC.step!(
     kwargs...
 )
     if init_params === nothing
-        return propose(spl, model)
+        return propose(rng, spl, model)
     else
         return Transition(model, init_params)
     end
@@ -162,15 +171,15 @@ end
 # either a new proposal (if accepted) or the previous proposal 
 # (if not accepted).
 function AbstractMCMC.step!(
-    rng::AbstractRNG,
+    rng::Random.AbstractRNG,
     model::DensityModel,
     spl::MetropolisHastings,
     ::Integer,
     params_prev::Transition;
     kwargs...
 )
     # Generate a new proposal.
-    params = propose(spl, model, params_prev)
+    params = propose(rng, spl, model, params_prev)
 
     # Calculate the log acceptance probability.
     logα = logdensity(model, params) - logdensity(model, params_prev) + 

src/proposal.jl

@@ -9,8 +9,11 @@ struct RandomWalkProposal{P} <: Proposal{P}
 end
 
 # Random draws
-Base.rand(p::Proposal{<:Distribution}) = rand(p.proposal)
-Base.rand(p::Proposal{<:AbstractArray}) = map(rand, p.proposal)
+Base.rand(p::Proposal, args...) = rand(Random.GLOBAL_RNG, p, args...)
+Base.rand(rng::Random.AbstractRNG, p::Proposal{<:Distribution}) = rand(rng, p.proposal)
+function Base.rand(rng::Random.AbstractRNG, p::Proposal{<:AbstractArray})
+    return map(x -> rand(rng, x), p.proposal)
+end
 
 # Densities
 Distributions.logpdf(p::Proposal{<:Distribution}, v) = logpdf(p.proposal, v)
@@ -23,16 +26,17 @@ end
 # Random Walk #
 ###############
 
-function propose(p::RandomWalkProposal, m::DensityModel)
-    return propose(StaticProposal(p.proposal), m)
+function propose(rng::Random.AbstractRNG, p::RandomWalkProposal, m::DensityModel)
+    return propose(rng, StaticProposal(p.proposal), m)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     proposal::RandomWalkProposal{<:Union{Distribution,AbstractArray}}, 
     model::DensityModel, 
     t
 )
-    return t + rand(proposal)
+    return t + rand(rng, proposal)
 end
 
 function q(
@@ -48,11 +52,12 @@ end
 ##########
 
 function propose(
+    rng::Random.AbstractRNG,
     proposal::StaticProposal{<:Union{Distribution,AbstractArray}},
     model::DensityModel,
     t=nothing
 )
-    return rand(proposal)
+    return rand(rng, proposal)
 end
 
 function q(
@@ -76,18 +81,20 @@ for T in (StaticProposal, RandomWalkProposal)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     proposal::Proposal{<:Function},
     model::DensityModel
 )
-    return propose(proposal(), model)
+    return propose(rng, proposal(), model)
 end
 
 function propose(
+    rng::Random.AbstractRNG,
     proposal::Proposal{<:Function}, 
     model::DensityModel,
     t
 )
-    return propose(proposal(t), model)
+    return propose(rng, proposal(t), model)
 end
 
 function q(

src/structarray-connect.jl

@@ -2,7 +2,7 @@ import .StructArrays: StructArray
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    rng::AbstractRNG, 
+    rng::Random.AbstractRNG, 
     model::DensityModel, 
     s::Metropolis, 
     N::Integer,