Addition of step_warmup

docs/src/api.md

@@ -71,9 +71,13 @@ Common keyword arguments for regular and parallel sampling are:
 - `progress` (default: `AbstractMCMC.PROGRESS[]` which is `true` initially):  toggles progress logging
 - `chain_type` (default: `Any`): determines the type of the returned chain
 - `callback` (default: `nothing`): if `callback !== nothing`, then
-  `callback(rng, model, sampler, sample, state, iteration)` is called after every sampling step,
-  where `sample` is the most recent sample of the Markov chain and `state` and `iteration` are the current state and iteration of the sampler
-- `discard_initial` (default: `0`): number of initial samples that are discarded
+  `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
+- `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.
 - `initial_state` (default: `nothing`): if `initial_state !== nothing`, the first call to [`AbstractMCMC.step`](@ref)
   is passed `initial_state` as the `state` argument.

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

@@ -73,6 +73,23 @@
 """
 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...)
+function step_warmup(rng, model, sampler, state; kwargs...)
+    return step(rng, model, sampler, state; kwargs...)
+end
+
 """
     samples(sample, model, sampler[, N; kwargs...])
 

src/sample.jl

@@ -40,6 +40,11 @@
 ```
 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 @@
 
 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 @@
 end
 
 # Default implementations of regular and parallel sampling.
-
 function mcmcsample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
@@ -100,15 +109,28 @@
     progress=PROGRESS[],
     progressname="Sampling",
     callback=nothing,
-    discard_initial=0,
+    num_warmup::Int=0,
+    discard_initial::Int=num_warmup,
     thinning=1,
     chain_type::Type=Any,
     initial_state=nothing,
     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()
@@ -123,22 +145,41 @@
         end
 
         # Obtain the initial sample and state.
-        sample, state = if initial_state === nothing
-            step(rng, model, sampler; kwargs...)
+        sample, state = if num_warmup > 0
+            if initial_state === nothing
+                step_warmup(rng, model, sampler; kwargs...)
+            else
+                step_warmup(rng, model, sampler, initial_state; kwargs...)
+            end
         else
-            step(rng, model, sampler, initial_state; kwargs...)
+            if initial_state === nothing
+                step(rng, model, sampler; kwargs...)
+            else
+                step(rng, model, sampler, initial_state; kwargs...)
+            end
+        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
 
         # Run callback.
@@ -148,19 +189,16 @@
         samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)
         samples = save!!(samples, sample, 1, 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 the sampler.
         for i in 2: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
@@ -170,7 +208,11 @@
             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 ||
@@ -179,6 +221,9 @@
             # 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
@@ -214,28 +259,45 @@
     progress=PROGRESS[],
     progressname="Convergence sampling",
     callback=nothing,
-    discard_initial=0,
+    num_warmup=0,
+    discard_initial=num_warmup,
     thinning=1,
     initial_state=nothing,
     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 = if initial_state === nothing
-            step(rng, model, sampler; kwargs...)
+        sample, state = if num_warmup > 0
+            if initial_state === nothing
+                step_warmup(rng, model, sampler; kwargs...)
+            else
+                step_warmup(rng, model, sampler, initial_state; kwargs...)
+            end
         else
-            step(rng, model, sampler, state; kwargs...)
+            if initial_state === nothing
+                step(rng, model, sampler; kwargs...)
+            else
+                step(rng, model, sampler, initial_state; kwargs...)
+            end
         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_from_warmup
+                step_warmup(rng, model, sampler, state; kwargs...)
+            else
+                step(rng, model, sampler, state; kwargs...)
+            end
         end
 
         # Run callback.
@@ -247,16 +309,23 @@
 
         # Step through the sampler until stopping.
         i = 2
-
         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 ||