Reporter (#18)

* Add reporter to DynamicHMCChain

* Report mcmc steps

* Test reporting

* Increment version number

Project.toml

@@ -1,7 +1,7 @@
 name = "SampleChainsDynamicHMC"
 uuid = "6d9fd711-e8b2-4778-9c70-c1dfb499d4c4"
 authors = ["Chad Scherrer <[email protected]> and contributors"]
-version = "0.3.3"
+version = "0.3.4"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
@@ -33,8 +33,9 @@ julia = "1.5"
 
 [extras]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["LinearAlgebra", "ReverseDiff", "Test"]
+test = ["LinearAlgebra", "Logging", "ReverseDiff", "Test"]

src/SampleChainsDynamicHMC.jl

@@ -24,9 +24,10 @@ export dynamichmc
     meta        # Metadata associated with the sample as a whole
     state       
     transform
+    reporter
 end
 
-function DynamicHMCChain(t::TransformVariables.TransformTuple, Q::DynamicHMC.EvaluatedLogDensity, tree_stats, steps)
+function DynamicHMCChain(t::TransformVariables.TransformTuple, Q::DynamicHMC.EvaluatedLogDensity, tree_stats, steps, reporter = NoProgressReport())
     tQq = TransformVariables.transform(t, Q.q)
     T = typeof(tQq)
     samples = chainvec(tQq)
@@ -35,7 +36,7 @@ function DynamicHMCChain(t::TransformVariables.TransformTuple, Q::DynamicHMC.Eva
     meta = steps
     transform = t
 
-    return DynamicHMCChain{T}(samples, logq, info, meta, Q, transform)
+    return DynamicHMCChain{T}(samples, logq, info, meta, Q, transform, reporter)
 end
 
 TupleVectors.summarize(ch::DynamicHMCChain) = summarize(samples(ch))
@@ -127,13 +128,13 @@ end
 function SampleChains.newchain(rng::Random.AbstractRNG, config::DynamicHMCConfig, ℓ, tr)
     P = LogDensityProblems.TransformedLogDensity(tr, ℓ)
     ∇P = LogDensityProblems.ADgradient(config.ad_backend, P)
-    reporter = DynamicHMC.NoProgressReport()
+    reporter = config.reporter
 
     results = DynamicHMC.mcmc_keep_warmup(rng, ∇P, 0; 
           initialization = config.init          
         , warmup_stages  = config.warmup_stages 
         , algorithm      = config.algorithm     
-        , reporter       = config.reporter     
+        , reporter       = reporter     
         )
 
     steps = DynamicHMC.mcmc_steps(
@@ -143,17 +144,20 @@ function SampleChains.newchain(rng::Random.AbstractRNG, config::DynamicHMCConfig
 
     Q = results.final_warmup_state.Q
     Q, tree_stats = DynamicHMC.mcmc_next_step(steps, Q)
-    chain = DynamicHMCChain(tr, Q, tree_stats, steps)
+    chain = DynamicHMCChain(tr, Q, tree_stats, steps, reporter)
 end
 
 function SampleChains.newchain(config::DynamicHMCConfig, ℓ, tr)
     return newchain(Random.GLOBAL_RNG, config, ℓ, tr)
 end
 
 function SampleChains.sample!(chain::DynamicHMCChain, n::Int=1000)
+    reporter = getfield(chain, :reporter)
+    mcmc_reporter = DynamicHMC.make_mcmc_reporter(reporter, n; currently_warmup = false)
     @cleanbreak for j in 1:n
         Q, tree_stats = step!(chain)
         pushsample!(chain, Q, tree_stats)
+        DynamicHMC.report(mcmc_reporter, j)
     end 
     return chain
 end

test/notebook.jl

@@ -2,6 +2,7 @@
 
 using SampleChainsDynamicHMC
 using TransformVariables
+using Logging
 using Test
 
 function ℓ(nt)
@@ -65,3 +66,19 @@ using SampleChainsDynamicHMC.LogDensityProblems
 	sample!(chain, 90)
 	@test length(chain) == 100
 end
+
+@testset "reporting" begin
+	io = IOBuffer()
+	chains = with_logger(SimpleLogger(io)) do
+		newchain(dynamichmc(reporter=LogProgressReport()), ℓ, t)
+	end
+	warmup_log = String(take!(io))
+	@test !isempty(warmup_log)
+
+	io = IOBuffer()
+	with_logger(SimpleLogger(io)) do
+		sample!(chains, 10)
+	end
+	log = String(take!(io))
+	@test !isempty(log)
+end