Make log density function generation on-demand (#416)

This PR introduces on-demand generation of log density functions for
JuliaBUGS models.

The `skip_source_generation` parameter has been removed from `compile()`
and `BUGSModel()`. Instead, log density functions are now generated
on-demand when `set_evaluation_mode(model,
UseGeneratedLogDensityFunction())` is called. All models start with
`UseGraph()` mode.

---------

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

Author: sunxd3

JuliaBUGS/benchmark/juliabugs.jl

@@ -74,7 +74,8 @@ end
 # ! writing a _function_ to benchmark all models won't work because of world-age error
 
 function benchmark_JuliaBUGS_model_with_Mooncake(model::JuliaBUGS.BUGSModel)
-    # p = Base.Fix1(LogDensityProblems.logdensity, model)
+    # Use generated log density function for Mooncake
+    model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGeneratedLogDensityFunction())
     p = Base.Fix1(model.log_density_computation_function, model.evaluation_env)
     backend = AutoMooncake(; config=nothing)
     dim = LogDensityProblems.dimension(model)

JuliaBUGS/benchmark/run_benchmarks.jl

@@ -29,7 +29,8 @@ output_file = get(ENV, "BENCHMARK_OUTPUT", "benchmark_results.md")
 open(output_file, "w") do io
     println(io, "## Benchmark Results\n")
     cpu_info = first(Sys.cpu_info())
-    println(io, "**Julia $(VERSION)** on $(cpu_info.model)\n")
+    os_info = Sys.KERNEL
+    println(io, "**Julia $(VERSION)** on $(cpu_info.model) ($(os_info))\n")
     println(io, "Ratio = JuliaBUGS/Stan (lower is better for JuliaBUGS)\n")
     println(io, "| Model | Stan Params | JBUGS Params | LD Ratio | Grad Ratio |")
     println(io, "|:------|------------:|-------------:|---------:|-----------:|")
@@ -52,6 +53,6 @@ open(output_file, "w") do io
     end
     println(
         io,
-        "\n*Note: Performance comparison may not be apples-to-apples as parameter counts can differ due to different model parameterizations.*",
+        "\n*Note: Stan benchmarks use hand-optimized Stan models, not direct BUGS translations. Comparison is illustrative only.*",
     )
 end

JuliaBUGS/examples/gp.jl

@@ -120,8 +120,8 @@ model = gp_golf_putting(
     data.jitter, # Numerical stability term
 )
 
-# Optionally, set the evaluation mode. Using generated functions can be faster.
-# model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGeneratedLogDensityFunction())
+# Generate the log density function for optimal performance
+model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGeneratedLogDensityFunction())
 
 # --- MCMC Setup with Custom LogDensityProblems Wrapper ---
 

JuliaBUGS/src/JuliaBUGS.jl

@@ -239,25 +239,22 @@ function validate_bugs_expression(expr, line_num)
 end
 
 """
-    compile(model_def, data[, initial_params]; skip_validation=false, skip_source_generation=false)
+    compile(model_def, data[, initial_params]; skip_validation=false)
 
 Compile the model with model definition and data. Optionally, initializations can be provided.
 If initializations are not provided, values will be sampled from the prior distributions.
 
 By default, validates that all functions in the model are in the BUGS allowlist (suitable for @bugs macro).
 Set `skip_validation=true` to skip validation (for @model macro usage).
 
-Set `skip_source_generation=true` to bypass generating optimized log-density functions. This produces
-models with stable type signatures (no anonymous function types in type parameters), which is essential
-for serialization. When enabled, the model will use graph traversal for evaluation (UseGraph mode),
-which may be slower but ensures type stability across sessions.
+The compiled model uses `UseGraph` evaluation mode by default. To use the optimized generated
+log-density function, call `set_evaluation_mode(model, UseGeneratedLogDensityFunction())`.
 """
 function compile(
     model_def::Expr,
     data::NamedTuple,
     initial_params::NamedTuple=NamedTuple();
     skip_validation::Bool=false,
-    skip_source_generation::Bool=false,
     eval_module::Module=@__MODULE__,
 )
     # Validate functions by default (for @bugs macro usage)
@@ -287,15 +284,7 @@ function compile(
             values(eval_env),
         ),
     )
-    return BUGSModel(
-        g,
-        nonmissing_eval_env,
-        model_def,
-        data,
-        initial_params,
-        true,
-        skip_source_generation,
-    )
+    return BUGSModel(g, nonmissing_eval_env, model_def, data, initial_params, true)
 end
 # function compile(
 #     model_str::String,

JuliaBUGS/src/gibbs.jl

@@ -215,13 +215,13 @@ julia> using JuliaBUGS: expand_variables, @varname
 julia> model_parameters = [@varname(x[1]), @varname(x[2]), @varname(x[3]), @varname(y)];
 
 julia> expand_variables([@varname(x)], model_parameters)
-3-element Vector{AbstractPPL.VarName}:
+3-element Vector{VarName}:
  x[1]
  x[2]
  x[3]
 
 julia> expand_variables([@varname(x[1]), @varname(y)], model_parameters)
-2-element Vector{AbstractPPL.VarName}:
+2-element Vector{VarName}:
  x[1]
  y
 ```

JuliaBUGS/src/model/abstractppl.jl

@@ -20,9 +20,11 @@ import AbstractPPL: condition, decondition, evaluate!!
 #######################
 
 """
-    condition(model::BUGSModel, conditioning_spec; regenerate_log_density::Bool=true)
+    condition(model::BUGSModel, conditioning_spec)
 
 Create a new model by conditioning on specified variables with given values.
+The returned model uses `UseGraph` evaluation mode. To use optimized log density
+computation, call `set_evaluation_mode(model, UseGeneratedLogDensityFunction())`.
 
 # Arguments
 - `model::BUGSModel`: The model to condition
@@ -45,12 +47,6 @@ New `BUGSModel` with:
 
 # Examples
 ```jldoctest condition
-julia> using JuliaBUGS: @bugs, compile, @varname, initialize!
-
-julia> using JuliaBUGS.Model: condition, parameters
-
-julia> using Test
-
 julia> model_def = @bugs begin
            for i in 1:3
                x[i] ~ Normal(0, 1)
@@ -69,7 +65,7 @@ julia> model_cond.evaluation_env.x[1:2]
  2.0
 
 julia> parameters(model_cond)
-2-element Vector{AbstractPPL.VarName}:
+2-element Vector{VarName}:
  x[3]
  y
 
@@ -86,7 +82,7 @@ julia> model_cond2.evaluation_env.x
  7.0
 
 julia> parameters(model_cond2)  # All x[i] removed, only y remains
-1-element Vector{AbstractPPL.VarName}:
+1-element Vector{VarName}:
  y
 
 julia> # Check parameter lengths
@@ -105,8 +101,8 @@ julia> # NamedTuple syntax
 julia> model_cond3.evaluation_env.y
 10.0
 
-julia> parameters(model_cond3)  # y removed, only x[i] remain
-3-element Vector{AbstractPPL.VarName}:
+julia> sort(parameters(model_cond3); by=string)  # y removed, only x[i] remain
+3-element Vector{VarName}:
  x[1]
  x[2]
  x[3]
@@ -130,12 +126,12 @@ julia> model_cond4.evaluation_env.x[[1, 3]]
  3.0
 
 julia> parameters(model_cond4)
-2-element Vector{AbstractPPL.VarName}:
+2-element Vector{VarName}:
  x[2]
  y
 ```
 """
-function condition(model::BUGSModel, conditioning_spec; regenerate_log_density::Bool=true)
+function condition(model::BUGSModel, conditioning_spec)
     # Parse and validate conditioning specification
     var_values = _parse_conditioning_spec(conditioning_spec, model)::Dict{<:VarName,<:Any}
     vars_to_condition = collect(keys(var_values))::Vector{<:VarName}
@@ -152,12 +148,12 @@ function condition(model::BUGSModel, conditioning_spec; regenerate_log_density::
     new_graph = _mark_as_observed(model.g, expanded_vars)
 
     # Create updated model with conditioned variables
+    # Log density function will be generated when set_evaluation_mode is called with UseGeneratedLogDensityFunction
     return _create_modified_model(
         model,
         new_graph,
         new_evaluation_env;
         base_model=isnothing(model.base_model) ? model : model.base_model,
-        regenerate_log_density=regenerate_log_density,
     )
 end
 
@@ -256,14 +252,6 @@ For base_model restoration (no args):
 
 # Examples
 ```jldoctest decondition
-julia> using JuliaBUGS: @bugs, compile
-
-julia> using JuliaBUGS.Model: condition, parameters, decondition
-
-julia> using AbstractPPL: @varname
-
-julia> using Test
-
 julia> model_def = @bugs begin
            x ~ dnorm(0, 1)
            y ~ dnorm(x, 1) 
@@ -276,13 +264,13 @@ julia> # Condition model
        model_cond = condition(model, (; x = 1.0, y = 1.5));
 
 julia> parameters(model_cond)
-AbstractPPL.VarName[]
+VarName[]
 
 julia> # Partial deconditioning with specified variables
        model_d1 = decondition(model_cond, [@varname(y)]);
 
 julia> parameters(model_d1)
-1-element Vector{AbstractPPL.VarName}:
+1-element Vector{VarName}:
  y
 
 julia> # Full restoration to base model (no arguments)
@@ -351,8 +339,8 @@ julia> # Decondition with subsumption
            decondition(model_arr_cond, [@varname(v)])
        );
 
-julia> parameters(model_arr_decon)
-3-element Vector{AbstractPPL.VarName}:
+julia> sort(parameters(model_arr_decon); by=string)
+3-element Vector{VarName}:
  v[1]
  v[2]
  v[3]
@@ -552,7 +540,6 @@ function _create_modified_model(
     new_graph::BUGSGraph,
     new_evaluation_env::NamedTuple;
     base_model=nothing,
-    regenerate_log_density::Bool=true,
 )
     # Create new graph evaluation data
     new_graph_evaluation_data = GraphEvaluationData(new_graph)
@@ -563,42 +550,29 @@ function _create_modified_model(
         model, new_parameters
     )
 
-    # Generate new log density function and update graph evaluation data
-    new_log_density_computation_function, updated_graph_evaluation_data =
-        if regenerate_log_density
-            _regenerate_log_density_function(
-                model.model_def, new_graph, new_evaluation_env, new_graph_evaluation_data
-            )
-        else
-            # Skip regeneration (fast path): ensure stale code isn't used
-            nothing, new_graph_evaluation_data
-        end
-
     # Recompute mutable symbols for the new graph
-    new_mutable_symbols = get_mutable_symbols(updated_graph_evaluation_data)
+    new_mutable_symbols = get_mutable_symbols(new_graph_evaluation_data)
 
     # Create the new model with all updated fields
+    # Log density function is NOT generated here - it will be generated on-demand
+    # when set_evaluation_mode(model, UseGeneratedLogDensityFunction()) is called
     kwargs = Dict{Symbol,Any}(
         :untransformed_param_length => new_untransformed_param_length,
         :transformed_param_length => new_transformed_param_length,
         :evaluation_env => new_evaluation_env,
-        :graph_evaluation_data => updated_graph_evaluation_data,
+        :graph_evaluation_data => new_graph_evaluation_data,
         :g => new_graph,
-        :log_density_computation_function => new_log_density_computation_function,
+        :log_density_computation_function => nothing,
         :mutable_symbols => new_mutable_symbols,
+        :evaluation_mode => UseGraph(),
     )
 
-    # Force graph evaluation mode when skipping regeneration to avoid stale compiled code
-    if !regenerate_log_density
-        kwargs[:evaluation_mode] = UseGraph()
-        kwargs[:log_density_computation_function] = nothing
-    end
-
     # Add base_model if provided
     if !isnothing(base_model)
         kwargs[:base_model] = base_model
     end
 
+    # Return model without precomputing log density function (on-demand generation)
     return BUGSModel(model; kwargs...)
 end