Merge pull request #241 from ReactiveBayes/dev-stack-constraints

Do not use the plus for stacked constraints

Author: bvdmitri

docs/src/plugins/constraint_specification.md

@@ -69,6 +69,19 @@ We can specify constraints over the first `toy_model` submodel using the followi
 end
 ```
 
+## Stacked functional form constraints
+In the constraints macro, we can specify multiple functional form constraints over the same variable. For example, suppose we have the following model:
+```@example constraints
+@constraints begin 
+    q(x) :: Normal :: Beta
+end
+```
+In this constraint the posterior over `x` will first be constrained to be a normal distribution, and then the result with be constrained to be a beta distribution.
+This might be useful to create a chain of constraints that are applied in order. The resulting constraint is a tuple of constraints.
+
+!!! note 
+    The inference backend must support stacked constraints for this feature to work. Some combinations of stacked constraints might not be supported or theoretically sound.
+
 ## Default constraints
 While we can specify constraints over all instances of a submodel at a specific layer of the hierarchy, we're not guaranteed to have all instances of a submodel at a specific layer of the hierarchy. To this extent, we can specify default constraints that apply to all instances of a specific submodel. For example, we can define the following model, where we have a `recursive_model` instance at every layer of the hierarchy:
 ```@example constraints

src/plugins/variational_constraints/variational_constraints_macro.jl

@@ -66,9 +66,14 @@ function create_submodel_constraints(e::Expr)
     end
 end
 
+stack_constraints(a, b) = (a, b)
+stack_constraints(a, b::Tuple) = (a, b...)
+stack_constraints(a::Tuple, b) = (a..., b)
+stack_constraints(a::Tuple, b::Tuple) = (a..., b...)
+
 function rewrite_stacked_constraints(e::Expr)
     if @capture(e, (a_::b_::c_))
-        return :($a::($b + $c))
+        return :($a::(GraphPPL.stack_constraints($b, $c)))
     else
         return e
     end

test/plugins/variational_constraints/variational_constraints_macro_tests.jl

@@ -170,7 +170,7 @@ end
     end
     output = quote
         q(x, y) = q(x)q(y)
-        q(x)::(PointMass() + SampleList())
+        q(x)::GraphPPL.stack_constraints(PointMass(), SampleList())
     end
     @test_expression_generating apply_pipeline(input, rewrite_stacked_constraints) output
 
@@ -181,11 +181,20 @@ end
     end
     output = quote
         q(x, y) = q(x)q(y)
-        q(x)::((PointMass() + Sample) + SampleList())
+        q(x)::GraphPPL.stack_constraints(GraphPPL.stack_constraints(PointMass(), Sample), SampleList())
     end
     @test_expression_generating apply_pipeline(input, rewrite_stacked_constraints) output
 end
 
+@testitem "stack_constraints" begin 
+    import GraphPPL: stack_constraints
+
+    @test stack_constraints(1, 2) == (1, 2)
+    @test stack_constraints(1, (2, 1)) == (1, 2, 1)
+    @test stack_constraints((1, 2), 3) == (1, 2, 3)
+    @test stack_constraints((1, 3), (2, 1)) == (1, 3, 2, 1)
+end
+
 @testitem "replace_begin_end" begin
     import GraphPPL: replace_begin_end, apply_pipeline