Ensure only model parameters are included in the sorted parameter list (#292)

If a Julia code to compute log density can be generated for a model,
then the topo order of nodes in the model need to be updated to sync
with the generated Julia source.

The original code used the topologically sorted list of all nodes
(`pass.sorted_nodes`) to determine the final ordered list of parameters
for the `Model`. However, `pass.sorted_nodes` can include variables that
are purely "transformed data" - deterministic nodes computed entirely
from provided data constants at compile time (as described in
`docs/src/source_gen.md` under "Handling Mixed Data Transformation and
Deterministic Assignments"). For example, if `y = [1, 2, missing,
missing, 2]` and the model has `x[i] = y[i] + 1`, then `x[1]`, `x[2]`,
and `x[5]` are transformed data, while only `x[3]` and `x[4]` (dependent
on the missing `y` values) are actual deterministic parameters within
the model graph. Including transformed data nodes in the final
`parameters` list is incorrect, as they are constants and not part of
the model's parameter space to be evaluated or sampled.

This PR is a simple fix by filtering the `pass.sorted_nodes` list,
keeping only those `VarName`s that are already present in the
`parameters` set. This ensures the final `parameters` field of the
`Model` struct contains only true model parameters, correctly ordered
according to the graph's topology.

---------

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

Author: sunxd3

.github/workflows/TestsMacOS.yml

@@ -12,11 +12,11 @@ jobs:
   test:
     name: Julia ${{ matrix.version }} on ${{ matrix.os }} (${{ matrix.arch }})
     runs-on: ${{ matrix.os }}
-    continue-on-error: ${{ matrix.version == 'nightly' }}
+    continue-on-error: ${{ matrix.version == 'pre' }}
     strategy:
       fail-fast: false
       matrix:
-        version: ['1', '1.10', 'nightly']
+        version: ['1', 'min', 'pre']
         os: [macOS-latest]
         arch: [aarch64]
 
@@ -55,6 +55,11 @@ jobs:
         env:
           TEST_GROUP: "log_density"
 
+      - name: Running `source_gen` tests
+        uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "source_gen"
+
       - name: Running `gibbs` tests
         uses: nick-fields/retry@v3
         with:
@@ -77,17 +82,4 @@ jobs:
         uses: julia-actions/julia-runtest@v1
         env:
           TEST_GROUP: "experimental"
-      
-      - uses: julia-actions/julia-processcoverage@v1
-        if: matrix.coverage
 
-      - uses: codecov/codecov-action@v4
-        if: matrix.coverage
-        with:
-          file: lcov.info
-      
-      - uses: coverallsapp/github-action@master
-        if: matrix.coverage
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          path-to-lcov: lcov.info

docs/src/source_gen.md

@@ -457,6 +457,8 @@ end
 
 We made a simple change to the program to prepare for lowering: we need to distinguish between observations and model parameters (because they correspond to different code). We introduce a new operator into the program `\eqsim` to indicate that the left hand side is an observation.
 
+### Handling Mixed Observations and Parameters
+
 BUGS supports mixing observations and model parameters for different elements of the same array variable.
 To support this, we introduce a guard to use conditional logic to decide what computation to do for different iteration of the same statement.
 
@@ -487,3 +489,27 @@ begin
     end
 end
 ```
+
+### Handling Mixed Data Transformation and Deterministic Assignments
+
+For instance
+
+```julia
+for i in 1:5
+    x[i] = y[i] + 1
+end
+```
+
+if the data is
+
+```julia
+y = [1, 2, missing, missing, 2]
+```
+
+this is generally allowed in BUGS.
+
+`x[1], x[2], x[5]` can be computed at compile time, so these are "transformed data".
+`x[3], x[4]` need to be computed at evaluation time.
+And only `x[3]` and `x[4]` are in the compiled graph.
+
+For generated Julia program, if a statement can be eliminated because all the variables stemmed from this statement are "transformed data". While in the above case, where a statements corresponds to both transformed data and deterministic variables. It will be left in the generated program as is. In this case, there will be redundant computation.

src/model.jl

@@ -233,6 +233,9 @@ function BUGSModel(
         sorted_nodes = pass.sorted_nodes
         original_parameters_length = length(parameters)
         parameters = VarName[vn for vn in sorted_nodes if vn in parameters]
+        sorted_nodes = [
+            vn for vn in sorted_nodes if vn in flattened_graph_node_data.sorted_nodes
+        ]
         @assert length(parameters) == original_parameters_length "there are less parameters in the generated log density function than in the original model"
         flattened_graph_node_data = FlattenedGraphNodeData(g, sorted_nodes)
     else

test/source_gen.jl

@@ -1,13 +1,10 @@
 using BangBang
 using Bijectors
 using JuliaBUGS
-using JuliaBUGS: _generate_lowered_model_def, _gen_log_density_computation_function_expr
-using JuliaBUGS: CollectSortedNodes
 using JuliaBUGS.BUGSPrimitives
 using LogDensityProblems
 using OrderedCollections
 
-# bones, mice, kidney have missings in data
 test_examples = [
     :rats,
     :pumps,
@@ -32,52 +29,21 @@ test_examples = [
     :air,
     :birats,
     :schools,
+    :cervix,
 ]
 
-function _create_model(model_name::Symbol)
-    (; model_def, data, inits) = getfield(JuliaBUGS.BUGSExamples, model_name)
-    model = compile(model_def, data, inits)
-    evaluation_env = deepcopy(model.evaluation_env)
-    return model, evaluation_env
-end
-
-function _create_bugsmodel_with_consistent_sorted_nodes(
-    model::JuliaBUGS.BUGSModel, reconstructed_model_def
-)
-    pass = CollectSortedNodes(model.evaluation_env)
-    JuliaBUGS.analyze_block(pass, reconstructed_model_def)
-    sorted_nodes = pass.sorted_nodes
-    sorted_parameters = [vn for vn in sorted_nodes if vn in model.parameters]
-    new_flattened_graph_node_data = JuliaBUGS.FlattenedGraphNodeData(model.g, sorted_nodes)
-    new_model = BangBang.setproperty!!(model, :parameters, sorted_parameters)
-    new_model = BangBang.setproperty!!(
-        new_model, :flattened_graph_node_data, new_flattened_graph_node_data
-    )
-    return new_model
-end
-
 @testset "source_gen: $example_name" for example_name in test_examples
-    model, evaluation_env = _create_model(example_name)
-    lowered_model_def, reconstructed_model_def = _generate_lowered_model_def(
-        model.model_def, model.g, evaluation_env
-    )
-    log_density_computation_expr = _gen_log_density_computation_function_expr(
-        lowered_model_def, evaluation_env, gensym(example_name)
-    )
-    log_density_computation_function = eval(log_density_computation_expr)
-
-    model_with_consistent_sorted_nodes = _create_bugsmodel_with_consistent_sorted_nodes(
-        model, reconstructed_model_def
-    )
-    result_with_old_model = JuliaBUGS.evaluate!!(model)[2]
-    params = JuliaBUGS.getparams(model_with_consistent_sorted_nodes)
-    result_with_bugsmodel = JuliaBUGS.evaluate!!(
-        model_with_consistent_sorted_nodes, params
-    )[2]
-    result_with_log_density_computation_function = log_density_computation_function(
-        evaluation_env, params
-    )
-    @test result_with_old_model ≈ result_with_bugsmodel
+    (; model_def, data, inits) = getfield(JuliaBUGS.BUGSExamples, example_name)
+    model = compile(model_def, data, inits)
+    params = JuliaBUGS.getparams(model)
+    result_with_bugsmodel = begin
+        model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGraph())
+        LogDensityProblems.logdensity(model, params)
+    end
+    result_with_log_density_computation_function = begin
+        model = JuliaBUGS.set_evaluation_mode(model, JuliaBUGS.UseGeneratedLogDensityFunction())
+        LogDensityProblems.logdensity(model, params)
+    end
     @test result_with_log_density_computation_function ≈ result_with_bugsmodel
 end
 
@@ -88,3 +54,17 @@ end
         end
     )
 end
+
+@testset "mixed data transformation and deterministic assignments" begin
+    model_def = @bugs begin
+        for i in 1:5
+            y[i] ~ Normal(0, 1)
+        end
+        for i in 1:5
+            x[i] = y[i] + 1
+        end
+    end
+    data = (; y=[1, 2, missing, missing, 2])
+
+    model = compile(model_def, data)
+end