Merge pull request #13 from TuringLang/csp/iterator

Add iterator interface

Author: cpfiffer

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 = "0.3.0"
+version = "0.4.0"
 
 [deps]
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"

src/AbstractMCMC.jl

@@ -406,4 +406,56 @@ function psample(
     return reduce(chainscat, chains)
 end
 
+
+##################
+# Iterator tools #
+##################
+struct Stepper{A<:AbstractRNG, ModelType<:AbstractModel, SamplerType<:AbstractSampler, K}
+    rng::A
+    model::ModelType
+    s::SamplerType
+    kwargs::K
+end
+
+function Base.iterate(stp::Stepper, state=nothing)
+    t = step!(stp.rng, stp.model, stp.s, 1, state; stp.kwargs...)
+    return t, t
+end
+
+Base.IteratorSize(::Type{<:Stepper}) = Base.IsInfinite()
+Base.IteratorEltype(::Type{<:Stepper}) = Base.EltypeUnknown()
+
+"""
+    steps!([rng::AbstractRNG, ]model::AbstractModel, s::AbstractSampler, kwargs...)
+
+`steps!` returns an iterator that returns samples continuously, after calling `sample_init!`.
+
+Usage:
+
+```julia
+for transition in steps!(MyModel(), MySampler())
+    println(transition)
+
+    # Do other stuff with transition below.
+end
+```
+"""
+function steps!(
+    model::AbstractModel,
+    s::AbstractSampler,
+    kwargs...
+)
+    return steps!(GLOBAL_RNG, model, s; kwargs...)
+end
+
+function steps!(
+    rng::AbstractRNG,
+    model::AbstractModel,
+    s::AbstractSampler,
+    kwargs...
+)
+    sample_init!(rng, model, s, 0)
+    return Stepper(rng, model, s, kwargs)
+end
+
 end # module AbstractMCMC

test/runtests.jl

@@ -1,5 +1,5 @@
 using AbstractMCMC
-using AbstractMCMC: sample, psample
+using AbstractMCMC: sample, psample, steps!
 
 using Random
 using Statistics
@@ -57,4 +57,30 @@ include("interface.jl")
         @test chain1 isa Vector{MyTransition}
         @test chain2 isa MyChain
     end
+
+    @testset "Iterator sampling" begin
+        Random.seed!(1234)
+        as = []
+        bs = []
+
+        iter = steps!(MyModel(), MySampler())
+
+        for (count, t) in enumerate(iter)
+            if count >= 1000
+                break
+            end
+
+            push!(as, t.a)
+            push!(bs, t.b)
+        end
+
+        @test mean(as) ≈ 0.5 atol=1e-2
+        @test var(as) ≈ 1 / 12 atol=5e-3
+        @test mean(bs) ≈ 0.0 atol=5e-2
+        @test var(bs) ≈ 1 atol=5e-2
+
+        println(eltype(iter))
+        @test Base.IteratorSize(iter) == Base.IsInfinite()
+        @test Base.IteratorEltype(iter) == Base.EltypeUnknown()
+    end
 end