Add some inter-op examples (#249)

partially address #245

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: sunxd3

examples/.JuliaFormatter.toml

@@ -0,0 +1,2 @@
+style = "blue"
+always_use_return = false

examples/Project.toml

@@ -0,0 +1,20 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"
+AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+JuliaBUGS = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

examples/README.md

@@ -0,0 +1,9 @@
+# JuliaBUGS Examples
+
+We adapted some examples to show how to use JuliaBUGS in this repo.
+
+## Sources
+
+* SIR: https://github.com/TuringLang/Turing-Workshop/tree/main/2023-MRC-BSU-and-UKHSA/Part-2-More-Julia-and-some-Bayesian-inference
+* GP: https://turinglang.org/docs/tutorials/gaussian-processes-introduction/
+* BNN: https://turinglang.org/docs/tutorials/bayesian-neural-networks/

examples/bnn.jl

@@ -0,0 +1,109 @@
+using JuliaBUGS
+
+using AbstractMCMC
+using ADTypes
+using AdvancedHMC
+using DifferentiationInterface
+using FillArrays
+using Functors
+using LinearAlgebra
+using LogDensityProblems
+using LogDensityProblemsAD
+using Lux
+using MCMCChains
+using Mooncake
+using Random
+
+## data simulation
+
+# Number of points to generate
+N = 80
+M = round(Int, N / 4)
+rng = Random.default_rng()
+Random.seed!(rng, 1234)
+
+# Generate artificial data
+x1s = rand(rng, Float32, M) * 4.5f0;
+x2s = rand(rng, Float32, M) * 4.5f0;
+xt1s = Array([[x1s[i] + 0.5f0; x2s[i] + 0.5f0] for i in 1:M])
+x1s = rand(rng, Float32, M) * 4.5f0;
+x2s = rand(rng, Float32, M) * 4.5f0;
+append!(xt1s, Array([[x1s[i] - 5.0f0; x2s[i] - 5.0f0] for i in 1:M]))
+
+x1s = rand(rng, Float32, M) * 4.5f0;
+x2s = rand(rng, Float32, M) * 4.5f0;
+xt0s = Array([[x1s[i] + 0.5f0; x2s[i] - 5.0f0] for i in 1:M])
+x1s = rand(rng, Float32, M) * 4.5f0;
+x2s = rand(rng, Float32, M) * 4.5f0;
+append!(xt0s, Array([[x1s[i] - 5.0f0; x2s[i] + 0.5f0] for i in 1:M]))
+
+# Store all the data for later
+xs = [xt1s; xt0s]
+xs_hcat = Float64.(reduce(hcat, xs))
+ts = [ones(2 * M); zeros(2 * M)]
+
+alpha = 0.09
+sigma = sqrt(1.0 / alpha)
+
+## 
+
+# Construct a neural network using Lux
+nn_initial = Chain(Dense(2 => 3, tanh), Dense(3 => 2, tanh), Dense(2 => 1, σ))
+
+# Initialize the model weights and state
+ps, st = Lux.setup(rng, nn_initial)
+
+Lux.parameterlength(nn_initial) # number of parameters in NN
+
+function vector_to_parameters(ps_new::AbstractVector, ps::NamedTuple)
+    @assert length(ps_new) == Lux.parameterlength(ps)
+    i = 1
+    function get_ps(x)
+        z = reshape(view(ps_new, i:(i + length(x) - 1)), size(x))
+        i += length(x)
+        return z
+    end
+    return fmap(get_ps, ps)
+end
+
+const nn = StatefulLuxLayer{true}(nn_initial, nothing, st)
+
+model_def = @bugs begin
+    parameters[1:nparameters] ~ parameter_distribution(nparameters, sigma)
+    predictions[1:N] = make_prediction(parameters[1:nparameters], xs[:, :])
+    for i in 1:N
+        ts[i] ~ Bernoulli(predictions[i])
+    end
+end
+
+JuliaBUGS.@register_primitive function parameter_distribution(nparameters, sigma)
+    return MvNormal(zeros(nparameters), Diagonal(abs2.(sigma .* ones(nparameters))))
+end
+
+JuliaBUGS.@register_primitive function make_prediction(parameters, xs; ps=ps, nn=nn)
+    return Lux.apply(nn, f32(xs), f32(vector_to_parameters(parameters, ps)))
+end
+
+@eval JuliaBUGS begin
+    ps = Main.ps
+    nn = Main.nn
+    Lux = Main.Lux
+    f32 = Main.f32
+    vector_to_parameters = Main.vector_to_parameters
+end
+
+data = (nparameters=Lux.parameterlength(nn), xs=xs_hcat, ts=ts, N=length(ts), sigma=sigma)
+
+model = compile(model_def, data)
+
+ad_model = ADgradient(AutoMooncake(; config=Mooncake.Config()), model)
+
+# sampling is slow, so sample 10 of them to verify that this can work
+samples_and_stats = AbstractMCMC.sample(
+    ad_model,
+    NUTS(0.65),
+    10;
+    chain_type=Chains,
+    # n_adapts=1000, 
+    # discard_initial=1000
+)

examples/gp.jl

@@ -0,0 +1,217 @@
+# Example demonstrating the use of Gaussian Processes (GPs) within JuliaBUGS
+# for modeling golf putting accuracy based on distance.
+# This example uses AbstractGPs.jl for the GP implementation and AdvancedHMC.jl
+# for sampling from the posterior distribution.
+
+using JuliaBUGS
+using JuliaBUGS: @model
+
+# Required packages for GP modeling and MCMC
+using AbstractGPs, Distributions, LogExpFunctions
+using LogDensityProblems, LogDensityProblemsAD
+using AbstractMCMC, AdvancedHMC, MCMCChains
+
+# Differentiation backend
+using DifferentiationInterface
+using Mooncake: Mooncake
+
+# --- Data Definition ---
+
+# Golf putting data from Gelman et al. (BDA3, Chapter 5)
+golf_data = (
+    distance=[2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], # Distance in feet
+    n=[ # Number of putts attempted
+        1443,
+        694,
+        455,
+        353,
+        272,
+        256,
+        240,
+        217,
+        200,
+        237,
+        202,
+        192,
+        174,
+        167,
+        201,
+        195,
+        191,
+        147,
+        152,
+    ],
+    y=[ # Number of successful putts
+        1346,
+        577,
+        337,
+        208,
+        149,
+        136,
+        111,
+        69,
+        67,
+        75,
+        52,
+        46,
+        54,
+        28,
+        27,
+        31,
+        33,
+        20,
+        24,
+    ],
+)
+
+# Prepare data in the format expected by the BUGS model
+data = (
+    d=golf_data.distance,
+    n=golf_data.n,
+    y=golf_data.y,
+    jitter=1e-6, # Small value added to GP kernel diagonal for numerical stability
+    N=length(golf_data.distance),
+)
+
+# --- BUGS Model Definition ---
+
+@model function gp_golf_putting((; v, l, f_latent, y), N, n, d, jitter)
+    # Priors for GP hyperparameters
+    v ~ Distributions.Gamma(2, 1) # Variance
+    l ~ Distributions.Gamma(4, 1) # Lengthscale
+
+    # Latent GP function values
+    # f_latent represents the underlying putting success probability (on logit scale)
+    # modeled by a GP.
+    f_latent[1:N] ~ gp_predict(v, l, d[1:N], jitter)
+
+    # Likelihood: Binomial distribution for observed successes
+    # The success probability for each distance is the logistic transformation of the latent GP value.
+    y[1:N] ~ y_distribution(n[1:N], f_latent[1:N])
+end
+
+# --- Custom Primitive Definitions for BUGS ---
+
+# Register the GP kernel type with JuliaBUGS
+# This allows using AbstractGPs types directly in the model definition.
+JuliaBUGS.@register_primitive GP with_lengthscale SEKernel
+
+# Define a function callable within the BUGS model to compute GP predictions.
+# BUGS requires functions to operate on basic numerical types, so this wraps the GP call.
+JuliaBUGS.@register_primitive function gp_predict(v, l, d, jitter)
+    # Create a GP with a Squared Exponential kernel using the provided hyperparameters
+    kernel = v * with_lengthscale(SEKernel(), l)
+    gp = GP(kernel)
+    # Return the distribution representing the GP evaluated at distances `d` with jitter
+    return gp(d, jitter)
+end
+
+# Define a function for the observation model (likelihood).
+# This creates a product distribution of Binomials, one for each distance.
+JuliaBUGS.@register_primitive function y_distribution(n, f_latent)
+    return product_distribution(Binomial.(n, logistic.(f_latent)))
+end
+
+# --- Model Instantiation ---
+
+# Create the JuliaBUGS model instance
+# Provide initial values (missing for parameters to be inferred) and observed data
+model = gp_golf_putting(
+    (; v=missing, l=missing, f_latent=fill(missing, data.N), y=data.y),
+    data.N,      # Number of observations
+    data.n,      # Observed attempts
+    data.d,      # Observed distances
+    data.jitter, # Numerical stability term
+)
+
+# Optionally, set the evaluation mode. Using generated functions can be faster.
+# model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGeneratedLogDensityFunction())
+
+# --- MCMC Setup with Custom LogDensityProblems Wrapper ---
+
+# We need a wrapper around the JuliaBUGS model to interface with LogDensityProblems
+# and utilize automatic differentiation (AD) via Mooncake.jl for gradient computation,
+# which is required by AdvancedHMC.
+
+struct BUGSMooncakeModel{T,P}
+    model::T # The JuliaBUGS model
+    prep::P  # Pre-allocated workspace for gradient computation using Mooncake
+end
+
+# Define the function to compute the log density using the JuliaBUGS model's internal function
+f(x) = model.log_density_computation_function(model.evaluation_env, x)
+
+# Prepare the differentiation backend (Mooncake)
+backend = AutoMooncake(; config=nothing)
+x_init = rand(LogDensityProblems.dimension(model)) # Initial point for testing/preparation
+prep = prepare_gradient(f, backend, x_init)
+
+# Create the wrapped model instance
+bugsmooncake = BUGSMooncakeModel(model, prep)
+
+# --- LogDensityProblems Interface Implementation for the Wrapper ---
+
+# Define logdensity function for the wrapper
+function LogDensityProblems.logdensity(model::BUGSMooncakeModel, x::AbstractVector)
+    return f(x) # Calls the underlying JuliaBUGS log density function
+end
+
+# Define logdensity_and_gradient function using the prepared DifferentiationInterface setup
+function LogDensityProblems.logdensity_and_gradient(
+    model::BUGSMooncakeModel, x::AbstractVector
+)
+    # Computes both the log density and its gradient using Mooncake AD
+    return DifferentiationInterface.value_and_gradient(
+        f, model.prep, AutoMooncake(; config=nothing), x
+    )
+end
+
+# Define dimension function
+function LogDensityProblems.dimension(model::BUGSMooncakeModel)
+    return LogDensityProblems.dimension(model.model) # Delegates to the original model
+end
+
+# Define a custom bundle_samples function to convert the AdvancedHMC.Transition to a Chains object
+function AbstractMCMC.bundle_samples(
+    ts::Vector{<:AdvancedHMC.Transition},
+    logdensitymodel::AbstractMCMC.LogDensityModel{<:BUGSMooncakeModel},
+    sampler::AdvancedHMC.AbstractHMCSampler,
+    state,
+    chain_type::Type{Chains};
+    discard_initial=0,
+    thinning=1,
+    kwargs...,
+)
+    stats_names = collect(keys(merge((; lp=ts[1].z.ℓπ.value), AdvancedHMC.stat(ts[1]))))
+    stats_values = [
+        vcat([ts[i].z.ℓπ.value..., collect(values(AdvancedHMC.stat(ts[i])))...]) for
+        i in eachindex(ts)
+    ]
+
+    return JuliaBUGS.gen_chains(
+        logdensitymodel.logdensity.model,
+        [t.z.θ for t in ts],
+        stats_names,
+        stats_values;
+        discard_initial=discard_initial,
+        thinning=thinning,
+        kwargs...,
+    )
+end
+
+# Specify capabilities (indicates gradient availability)
+function LogDensityProblems.capabilities(::Type{<:BUGSMooncakeModel})
+    return LogDensityProblems.LogDensityOrder{1}() # Can compute up to the gradient
+end
+
+# --- MCMC Sampling ---
+
+# Sample from the posterior distribution using AdvancedHMC's NUTS sampler
+samples_and_stats = AbstractMCMC.sample(
+    AbstractMCMC.LogDensityModel(bugsmooncake), # Wrap the model for AbstractMCMC
+    AdvancedHMC.NUTS(0.65), # No-U-Turn Sampler
+    1000;                   # Total number of samples
+    chain_type=Chains,      # Store results as MCMCChains object
+    n_adapts=500,           # Number of adaptation steps for NUTS
+    discard_initial=500,    # Number of initial samples (warmup) to discard;
+)