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/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`, 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,22 @@ 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 [`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 a "warmup"-stage initial stage that is different from the
+standard iteration.
+
+By default, this simply calls [`step`](@ref.)
+"""
+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

@@ -91,7 +91,29 @@ function StatsBase.sample(
 end
 
 # Default implementations of regular and parallel sampling.
-
+"""
+    mcmcsample(rng, model, sampler, N_or_is_done; kwargs...)
+
+Default implementation of `sample` for a `model` and `sampler`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: the random number generator to use.
+- `model::AbstractModel`: the model to sample from.
+- `sampler::AbstractSampler`: the sampler to use.
+- `N::Integer`: the number of samples to draw.
+
+# Keyword arguments
+- `progress`: whether to display a progress bar. Defaults to `true`.
+- `progressname`: the name of the progress bar. Defaults to `"Sampling"`.
+- `callback`: a function that is called after each [`AbstractMCMC.step`](@ref).
+   Defaults to `nothing`.
+- `num_warmup`: number of warmup samples to draw. Defaults to `0`.
+- `discard_initial`: number of initial samples to discard. Defaults to `num_warmup`.
+- `thinning`: number of samples to discard between samples. Defaults to `1`.
+- `chain_type`: the type to pass to [`AbstractMCMC.bundle_samples`](@ref) at the
+  end of sampling to wrap up the resulting samples nicely. Defaults to `Any`.
+- `kwargs...`: Additional keyword arguments to pass on to [`AbstractMCMC.step`](@ref).
+"""
 function mcmcsample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
@@ -100,14 +122,21 @@ function mcmcsample(
     progress=PROGRESS[],
     progressname="Sampling",
     callback=nothing,
-    discard_initial=0,
+    num_warmup=0,
+    discard_initial=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")
-    Ntotal = thinning * (N - 1) + discard_initial + 1
+    Ntotal = thinning * (N - 1) + discard_initial + num_warmup + 1
+
+    # 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
+    discard_from_sample = max(discard_initial - discard_from_warmup, 0)
 
     # Start the timer
     start = time()
@@ -124,34 +153,76 @@ function mcmcsample(
         # Obtain the initial sample and state.
         sample, state = step(rng, model, sampler; kwargs...)
 
-        # 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
-            end
-
+        # Warmup sampling.
+        for _ in 1:discard_from_warmup
             # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            sample, state = step_warmup(rng, model, sampler, state; kwargs...)
         end
 
-        # Run callback.
-        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)
+        i = 1
+        if keep_from_warmup > 0
+            # Run callback.
+            callback === nothing ||
+                callback(rng, model, sampler, sample, state, i; kwargs...)
+
+            # Save the sample.
+            samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
+            samples = save!!(samples, sample, i, model, sampler; kwargs...)
+
+            # Step through remainder of warmup iterations and save.
+            i += 1
+            for _ in (discard_from_warmup + 1):num_warmup
+                # Update the progress bar.
+                if progress && i >= next_update
+                    ProgressLogging.@logprogress i / Ntotal
+                    next_update = i + threshold
+                end
+
+                # Obtain the next sample and state.
+                sample, state = step_warmup(rng, model, sampler, state; kwargs...)
 
-        # Save the sample.
-        samples = AbstractMCMC.samples(sample, model, sampler, N; kwargs...)
-        samples = save!!(samples, sample, 1, model, sampler, N; kwargs...)
+                # Run callback.
+                callback === nothing ||
+                    callback(rng, model, sampler, sample, state, i; kwargs...)
+
+                # Save the sample.
+                samples = save!!(samples, sample, i, model, sampler; kwargs...)
+                i += 1
+            end
+        else
+            # Discard additional initial samples, if needed.
+            for _ in 1:discard_from_sample
+                # Update the progress bar.
+                if progress && i >= next_update
+                    ProgressLogging.@logprogress i / Ntotal
+                    next_update = i + threshold
+                end
+
+                # Obtain the next sample and state.
+                sample, state = step(rng, model, sampler, state; kwargs...)
+            end
+
+            # Run callback.
+            callback === nothing ||
+                callback(rng, model, sampler, sample, state, i; kwargs...)
+
+            # Save the sample.
+            samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
+            samples = save!!(samples, sample, i, model, sampler; kwargs...)
+
+            # Increment iteration number.
+            i += 1
+        end
 
         # Update the progress bar.
-        itotal = 1 + discard_initial
+        itotal = i
         if progress && itotal >= next_update
             ProgressLogging.@logprogress itotal / Ntotal
             next_update = itotal + threshold
         end
 
         # 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.
@@ -174,6 +245,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,10 +283,16 @@ 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
+    discard_from_sample = max(discard_initial - discard_from_warmup, 0)
 
     # Start the timer
     start = time()
@@ -222,21 +302,54 @@ function mcmcsample(
         # Obtain the initial sample and state.
         sample, state = step(rng, model, sampler; kwargs...)
 
-        # Discard initial samples.
-        for _ in 1:discard_initial
+        # Warmup sampling.
+        for _ in 1:discard_from_warmup
             # Obtain the next sample and state.
-            sample, state = step(rng, model, sampler, state; kwargs...)
+            sample, state = step_warmup(rng, model, sampler, state; kwargs...)
         end
 
-        # Run callback.
-        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)
+        i = 1
+        if keep_from_warmup > 0
+            # Run callback.
+            callback === nothing ||
+                callback(rng, model, sampler, sample, state, i; kwargs...)
+
+            # Save the sample.
+            samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
+            samples = save!!(samples, sample, i, model, sampler; kwargs...)
 
-        # Save the sample.
-        samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
-        samples = save!!(samples, sample, 1, model, sampler; kwargs...)
+            # Step through remainder of warmup iterations and save.
+            i += 1
+            for _ in (discard_from_warmup + 1):num_warmup
+                # Obtain the next sample and state.
+                sample, state = step_warmup(rng, model, sampler, state; kwargs...)
 
-        # Step through the sampler until stopping.
-        i = 2
+                # Run callback.
+                callback === nothing ||
+                    callback(rng, model, sampler, sample, state, i; kwargs...)
+
+                # Save the sample.
+                samples = save!!(samples, sample, i, model, sampler; kwargs...)
+                i += 1
+            end
+        else
+            # Discard additional initial samples, if needed.
+            for _ in 1:discard_from_sample
+                # Obtain the next sample and state.
+                sample, state = step(rng, model, sampler, state; kwargs...)
+            end
+
+            # Run callback.
+            callback === nothing ||
+                callback(rng, model, sampler, sample, state, i; kwargs...)
+
+            # Save the sample.
+            samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
+            samples = save!!(samples, sample, i, model, sampler; kwargs...)
+
+            # Increment iteration number.
+            i += 1
+        end
 
         while !isdone(rng, model, sampler, samples, state, i; progress=progress, kwargs...)
             # Discard thinned samples.