Merge 471d9bb into caeade2

TuringLang · Oct 1, 2023 · d969040 · d969040
2 parents caeade2 + 471d9bb
commit d969040
Show file tree

Hide file tree

Showing 5 changed files with 190 additions and 82 deletions.
diff --git a/Project.toml b/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.2"
+version = "4.6.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
@@ -30,9 +30,10 @@ Transducers = "0.4.30"
 julia = "1.6"
 
 [extras]
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["IJulia", "Statistics", "Test"]
+test = ["FillArrays", "IJulia", "Statistics", "Test"]
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -75,14 +75,16 @@ Common keyword arguments for regular and parallel sampling are:
   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
 - `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.
 
 !!! info
     The common keyword arguments `progress`, `chain_type`, and `callback` are not supported by the iterator [`AbstractMCMC.steps`](@ref) and the transducer [`AbstractMCMC.Sample`](@ref).
 
 There is no "official" way for providing initial parameter values yet.
-However, multiple packages such as [EllipticalSliceSampling.jl](https://github.com/TuringLang/EllipticalSliceSampling.jl) and [AdvancedMH.jl](https://github.com/TuringLang/AdvancedMH.jl) support an `init_params` keyword argument for setting the initial values when sampling a single chain.
-To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, [we decided to support `init_params` in the default implementations of the ensemble methods](https://github.com/TuringLang/AbstractMCMC.jl/pull/94):
-- `init_params` (default: `nothing`): if set to `init_params !== nothing`, then the `i`th element of `init_params` is used as initial parameters of the `i`th chain. If one wants to use the same initial parameters `x` for every chain, one can specify e.g. `init_params = Iterators.repeated(x)` or `init_params = FillArrays.Fill(x, N)`.
+However, multiple packages such as [EllipticalSliceSampling.jl](https://github.com/TuringLang/EllipticalSliceSampling.jl) and [AdvancedMH.jl](https://github.com/TuringLang/AdvancedMH.jl) support an `initial_params` keyword argument for setting the initial values when sampling a single chain.
+To ensure that sampling multiple chains "just works" when sampling of a single chain is implemented, [we decided to support `initial_params` in the default implementations of the ensemble methods](https://github.com/TuringLang/AbstractMCMC.jl/pull/94):
+- `initial_params` (default: `nothing`): if `initial_params isa AbstractArray`, then the `i`th element of `initial_params` is used as initial parameters of the `i`th chain. If one wants to use the same initial parameters `x` for every chain, one can specify e.g. `initial_params = FillArrays.Fill(x, N)`.
 
 Progress logging can be enabled and disabled globally with `AbstractMCMC.setprogress!(progress)`.
 

diff --git a/src/sample.jl b/src/sample.jl
@@ -103,6 +103,7 @@ function mcmcsample(
     discard_initial=0,
     thinning=1,
     chain_type::Type=Any,
+    initial_state=nothing,
     kwargs...,
 )
     # Check the number of requested samples.
@@ -122,7 +123,11 @@ function mcmcsample(
         end
 
         # Obtain the initial sample and state.
-        sample, state = step(rng, model, sampler; kwargs...)
+        sample, state = if initial_state === nothing
+            step(rng, model, sampler; kwargs...)
+        else
+            step(rng, model, sampler, state; kwargs...)
+        end
 
         # Discard initial samples.
         for i in 1:discard_initial
@@ -211,6 +216,7 @@ function mcmcsample(
     callback=nothing,
     discard_initial=0,
     thinning=1,
+    initial_state=nothing,
     kwargs...,
 )
 
@@ -220,7 +226,11 @@ function mcmcsample(
 
     @ifwithprogresslogger progress name = progressname begin
         # Obtain the initial sample and state.
-        sample, state = step(rng, model, sampler; kwargs...)
+        sample, state = if initial_state === nothing
+            step(rng, model, sampler; kwargs...)
+        else
+            step(rng, model, sampler, state; kwargs...)
+        end
 
         # Discard initial samples.
         for _ in 1:discard_initial
@@ -288,7 +298,8 @@ function mcmcsample(
     nchains::Integer;
     progress=PROGRESS[],
     progressname="Sampling ($(min(nchains, Threads.nthreads())) threads)",
-    init_params=nothing,
+    initial_params=nothing,
+    initial_state=nothing,
     kwargs...,
 )
     # Check if actually multiple threads are used.
@@ -312,8 +323,9 @@ function mcmcsample(
     # Create a seed for each chain using the provided random number generator.
     seeds = rand(rng, UInt, nchains)
 
-    # Ensure that initial parameters are `nothing` or indexable
-    _init_params = _first_or_nothing(init_params, nchains)
+    # Ensure that initial parameters and states are `nothing` or of the correct length
+    check_initial_params(initial_params, nchains)
+    check_initial_state(initial_state, nchains)
 
     # Set up a chains vector.
     chains = Vector{Any}(undef, nchains)
@@ -364,10 +376,15 @@ function mcmcsample(
                                 _sampler,
                                 N;
                                 progress=false,
-                                init_params=if _init_params === nothing
+                                initial_params=if initial_params === nothing
+                                    nothing
+                                else
+                                    initial_params[chainidx]
+                                end,
+                                initial_state=if initial_state === nothing
                                     nothing
                                 else
-                                    _init_params[chainidx]
+                                    initial_state[chainidx]
                                 end,
                                 kwargs...,
                             )
@@ -397,7 +414,8 @@ function mcmcsample(
     nchains::Integer;
     progress=PROGRESS[],
     progressname="Sampling ($(Distributed.nworkers()) processes)",
-    init_params=nothing,
+    initial_params=nothing,
+    initial_state=nothing,
     kwargs...,
 )
     # Check if actually multiple processes are used.
@@ -410,6 +428,10 @@ function mcmcsample(
         @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
     end
 
+    # Ensure that initial parameters and states are `nothing` or of the correct length
+    check_initial_params(initial_params, nchains)
+    check_initial_state(initial_state, nchains)
+
     # Create a seed for each chain using the provided random number generator.
     seeds = rand(rng, UInt, nchains)
 
@@ -445,7 +467,7 @@ function mcmcsample(
 
             Distributed.@async begin
                 try
-                    function sample_chain(seed, init_params=nothing)
+                    function sample_chain(seed, initial_params, initial_state)
                         # Seed a new random number generator with the pre-made seed.
                         Random.seed!(rng, seed)
 
@@ -456,7 +478,8 @@ function mcmcsample(
                             sampler,
                             N;
                             progress=false,
-                            init_params=init_params,
+                            initial_params=initial_params,
+                            initial_state=initial_state,
                             kwargs...,
                         )
 
@@ -466,11 +489,9 @@ function mcmcsample(
                         # Return the new chain.
                         return chain
                     end
-                    chains = if init_params === nothing
-                        Distributed.pmap(sample_chain, pool, seeds)
-                    else
-                        Distributed.pmap(sample_chain, pool, seeds, init_params)
-                    end
+                    chains = Distributed.pmap(
+                        sample_chain, pool, seeds, initial_params, initial_state
+                    )
                 finally
                     # Stop updating the progress bar.
                     progress && put!(channel, false)
@@ -491,19 +512,24 @@ function mcmcsample(
     N::Integer,
     nchains::Integer;
     progressname="Sampling",
-    init_params=nothing,
+    initial_params=nothing,
+    initial_state=nothing,
     kwargs...,
 )
     # Check if the number of chains is larger than the number of samples
     if nchains > N
         @warn "Number of chains ($nchains) is greater than number of samples per chain ($N)"
     end
 
+    # Ensure that initial parameters and states are `nothing` or of the correct length
+    check_initial_params(initial_params, nchains)
+    check_initial_state(initial_state, nchains)
+
     # Create a seed for each chain using the provided random number generator.
     seeds = rand(rng, UInt, nchains)
 
     # Sample the chains.
-    function sample_chain(i, seed, init_params=nothing)
+    function sample_chain(i, seed, initial_params=nothing)
         # Seed a new random number generator with the pre-made seed.
         Random.seed!(rng, seed)
 
@@ -514,15 +540,15 @@ function mcmcsample(
             sampler,
             N;
             progressname=string(progressname, " (Chain ", i, " of ", nchains, ")"),
-            init_params=init_params,
+            initial_params=initial_params,
             kwargs...,
         )
     end
 
-    chains = if init_params === nothing
+    chains = if initial_params === nothing
         map(sample_chain, 1:nchains, seeds)
     else
-        map(sample_chain, 1:nchains, seeds, init_params)
+        map(sample_chain, 1:nchains, seeds, initial_params)
     end
 
     # Concatenate the chains together.
@@ -532,31 +558,38 @@ end
 tighten_eltype(x) = x
 tighten_eltype(x::Vector{Any}) = map(identity, x)
 
-"""
-    _first_or_nothing(x, n::Int)
-
-Return the first `n` elements of collection `x`, or `nothing` if `x === nothing`.
+@nospecialize check_initial_params(x, n) = throw(
+    ArgumentError(
+        "initial parameters must be specified as a vector of length equal to the number of chains or `nothing`",
+    ),
+)
+check_initial_params(::Nothing, n) = nothing
+function check_initial_params(x::AbstractArray, n)
+    if length(x) != n
+        throw(
+            ArgumentError(
+                "incorrect number of initial parameters (expected $n, received $(length(x))"
+            ),
+        )
+    end
 
-If `x !== nothing`, then `x` has to contain at least `n` elements.
-"""
-function _first_or_nothing(x, n::Int)
-    y = _first(x, n)
-    length(y) == n || throw(
-        ArgumentError("not enough initial parameters (expected $n, received $(length(y))"),
-    )
-    return y
+    return nothing
 end
-_first_or_nothing(::Nothing, ::Int) = nothing
 
-# `first(x, n::Int)` requires Julia 1.6
-function _first(x, n::Int)
-    @static if VERSION >= v"1.6.0-DEV.431"
-        first(x, n)
-    else
-        if x isa AbstractVector
-            @inbounds x[firstindex(x):min(firstindex(x) + n - 1, lastindex(x))]
-        else
-            collect(Iterators.take(x, n))
-        end
+@nospecialize check_initial_state(x, n) = throw(
+    ArgumentError(
+        "initial states must be specified as a vector of length equal to the number of chains or `nothing`",
+    ),
+)
+check_initial_state(::Nothing, n) = nothing
+function check_initial_state(x::AbstractArray, n)
+    if length(x) != n
+        throw(
+            ArgumentError(
+                "incorrect number of initial states (expected $n, received $(length(x))"
+            ),
+        )
     end
+
+    return nothing
 end