@mcmcjs/julia 0.11.3 → 0.13.0

package/dist/fitlib.jl

@@ -18,6 +18,8 @@ import AdvancedHMC
 import ForwardDiff
 # Loading Mooncake activates AbstractPPL's native extension for AutoMooncake.
 import Mooncake
+# Loading ReverseDiff activates the DynamicPPL extension for AutoReverseDiff.
+import ReverseDiff
 import FlexiChains
 import DimensionalData
 
@@ -86,13 +88,50 @@ function Logging.handle_message(
     Logging.handle_message(logger.fallback, level, message, _module, group, id, file, line; kwargs...)
 end
 
+function build_adtype(name)
+    name == "forwarddiff" && return Turing.ADTypes.AutoForwardDiff()
+    name == "reversediff" && return Turing.ADTypes.AutoReverseDiff()
+    name == "mooncake" && return Turing.ADTypes.AutoMooncake(; config = nothing)
+    error("unsupported adtype: $name")
+end
+
 function build_sampler(sampler, backend)
     algorithm = get(sampler, "algorithm", "NUTS")
-    algorithm == "NUTS" || error("unsupported sampler algorithm: $algorithm")
+    # Prior draws are iid: no warmup, no adaptation. The JuliaBUGS path samples
+    # ancestrally (sample_bugs_prior) instead of going through a sampler object.
+    algorithm == "Prior" && return Turing.Prior(), 0
     warmup = Int(get(sampler, "warmup", 1000))
     adapt_delta = Float64(get(sampler, "adapt_delta", 0.8))
-    nuts = backend == "juliabugs" ? AdvancedHMC.NUTS(adapt_delta) : Turing.NUTS(warmup, adapt_delta)
-    return nuts, warmup
+    if backend == "juliabugs"
+        algorithm == "NUTS" || error("the juliabugs backend supports NUTS and Prior only")
+        return AdvancedHMC.NUTS(adapt_delta), warmup
+    end
+    adkw = haskey(sampler, "adtype") ? (; adtype = build_adtype(sampler["adtype"])) : (;)
+    algorithm == "NUTS" && return Turing.NUTS(warmup, adapt_delta; adkw...), warmup
+    algorithm == "HMC" &&
+        return Turing.HMC(Float64(sampler["step_size"]), Int(sampler["leapfrog_steps"]); adkw...),
+        warmup
+    algorithm == "HMCDA" &&
+        return Turing.HMCDA(warmup, adapt_delta, Float64(sampler["lambda"]); adkw...), warmup
+    algorithm == "MH" && return Turing.MH(), warmup
+    error("unsupported sampler algorithm: $algorithm")
+end
+
+# Extra keyword arguments for AbstractMCMC.sample, from the request's sampler
+# table: thinning, burn-in for the non-adaptive samplers (NUTS and HMCDA discard
+# their own adaptation), and named starting values replicated across chains.
+function sampling_kwargs(sampler, warmup, chains)
+    kw = (;)
+    thin = Int(get(sampler, "thin", 1))
+    thin > 1 && (kw = merge(kw, (; thinning = thin)))
+    algorithm = get(sampler, "algorithm", "NUTS")
+    algorithm in ("MH", "HMC") && warmup > 0 && (kw = merge(kw, (; discard_initial = warmup)))
+    if haskey(sampler, "initial_params")
+        nt = (; (Symbol(k) => narrow(v) for (k, v) in sampler["initial_params"])...)
+        strategy = Turing.DynamicPPL.InitFromParams(nt)
+        kw = merge(kw, (; initial_params = fill(strategy, chains)))
+    end
+    return kw
 end
 
 to_namedtuple(data) = (; (Symbol(k) => v for (k, v) in data)...)
@@ -288,9 +327,9 @@ end
 # 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. With a callback, draws stream.
-function build_and_sample(entry, data, sampler, draws, chains, rng; callback = nothing)
+function build_and_sample(entry, data, sampler, draws, chains, rng; callback = nothing, extra = (;))
     model = condition_on_data(entry(data), data)
-    kw = callback === nothing ? NamedTuple() : (; callback)
+    kw = callback === nothing ? extra : merge(extra, (; callback))
     return Logging.with_logger(JSONProgressLogger()) do
         sample(rng, model, sampler, MCMCSerial(), draws, chains; progress = true, kw...)
     end
@@ -310,7 +349,9 @@ function vnchain_to_wire(chn)
     total = nParams + length(extras)
 
     flat = Vector{Union{Float64,Nothing}}(undef, nIter * total * nChains)
-    cell(v) = v === missing ? nothing : Float64(v)
+    # JSON has no Inf/NaN; a non-finite draw (e.g. 1/sqrt(tau) under a diffuse
+    # prior) becomes null, which the samples parser reads back as NaN.
+    cell(v) = v === missing || !isfinite(v) ? nothing : Float64(v)
     for c in 1:nChains, p in 1:nParams, i in 1:nIter
         flat[i + (p - 1) * nIter + (c - 1) * nIter * total] = cell(arr[i, c, p])
     end
@@ -335,8 +376,9 @@ end
 # 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)
+function sample_bugs(model, sampler, warmup, draws, chains, rng; callback = nothing, thin = 1)
+    kw = callback === nothing ? (;) : (; callback)
+    thin > 1 && (kw = merge(kw, (; thinning = thin)))
     return Logging.with_logger(JSONProgressLogger()) do
         JuliaBUGS.AbstractMCMC.sample(
             rng, model, sampler, JuliaBUGS.AbstractMCMC.MCMCSerial(), draws, chains;
@@ -346,6 +388,31 @@ function sample_bugs(model, sampler, warmup, draws, chains, rng; callback = noth
     end
 end
 
+# JuliaBUGS prior sampling: each draw is one ancestral pass (evaluate!! with an
+# rng samples every unobserved stochastic node from its prior and recomputes the
+# deterministic ones; observed data stays fixed). The collected variables are the
+# model parameters plus generated quantities, the same columns a posterior fit
+# produces through gen_chains.
+function sample_bugs_prior(model, draws, chains, rng)
+    model isa JuliaBUGS.BUGSModelWithGradient && (model = model.base_model)
+    APPL = JuliaBUGS.AbstractPPL
+    vars = vcat(JuliaBUGS.model_parameters(model), JuliaBUGS.generated_quantities(model))
+    isempty(vars) && error("the model has no parameters or generated quantities to sample")
+    grab(env, vn) = (v = APPL.getvalue(env, vn); v isa AbstractArray ? copy(v) : v)
+    columns = Dict(vn => Vector{Any}(undef, draws * chains) for vn in vars)
+    for c in 1:chains, i in 1:draws
+        env, _ = APPL.evaluate!!(rng, model)
+        for vn in vars
+            columns[vn][i + (c - 1) * draws] = grab(env, vn)
+        end
+    end
+    dict = Dict{FlexiChains.ParameterOrExtra{<:VarName},Array}(
+        FlexiChains.Parameter(vn) =>
+            [columns[vn][i + (c - 1) * draws] for i in 1:draws, c in 1:chains] for vn in vars
+    )
+    return FlexiChains.FlexiChain{VarName}(draws, chains, dict)
+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
@@ -412,6 +479,69 @@ function bugs_draw_streamer(model, draws_per_chain::Int, batch_size::Int)
     end
 end
 
+# JuliaBUGS pointwise log-likelihood: restore each posterior draw into the
+# evaluation environment, recompute the deterministic nodes in topological
+# order, and record logpdf per observed node. Every unobserved stochastic node
+# must be covered by the posterior columns, or log p(y | theta) would not be a
+# function of the draw.
+function loglik_bugs(model, posterior)
+    model isa JuliaBUGS.BUGSModelWithGradient && (model = model.base_model)
+    APPL = JuliaBUGS.AbstractPPL
+    gd = model.graph_evaluation_data
+
+    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
+
+    restore = Tuple{APPL.VarName,Int}[]
+    for vn in JuliaBUGS.model_parameters(model)
+        leaves = collect(APPL.varname_leaves(vn, APPL.getvalue(model.evaluation_env, vn)))
+        found = count(l -> string(l) in posterior_names, leaves)
+        found == length(leaves) ||
+            error("the posterior samples do not cover $(vn); were they produced by fitting this spec?")
+        append!(restore, [(l, col[string(l)]) for l in leaves])
+    end
+    observed = [
+        (vn, i) for (i, vn) in enumerate(gd.sorted_nodes) if
+        gd.is_stochastic_vals[i] && gd.is_observed_vals[i]
+    ]
+    isempty(observed) && error("the model has no observed variables to compute a log-likelihood for")
+
+    pnames = string.(first.(observed))
+    n = length(pnames)
+    flat_out = Vector{Union{Float64,Nothing}}(undef, nIter * n * nChains)
+    for c in 1:nChains, i in 1:nIter
+        env = model.evaluation_env
+        for (leaf, p) in restore
+            env = JuliaBUGS.BangBang.setindex!!(env, cell(i, p, c), leaf)
+        end
+        for (k, node) in enumerate(gd.sorted_nodes)
+            if !gd.is_stochastic_vals[k]
+                env = JuliaBUGS.BangBang.setindex!!(
+                    env, gd.node_function_vals[k](env, gd.loop_vars_vals[k]), node,
+                )
+            end
+        end
+        for (k, (vn, idx)) in enumerate(observed)
+            dist = gd.node_function_vals[idx](env, gd.loop_vars_vals[idx])
+            v = logpdf(dist, APPL.getvalue(env, vn))
+            flat_out[i + (k - 1) * nIter + (c - 1) * nIter * n] = isfinite(v) ? Float64(v) : nothing
+        end
+    end
+    return Dict(
+        "size" => [nIter, n, nChains],
+        "value_flat" => flat_out,
+        "parameters" => pnames,
+        "name_map" => Dict("parameters" => pnames, "internals" => String[]),
+    )
+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.
@@ -479,7 +609,8 @@ function predict_bugs(model, posterior, targets, rng)
     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]
+        v = arr[i, p, c]
+        flat_out[i + (p - 1) * nIter + (c - 1) * nIter * n] = isfinite(v) ? v : nothing
     end
     return Dict(
         "size" => [nIter, n, nChains],
@@ -543,7 +674,26 @@ function handle_request(request)
         modelmod = load_model_module(request["model"]["file"])
         entry = resolve_entry(modelmod, get(request["model"], "entry", nothing))
 
-        wire = if mode == "predict"
+        wire = if mode == "loglik"
+            if backend == "juliabugs"
+                model = Base.invokelatest(entry, bugs_namedtuple(request["data"]))
+                stage = "load_samples"
+                posterior = JSON.parsefile(request["samples"])
+                stage = "loglik"
+                Base.invokelatest(loglik_bugs, model, posterior)
+            else
+                data = ModelData(to_namedtuple(request["data"]))
+                model = Base.invokelatest(entry, data)
+                model = Base.invokelatest(condition_on_data, model, data)
+                stage = "load_samples"
+                chn = wire_to_vnchain(JSON.parsefile(request["samples"]), String[])
+                stage = "loglik"
+                pll = Base.invokelatest(Turing.DynamicPPL.pointwise_loglikelihoods, model, chn)
+                isempty(FlexiChains.parameters(pll)) &&
+                    error("the model has no observed variables to compute a log-likelihood for")
+                vnchain_to_wire(pll)
+            end
+        elseif mode == "predict"
             if backend == "juliabugs"
                 model = Base.invokelatest(entry, bugs_predict_namedtuple(request["data"]))
                 stage = "load_samples"
@@ -568,19 +718,26 @@ function handle_request(request)
             stage = "sample"
             if backend == "juliabugs"
                 model = Base.invokelatest(entry, data)
-                # 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,
-                )
+                chn = if get(request["sampler"], "algorithm", "NUTS") == "Prior"
+                    Base.invokelatest(sample_bugs_prior, model, draws, chains, rng)
+                else
+                    # 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
+                    Base.invokelatest(
+                        sample_bugs, model, sampler, warmup, draws, chains, rng;
+                        callback = cb, thin = Int(get(request["sampler"], "thin", 1)),
+                    )
+                end
                 vnchain_to_wire(chn)
             else
                 cb = get(request, "stream_draws", false) ?
                     draw_streamer(draws, Int(get(request, "draw_batch_size", 25))) : nothing
+                extra = sampling_kwargs(request["sampler"], warmup, chains)
                 chn = Base.invokelatest(
-                    build_and_sample, entry, data, sampler, draws, chains, rng; callback = cb,
+                    build_and_sample, entry, data, sampler, draws, chains, rng;
+                    callback = cb, extra,
                 )
                 vnchain_to_wire(chn)
             end