fix: support Colon() in middle of indices for NodeLabel (#275)

CHANGELOG.md

@@ -0,0 +1,8 @@
+\## \[Unreleased]
+
+
+
+\### Fixed
+
+\- Fixed `Colon()` not working in the middle of indices for `NodeLabel` getindex (\[#303](https://github.com/ReactiveBayes/GraphPPL.jl/pull/303))
+

src/graph_engine.jl

@@ -266,6 +266,7 @@ end
 Base.length(label::NodeLabel) = 1
 Base.size(label::NodeLabel) = ()
 Base.getindex(label::NodeLabel, any) = label
+Base.getindex(label::NodeLabel, i1, is...) = label
 Base.:(<)(left::NodeLabel, right::NodeLabel) = left.global_counter < right.global_counter
 Base.broadcastable(label::NodeLabel) = Ref(label)
 

src/plugins/variational_constraints/variational_constraints_engine.jl

@@ -646,7 +646,7 @@ function materialize_constraints!(model::Model, node_label::NodeLabel, node_data
 
     if !is_valid_partition(constraint_set)
         error(
-            lazy"Factorization constraint set at node $node_label is not a valid constraint set. Please check your model definition and constraint specification. (Constraint set: $constraint_bitset)"
+            lazy"Factorization constraint set at node $node_label is not a valid constraint set. Please check your model definition and constraint specification. (Constraint set: $constraint_bitset)\n\nHint: A valid constraint set requires that each interface belongs to exactly one factorization group. In the constraint matrix, each column must have exactly one non-zero entry per group. Make sure all variables appearing in your model are included in your factorization constraint."
         )
     end
 

test/plugins/variational_constraints/variational_constraints_engine_tests.jl

@@ -1596,3 +1596,42 @@ end
     @test map(rhs, collect(stack)) == [:a]
 end
 
+@testitem "invalid factorization constraint set throws a descriptive error with a hint" begin
+    using BitSetTuples
+    import GraphPPL:
+        create_model,
+        with_plugins,
+        materialize_constraints!,
+        getcontext,
+        setextra!,
+        VariationalConstraintsPlugin,
+        PluginsCollection,
+        VariationalConstraintsFactorizationBitSetKey
+
+    include("../../testutils.jl")
+
+    using .TestUtils.ModelZoo
+
+    # A constraint set where the same interface appears in more than one group is invalid:
+    # `((1,), (3,), (1, 3))` puts interface `1` in both the first and the third group.
+    model = create_model(with_plugins(simple_model(), PluginsCollection(VariationalConstraintsPlugin())))
+    ctx = getcontext(model)
+    node = ctx[NormalMeanVariance, 2]
+
+    setextra!(model[node], VariationalConstraintsFactorizationBitSetKey, BoundedBitSetTuple(((1,), (3,), (1, 3))))
+
+    error = try
+        materialize_constraints!(model, node)
+        nothing
+    catch e
+        e
+    end
+
+    @test error !== nothing
+    @test error isa ErrorException
+    message = sprint(showerror, error)
+    @test occursin("not a valid constraint set", message)
+    @test occursin("valid constraint set requires", message)
+    @test occursin("each interface belongs to exactly one factorization group", message)
+    @test occursin("Hint", message)
+end

test/test_colon_index.jl

@@ -0,0 +1,17 @@
+using Test
+using GraphPPL
+using Distributions
+
+@model function test_colon_bug(media)
+    for g in 1:2
+        for i in 1:3
+            media_effect[g, i] := media[g, :, i]
+        end
+    end
+end
+
+@testset "Colon index in middle of indices" begin
+    media_data = rand(2, 5, 3)
+    model = create_model(test_colon_bug(media = media_data))
+    @test model !== nothing
+end