@mcmcjs/julia 0.6.1 → 0.7.0

package/dist/fitlib.jl

@@ -96,6 +96,21 @@ end
 
 to_namedtuple(data) = (; (Symbol(k) => v for (k, v) in data)...)
 
+# Binds the data so a model can read a variable as a property (`data.y`) or by
+# index (`data["y"]` / `data[:y]`), with `haskey`/`keys` supported, so models
+# written in either idiom run unchanged. The underlying NamedTuple is reached
+# via getfield to avoid colliding with a data variable also named `nt`.
+struct ModelData
+    nt::NamedTuple
+end
+Base.getproperty(d::ModelData, k::Symbol) = getproperty(getfield(d, :nt), k)
+Base.getindex(d::ModelData, k::Symbol) = getproperty(getfield(d, :nt), k)
+Base.getindex(d::ModelData, k::AbstractString) = getproperty(getfield(d, :nt), Symbol(k))
+Base.haskey(d::ModelData, k::Symbol) = haskey(getfield(d, :nt), k)
+Base.haskey(d::ModelData, k::AbstractString) = haskey(getfield(d, :nt), Symbol(k))
+Base.keys(d::ModelData) = keys(getfield(d, :nt))
+Base.propertynames(d::ModelData) = propertynames(getfield(d, :nt))
+
 # JuliaBUGS requires dense arrays with a concrete numeric eltype, but JSON parses
 # them as Vector{Any} (and nested arrays as vectors of vectors). narrow promotes a
 # numeric array to a concrete eltype and stacks equal-sized nested arrays into a
@@ -140,14 +155,83 @@ function chains_from_wire(wire, targets)
     return MCMCChains.Chains(arr, syms, Dict(:parameters => syms); info = (; varname_to_symbol = vmap))
 end
 
+# Expand one sampler transition to canonical scalar leaves (theta -> theta[1], ...),
+# the same names the final samples file uses, so streamed batches reconstruct it.
+function flatten_draw(transition)
+    out = Pair{String,Float64}[]
+    for (vn, val) in pairs(transition.params)
+        for (leaf, leafval) in Turing.DynamicPPL.varname_and_value_leaves(vn, val)
+            push!(out, string(leaf) => Float64(leafval))
+        end
+    end
+    return out
+end
+
+# A per-draw callback that emits {"mcmcjs":"draws",...} batches on PROGRESS_IO as
+# sampling proceeds. AbstractMCMC passes the 1-based chain index as `chain_number`;
+# each batch carries a per-chain monotonic seq, and a chain's batches concatenate
+# (by leaf name) to the final samples file. Assumes the callback is not invoked
+# concurrently across chains (true for the MCMCSerial path build_and_sample uses).
+function draw_streamer(draws_per_chain::Int, batch_size::Int)
+    chain = Ref(-1)
+    seq = Ref(0)
+    names = String[]
+    buffer = Dict{String,Vector{Float64}}()
+    filled = Ref(0)
+    function emit(iteration)
+        filled[] == 0 && return
+        cols = Dict{String,Vector{Float64}}(n => copy(buffer[n]) for n in names)
+        println(
+            PROGRESS_IO[],
+            JSON.json(
+                Dict(
+                    "mcmcjs" => "draws",
+                    "chain" => chain[],
+                    "seq" => seq[],
+                    "iteration" => iteration,
+                    "draws" => cols,
+                ),
+            ),
+        )
+        flush(PROGRESS_IO[])
+        seq[] += 1
+        for n in names
+            empty!(buffer[n])
+        end
+        filled[] = 0
+    end
+    return function (_rng, _model, _sampler, transition, _state, iteration; chain_number = 1, _kwargs...)
+        c = Int(chain_number) - 1
+        if c != chain[]
+            chain[] = c
+            seq[] = 0
+        end
+        leaves = flatten_draw(transition)
+        if isempty(names)
+            for (n, _) in leaves
+                push!(names, n)
+                buffer[n] = Float64[]
+            end
+        end
+        for (n, v) in leaves
+            push!(get!(buffer, n, Float64[]), v)
+        end
+        filled[] += 1
+        if filled[] >= batch_size || iteration >= draws_per_chain
+            emit(iteration)
+        end
+    end
+end
+
 # Build the model and sample in one call so both run in the latest world age:
 # the model methods are defined by `include` at runtime, so a plain call from
 # this (older) world would fail with "method too new". invokelatest fixes that.
-# Returns Turing's default chain type, a FlexiChain.
-function build_and_sample(entry, data, sampler, draws, chains, rng)
+# Returns Turing's default chain type, a FlexiChain. With a callback, draws stream.
+function build_and_sample(entry, data, sampler, draws, chains, rng; callback = nothing)
     model = entry(data)
+    kw = callback === nothing ? NamedTuple() : (; callback)
     return Logging.with_logger(JSONProgressLogger()) do
-        sample(rng, model, sampler, MCMCSerial(), draws, chains; progress = true)
+        sample(rng, model, sampler, MCMCSerial(), draws, chains; progress = true, kw...)
     end
 end
 
@@ -251,7 +335,7 @@ function handle_request(request)
 
         wire = if mode == "predict"
             backend == "turing" || error("predict is not yet supported for backend $backend")
-            model = Base.invokelatest(entry, predict_namedtuple(request["data"]))
+            model = Base.invokelatest(entry, ModelData(predict_namedtuple(request["data"])))
             stage = "load_samples"
             rchn = chains_from_wire(JSON.parsefile(request["samples"]), request["predict"]["targets"])
             stage = "predict"
@@ -260,7 +344,7 @@ function handle_request(request)
                 error("no variables predicted; check that [predict].targets name unconditioned outcomes")
             chains_to_wire(pp)
         else
-            data = backend == "juliabugs" ? bugs_namedtuple(request["data"]) : to_namedtuple(request["data"])
+            data = backend == "juliabugs" ? bugs_namedtuple(request["data"]) : ModelData(to_namedtuple(request["data"]))
             sampler, warmup = build_sampler(request["sampler"], backend)
             draws = Int(request["sampler"]["draws"])
             chains = Int(get(request["sampler"], "chains", 4))
@@ -270,7 +354,11 @@ function handle_request(request)
                 chn = Base.invokelatest(sample_bugs, model, sampler, warmup, draws, chains, rng)
                 chains_to_wire(chn)
             else
-                chn = Base.invokelatest(build_and_sample, entry, data, sampler, draws, chains, rng)
+                cb = get(request, "stream_draws", false) ?
+                    draw_streamer(draws, Int(get(request, "draw_batch_size", 25))) : nothing
+                chn = Base.invokelatest(
+                    build_and_sample, entry, data, sampler, draws, chains, rng; callback = cb,
+                )
                 vnchain_to_wire(chn)
             end
         end