Addition of step_warmup

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 = "4.4.0"
+version = "4.4.1"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

docs/src/api.md

@@ -73,7 +73,11 @@ Common keyword arguments for regular and parallel sampling are:
 - `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
+- `num_warmup` (default: `0`): number of "warm-up" steps to take before the first "regular" step, 
+   i.e. number of times to call [`AbstractMCMC.step_warmup`](@ref) before the first call to 
+   [`AbstractMCMC.step`](@ref).
+- `discard_initial` (default: `num_warmup`): number of initial samples that are discarded. Note that
+  if `discard_initial < num_warmup`, warm-up samples will also be included in the resulting samples.
 - `thinning` (default: `1`): factor by which to thin samples.
 
 !!! info

docs/src/design.md

@@ -63,6 +63,15 @@ the sampling step of the inference method.
 AbstractMCMC.step
 ```
 
+If one also has some special handling of the warmup-stage of sampling, then this can be specified by overloading
+
+```@docs
+AbstractMCMC.step_warmup
+```
+
+which will be used for the first `num_warmup` iterations, as specified as a keyword argument to [`AbstractMCMC.sample`](@ref). 
+Note that this is optional; by default it simply calls [`AbstractMCMC.step`](@ref) from above.
+
 ## Collecting samples
 
 !!! note

src/interface.jl

@@ -56,6 +56,21 @@ current `state` of the sampler.
 """
 function step end
 
+"""
+    step_warmup(rng, model, sampler[, state; kwargs...])
+
+Return a 2-tuple of the next sample and the next state of the MCMC `sampler` for `model`.
+
+When sampling using [`sample`](@ref), this takes the place of [`AbstractMCMC.step`](@ref) in the first
+`num_warmup` number of iterations, as specified by the `num_warmup` keyword to [`sample`](@ref).
+This is useful if the sampler has an initial "warmup"-stage that is different from the
+standard iteration.
+
+By default, this simply calls [`AbstractMCMC.step`](@ref).
+"""
+step_warmup(rng, model, sampler; kwargs...) = step(rng, model, sampler; kwargs...)
+step_warmup(rng, model, sampler, state; kwargs...) = step(rng, model, sampler, state; kwargs...)
+
 """
     samples(sample, model, sampler[, N; kwargs...])
 

src/sample.jl

@@ -40,6 +40,11 @@ isdone(rng, model, sampler, samples, state, iteration; kwargs...)
 ```
 where `state` and `iteration` are the current state and iteration of the sampler, respectively.
 It should return `true` when sampling should end, and `false` otherwise.
+
+# Keyword arguments
+
+See https://turinglang.org/AbstractMCMC.jl/dev/api/#Common-keyword-arguments for common keyword
+arguments.
 """
 function StatsBase.sample(
     rng::Random.AbstractRNG,
@@ -77,6 +82,11 @@ end
 
 Sample `nchains` Monte Carlo Markov chains from the `model` with the `sampler` in parallel
 using the `parallel` algorithm, and combine them into a single chain.
+
+# Keyword arguments
+
+See https://turinglang.org/AbstractMCMC.jl/dev/api/#Common-keyword-arguments for common keyword
+arguments.
 """
 function StatsBase.sample(
     rng::Random.AbstractRNG,
@@ -91,7 +101,6 @@ function StatsBase.sample(
 end
 
 # Default implementations of regular and parallel sampling.
-
 function mcmcsample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
@@ -100,14 +109,27 @@ function mcmcsample(
     progress=PROGRESS[],
     progressname="Sampling",
     callback=nothing,
-    discard_initial=0,
+    num_warmup::Int=0,
+    discard_initial::Int=num_warmup,
     thinning=1,
     chain_type::Type=Any,
     kwargs...,
 )
     # Check the number of requested samples.
     N > 0 || error("the number of samples must be ≥ 1")
+    discard_initial >= 0 ||
+        throw(ArgumentError("number of discarded samples must be non-negative"))
+    num_warmup >= 0 ||
+        throw(ArgumentError("number of warm-up samples must be non-negative"))
     Ntotal = thinning * (N - 1) + discard_initial + 1
+    Ntotal >= num_warmup || throw(
+        ArgumentError("number of warm-up samples exceeds the total number of samples")
+    )
+
+    # Determine how many samples to drop from `num_warmup` and the
+    # main sampling process before we start saving samples.
+    discard_from_warmup = min(num_warmup, discard_initial)
+    keep_from_warmup = num_warmup - discard_from_warmup
 
     # Start the timer
     start = time()
@@ -122,40 +144,58 @@ function mcmcsample(
         end
 
         # Obtain the initial sample and state.
-        sample, state = step(rng, model, sampler; kwargs...)
+        sample, state = if num_warmup > 0
+            step_warmup(rng, model, sampler; kwargs...)
+        else
+            step(rng, model, sampler; kwargs...)
+        end
+
+        # Update the progress bar.
+        itotal = 1
+        if progress && itotal >= next_update
+            ProgressLogging.@logprogress itotal / Ntotal
+            next_update = itotal + threshold
+        end
 
         # Discard initial samples.
-        for i in 1:discard_initial
-            # Update the progress bar.
-            if progress && i >= next_update
-                ProgressLogging.@logprogress i / Ntotal
-                next_update = i + threshold
+        for j in 1:discard_initial
+            # Obtain the next sample and state.
+            sample, state = if j ≤ num_warmup
+                step_warmup(rng, model, sampler, state; kwargs...)
+            else
+                step(rng, model, sampler, state; kwargs...)
             end
 
-            # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            # Update the progress bar.
+            if progress && (itotal += 1) >= next_update
+                ProgressLogging.@logprogress itotal / Ntotal
+                next_update = itotal + threshold
+            end
         end
 
+        # Initialize iteration counter.
+        i = 1
+
         # Run callback.
-        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)
+        callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)
 
         # Save the sample.
         samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)
-        samples = save!!(samples, sample, 1, model, sampler, N; kwargs...)
+        samples = save!!(samples, sample, i, model, sampler, N; kwargs...)
 
-        # Update the progress bar.
-        itotal = 1 + discard_initial
-        if progress && itotal >= next_update
-            ProgressLogging.@logprogress itotal / Ntotal
-            next_update = itotal + threshold
-        end
+        # Step through remainder of warmup iterations and save.
+        i += 1
 
         # Step through the sampler.
-        for i in 2:N
+        while i ≤ N
             # Discard thinned samples.
             for _ in 1:(thinning - 1)
                 # Obtain the next sample and state.
-                sample, state = step(rng, model, sampler, state; kwargs...)
+                sample, state = if i ≤ keep_from_warmup
+                    step_warmup(rng, model, sampler, state; kwargs...)
+                else
+                    step(rng, model, sampler, state; kwargs...)
+                end
 
                 # Update progress bar.
                 if progress && (itotal += 1) >= next_update
@@ -165,7 +205,11 @@ function mcmcsample(
             end
 
             # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            sample, state = if i ≤ keep_from_warmup
+                step_warmup(rng, model, sampler, state; kwargs...)
+            else
+                step(rng, model, sampler, state; kwargs...)
+            end
 
             # Run callback.
             callback === nothing ||
@@ -174,6 +218,9 @@ function mcmcsample(
             # Save the sample.
             samples = save!!(samples, sample, i, model, sampler, N; kwargs...)
 
+            # Increment iteration counter.
+            i += 1
+
             # Update the progress bar.
             if progress && (itotal += 1) >= next_update
                 ProgressLogging.@logprogress itotal / Ntotal
@@ -209,44 +256,67 @@ function mcmcsample(
     progress=PROGRESS[],
     progressname="Convergence sampling",
     callback=nothing,
-    discard_initial=0,
+    num_warmup=0,
+    discard_initial=num_warmup,
     thinning=1,
     kwargs...,
 )
+    # Determine how many samples to drop from `num_warmup` and the
+    # main sampling process before we start saving samples.
+    discard_from_warmup = min(num_warmup, discard_initial)
+    keep_from_warmup = num_warmup - discard_from_warmup
 
     # Start the timer
     start = time()
     local state
 
     @ifwithprogresslogger progress name = progressname begin
         # Obtain the initial sample and state.
-        sample, state = step(rng, model, sampler; kwargs...)
+        sample, state = if num_warmup > 0
+            step_warmup(rng, model, sampler; kwargs...)
+        else
+            step(rng, model, sampler; kwargs...)
+        end
 
         # Discard initial samples.
-        for _ in 1:discard_initial
+        for j in 1:discard_initial
             # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            sample, state = if j ≤ discard_num_warmup
+                step_warmup(rng, model, sampler, state; kwargs...)
+            else
+                step(rng, model, sampler, state; kwargs...)
+            end
         end
 
+        # Initialize iteration counter.
+        i = 1
+
         # Run callback.
-        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)
+        callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)
 
         # Save the sample.
         samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
-        samples = save!!(samples, sample, 1, model, sampler; kwargs...)
-
-        # Step through the sampler until stopping.
-        i = 2
+        samples = save!!(samples, sample, i, model, sampler; kwargs...)
+
+        i += 1
 
         while !isdone(rng, model, sampler, samples, state, i; progress=progress, kwargs...)
             # Discard thinned samples.
             for _ in 1:(thinning - 1)
                 # Obtain the next sample and state.
-                sample, state = step(rng, model, sampler, state; kwargs...)
+                sample, state = if i ≤ keep_from_warmup
+                    step_warmup(rng, model, sampler, state; kwargs...)
+                else
+                    step(rng, model, sampler, state; kwargs...)
+                end
             end
 
             # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            sample, state = if i ≤ keep_from_warmup
+                step_warmup(rng, model, sampler, state; kwargs...)
+            else
+                step(rng, model, sampler, state; kwargs...)
+            end
 
             # Run callback.
             callback === nothing ||