Add an argument to condition function to skip generating and evaling log density function (#394)

## Motivation
`condition` currently regenerates the log-density function on every
call, invoking MacroTools transforms and `eval`. This is too slow for
hot loops.

We want a fast path for repeated conditioning that:
- Avoids regeneration
- Lets users update observed values in-place without changing graph
structure
- Allows explicit regeneration later to regain compiled performance when
desired

## Summary of Changes
- `condition` now supports a fast path via a keyword flag:
- `condition(model, conditioning_spec;
regenerate_log_density::Bool=true)`
- When `false`: skips regeneration, nulls the compiled function, and
forces graph evaluation mode
- New helper for value updates without reconditioning:
  - `set_observed_values!(model, obs::Dict{<:VarName,<:Any})`
- Updates evaluation environment values of already-observed stochastic
variables in-place
  - Validates existence, stochasticity, and observation status
- New function to regenerate the compiled log-density function without
changing the evaluation mode:
  - `regenerate_log_density_function(model; force=false)`
- Re-generates for the model’s current graph + environment and refreshes
graph evaluation data
- Exports added:
- `set_observed_values!` and `regenerate_log_density_function` (from
`Model` module)

## Tiny Benchmark for Sanity Check
with 
```julia
@bugs begin
    for i in 1:N
        x[i] ~ Normal(0, 1)
    end
    y ~ Normal(sum(x[:]), 1)
end
```

```
BenchmarkTools results (N=100, iters=200):
  Single condition: regular = 13.96 ms, fast = 0.15 ms, speedup ≈ 92.3x
  Loop conditioning: regular = 28139.5 ms, fast(update only) = 0.2 ms, speedup ≈ 112840.1x
```

Interpretation:
- Single call to `condition`: ~90x faster when skipping regeneration
- Hot loop: one fast `condition` + repeated `set_observed_values!` is
~1e5x faster than repeatedly reconditioning

Author: sunxd3

JuliaBUGS/Project.toml

@@ -1,6 +1,6 @@
 name = "JuliaBUGS"
 uuid = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
-version = "0.10.1"
+version = "0.10.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

JuliaBUGS/src/model/Model.jl

@@ -19,6 +19,7 @@ include("abstractppl.jl")
 include("logdensityproblems.jl")
 
 export parameters, variables, initialize!, getparams, settrans, set_evaluation_mode
+export regenerate_log_density_function, set_observed_values!
 export evaluate_with_rng!!, evaluate_with_env!!, evaluate_with_values!!
 
 end # Model

JuliaBUGS/src/model/abstractppl.jl

@@ -20,7 +20,7 @@ import AbstractPPL: condition, decondition, evaluate!!
 #######################
 
 """
-    condition(model::BUGSModel, conditioning_spec)
+    condition(model::BUGSModel, conditioning_spec; regenerate_log_density::Bool=true)
 
 Create a new model by conditioning on specified variables with given values.
 
@@ -135,7 +135,7 @@ julia> parameters(model_cond4)
  y
 ```
 """
-function condition(model::BUGSModel, conditioning_spec)
+function condition(model::BUGSModel, conditioning_spec; regenerate_log_density::Bool=true)
     # 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}
@@ -157,6 +157,7 @@ function condition(model::BUGSModel, conditioning_spec)
         new_graph,
         new_evaluation_env;
         base_model=isnothing(model.base_model) ? model : model.base_model,
+        regenerate_log_density=regenerate_log_density,
     )
 end
 
@@ -551,6 +552,7 @@ 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)
@@ -562,9 +564,15 @@ function _create_modified_model(
     )
 
     # Generate new log density function and update graph evaluation data
-    new_log_density_computation_function, updated_graph_evaluation_data = _regenerate_log_density_function(
-        model.model_def, new_graph, new_evaluation_env, new_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)
@@ -580,6 +588,12 @@ function _create_modified_model(
         :mutable_symbols => new_mutable_symbols,
     )
 
+    # 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
@@ -623,6 +637,62 @@ function _regenerate_log_density_function(
     end
 end
 
+#######################
+# Observed Value Updates
+#######################
+
+"""
+    set_observed_values!(model::BUGSModel, obs::Dict{<:VarName,<:Any})
+
+Update values of observed stochastic variables without reconditioning or regenerating code.
+
+Validates that each variable exists in the model, is stochastic, and is currently observed.
+Updates the evaluation environment in place and returns the updated model.
+"""
+function set_observed_values!(model::BUGSModel, obs::Dict{<:VarName,<:Any})
+    new_env = model.evaluation_env
+    for (vn, val) in obs
+        if vn ∉ labels(model.g)
+            throw(ArgumentError("Variable $vn does not exist in the model"))
+        end
+        node_info = model.g[vn]
+        if !node_info.is_stochastic
+            throw(ArgumentError("Cannot update $vn: it is deterministic (logical)"))
+        end
+        if !node_info.is_observed
+            throw(ArgumentError("Cannot update $vn: it is not observed"))
+        end
+        new_env = BangBang.setindex!!(new_env, val, vn)
+    end
+    return BangBang.setproperty!!(model, :evaluation_env, new_env)
+end
+
+"""
+    regenerate_log_density_function(model::BUGSModel; force::Bool=false)
+
+Generate and attach a compiled log-density function for the model's current graph and evaluation environment.
+
+Does not change the evaluation mode. When `force=false`, preserves an existing compiled function; when `force=true`,
+overwrites it if a new one can be generated. Returns the updated model (or the original if generation is not possible).
+"""
+function regenerate_log_density_function(model::BUGSModel; force::Bool=false)
+    new_fn, updated_graph_eval_data = _regenerate_log_density_function(
+        model.model_def, model.g, model.evaluation_env, model.graph_evaluation_data
+    )
+    # Always refresh graph_evaluation_data from regeneration helper (it may refine ordering)
+    model = BangBang.setproperty!!(model, :graph_evaluation_data, updated_graph_eval_data)
+
+    if isnothing(new_fn)
+        # Cannot generate compiled function; leave as-is
+        return model
+    end
+
+    if force || isnothing(model.log_density_computation_function)
+        model = BangBang.setproperty!!(model, :log_density_computation_function, new_fn)
+    end
+    return model
+end
+
 #######################
 # Evaluation API
 #######################

JuliaBUGS/test/model/abstractppl.jl

@@ -5,6 +5,8 @@ using JuliaBUGS.Model:
     decondition,
     parameters,
     set_evaluation_mode,
+    set_observed_values!,
+    regenerate_log_density_function,
     UseGeneratedLogDensityFunction,
     UseGraph
 using LogDensityProblems
@@ -149,6 +151,111 @@ JuliaBUGS.@bugs_primitive Normal Gamma
             logp = LogDensityProblems.logdensity(model_cond_gen, Float64[])
             @test isfinite(logp)
         end
+
+        @testset "Fast conditioning path and observed value updates" begin
+            @testset "Fast conditioning skips regeneration and forces graph mode" begin
+                model_def = @bugs begin
+                    x ~ Normal(0, 1)
+                    y ~ Normal(x, 1)
+                end
+
+                model = compile(model_def, (;))
+
+                # Regular conditioning (default regenerates)
+                model_reg = condition(model, Dict(@varname(x) => 1.0))
+                @test !isnothing(model_reg.log_density_computation_function)
+
+                # Fast conditioning (no regeneration)
+                model_fast = condition(
+                    model, Dict(@varname(x) => 1.0); regenerate_log_density=false
+                )
+                @test model_fast.log_density_computation_function === nothing
+                @test model_fast.evaluation_mode isa UseGraph
+
+                # Same parameters
+                @test parameters(model_reg) == parameters(model_fast)
+
+                # Compare log density via graph evaluation
+                params = zeros(length(parameters(model_fast)))
+                logp_fast = LogDensityProblems.logdensity(
+                    set_evaluation_mode(model_fast, UseGraph()), params
+                )
+                logp_reg = LogDensityProblems.logdensity(
+                    set_evaluation_mode(model_reg, UseGraph()), params
+                )
+                @test logp_fast ≈ logp_reg
+            end
+
+            @testset "set_observed_values! updates values and validates" begin
+                # Model with a deterministic node to hit validation
+                model_def = @bugs begin
+                    x ~ Normal(0, 1)
+                    y = x^2              # deterministic
+                    z ~ Normal(y, 1)
+                end
+
+                model = compile(model_def, (; z=2.0))
+
+                # Fast condition on x
+                m = condition(model, Dict(@varname(x) => 1.0); regenerate_log_density=false)
+                @test m.evaluation_mode isa UseGraph
+
+                # Update observed x value without reconditioning
+                m2 = set_observed_values!(m, Dict(@varname(x) => 2.0))
+                @test m2.evaluation_env.x == 2.0
+
+                # Structure unchanged (x observed; only z observed from data; no parameters)
+                # Here parameters is empty because y is deterministic and z is observed
+                @test parameters(m2) == parameters(m)
+
+                # Errors on invalid updates
+                @test_throws ArgumentError set_observed_values!(
+                    m2, Dict(@varname(y) => 3.0)
+                )  # deterministic
+
+                # Updating originally observed data should be allowed
+                m3 = set_observed_values!(m2, Dict(@varname(z) => 1.0))
+                @test m3.evaluation_env.z == 1.0
+                # To test non-observed error, try updating a parameter in a different model.
+            end
+
+            @testset "set_observed_values! errors on non-observed variables" begin
+                model_def = @bugs begin
+                    x ~ Normal(0, 1)
+                    y ~ Normal(x, 1)
+                end
+                model = compile(model_def, (;))
+                m = condition(model, Dict(@varname(x) => 1.0); regenerate_log_density=false)
+                # y is not observed
+                @test_throws ArgumentError set_observed_values!(m, Dict(@varname(y) => 0.0))
+            end
+
+            @testset "Regeneration after fast conditioning (no mode change)" begin
+                model_def = @bugs begin
+                    x ~ Normal(0, 1)
+                    y ~ Normal(x, 1)
+                end
+
+                model = compile(model_def, (;))
+                m = condition(model, Dict(@varname(x) => 1.0); regenerate_log_density=false)
+                @test m.log_density_computation_function === nothing
+
+                # Regenerate compiled function without changing mode
+                m2 = regenerate_log_density_function(m)
+                @test !isnothing(m2.log_density_computation_function)
+                @test m2.evaluation_mode isa UseGraph
+
+                # Can switch to generated mode explicitly and match graph
+                params = zeros(length(parameters(m2)))
+                logp_gen = LogDensityProblems.logdensity(
+                    set_evaluation_mode(m2, UseGeneratedLogDensityFunction()), params
+                )
+                logp_graph = LogDensityProblems.logdensity(
+                    set_evaluation_mode(m2, UseGraph()), params
+                )
+                @test logp_gen ≈ logp_graph
+            end
+        end
     end
 
     @testset "decondition" begin