Add auto-marginalization implementation to the main package

JuliaBUGS/.JuliaFormatter.toml

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

JuliaBUGS/src/BUGSExamples/Volume_2/08_Cervix.jl

@@ -47,6 +47,7 @@ model
 }
 """
 
+#! format: off
 data = (N = 2044,
     d = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -442,6 +443,7 @@ data = (N = 2044,
 
 inits = (beta0C = 0, beta = 0, q = 0.5, phi = [0.5 0.5; 0.5 0.5])
 inits_alternative = (beta0C = 1.0, beta = 1.0, q = 0.75, phi = [0.15 0.15; 0.15 0.15])
+#! format: on
 
 reference_results = (
     var"beta" = (mean = 0.6156, std = 0.3406),

JuliaBUGS/src/BUGSExamples/Volume_2/11_Schools.jl

@@ -70,6 +70,7 @@ model
 }
 """
 
+#! format: off
 data = (N = 1978, M = 38, mn = [0, 0, 0],
     prec = [0.0001 0 0; 0 0.0001 0; 0 0 0.0001],
     R = [0.1 0.005 0.005; 0.005 0.1 0.005; 0.005 0.005 0.1],
@@ -1573,6 +1574,7 @@ inits_alternative = (theta = 0.1, phi = 0, gamma = [1.0, 1.0, 1.0],
              0 0 0
              0 0 0
              0 0 0])
+#! format: on
 
 reference_results = (
     var"beta[1]" = (mean = 2.589E-4, std = 9.8E-5),

JuliaBUGS/src/BUGSExamples/Volume_2/16_Stagnant.jl

@@ -52,4 +52,4 @@ reference_results = (
 )
 
 stagnant = Example(
-    name, model_def, original, data, inits, inits_alternative, reference_results)
+    name, model_def, original, data, inits, inits_alternative, reference_results)

JuliaBUGS/src/BUGSExamples/Volume_2/18_Pigs.jl

@@ -0,0 +1 @@
+

JuliaBUGS/src/BUGSExamples/Volume_2/19_Simulating_data.jl

@@ -0,0 +1 @@
+

JuliaBUGS/src/BUGSExamples/Volume_3/09_Hips4.jl

@@ -92,7 +92,7 @@ model_def = @bugs begin
             for t in 2:N
                 for s in 1:S
                     var"pi"[
-                        n, k, t, s] = inprod(
+                    n, k, t, s] = inprod(
                         var"pi"[n, k, t - 1, :], Lambda[n, k, t, :, s])
                 end
             end

JuliaBUGS/src/BUGSExamples/Volume_4/11_Methadone.jl

@@ -51,7 +51,8 @@ end
 data_dict = JSON.parsefile(joinpath(readdir(), "methadone_data.json"))
 data = NamedTuple{Tuple([Symbol(key) for key in keys(data_dict)])}(Tuple([map(identity,
                                                                               val)
-                                                                          for val in values(data_dict)]))
+                                                                          for val in
+                                                                              values(data_dict)]))
 
 inits = (
     lambda = 0,

JuliaBUGS/src/JuliaBUGS.jl

@@ -42,7 +42,8 @@ using .Model:
     BUGSModel,
     evaluate_with_values!!,
     UseGraph,
-    UseGeneratedLogDensityFunction
+    UseGeneratedLogDensityFunction,
+    UseAutoMarginalization
 
 include("independent_mh.jl")
 include("gibbs.jl")

JuliaBUGS/src/model/Model.jl

@@ -9,6 +9,7 @@ using Graphs
 using LinearAlgebra
 using JuliaBUGS: JuliaBUGS, BUGSGraph
 using JuliaBUGS.BUGSPrimitives
+using LogExpFunctions
 using MetaGraphsNext
 using Random
 
@@ -20,5 +21,8 @@ include("logdensityproblems.jl")
 
 export parameters, variables, initialize!, getparams, settrans, set_evaluation_mode
 export evaluate_with_rng!!, evaluate_with_env!!, evaluate_with_values!!
+export evaluate_with_marginalization_rng!!,
+    evaluate_with_marginalization_env!!, evaluate_with_marginalization_values!!
+export UseAutoMarginalization, enumerate_discrete_values
 
 end # Model

JuliaBUGS/src/model/abstractppl.jl

@@ -569,7 +569,7 @@ function _create_modified_model(
     # Recompute mutable symbols for the new graph
     new_mutable_symbols = get_mutable_symbols(updated_graph_evaluation_data)
 
-    # Create the new model with all updated fields
+    # Create the new model with all updated fields (without auto-marg caches yet)
     kwargs = Dict{Symbol,Any}(
         :untransformed_param_length => new_untransformed_param_length,
         :transformed_param_length => new_transformed_param_length,
@@ -585,7 +585,33 @@ function _create_modified_model(
         kwargs[:base_model] = base_model
     end
 
-    return BUGSModel(model; kwargs...)
+    new_model = BUGSModel(model; kwargs...)
+
+    # Compute and attach auto-marg caches once for the new graph
+    try
+        order = JuliaBUGS.Model._compute_marginalization_order(new_model)
+        keys = JuliaBUGS.Model._precompute_minimal_cache_keys(new_model, order)
+
+        gd = new_model.graph_evaluation_data
+        gd_cached = GraphEvaluationData{
+            typeof(gd.node_function_vals),typeof(gd.loop_vars_vals)
+        }(
+            gd.sorted_nodes,
+            gd.sorted_parameters,
+            gd.is_stochastic_vals,
+            gd.is_observed_vals,
+            gd.node_function_vals,
+            gd.loop_vars_vals,
+            gd.node_types,
+            gd.is_discrete_finite_vals,
+            keys,
+            order,
+        )
+        return BUGSModel(new_model; graph_evaluation_data=gd_cached)
+    catch
+        # If caches cannot be computed (e.g., unsupported model), return model as-is.
+        return new_model
+    end
 end
 
 # Common helper function to regenerate log density function
@@ -698,8 +724,14 @@ function evaluate!!(
     temperature=1.0,
     transformed=model.transformed,
 )
-    evaluation_env, log_densities = evaluate_with_values!!(
-        model, flattened_values; temperature=temperature, transformed=transformed
-    )
+    if model.evaluation_mode isa UseAutoMarginalization
+        evaluation_env, log_densities = evaluate_with_marginalization_values!!(
+            model, flattened_values; temperature=temperature, transformed=transformed
+        )
+    else
+        evaluation_env, log_densities = evaluate_with_values!!(
+            model, flattened_values; temperature=temperature, transformed=transformed
+        )
+    end
     return evaluation_env, log_densities.tempered_logjoint
 end