@mcmcjs/julia 0.9.1 → 0.11.0

package/dist/fitlib.jl

@@ -17,6 +17,8 @@ using Pkg
 import JuliaBUGS
 import AdvancedHMC
 import ForwardDiff
+# Loading Mooncake activates AbstractPPL's native extension for AutoMooncake.
+import Mooncake
 import FlexiChains
 import DimensionalData
 
@@ -131,6 +133,9 @@ bugs_namedtuple(data) = (; (Symbol(k) => narrow(v) for (k, v) in data)...)
 to_missing(x) = x === nothing ? missing : (x isa AbstractVector ? map(to_missing, x) : x)
 predict_namedtuple(data) = (; (Symbol(k) => to_missing(v) for (k, v) in data)...)
 
+# A blanked target compiles to a latent node; JuliaBUGS drops scalar missings.
+bugs_predict_namedtuple(data) = (; (Symbol(k) => narrow(to_missing(v)) for (k, v) in data)...)
+
 # Rebuild a posterior Chains (latents only) from our MCMCChains-JSON wire object so
 # it can feed Turing.predict. The varname_to_symbol map is required for VarName
 # indexing; target outcome columns are excluded.
@@ -297,17 +302,164 @@ function vnchain_to_wire(chn)
 end
 
 # JuliaBUGS compiles the model at runtime, so building and sampling must run in
-# separate latest-world frames; merging them reintroduces a world-age error. The
-# result is already an MCMCChains.Chains, so it is not re-wrapped.
-function sample_bugs(model, sampler, warmup, draws, chains, rng)
+# separate latest-world frames; merging them reintroduces a world-age error.
+# Returns a FlexiChain like the Turing path; gen_chains recovers generated
+# quantities (unobserved nodes with no observed descendants) after sampling.
+# With a callback, draws stream (see bugs_draw_streamer).
+function sample_bugs(model, sampler, warmup, draws, chains, rng; callback = nothing)
+    kw = callback === nothing ? NamedTuple() : (; callback)
     return Logging.with_logger(JSONProgressLogger()) do
         JuliaBUGS.AbstractMCMC.sample(
             rng, model, sampler, JuliaBUGS.AbstractMCMC.MCMCSerial(), draws, chains;
-            chain_type = MCMCChains.Chains, n_adapts = warmup, discard_initial = warmup, progress = true,
+            chain_type = FlexiChains.VNChain, n_adapts = warmup, discard_initial = warmup,
+            progress = true, kw...,
         )
     end
 end
 
+# The JuliaBUGS counterpart to draw_streamer. A JuliaBUGS transition is a raw
+# AdvancedHMC.Transition whose params are an unconstrained vector, so (unlike the
+# Turing path) draws cannot be read off it directly. Per batch we rebuild a small
+# FlexiChain through the same gen_chains + vnchain_to_wire path the final file
+# uses, so the streamed leaf names (parameters, generated quantities, and sampler
+# stats) match the file exactly. Reconstruction runs on a copy of the model so it
+# never perturbs the sampler's state. AbstractMCMC passes the 1-based chain index
+# as chain_number; each chain's batches carry a per-chain monotonic seq.
+# TODO: once the pinned JuliaBUGS carries TuringLang/JuliaBUGS.jl#517 (with
+# AbstractMCMC#212 and AbstractPPL#178), a transition exposes named constrained
+# params (a VarNamedTuple) plus generated quantities, so this collapses to a
+# per-draw flatten like draw_streamer, dropping the deepcopy and reconstruction.
+function bugs_draw_streamer(model, draws_per_chain::Int, batch_size::Int)
+    recon = deepcopy(model)
+    side_rng = StableRNG(0)
+    chain = Ref(-1)
+    seq = Ref(0)
+    θs = Vector{Vector{Float64}}()
+    lps = Float64[]
+    stats = Any[]
+    function emit(iteration)
+        isempty(θs) && return
+        stats_names = collect(keys(merge((; lp = lps[1]), stats[1])))
+        stats_values = [vcat(lps[i], collect(values(stats[i]))) for i in eachindex(θs)]
+        chn = JuliaBUGS.gen_chains(FlexiChains.VNChain, recon, θs, stats_names, stats_values; rng = side_rng)
+        wire = vnchain_to_wire(chn)
+        nIter, total, _ = wire["size"]
+        flat = wire["value_flat"]
+        params = wire["parameters"]
+        cols = Dict(params[p] => [flat[i + (p - 1) * nIter] for i in 1:nIter] for p in 1:total)
+        println(
+            PROGRESS_IO[],
+            JSON.json(
+                Dict(
+                    "mcmcjs" => "draws",
+                    "chain" => chain[],
+                    "seq" => seq[],
+                    "iteration" => iteration,
+                    "draws" => cols,
+                ),
+            ),
+        )
+        flush(PROGRESS_IO[])
+        seq[] += 1
+        empty!(θs)
+        empty!(lps)
+        empty!(stats)
+    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
+            empty!(θs)
+            empty!(lps)
+            empty!(stats)
+        end
+        push!(θs, copy(transition.z.θ))
+        push!(lps, transition.z.ℓπ.value)
+        push!(stats, AdvancedHMC.stat(transition))
+        if length(θs) >= batch_size || iteration >= draws_per_chain
+            emit(iteration)
+        end
+    end
+end
+
+# JuliaBUGS posterior prediction: condition the (target-blanked) model on the
+# posterior parameter columns, then evaluate!! ancestral-samples the remaining
+# latents once per posterior draw. No sampler or gradient runs.
+function predict_bugs(model, posterior, targets, rng)
+    model isa JuliaBUGS.BUGSModelWithGradient && (model = model.base_model)
+    APPL = JuliaBUGS.AbstractPPL
+    target_syms = Set(Symbol.(targets))
+    is_target(vn) = APPL.getsym(vn) in target_syms
+
+    sz = posterior["size"]
+    nIter, nCols, nChains = Int(sz[1]), Int(sz[2]), Int(sz[3])
+    flat = posterior["value_flat"]
+    col = Dict(n => i for (i, n) in enumerate(posterior["parameters"]))
+    posterior_names = Set(posterior["name_map"]["parameters"])
+    cell(i, p, c) = begin
+        v = flat[i + (p - 1) * nIter + (c - 1) * nIter * nCols]
+        v === nothing ? NaN : Float64(v)
+    end
+
+    # Nodes fully covered by posterior columns are conditioned; anything
+    # uncovered stays latent and is forward-sampled, like Turing.predict.
+    gd = model.graph_evaluation_data
+    conditioned = APPL.VarName[]
+    restore = Tuple{APPL.VarName,Int}[]
+    for (i, vn) in enumerate(gd.sorted_nodes)
+        (gd.is_stochastic_vals[i] && !gd.is_observed_vals[i] && !is_target(vn)) || continue
+        leaves = collect(APPL.varname_leaves(vn, APPL.getvalue(model.evaluation_env, vn)))
+        found = count(l -> string(l) in posterior_names, leaves)
+        if found == length(leaves)
+            push!(conditioned, vn)
+            append!(restore, [(l, col[string(l)]) for l in leaves])
+        elseif found > 0
+            error("the posterior samples cover $(vn) only partially; refit before predicting")
+        end
+    end
+    isempty(conditioned) &&
+        error("the posterior samples share no parameters with the model; were they produced by fitting this spec?")
+    cond = APPL.condition(model, conditioned)
+
+    any(is_target, cond.graph_evaluation_data.sorted_nodes) ||
+        error("no variables predicted; check that [predict].targets name model variables")
+
+    out_leaves = APPL.VarName[]
+    arr = Array{Float64,3}(undef, 0, 0, 0)
+    for c in 1:nChains, i in 1:nIter
+        env = cond.evaluation_env
+        for (leaf, p) in restore
+            env = JuliaBUGS.BangBang.setindex!!(env, cell(i, p, c), leaf)
+        end
+        cond = JuliaBUGS.BangBang.setproperty!!(cond, :evaluation_env, env)
+        drawn, _ = APPL.evaluate!!(rng, cond)
+        if isempty(out_leaves)
+            out_leaves = [
+                l for vn in cond.graph_evaluation_data.sorted_nodes if is_target(vn)
+                for l in APPL.varname_leaves(vn, APPL.getvalue(drawn, vn))
+            ]
+            arr = Array{Float64,3}(undef, nIter, length(out_leaves), nChains)
+        end
+        for (k, leaf) in enumerate(out_leaves)
+            arr[i, k, c] = Float64(APPL.getvalue(drawn, leaf))
+        end
+    end
+
+    pnames = string.(out_leaves)
+    n = length(pnames)
+    flat_out = Vector{Union{Float64,Nothing}}(undef, nIter * n * nChains)
+    for c in 1:nChains, p in 1:n, i in 1:nIter
+        flat_out[i + (p - 1) * nIter + (c - 1) * nIter * n] = arr[i, p, c]
+    end
+    return Dict(
+        "size" => [nIter, n, nChains],
+        "value_flat" => flat_out,
+        "parameters" => pnames,
+        "name_map" => Dict("parameters" => pnames, "internals" => String[]),
+    )
+end
+
 # Build the exact wire object parseMCMCChainsJson consumes: value_flat is indexed
 # i + p*nIter + c*nIter*nParams (column-major over [iteration, parameter, chain]).
 function chains_to_wire(chn)
@@ -333,7 +485,7 @@ function provenance()
     packages = Dict{String,String}()
     for (_, info) in Pkg.dependencies()
         if info.name in
-           ("Turing", "FlexiChains", "MCMCChains", "DynamicPPL", "AbstractMCMC", "JuliaBUGS", "AdvancedHMC", "ForwardDiff", "ADTypes") &&
+           ("Turing", "FlexiChains", "MCMCChains", "DynamicPPL", "AbstractMCMC", "JuliaBUGS", "AdvancedHMC", "ForwardDiff", "Mooncake", "ADTypes") &&
            info.version !== nothing
             packages[info.name] = string(info.version)
         end
@@ -372,15 +524,22 @@ function handle_request(request)
         entry = resolve_entry(Main, get(request["model"], "entry", nothing))
 
         wire = if mode == "predict"
-            backend == "turing" || error("predict is not yet supported for backend $backend")
-            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"
-            pp = Base.invokelatest(predict, rng, model, rchn)
-            isempty(names(pp)) &&
-                error("no variables predicted; check that [predict].targets name unconditioned outcomes")
-            chains_to_wire(pp)
+            if backend == "juliabugs"
+                model = Base.invokelatest(entry, bugs_predict_namedtuple(request["data"]))
+                stage = "load_samples"
+                posterior = JSON.parsefile(request["samples"])
+                stage = "predict"
+                Base.invokelatest(predict_bugs, model, posterior, request["predict"]["targets"], rng)
+            else
+                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"
+                pp = Base.invokelatest(predict, rng, model, rchn)
+                isempty(names(pp)) &&
+                    error("no variables predicted; check that [predict].targets name unconditioned outcomes")
+                chains_to_wire(pp)
+            end
         else
             data = backend == "juliabugs" ? bugs_namedtuple(request["data"]) : ModelData(to_namedtuple(request["data"]))
             sampler, warmup = build_sampler(request["sampler"], backend)
@@ -389,8 +548,14 @@ function handle_request(request)
             stage = "sample"
             if backend == "juliabugs"
                 model = Base.invokelatest(entry, data)
-                chn = Base.invokelatest(sample_bugs, model, sampler, warmup, draws, chains, rng)
-                chains_to_wire(chn)
+                # Reconstruction reads the plain BUGSModel; entry may return a gradient wrapper.
+                base = hasproperty(model, :base_model) ? model.base_model : model
+                cb = get(request, "stream_draws", false) ?
+                    bugs_draw_streamer(base, draws, Int(get(request, "draw_batch_size", 25))) : nothing
+                chn = Base.invokelatest(
+                    sample_bugs, model, sampler, warmup, draws, chains, rng; callback = cb,
+                )
+                vnchain_to_wire(chn)
             else
                 cb = get(request, "stream_draws", false) ?
                     draw_streamer(draws, Int(get(request, "draw_batch_size", 25))) : nothing

package/dist/index.cjs

@@ -268,6 +268,7 @@ async function runFit(spec, resolved, io) {
     ({ stdout, stderr, code, cancelled } = await io.spawn(resolved.command, args, {
       onProgress: io.onProgress,
       onDraws: io.onDraws,
+      onLog: io.onLog,
       signal: io.signal
     }));
   } finally {
@@ -397,6 +398,7 @@ var PACKAGES = [
   "JuliaBUGS",
   "AdvancedHMC",
   "ForwardDiff",
+  "Mooncake",
   "JSON",
   "StableRNGs",
   "StatsFuns"