Pass num_warmup as kwarg to step_warmup (#178)

* Pass `num_warmup` kwarg to `step_warmup()`

* Mention the num_warmup kwarg in step_warmup's docstring

Author: penelopeysm

Project.toml

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

src/interface.jl

@@ -83,6 +83,10 @@ When sampling using [`sample`](@ref), this takes the place of [`AbstractMCMC.ste
 This is useful if the sampler has an initial "warmup"-stage that is different from the
 standard iteration.
 
+The total number of warmup steps requested in sampling will be passed to the `step_warmup`
+function as the `num_warmup` keyword argument. This allows implementations of `step_warmup`
+to customise their behavior based on this information.
+
 By default, this simply calls [`AbstractMCMC.step`](@ref).
 """
 step_warmup(rng, model, sampler; kwargs...) = step(rng, model, sampler; kwargs...)

src/sample.jl

@@ -179,9 +179,9 @@ function mcmcsample(
             # Obtain the initial sample and state.
             sample, state = if num_warmup > 0
                 if initial_state === nothing
-                    step_warmup(rng, model, sampler; kwargs...)
+                    step_warmup(rng, model, sampler; num_warmup, kwargs...)
                 else
-                    step_warmup(rng, model, sampler, initial_state; kwargs...)
+                    step_warmup(rng, model, sampler, initial_state; num_warmup, kwargs...)
                 end
             else
                 if initial_state === nothing
@@ -202,7 +202,7 @@ function mcmcsample(
             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...)
+                    step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
                 else
                     step(rng, model, sampler, state; kwargs...)
                 end
@@ -229,7 +229,7 @@ function mcmcsample(
                 for _ in 1:(thinning - 1)
                     # Obtain the next sample and state.
                     sample, state = if i ≤ keep_from_warmup
-                        step_warmup(rng, model, sampler, state; kwargs...)
+                        step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
                     else
                         step(rng, model, sampler, state; kwargs...)
                     end
@@ -244,7 +244,7 @@ function mcmcsample(
 
                 # Obtain the next sample and state.
                 sample, state = if i ≤ keep_from_warmup
-                    step_warmup(rng, model, sampler, state; kwargs...)
+                    step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
                 else
                     step(rng, model, sampler, state; kwargs...)
                 end
@@ -328,9 +328,9 @@ function mcmcsample(
         # Obtain the initial sample and state.
         sample, state = if num_warmup > 0
             if initial_state === nothing
-                step_warmup(rng, model, sampler; kwargs...)
+                step_warmup(rng, model, sampler; num_warmup, kwargs...)
             else
-                step_warmup(rng, model, sampler, initial_state; kwargs...)
+                step_warmup(rng, model, sampler, initial_state; num_warmup, kwargs...)
             end
         else
             if initial_state === nothing
@@ -344,7 +344,7 @@ function mcmcsample(
         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...)
+                step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
             else
                 step(rng, model, sampler, state; kwargs...)
             end
@@ -364,15 +364,15 @@ function mcmcsample(
             for _ in 1:(thinning - 1)
                 # Obtain the next sample and state.
                 sample, state = if i ≤ keep_from_warmup
-                    step_warmup(rng, model, sampler, state; kwargs...)
+                    step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
                 else
                     step(rng, model, sampler, state; kwargs...)
                 end
             end
 
             # Obtain the next sample and state.
             sample, state = if i ≤ keep_from_warmup
-                step_warmup(rng, model, sampler, state; kwargs...)
+                step_warmup(rng, model, sampler, state; num_warmup, kwargs...)
             else
                 step(rng, model, sampler, state; kwargs...)
             end

test/utils.jl

@@ -26,8 +26,11 @@ function AbstractMCMC.step_warmup(
     state::Union{Nothing,Integer}=nothing;
     loggers=false,
     initial_params=nothing,
+    num_warmup,
     kwargs...,
 )
+    num_warmup isa Integer ||
+        error("num_warmup should have been passed as a keyword argument to step_warmup")
     transition, state = AbstractMCMC.step(
         rng, model, sampler, state; loggers, initial_params, kwargs...
     )