fix performance by moving more computation to the construction

Author: sunxd3

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

JuliaBUGS/src/model/bugsmodel.jl

@@ -31,9 +31,8 @@ end
 Return the finite support for a discrete univariate distribution.
 Relies on Distributions.support to provide an iterable, finite range.
 """
-enumerate_discrete_values(dist::Distributions.DiscreteUnivariateDistribution) = Distributions.support(
-    dist
-)
+enumerate_discrete_values(dist::Distributions.DiscreteUnivariateDistribution) =
+    Distributions.support(dist)
 
 """
     classify_node_type(dist)
@@ -77,8 +76,15 @@ struct GraphEvaluationData{TNF,TV}
     node_types::Vector{Symbol}
     is_discrete_finite_vals::Vector{Bool}
     minimal_cache_keys::Dict{Int,Vector{Int}}
+    marginalization_order::Vector{Int}
 end
 
+"""
+    GraphEvaluationData(compat constructor)
+
+Backward-compatible constructor that fills new caching fields with defaults
+when older call sites provide only the first nine fields.
+"""
 function GraphEvaluationData(
     g::BUGSGraph,
     sorted_nodes::Vector{<:VarName}=VarName[
@@ -114,7 +120,7 @@ function GraphEvaluationData(
         end
     end
 
-    return GraphEvaluationData(
+    return GraphEvaluationData{typeof(node_function_vals),typeof(loop_vars_vals)}(
         sorted_nodes,
         sorted_parameters,
         is_stochastic_vals,
@@ -124,6 +130,7 @@ function GraphEvaluationData(
         node_types,
         is_discrete_finite_vals,
         Dict{Int,Vector{Int}}(),
+        Int[],
     )
 end
 
@@ -233,7 +240,7 @@ function BUGSModel(
         JuliaBUGS.Model._precompute_minimal_cache_keys(m, collect(1:n))
     end
     # Attach minimal cache keys to GraphEvaluationData (order remains default)
-    gd2 = GraphEvaluationData(
+    gd2 = GraphEvaluationData{typeof(gd.node_function_vals),typeof(gd.loop_vars_vals)}(
         gd.sorted_nodes,
         gd.sorted_parameters,
         gd.is_stochastic_vals,
@@ -243,6 +250,7 @@ function BUGSModel(
         gd.node_types,
         gd.is_discrete_finite_vals,
         minimal_keys,
+        gd.marginalization_order,
     )
     # Return final model with cached minimal keys
     return BUGSModel(
@@ -368,7 +376,10 @@ function BUGSModel(
     end
 
     # Update graph_evaluation_data with the computed node types
-    graph_evaluation_data = GraphEvaluationData(
+    graph_evaluation_data = GraphEvaluationData{
+        typeof(graph_evaluation_data.node_function_vals),
+        typeof(graph_evaluation_data.loop_vars_vals),
+    }(
         graph_evaluation_data.sorted_nodes,
         graph_evaluation_data.sorted_parameters,
         graph_evaluation_data.is_stochastic_vals,
@@ -378,6 +389,7 @@ function BUGSModel(
         node_types,
         is_discrete_finite_vals,
         Dict{Int,Vector{Int}}(),
+        Int[],
     )
 
     lowered_model_def, reconstructed_model_def = JuliaBUGS._generate_lowered_model_def(
@@ -429,7 +441,9 @@ function BUGSModel(
         end
 
         # Reconstruct GraphEvaluationData while preserving classification
-        graph_evaluation_data = GraphEvaluationData(
+        graph_evaluation_data = GraphEvaluationData{
+            typeof(new_gd.node_function_vals),typeof(new_gd.loop_vars_vals)
+        }(
             new_gd.sorted_nodes,
             new_gd.sorted_parameters,
             new_gd.is_stochastic_vals,
@@ -439,6 +453,7 @@ function BUGSModel(
             new_node_types,
             new_is_discrete_finite_vals,
             Dict{Int,Vector{Int}}(),
+            Int[],
         )
     else
         log_density_computation_function = nothing
@@ -464,14 +479,17 @@ function BUGSModel(
         mutable_symbols,
         nothing,
     )
-    # Precompute minimal cache keys for the default order (1:n)
+    # Precompute marginalization order and minimal cache keys once
     n = length(graph_evaluation_data.sorted_nodes)
-    sorted_indices = collect(1:n)
+    sorted_indices = JuliaBUGS.Model._compute_marginalization_order(model_without_min_keys)
     minimal_keys = JuliaBUGS.Model._precompute_minimal_cache_keys(
         model_without_min_keys, sorted_indices
     )
-    # Attach minimal keys to GraphEvaluationData
-    graph_evaluation_data_with_keys = GraphEvaluationData(
+    # Attach cached order and keys to GraphEvaluationData
+    graph_evaluation_data_with_keys = GraphEvaluationData{
+        typeof(graph_evaluation_data.node_function_vals),
+        typeof(graph_evaluation_data.loop_vars_vals),
+    }(
         graph_evaluation_data.sorted_nodes,
         graph_evaluation_data.sorted_parameters,
         graph_evaluation_data.is_stochastic_vals,
@@ -481,6 +499,7 @@ function BUGSModel(
         graph_evaluation_data.node_types,
         graph_evaluation_data.is_discrete_finite_vals,
         minimal_keys,
+        sorted_indices,
     )
 
     # Return final model with cached minimal keys

JuliaBUGS/src/model/evaluation.jl

@@ -731,14 +731,17 @@ function evaluate_with_marginalization_values!!(
         )
     end
 
-    # Compute an order that minimizes frontier growth (interleave discrete parents before observed children)
+    # Use cached marginalization order and minimal frontier keys when available
     gd = model.graph_evaluation_data
     n = length(gd.sorted_nodes)
-    sorted_indices = _compute_marginalization_order(model)
-
-    # Compute minimal cache keys for this specific order
-    # (do not reuse cached keys if they were built for a different order)
-    minimal_keys = _precompute_minimal_cache_keys(model, sorted_indices)
+    # Strictly require caches to be present for performance
+    if isempty(gd.marginalization_order) || isempty(gd.minimal_cache_keys)
+        error(
+            "Auto marginalization cache missing. This model was not prepared for UseAutoMarginalization.",
+        )
+    end
+    sorted_indices = gd.marginalization_order
+    minimal_keys = gd.minimal_cache_keys
 
     # Initialize memoization cache
     # Size hint: at most 2^|discrete_finite| * |nodes| entries