Add discard_initial keyword argument (#47)

* Add `discard_initial` keyword argument

* Update README

* Update Transducer syntax

* Bump version

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probablistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "2.0.0"
+version = "2.1.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

README.md

@@ -68,9 +68,10 @@ Common keyword arguments for regular and parallel sampling (not supported by the
 are:
 - `progress` (default: `true`):  toggles progress logging
 - `chain_type` (default: `Any`): determines the type of the returned chain
-- `callback` (default: `nothing`): if `callback !== nohting`, then
+- `callback` (default: `nothing`): if `callback !== nothing`, then
   `callback(rng, model, sampler, sample, iteration)` is called after every sampling step,
   where `sample` is the most recent sample of the Markov chain and `iteration` is the current iteration
+- `discard_initial` (default: `0`): number of initial samples that are discarded
 
 Additionally, AbstractMCMC defines the abstract type `AbstractChains` for Markov chains and the
 method `AbstractMCMC.chainscat(::AbstractChains...)` for concatenating multiple chains.

src/sample.jl

@@ -64,16 +64,27 @@ function mcmcsample(
     progress = true,
     progressname = "Sampling",
     callback = nothing,
+    discard_initial = 0,
     chain_type::Type=Any,
     kwargs...
 )
     # Check the number of requested samples.
     N > 0 || error("the number of samples must be ≥ 1")
+    Ntotal = N + discard_initial
 
     @ifwithprogresslogger progress name=progressname begin
         # Obtain the initial sample and state.
         sample, state = step(rng, model, sampler; kwargs...)
 
+        # Discard initial samples.
+        for i in 1:(discard_initial - 1)
+            # Update the progress bar.
+            progress && ProgressLogging.@logprogress i/Ntotal
+
+            # Obtain the next sample and state.
+            sample, state = step(rng, model, sampler, state; kwargs...)
+        end
+
         # Run callback.
         callback === nothing || callback(rng, model, sampler, sample, 1)
 
@@ -82,7 +93,7 @@ function mcmcsample(
         samples = save!!(samples, sample, 1, model, sampler, N; kwargs...)
 
         # Update the progress bar.
-        progress && ProgressLogging.@logprogress 1/N
+        progress && ProgressLogging.@logprogress (1 + discard_initial) / Ntotal
 
         # Step through the sampler.
         for i in 2:N
@@ -96,7 +107,7 @@ function mcmcsample(
             samples = save!!(samples, sample, i, model, sampler, N; kwargs...)
 
             # Update the progress bar.
-            progress && ProgressLogging.@logprogress i/N
+            progress && ProgressLogging.@logprogress (i + discard_initial) / Ntotal
         end
     end
 
@@ -129,12 +140,19 @@ function mcmcsample(
     progress = true,
     progressname = "Convergence sampling",
     callback = nothing,
+    discard_initial = 0,
     kwargs...
 )
     @ifwithprogresslogger progress name=progressname begin
         # Obtain the initial sample and state.
         sample, state = step(rng, model, sampler; kwargs...)
 
+        # Discard initial samples.
+        for _ in 2:discard_initial
+            # Obtain the next sample and state.
+            sample, state = step(rng, model, sampler, state; kwargs...)
+        end
+
         # Run callback.
         callback === nothing || callback(rng, model, sampler, sample, 1)
 

test/sample.jl

@@ -220,10 +220,23 @@
         @test chain2 isa MyChain
     end
 
+    @testset "Discard initial samples" begin
+        chain = sample(MyModel(), MySampler(), 100; sleepy = true, discard_initial = 50)
+        @test length(chain) == 100
+        @test !ismissing(chain[1].a)
+    end
+
     @testset "Sample without predetermined N" begin
         Random.seed!(1234)
         chain = sample(MyModel(), MySampler())
         bmean = mean(x.b for x in chain)
+        @test ismissing(chain[1].a)
+        @test abs(bmean) <= 0.001 && length(chain) < 10_000
+
+        # Discard initial samples.
+        chain = sample(MyModel(), MySampler(); discard_initial = 50)
+        bmean = mean(x.b for x in chain)
+        @test !ismissing(chain[1].a)
         @test abs(bmean) <= 0.001 && length(chain) < 10_000
     end
 

test/transducer.jl

@@ -7,7 +7,7 @@
         Logging.with_logger(TerminalLogger()) do
             xf = AbstractMCMC.Sample(MyModel(), MySampler();
                                      sleepy = true, logger = true)
-            chain = collect(xf, withprogress(1:N; interval=1e-3))
+            chain = withprogress(1:N; interval=1e-3) |> xf |> collect
         end
 
         # test output type and size
@@ -24,16 +24,19 @@
 
     @testset "drop" begin
         xf = AbstractMCMC.Sample(MyModel(), MySampler())
-        chain = collect(xf |> Drop(1), 1:10)
+        chain = 1:10 |> xf |> Drop(1) |> collect
         @test chain isa Vector{MySample{Float64,Float64}}
         @test length(chain) == 9
     end
 
     # Reproduce iterator example
     @testset "iterator example" begin
         # filter missing values and split transitions
-        xf = AbstractMCMC.Sample(MyModel(), MySampler()) |>
-            OfType(MySample{Float64,Float64}) |> Map(x -> (x.a, x.b))
+        xf = opcompose(
+            AbstractMCMC.Sample(MyModel(), MySampler()),
+            OfType(MySample{Float64,Float64}),
+            Map(x -> (x.a, x.b)),
+        )
         as, bs = foldl(xf, 1:999; init = (Float64[], Float64[])) do (as, bs), (a, b)
             push!(as, a)
             push!(bs, b)