add experiment package

Author: sunxd3

experiments/AutoMarginalizationExperiments/Project.toml

@@ -0,0 +1,45 @@
+name = "AutoMarginalizationExperiments"
+uuid = "7a1de1b0-2fb5-4cf5-9df0-9a8847935917"
+version = "0.1.0"
+
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
+AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RDatasets = "ce6b1742-4840-55fa-b093-852dadbb1d8b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+
+[compat]
+julia = "1.10, 1.11"
+Distributions = "0.25"
+ADTypes = "1"
+ForwardDiff = "0.10, 0.11"
+LogDensityProblems = "2"
+LogDensityProblemsAD = "1"
+OrdinaryDiffEq = "6"
+RDatasets = "0.7"
+BenchmarkTools = "1"
+Graphs = "1"
+MetaGraphsNext = "0.6, 0.7"
+StaticArrays = "1"
+AdvancedHMC = "0.6, 0.7, 0.8"
+AbstractMCMC = "5"
+MCMCChains = "6, 7"
+AbstractPPL = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13"
+LogExpFunctions = "0.3"

experiments/AutoMarginalizationExperiments/README.md

@@ -0,0 +1,27 @@
+# AutoMarginalizationExperiments
+
+A lightweight experiment harness (as a package) to showcase finite‑support discrete auto‑marginalization in JuliaBUGS combined with HMC.
+
+Important: Auto-marginalization is on the current JuliaBUGS branch, not the latest release. Develop JuliaBUGS into this environment first:
+
+```
+julia --project=experiments/AutoMarginalizationExperiments -e 'using Pkg; Pkg.develop(path="JuliaBUGS"); Pkg.instantiate()'
+```
+
+Then try a quick run:
+
+```
+julia --project=experiments/AutoMarginalizationExperiments -e 'using Pkg; Pkg.instantiate()'
+julia --project=experiments/AutoMarginalizationExperiments -e 'using AutoMarginalizationExperiments; AutoMarginalizationExperiments.run_gmm_autmarg_nuts(2000, 3)'
+```
+
+Goals (aligned with experiments/plan.md):
+- Exactness/gradient checks on small models (GMM/HMM).
+- Scaling vs weighted frontier width and order selection.
+- ODE + finite discrete noise (PK Theoph) with reuse ablations.
+- Single changepoint over a finite grid.
+- “Must‑have” demos: GMM and HMM with NUTS on the marginalized target.
+
+Folders:
+- `src/` — package modules: metrics, ordering helpers, synthetic GMM/HMM, NUTS harness.
+- Future: `scripts/` for CLI drivers and CSV logging; `pk_theoph/`, `changepoint_step/` subfolders.

experiments/AutoMarginalizationExperiments/scripts/benchmarks/gmm_hmm_logdensity.jl

@@ -0,0 +1,37 @@
+#!/usr/bin/env julia
+using AutoMarginalizationExperiments
+using JuliaBUGS
+using LogDensityProblems
+using BenchmarkTools
+using Printf
+
+function bench_gmm(; N_list=[1_000, 5_000, 10_000], K=3, reps=10, seed=1)
+    @printf "GMM logdensity benchmark (auto-marg), K=%d\n" K
+    for N in N_list
+        data, _ = AutoMarginalizationExperiments.synth_gmm(N; seed=seed, weights=fill(1/K,K), mus=collect(range(-2,2; length=K)), sigmas=fill(1.0,K))
+        dataK = (data..., K=K)
+        mdef = AutoMarginalizationExperiments.build_gmm_model(K)
+        model, θ = AutoMarginalizationExperiments.compile_autmarg(mdef, dataK)
+        # warmup
+        LogDensityProblems.logdensity(model, θ)
+        b = @benchmark LogDensityProblems.logdensity($model, $θ) samples=$reps evals=1
+        @printf "  N=%6d  median=%.3f ms  mean=%.3f ms  allocs=%d  bytes=%d\n" N (median(b).time/1e6) (mean(b).time/1e6) median(b).allocs median(b).memory
+    end
+end
+
+function bench_hmm(; T_list=[200, 500, 1000], reps=10, seed=1)
+    @printf "HMM logdensity benchmark (auto-marg), S=2\n"
+    for T in T_list
+        data, _ = AutoMarginalizationExperiments.synth_hmm_binary(T; seed=seed)
+        mdef = AutoMarginalizationExperiments.build_hmm2_model()
+        model, θ = AutoMarginalizationExperiments.compile_autmarg(mdef, data)
+        LogDensityProblems.logdensity(model, θ)
+        b = @benchmark LogDensityProblems.logdensity($model, $θ) samples=$reps evals=1
+        @printf "  T=%6d  median=%.3f ms  mean=%.3f ms  allocs=%d  bytes=%d\n" T (median(b).time/1e6) (mean(b).time/1e6) median(b).allocs median(b).memory
+    end
+end
+
+bench_gmm()
+println()
+bench_hmm()
+

experiments/AutoMarginalizationExperiments/scripts/exactness_gmm.jl

@@ -0,0 +1,88 @@
+#!/usr/bin/env julia
+using AutoMarginalizationExperiments
+using JuliaBUGS
+using ADTypes
+using Distributions
+using ForwardDiff
+using LogExpFunctions
+using LogDensityProblems
+using LogDensityProblemsAD
+using Printf
+
+function main(; N=12, K=2, seed=1)
+    data, truth = AutoMarginalizationExperiments.synth_gmm(N; seed=seed, weights=fill(1/K,K), mus=collect(range(-1.0,1.0; length=K)), sigmas=fill(0.8, K))
+    dataK = (data..., K=K)
+    model_def = AutoMarginalizationExperiments.build_gmm_model(K)
+    model, θ = AutoMarginalizationExperiments.compile_autmarg(model_def, dataK)
+    ad_model = ADgradient(AutoForwardDiff(), model)
+
+    # Build mapping from θ to variable values
+    gd = model.graph_evaluation_data
+    # Continuous parameters only
+    cont_vars = JuliaBUGS.Model.VarName[]
+    for vn in gd.sorted_parameters
+        idx = findfirst(==(vn), gd.sorted_nodes)
+        if idx !== nothing && gd.node_types[idx] == :continuous
+            push!(cont_vars, vn)
+        end
+    end
+    var_lengths = Dict{JuliaBUGS.Model.VarName,Int}()
+    for vn in cont_vars
+        var_lengths[vn] = model.transformed_var_lengths[vn]
+    end
+    offsets = Dict{JuliaBUGS.Model.VarName,Int}()
+    start = 1
+    for vn in cont_vars
+        offsets[vn] = start
+        start += var_lengths[vn]
+    end
+
+    function unpack(θvec)
+        T = eltype(θvec)
+        mus = fill(zero(T), K)
+        sigmas = fill(zero(T), K)
+        for vn in cont_vars
+            name = string(vn)
+            s = offsets[vn]
+            # parse index inside brackets
+            idx = try parse(Int, name[findfirst('[', name)+1:findfirst(']', name)-1]) catch; 0 end
+            if startswith(name, "mu[") && idx ≥ 1 && idx ≤ K
+                mus[idx] = θvec[s]
+            elseif startswith(name, "sigma[") && idx ≥ 1 && idx ≤ K
+                sigmas[idx] = exp(θvec[s])
+            end
+        end
+        return mus, sigmas
+    end
+
+    function logjoint_closed(θvec)
+        mus, sigmas = unpack(θvec)
+        @assert length(mus) == K && length(sigmas) == K
+        # Priors: mu ~ Normal(0,5), sigma ~ Exponential(1) with log-Jacobian from exp transform
+        lp = 0.0
+        for k in 1:K
+            lp += logpdf(Distributions.Normal(0,5), mus[k])
+            lp += logpdf(Distributions.Exponential(1.0), sigmas[k]) + log(sigmas[k]) # jacobian of exp
+        end
+        # Likelihood: product over i of sum_k w_k N(y_i | mu_k, sigma_k)
+        w = 1.0 / K
+        for yi in data.y
+            terms = similar(mus)
+            @inbounds for k in 1:K
+                terms[k] = log(w) + logpdf(Distributions.Normal(mus[k], sigmas[k]), yi)
+            end
+            lp += LogExpFunctions.logsumexp(terms)
+        end
+        return lp
+    end
+
+    val_ad, grad_ad = LogDensityProblems.logdensity_and_gradient(ad_model, θ)
+    val_cf = logjoint_closed(θ)
+    grad_cf = ForwardDiff.gradient(logjoint_closed, θ)
+
+    @printf "N=%d, K=%d\n" N K
+    @printf "value: engine=%.8f, closed=%.8f, absdiff=%.3e\n" val_ad val_cf abs(val_ad - val_cf)
+    @printf "grad max-abs-diff: %.3e\n" maximum(abs.(grad_ad .- grad_cf))
+end
+
+main()

experiments/AutoMarginalizationExperiments/scripts/ordering_ablation.jl

@@ -0,0 +1,78 @@
+#!/usr/bin/env julia
+using AutoMarginalizationExperiments
+using JuliaBUGS
+using LogDensityProblems
+using Printf
+using Random
+
+function hmm_logdensity_with_order(model, θ, order)
+    # Prepare caches and offsets as in the benchmark script
+    gd = model.graph_evaluation_data
+    minimal_keys = AutoMarginalizationExperiments.prepare_minimal_cache_keys(model, order)
+    # Build continuous-only param order and offsets
+    cont_vars = JuliaBUGS.Model.VarName[]
+    var_lengths = Dict{JuliaBUGS.Model.VarName,Int}()
+    for vn in gd.sorted_parameters
+        idx = findfirst(==(vn), gd.sorted_nodes)
+        if idx !== nothing && gd.node_types[idx] == :continuous
+            push!(cont_vars, vn)
+            var_lengths[vn] = model.transformed_var_lengths[vn]
+        end
+    end
+    offsets = Dict{JuliaBUGS.Model.VarName,Int}()
+    start = 1
+    for vn in cont_vars
+        offsets[vn] = start
+        start += var_lengths[vn]
+    end
+    env = JuliaBUGS.Model.smart_copy_evaluation_env(model.evaluation_env, model.mutable_symbols)
+    memo = Dict{Tuple{Int,UInt64},Any}()
+    return JuliaBUGS.Model._marginalize_recursive(
+        model, env, order, θ, offsets, var_lengths, memo, minimal_keys,
+    )
+end
+
+function peak_frontier_size(minimal_keys)
+    isempty(minimal_keys) && return 0
+    return maximum((length(v) for v in values(minimal_keys)))
+end
+
+function main(; T=300, reps=10, seed=1)
+    rng = MersenneTwister(seed)
+    data, _ = AutoMarginalizationExperiments.synth_hmm_binary(T; seed=seed)
+    model_def = AutoMarginalizationExperiments.build_hmm2_model()
+    model, θ = AutoMarginalizationExperiments.compile_autmarg(model_def, data)
+
+    gd = model.graph_evaluation_data
+    default_order = isempty(gd.marginalization_order) ? collect(1:length(gd.sorted_nodes)) : gd.marginalization_order
+    interleaved = AutoMarginalizationExperiments.build_interleaved_order(model)
+
+    # Warmup
+    hmm_logdensity_with_order(model, θ, default_order)
+    hmm_logdensity_with_order(model, θ, interleaved)
+
+    # Measure
+    function timeit(order)
+        t = @elapsed begin
+            for _ in 1:reps
+                hmm_logdensity_with_order(model, θ, order)
+            end
+        end
+        mk = AutoMarginalizationExperiments.prepare_minimal_cache_keys(model, order)
+        return t, peak_frontier_size(mk)
+    end
+
+    t_def, w_def = timeit(default_order)
+    t_int, w_int = timeit(interleaved)
+
+    l_def = hmm_logdensity_with_order(model, θ, default_order)
+    l_int = hmm_logdensity_with_order(model, θ, interleaved)
+
+    @printf "HMM ordering ablation (T=%d, reps=%d)\n" T reps
+    @printf " default: time=%.4f s, peak_frontier=%d, logp=%.6f\n" t_def w_def l_def
+    @printf " interlv: time=%.4f s, peak_frontier=%d, logp=%.6f\n" t_int w_int l_int
+    @printf " abs diff in logp = %.3e\n" abs(l_def - l_int)
+end
+
+main()
+

experiments/AutoMarginalizationExperiments/scripts/setup.jl

@@ -0,0 +1,11 @@
+#!/usr/bin/env julia
+using Pkg
+using Printf
+
+root = normpath(joinpath(@__DIR__, "..", "..", ".."))
+jbugs = joinpath(root, "JuliaBUGS")
+@printf "Developing JuliaBUGS from %s\n" jbugs
+Pkg.develop(path=jbugs)
+Pkg.instantiate()
+println("OK")
+

experiments/AutoMarginalizationExperiments/src/AutoMarginalizationExperiments.jl

@@ -0,0 +1,33 @@
+module AutoMarginalizationExperiments
+
+using Random
+using LinearAlgebra
+using Statistics
+using Printf
+
+using Distributions
+using ADTypes
+using LogDensityProblems
+using LogDensityProblemsAD
+
+using JuliaBUGS
+using JuliaBUGS: @bugs, compile, settrans
+import JuliaBUGS.Model
+
+include("metrics.jl")
+include("ordering.jl")
+include("synth_gmm.jl")
+include("synth_hmm.jl")
+include("harness.jl")
+
+export 
+    # Metrics
+    Metrics,
+    # GMM
+    synth_gmm, build_gmm_model, run_gmm_autmarg_nuts,
+    # HMM
+    synth_hmm_binary, build_hmm2_model, run_hmm_autmarg_nuts,
+    # Ordering helpers
+    build_interleaved_order, prepare_minimal_cache_keys
+
+end # module