Merge pull request #403 from ReactiveBayes/binomial_regression

Add Binomial Regression Tests and Example Notebook

Project.toml

@@ -1,7 +1,7 @@
 name = "RxInfer"
 uuid = "86711068-29c9-4ff7-b620-ae75d7495b3d"
 authors = ["Bagaev Dmitry <[email protected]> and contributors"]
-version = "3.8.4"
+version = "3.9.0"
 
 [deps]
 BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e"
@@ -42,7 +42,7 @@ MacroTools = "0.5.6"
 Optim = "1.0.0"
 ProgressMeter = "1.0.0"
 Random = "1.9"
-ReactiveMP = "~4.4.4"
+ReactiveMP = "~4.5.0"
 Reexport = "1.2.0"
 Rocket = "1.8.0"
 Static = "0.8.10, 1"

codemeta.json

@@ -9,12 +9,12 @@
     "downloadUrl": "https://github.com/reactivebayes/RxInfer.jl/releases",
     "issueTracker": "https://github.com/reactivebayes/RxInfer.jl/issues",
     "name": "RxInfer.jl",
-    "version": "3.8.4",
+    "version": "3.9.0",
     "description": "Julia package for automated, scalable and efficient Bayesian inference on factor graphs with reactive message passing. ",
     "applicationCategory": "Statistics",
     "developmentStatus": "active",
     "readme": "https://reactivebayes.github.io/RxInfer.jl/stable/",
-    "softwareVersion": "3.8.4",
+    "softwareVersion": "3.9.0",
     "keywords": [
         "Bayesian inference",
         "message passing",

scripts/format.jl

@@ -10,7 +10,7 @@ s = ArgParseSettings()
 end
 
 commandline_args = parse_args(s)
-folders_to_format = ["scripts", "src", "test"]
+folders_to_format = ["scripts", "src", "test", "ext"]
 
 overwrite = commandline_args["overwrite"]
 formatted = all(map(folder -> JuliaFormatter.format(folder, overwrite = overwrite, verbose = true), folders_to_format))

test/models/regression/binomialreg_tests.jl

@@ -0,0 +1,86 @@
+@testitem "Linear regression with BinomialPolya node" begin
+    using BenchmarkTools, Plots, Dates, LinearAlgebra, StableRNGs
+
+    include(joinpath(@__DIR__, "..", "..", "utiltests.jl"))
+
+    function generate_synthetic_binomial_data(n_samples::Int, true_beta::Vector{Float64}; seed::Int = 42)
+        rng = StableRNG(seed)
+        n_features = length(true_beta)
+        # Generate design matrix X
+        X = randn(rng, n_samples, n_features)
+
+        # Generate number of trials for each observation
+        n_trials = rand(rng, 5:20, n_samples)
+
+        # Compute logits and probabilities
+        logits = X * true_beta
+        probs = 1 ./ (1 .+ exp.(-logits))
+
+        # Generate binomial outcomes
+        y = [rand(rng, Binomial(n_trials[i], probs[i])) for i in 1:n_samples]
+
+        return X, y, n_trials
+    end
+    n_samples = 1000
+    n_features = 2
+    true_beta = [-1.0, 0.6]
+    n_iterations = 100
+    n_sims = 20
+
+    @model function binomial_model(prior_xi, prior_precision, n_trials, X, y)
+        β ~ MvNormalWeightedMeanPrecision(prior_xi, prior_precision)
+        for i in eachindex(y)
+            y[i] ~ BinomialPolya(X[i], n_trials[i], β) where {dependencies = RequireMessageFunctionalDependencies(β = MvNormalWeightedMeanPrecision(prior_xi, prior_precision))}
+        end
+    end
+
+    function binomial_inference(binomial_model, iterations, X, y, n_trials, n_features)
+        return infer(
+            model = binomial_model(prior_xi = zeros(n_features), prior_precision = diageye(n_features)),
+            data = (X = X, y = y, n_trials = n_trials),
+            iterations = iterations,
+            free_energy = true,
+            options = (limit_stack_depth = 100,)
+        )
+    end
+
+    function run_simulation(n_sims::Int, n_samples::Int, true_beta::Vector{Float64}; iterations = n_iterations)
+        # Storage for results
+        n_features = length(true_beta)
+        coverage = Vector{Vector{Float64}}(undef, n_sims)
+        fes = Vector{Vector{Float64}}(undef, n_sims)
+        for sim in 1:n_sims
+            # Generate new dataset
+            X, y, n_trials = generate_synthetic_binomial_data(n_samples, true_beta, seed = sim)
+            X = [collect(row) for row in eachrow(X)]
+
+            # Run inference
+            results = binomial_inference(binomial_model, iterations, X, y, n_trials, n_features)
+            # Extract posterior parameters
+            post = results.posteriors[:β][end]
+            m = mean(post)
+            v = var(post)
+            estimates = map((x, y) -> Normal(x, sqrt(y)), m, v)
+            coverage[sim] = map((d, b) -> cdf(d, b), estimates, true_beta)
+            fes[sim] = results.free_energy
+        end
+
+        return coverage, fes
+    end
+
+    function in_credible_interval(x, lwr = 0.025, upr = 0.975)
+        return x >= lwr && x <= upr
+    end
+
+    coverage, fes = run_simulation(n_sims, n_samples, true_beta)
+    for i in 1:n_sims
+        @test fes[i][end] < fes[i][1]
+    end
+    coverages = Vector{Float64}(undef, n_features)
+    for i in 1:n_features
+        coverages[i] = sum(in_credible_interval.(getindex.(coverage, i))) / n_sims
+        @test coverages[i] >= 0.8
+    end
+
+    @test_benchmark "models" "binomialreg" binomial_inference(binomial_model, $n_iterations, $X, $y, $n_trials)
+end