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,14 @@ 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
+```
+
+Note that this is optional; by default it simply calls [`AbstractMCMC.step`](@ref) from above.
+
 ## Collecting samples
 
 !!! note

src/interface.jl

@@ -56,6 +56,19 @@ 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 `discard_initial`
+number of iterations. 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.)
+"""
+step_warmup(rng, model, sampler, state; kwargs...) = step(rng, model, sampler, state; kwargs...)
+
 """
     samples(sample, model, sampler[, N; kwargs...])
 

src/sample.jl

@@ -133,7 +133,7 @@ function mcmcsample(
             end
 
             # 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.