Merge pull request #191 from tpapp/tp/use-ohmythreads

Use OhMyThreads for threaded example.

Incidental: use explicit RNGs everywhere.

Author: tpapp

docs/Project.toml

@@ -3,9 +3,9 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-ThreadTools = "dbf13d8f-d36e-4350-8970-f3a99faba1a8"
 TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"
 TransformedLogDensities = "f9bc47f6-f3f8-4f3b-ab21-f8bc73906f26"
 

docs/src/worked_example.md

@@ -94,13 +94,13 @@ summarize_tree_statistics(results.tree_statistics)
 
 Usually one would run multiple chains and check convergence and mixing using generic MCMC diagnostics not specific to NUTS.
 
-The specifics of running multiple chains is up to the user: various forms of [parallel computing](https://docs.julialang.org/en/v1/manual/parallel-computing/) can be utilized depending on the problem scale and the hardware available. In the example below we use [multi-threading](https://docs.julialang.org/en/v1/manual/multi-threading/), using [ThreadTools.jl](https://github.com/baggepinnen/ThreadTools.jl); other excellent packages are available for threading.
+The specifics of running multiple chains is up to the user: various forms of [parallel computing](https://docs.julialang.org/en/v1/manual/parallel-computing/) can be utilized depending on the problem scale and the hardware available. In the example below we use [multi-threading](https://docs.julialang.org/en/v1/manual/multi-threading/), using [OhMyThreads.jl](https://github.com/JuliaFolds2/OhMyThreads.jl); other excellent packages are available for threading.
 
 It is easy to obtain posterior results for use with [MCMCDiagnosticsTools.jl](https://github.com/TuringLang/MCMCDiagnosticTools.jl/) with [`stack_posterior_matrices`](@ref):
 
 ```@example bernoulli
-using ThreadTools, MCMCDiagnosticTools
-results5 = tmap1(_ -> mcmc_with_warmup(Random.default_rng(), ∇P, 1000; reporter = NoProgressReport()), 1:5)
+using OhMyThreads, MCMCDiagnosticTools
+results5 = tcollect(mcmc_with_warmup(Random.default_rng(), ∇P, 1000; reporter = NoProgressReport()) for _ in 1:5)
 ess_rhat(stack_posterior_matrices(results5))
 ```
 

test/runtests.jl

@@ -14,7 +14,7 @@ using LogDensityTestSuite
 ### general test environment
 ###
 
-const RNG = Random.default_rng()   # shorthand
+const RNG = copy(Random.default_rng())   # shorthand
 Random.seed!(RNG, UInt32[0x23ef614d, 0x8332e05c, 0x3c574111, 0x121aa2f4])
 
 "Tolerant testing in a CI environment."

test/sample-correctness_tests.jl

@@ -11,9 +11,9 @@ const MCMC_ARGS2 = (warmup_stages = default_warmup_stages(; M = Symmetric),)
 
 @testset "NUTS tests with random normal" begin
     for _ in 1:10
-        K = rand(3:10)
-        μ = randn(K)
-        d = abs.(randn(K))
+        K = rand(RNG, 3:10)
+        μ = randn(RNG, K)
+        d = abs.(randn(RNG, K))
         C = rand_C(K)
         ℓ = multivariate_normal(μ, Diagonal(d) * C)
         title = "multivariate normal μ = $(μ) d = $(d) C = $(C)"
@@ -94,7 +94,7 @@ end
           0.0 0.0 0.7434985947205197]
     ℓ2 = multivariate_normal(ones(3), D2 * C2)
     ℓ = mix(0.2, ℓ1, ℓ2)
-    NUTS_tests(RNG, ℓ, "mixture of two normals", 1000; τ_alert = 0.15)
+    NUTS_tests(RNG, ℓ, "mixture of two normals", 1000; τ_alert = 0.15, p_alert = 0.005)
 end
 
 @testset "NUTS tests with heavier tails and skewness" begin

test/sample-correctness_utilities.jl

@@ -28,7 +28,7 @@ end
 "Random Cholesky factor for correlation matrix."
 function rand_C(K)
     t = TransformVariables.CorrCholeskyFactor(K)
-    TransformVariables.transform(t, randn(TransformVariables.dimension(t)) ./ 4)'
+    TransformVariables.transform(t, randn(RNG, TransformVariables.dimension(t)) ./ 4)'
 end
 
 """

test/test_NUTS.jl

@@ -1,6 +1,7 @@
 using DynamicHMC: TrajectoryNUTS, rand_bool_logprob, GeneralizedTurnStatistic,
     AcceptanceStatistic, leaf_acceptance_statistic, acceptance_rate, TreeStatisticsNUTS,
-    NUTS, sample_tree, combine_turn_statistics, combine_visited_statistics
+    NUTS, sample_tree, combine_turn_statistics, combine_visited_statistics, evaluate_ℓ,
+    Hamiltonian
 
 ###
 ### random booleans
@@ -87,13 +88,13 @@ end
     # A test for sample_tree with a fixed ϵ and κ, which is perfectly adapted and should
     # provide excellent mixing
     for _ in 1:10
-        K = rand(2:8)
+        K = rand(RNG, 2:8)
         N = 10000
-        μ = randn(K)
+        μ = randn(RNG, K)
         Σ = rand_Σ(K)
         L = cholesky(Σ).L
         ℓ = multivariate_normal(μ, L)
-        Q = evaluate_ℓ(ℓ, randn(K))
+        Q = evaluate_ℓ(ℓ, randn(RNG, K))
         H = Hamiltonian(GaussianKineticEnergy(Σ), ℓ)
         qs = Array{Float64}(undef, N, K)
         ϵ = 0.5
@@ -103,7 +104,8 @@ end
             qs[i, :] = Q.q
         end
         m, C = mean_and_cov(qs, 1)
-        @test vec(m) ≈ μ atol = 0.15 rtol = maximum(diag(C))*0.02 norm = x -> norm(x,1)
+        tol = maximum(diag(C)) / 50
+        @test vec(m) ≈ μ atol = tol rtol = tol  norm = x -> norm(x,1)
         @test cov(qs, dims = 1) ≈ L*L' atol = 0.1 rtol = 0.1
     end
 end

test/test_diagnostics.jl

@@ -6,16 +6,16 @@
     N = 1000
     directions = Directions(UInt32(0))
     function rand_invalidtree()
-        if rand() < 0.1
+        if rand(RNG) < 0.1
             REACHED_MAX_DEPTH
         else
-            left = rand(-5:5)
-            right = left + rand(0:5)
+            left = rand(RNG, -5:5)
+            right = left + rand(RNG, 0:5)
             InvalidTree(left, right)
         end
     end
-    tree_statistics = [TreeStatisticsNUTS(randn(), rand(0:5), rand_invalidtree(), rand(),
-                                          rand(1:30), directions) for _ in 1:N]
+    tree_statistics = [TreeStatisticsNUTS(randn(RNG), rand(RNG, 0:5), rand_invalidtree(),
+                                          rand(RNG), rand(RNG, 1:30), directions) for _ in 1:N]
     stats = summarize_tree_statistics(tree_statistics)
     # acceptance rates
     @test stats.N == N

test/test_hamiltonian.jl

@@ -19,21 +19,21 @@ end
 
 @testset "Gaussian KE full" begin
     for _ in 1:100
-        K = rand(2:10)
+        K = rand(RNG, 2:10)
         Σ = rand_Σ(Symmetric, K)
         κ = GaussianKineticEnergy(inv(Σ))
         (; M⁻¹, W) = κ
         @test W isa LowerTriangular
         @test M⁻¹ * W * W' ≈ Diagonal(ones(K))
         m, C = simulated_meancov(()->rand_p(RNG, κ), 10000)
         @test Matrix(Σ) ≈ C rtol = 0.1
-        test_KE_gradient(κ, randn(K))
+        test_KE_gradient(κ, randn(RNG, K))
     end
 end
 
 @testset "Gaussian KE diagonal" begin
     for _ in 1:100
-        K = rand(2:10)
+        K = rand(RNG, 2:10)
         Σ = rand_Σ(Diagonal, K)
         κ = GaussianKineticEnergy(inv(Σ))
         (; M⁻¹, W) = κ
@@ -42,7 +42,7 @@ end
         @test M⁻¹ * (W * W') ≈ Diagonal(ones(K))
         m, C = simulated_meancov(()->rand_p(RNG, κ), 10000)
         @test Matrix(Σ) ≈ C rtol = 0.1
-        test_KE_gradient(κ, randn(K))
+        test_KE_gradient(κ, randn(RNG, K))
     end
 end
 
@@ -57,7 +57,7 @@ end
         @test ℓ2 == ℓq
         @test ∇ℓ2 == ∇ℓq
     end
-    (; H, z, Σ ) = rand_Hz(rand(3:10))
+    (; H, z, Σ ) = rand_Hz(rand(RNG, 3:10))
     test_consistency(H, z)
     ϵ = find_stable_ϵ(H.κ, Σ)
     for _ in 1:10
@@ -81,10 +81,10 @@ end
     M = rand_Σ(Diagonal, n)
     m = diag(M)
     κ = GaussianKineticEnergy(inv(M))
-    q = randn(n)
-    p = randn(n)
+    q = randn(RNG, n)
+    p = randn(RNG, n)
     Σ = rand_Σ(n)
-    ℓ = multivariate_normal(randn(n), cholesky(Σ).L)
+    ℓ = multivariate_normal(randn(RNG, n), cholesky(Σ).L)
     H = Hamiltonian(κ, ℓ)
     ϵ = find_stable_ϵ(H.κ, Σ)
     z = PhasePoint(evaluate_ℓ(ℓ, q), p)
@@ -110,7 +110,7 @@ end
 
     @testset "invalid values" begin
         n = 3
-        ℓ = multivariate_normal(randn(n), I)
+        ℓ = multivariate_normal(randn(RNG, n), I)
         @test_throws DynamicHMCError evaluate_ℓ(ℓ, fill(NaN, n))
     end
 end
@@ -134,7 +134,7 @@ end
     end
 
     for _ in 1:100
-        (; H, z) = rand_Hz(rand(2:5))
+        (; H, z) = rand_Hz(rand(RNG, 2:5))
         ϵ = find_initial_stepsize(InitialStepsizeSearch(), local_log_acceptance_ratio(H, z))
         test_hamiltonian_invariance(H, z, 10, ϵ/100; atol = 0.5)
     end
@@ -226,7 +226,7 @@ function HMC_transition(H, z::PhasePoint, ϵ, L)
         z′ = leapfrog(H, z′, ϵ)
     end
     Δ = logdensity(H, z′) - π₀
-    accept = Δ > 0 || (rand() < exp(Δ))
+    accept = Δ > 0 || (rand(RNG) < exp(Δ))
     accept ? z′ : z
 end
 
@@ -249,7 +249,7 @@ end
     # Tests the leapfrog and Hamiltonian code with HMC.
     K = 2
     ℓ = multivariate_normal(zeros(K), Diagonal(ones(K)))
-    q = randn(K)
+    q = randn(RNG, K)
     H = Hamiltonian(GaussianKineticEnergy(Diagonal(ones(K))), ℓ)
     qs = HMC_sample(H, q, 10000, find_stable_ϵ(H.κ, Diagonal(ones(K))) / 5)
     m, C = mean_and_cov(qs, 1)

test/test_stepsize.jl

@@ -30,7 +30,7 @@ $(SIGNATURES)
 A parametric random acceptance rate that depends on the stepsize. For unit
 testing acceptance rate tuning.
 """
-dummy_acceptance_rate(ϵ, σ = 0.05) = min(1/ϵ * exp(randn()*σ - σ^2/2), 1)
+dummy_acceptance_rate(ϵ, σ = 0.05) = min(1/ϵ * exp(randn(RNG)*σ - σ^2/2), 1)
 
 mean_dummy_acceptance_rate(ϵ, σ = 0.05) = mean(dummy_acceptance_rate(ϵ, σ) for _ in 1:10000)
 
@@ -83,7 +83,7 @@ end
     p = InitialStepsizeSearch()
     _bkt(A, ϵ, C) = (A(ϵ) - p.log_threshold) * (A(ϵ * C) - p.log_threshold) ≤ 0
     for _ in 1:100
-        (; H, z) = rand_Hz(rand(3:5))
+        (; H, z) = rand_Hz(rand(RNG, 3:5))
         A = local_log_acceptance_ratio(H, z)
         ϵ = find_initial_stepsize(p, A)
         @test _bkt(A, ϵ, 0.5) || _bkt(A, ϵ, 2.0)

test/utilities.jl

@@ -4,11 +4,11 @@ $(SIGNATURES)
 Random positive definite matrix of size `n` x `n` (for testing).
 """
 function rand_Σ(::Type{Symmetric}, n)
-    A = randn(n, n)
+    A = randn(RNG, n, n)
     Symmetric(A'*A .+ 0.01)
 end
 
-rand_Σ(::Type{Diagonal}, n) = Diagonal(randn(n).^2 .+ 0.01)
+rand_Σ(::Type{Diagonal}, n) = Diagonal(randn(RNG, n).^2 .+ 0.01)
 
 rand_Σ(n::Int) = rand_Σ(Symmetric, n)
 
@@ -31,7 +31,7 @@ end
 @testset "simulated meancov" begin
     μ = [2, 1.2]
     D = [2.0, 0.7]
-    m, C = simulated_meancov(()-> randn(2) .* D .+ μ, 10000)
+    m, C = simulated_meancov(()-> randn(RNG, 2) .* D .+ μ, 10000)
     @test m ≈ μ atol = 0.05 rtol = 0.1
     @test C ≈ Diagonal(abs2.(D)) atol = 0.05 rtol = 0.1
 end