From 0371b10b56fc05a6fc0ab9b164e60ca043c0c465 Mon Sep 17 00:00:00 2001
From: Penelope Yong <penelopeysm@gmail.com>
Date: Tue, 5 Nov 2024 01:34:12 +0000
Subject: [PATCH] Deepcopy adaptor before starting sampling

This avoids the unintuitive behaviour seen in #379
---
 src/sampler.jl  |  2 ++
 test/sampler.jl | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/sampler.jl b/src/sampler.jl
index 0d898719..bb8edcdc 100644
--- a/src/sampler.jl
+++ b/src/sampler.jl
@@ -168,6 +168,8 @@ function sample(
     (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."
+    # Prevent adaptor from being mutated
+    adaptor = deepcopy(adaptor)
     # Prepare containers to store sampling results
     n_keep = n_samples - (drop_warmup ? n_adapts : 0)
     θs, stats = Vector{T}(undef, n_keep), Vector{NamedTuple}(undef, n_keep)
diff --git a/test/sampler.jl b/test/sampler.jl
index 109f01a4..22f88171 100644
--- a/test/sampler.jl
+++ b/test/sampler.jl
@@ -191,4 +191,31 @@ 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 in 1:5
+            samples, stats = sample(
+                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