Merge pull request #13 from biaslab/dev-broadcasting-syntax-v2

feat: add broadcasting syntax

Author: bvdmitri

docs/src/user-guide.md

@@ -177,6 +177,64 @@ node, y ~ NormalMeanVariance(mean, var)
 
 Having a node reference can be useful in case the user wants to return it from a model and to use it later on to specify initial joint marginal distributions.
 
+### Broadcasting syntax 
+
+!!! note 
+    Broadcasting syntax requires at least v2.1.0 of `GraphPPL.jl` 
+
+GraphPPL support broadcasting for `~` operator in the exact same way as Julia itself. A user is free to write an expression of the following form:
+
+```julia
+y = datavar(Float64, n)
+y .~ NormalMeanVariance(0.0, 1.0) # <- i.i.d observations
+```
+
+More complex expression are also allowed:
+
+```julia
+m ~ NormalMeanPrecision(0.0, 0.0001)
+t ~ Gamma(1.0, 1.0)
+
+y = randomvar(Float64, n)
+y .~ NormalMeanPrecision(m, t)
+```
+
+```julia
+A = constvar(...)
+x = randomvar(n)
+y = datavar(Vector{Float64}, n)
+
+w         ~ Wishart(3, diageye(2))
+x[1]      ~ MvNormalMeanPrecision(zeros(2), diageye(2))
+x[2:end] .~ A .* x[1:end-1] # <- State-space model with transition matrix A
+y        .~ MvNormalMeanPrecision(x, w) # <- Observations with unknown precision matrix
+```
+
+Note, however, that all variables that take part in the broadcasting operation must be defined before either with `randomvar` or `datavar`. The exception here is constants that are automatically converted to their `constvar` equivalent. If you want to prevent broadcasting for some constant (e.g. if you want to add a vector to a multivariate Gaussian distribution) use explicit `constvar` call:
+
+```julia
+# Suppose `x` is a 2-dimensional Gaussian distribution
+z .~ x .+ constvar([ 1, 1 ])
+# Which is equivalent to 
+for i in 1:n
+   z[i] ~ x[i] + constvar([ 1, 1 ])
+end
+```
+
+Without explicit `constvar` Julia's broadcasting machinery would instead attempt to unroll for loop in the following way:
+
+```julia
+# Without explicit `constvar`
+z .~ x .+ [ 1, 1 ]
+# Which is equivalent to 
+array = [1, 1]
+for i in 1:n
+   z[i] ~ x[i] + array[i] # This is wrong if `x[i]` is supposed to be a multivariate Gaussian 
+end
+```
+
+Read more about how broadcasting machinery works in Julia in [the official documentation](https://docs.julialang.org/en/v1/manual/arrays/#Broadcasting).
+
 ### Node creation options
 
 To pass optional arguments to the node creation constructor the user can use the `where { options...  }` options specification syntax.

docs/src/utils.md

@@ -4,4 +4,5 @@
 GraphPPL.ishead
 GraphPPL.isblock
 GraphPPL.iscall
+GraphPPL.isbroadcastedcall
 ```

src/backends/reactivemp.jl

@@ -83,23 +83,27 @@ function write_argument_guard(::ReactiveMPBackend, argument::Symbol)
 end
 
 function write_randomvar_expression(::ReactiveMPBackend, model, varexp, options, arguments)
-    return :($varexp = ReactiveMP.randomvar($model, $options, $(GraphPPL.fquote(varexp)), $(arguments...)))
+    return :($varexp = ReactiveMP.randomvar($model, $options, $(GraphPPL.fquote(varexp)), $(arguments...)); $varexp)
 end
 
 function write_datavar_expression(::ReactiveMPBackend, model, varexpr, options, type, arguments)
     errstr    = "The expression `$varexpr = datavar($(type))` is incorrect. datavar(::Type, [ dims... ]) requires `Type` as a first argument, but `$(type)` is not a `Type`."
     checktype = :(GraphPPL.ensure_type($(type)) || error($errstr))
-    return :($checktype; $varexpr = ReactiveMP.datavar($model, $options, $(GraphPPL.fquote(varexpr)), ReactiveMP.PointMass{ $type }, $(arguments...)))
+    return :($checktype; $varexpr = ReactiveMP.datavar($model, $options, $(GraphPPL.fquote(varexpr)), ReactiveMP.PointMass{ $type }, $(arguments...)); $varexpr)
 end
 
 function write_constvar_expression(::ReactiveMPBackend, model, varexpr, arguments)
-    return :($varexpr = ReactiveMP.constvar($model, $(GraphPPL.fquote(varexpr)), $(arguments...)))
+    return :($varexpr = ReactiveMP.constvar($model, $(GraphPPL.fquote(varexpr)), $(arguments...)); $varexpr)
 end
 
 function write_as_variable(::ReactiveMPBackend, model, varexpr)
     return :(ReactiveMP.as_variable($model, $varexpr))
 end
 
+function write_undo_as_variable(::ReactiveMPBackend, varexpr)
+    return :(ReactiveMP.undo_as_variable($varexpr))
+end
+
 function write_anonymous_variable(::ReactiveMPBackend, model, varexpr)
     return :(ReactiveMP.setanonymous!($varexpr, true))
 end
@@ -108,10 +112,18 @@ function write_make_node_expression(::ReactiveMPBackend, model, fform, variables
     return :($nodeexpr = ReactiveMP.make_node($model, $options, $fform, $varexpr, $(variables...)))
 end
 
+function write_broadcasted_make_node_expression(::ReactiveMPBackend, model, fform, variables, options, nodeexpr, varexpr)
+    return :($nodeexpr = ReactiveMP.make_node.($model, $options, $fform, $varexpr, $(variables...)))
+end
+
 function write_autovar_make_node_expression(::ReactiveMPBackend, model, fform, variables, options, nodeexpr, varexpr, autovarid)
     return :(($nodeexpr, $varexpr) = ReactiveMP.make_node($model, $options, $fform, ReactiveMP.AutoVar($(GraphPPL.fquote(autovarid))), $(variables...)))
 end
 
+function write_check_variable_existence(::ReactiveMPBackend, model, varid, errormsg)
+    return :(ReactiveMP.haskey($model, $(QuoteNode(varid))) || Base.error($errormsg))
+end
+
 function write_node_options(::ReactiveMPBackend, model, fform, variables, options)
     is_factorisation_option_present = false
     is_meta_option_present          = false

src/model.jl

@@ -69,7 +69,11 @@ function normalize_tilde_arguments(backend, model, args)
 end
 
 function __normalize_arg(backend, model, arg)
-    if @capture(arg, (f_(v__) where { options__ }) | (f_(v__)))
+    if @capture(arg, constvar(arguments__))
+        return write_constvar_expression(backend, model, gensym(:anonymous_constvar), arguments)
+    elseif @capture(arg, constvar.(arguments__))
+        return error("Broadcasting of `constvar` in the constvar.(...) expression is dissalowed. Use `constvar((i) -> ..., dims...)` form instead.")
+    elseif @capture(arg, (f_(v__) where { options__ }) | (f_(v__)) | (f_.(v__) where { options__ }) | (f_.(v__) ))
         if f === :(|>)
             @assert length(v) === 2 "Unsupported pipe syntax in model specification: $(arg)"
             f = v[2]
@@ -79,7 +83,31 @@ function __normalize_arg(backend, model, arg)
         nnodeexpr = gensym(:nnode)
         options  = options !== nothing ? options : []
         v = normalize_tilde_arguments(backend, model, v)
-        return :(($nnodeexpr, $nvarexpr) ~ $f($(v...); $(options...)); $(write_anonymous_variable(backend, model, nvarexpr)); $nvarexpr)
+        if isbroadcastedcall(arg)
+            # Strip dot call from broadcasting dot operators, like `.+` and define `BroadcastFunction` explicitly to avoid UndefVarError
+            f = first(string(f)) === '.' ? Symbol(string(f)[2:end]) : f 
+            # broadcasting variables
+            broadcasting_locals = map((_) -> gensym(:bv), v)
+            return quote 
+                # Here we manually unroll anonymous broadcasting calls
+                # Later on GraphPPL does not distinguish between local broadcasting `~` expression and a regular `~` expression
+                begin 
+                    Base.broadcast($(v...)) do $(broadcasting_locals...)
+                        # $initf
+                        ($nnodeexpr, $nvarexpr) ~ $f($(broadcasting_locals...); $(options...)); 
+                        $(write_anonymous_variable(backend, model, nvarexpr)); 
+                        $(write_undo_as_variable(backend, nvarexpr));
+                    end
+                end
+            end
+            
+        else
+            return quote 
+                ($nnodeexpr, $nvarexpr) ~ $f($(v...); $(options...)); 
+                $(write_anonymous_variable(backend, model, nvarexpr)); 
+                $(write_undo_as_variable(backend, nvarexpr));
+            end
+        end
     else
         return arg
     end
@@ -128,6 +156,11 @@ function write_constvar_expression end
 """
 function write_as_variable end
 
+"""
+    write_undo_as_variable(backend, varexpr)
+"""
+function write_undo_as_variable end
+
 """
     write_anonymous_variable(backend, model, varexpr)
 """
@@ -138,11 +171,21 @@ function write_anonymous_variable end
 """
 function write_make_node_expression end
 
+"""
+    write_broadcasted_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr)
+"""
+function write_broadcasted_make_node_expression end
+
 """
     write_autovar_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr, autovarid)
 """
 function write_autovar_make_node_expression end
 
+"""
+    write_check_variable_existence(backend, model, varid, errormsg)
+"""
+function write_check_variable_existence end
+
 """
     write_node_options(backend, model, fform, variables, options)
 """
@@ -237,7 +280,10 @@ function generate_model_expression(backend, model_options, model_specification)
 
     # Step 1: Probabilistic arguments normalisation
     ms_body = prewalk(ms_body) do expression
-        if @capture(expression, (varexpr_ ~ fform_(arguments__) where { options__ }) | (varexpr_ ~ fform_(arguments__)))
+        if @capture(expression, 
+            (varexpr_ ~ fform_(arguments__) where { options__ }) | (varexpr_ ~ fform_(arguments__)) |
+            (varexpr_ .~ fform_(arguments__) where { options__ }) | (varexpr_ .~ fform_(arguments__))
+        )
             options   = options === nothing ? [] : options
 
             # Filter out keywords arguments to options array
@@ -249,8 +295,9 @@ function generate_model_expression(backend, model_options, model_specification)
                 return !ifparameters
             end
 
-            varexpr   =  @capture(varexpr, (nodeid_, varid_)) ? varexpr : :(($(gensym(:nnode)), $varexpr))
-            return :($varexpr ~ $(fform)($((normalize_tilde_arguments(backend, model, arguments))...); $(options...)))
+            varexpr =  @capture(varexpr, (nodeid_, varid_)) ? varexpr : :(($(gensym(:nnode)), $varexpr))
+            operator = isbroadcastedcall(expression) ? Symbol(".~") : :(~)
+            return :($operator($varexpr, $(fform)($((normalize_tilde_arguments(backend, model, arguments))...); $(options...))))
         elseif @capture(expression, varexpr_ = randomvar(arguments__) where { options__ })
             return :($varexpr = randomvar($(arguments...); $(options...)))
         elseif @capture(expression, varexpr_ = datavar(arguments__) where { options__ })
@@ -263,6 +310,8 @@ function generate_model_expression(backend, model_options, model_specification)
             return :($varexpr = datavar($(arguments...); ))
         elseif @capture(expression, varexpr_ = constvar(arguments__))
             return :($varexpr = constvar($(arguments...)))
+        elseif @capture(expression, constvar.(arguments__))
+            error("Broadcasting of `constvar` in the constvar.(...) expression is dissalowed. Use `constvar((i) -> ..., dims...)` form instead.")
         else
             return expression
         end
@@ -295,11 +344,11 @@ function generate_model_expression(backend, model_options, model_specification)
         elseif @capture(expression, varexpr_ = randomvar(arguments__; options__))
             rvoptions = write_randomvar_options(backend, varexpr, options)
             return write_randomvar_expression(backend, model, varexpr, rvoptions, arguments)
-        # Step 2.3 Conver constvar calls 
+        # Step 2.3 Convert constvar calls 
         elseif @capture(expression, varexpr_ = constvar(arguments__))
             return write_constvar_expression(backend, model, varexpr, arguments)
         # Step 2.2 Convert tilde expressions
-        elseif @capture(expression, (nodeexpr_, varexpr_) ~ fform_(arguments__; kwarguments__))
+        elseif @capture(expression, ((nodeexpr_, varexpr_) ~ fform_(arguments__; kwarguments__)) | ((nodeexpr_, varexpr_) .~ fform_(arguments__; kwarguments__)))
 
             varexpr, short_id, full_id = parse_varexpr(varexpr)
 
@@ -310,13 +359,23 @@ function generate_model_expression(backend, model_options, model_specification)
             variables = map((argexpr) -> write_as_variable(backend, model, argexpr), arguments)
             options = write_node_options(backend, model, fform, [ varexpr, arguments... ], kwarguments)
 
-            # Indexed variables like `y[1]` cannot be created on the fly and should be pre-initialised with `y = randomvar(n)`
-            # Single variables like `y` can be created on the fly with the `AutoVar` marker
-            # In the second case if variable `y` has been initialised before `AutoVar` should simply return it
-            if isref(varexpr)
-                return write_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr)
+            if isbroadcastedcall(expression)
+                # Strip dot call from broadcasting dot operators, like `.+`
+                fform = first(string(fform)) === '.' ? Symbol(string(fform)[2:end]) : fform 
+                return quote 
+                    # In case of broadcasted call we assume that variable has been created before otherwise it should throw an error
+                    $(write_check_variable_existence(backend, model, short_id, "Cannot use variables named `$(short_id)` in the broadcasting call. `$(short_id)` sequence of variables must be created in advance."))
+                    $(write_broadcasted_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr))
+                end
             else
-                return write_autovar_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr, full_id)
+                # Indexed variables like `y[1]` cannot be created on the fly and should be pre-initialised with `y = randomvar(n)`
+                # Single variables like `y` can be created on the fly with the `AutoVar` marker
+                # In the second case if variable `y` has been initialised before `AutoVar` should simply return it
+                if isref(varexpr)
+                    return write_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr)
+                else
+                    return write_autovar_make_node_expression(backend, model, fform, variables, options, nodeexpr, varexpr, full_id)
+                end
             end
         else
             return expression

src/utils.jl

@@ -32,6 +32,32 @@ See also: [`ishead`](@ref)
 iscall(expr)       = ishead(expr, :call) && length(expr.args) >= 1
 iscall(expr, fsym) = iscall(expr) && first(expr.args) === fsym
 
+"""
+    isbroadcastedcall(expr) 
+    isbroadcastedcall(expr, fsym) 
+
+Checks if expression represents a broadcast call to some function. Optionally accepts `fsym` to check for exact function name match.
+
+See also: [`iscall`](@ref)
+"""
+function isbroadcastedcall(expr) 
+    if isblock(expr) # TODO add for other functions?
+        nextexpr = findnext(isexpr, expr.args, 1)
+        return nextexpr !== nothing ? isbroadcastedcall(expr.args[nextexpr]) : false
+    end
+    (iscall(expr) && length(expr.args) >= 1 && first(string(first(expr.args))) === '.') || # Checks for `:(a .+ b)` syntax
+        (ishead(expr, :(.))) # Checks for `:(f.(x))` syntax
+end
+
+function isbroadcastedcall(expr, fsym) 
+    if isblock(expr) # TODO add for other functions?
+        nextexpr = findnext(isexpr, expr.args, 1)
+        return nextexpr !== nothing ? isbroadcastedcall(expr.args[nextexpr], fsym) : false
+    end
+    (iscall(expr) && length(expr.args) >= 1 && first(string(first(expr.args))) === '.' && Symbol(string(first(expr.args))[2:end]) === fsym) || # Checks for `:(a .+ b)` syntax
+        (ishead(expr, :(.)) && first(expr.args) === fsym) # Checks for `:(f.(x))` syntax
+end
+
 """
     isref(expr)
 

test/utils.jl

@@ -54,6 +54,28 @@ end
 
 end
 
+@testset "isbroadcastedcall tests" begin 
+    import GraphPPL: isbroadcastedcall
+
+    @test isbroadcastedcall(:(f(1))) === false
+    @test isbroadcastedcall(:(f(1)), :f) === false
+    @test isbroadcastedcall(:(f(1)), :g) === false
+    @test isbroadcastedcall(:(if true 1 else 2 end)) === false
+    @test isbroadcastedcall(:(begin end)) === false
+
+    @test isbroadcastedcall(:(a .+ b)) === true
+    @test isbroadcastedcall(:(a .+ b), :(+)) === true
+    @test isbroadcastedcall(:(a .+ b), :(-)) === false
+
+    @test isbroadcastedcall(:(f.(a))) === true
+    @test isbroadcastedcall(:(f.(a, b))) === true
+    @test isbroadcastedcall(:(f.(a)), :f) === true
+    @test isbroadcastedcall(:(f.(a, b)), :f) === true
+    @test isbroadcastedcall(:(f.(a)), :g) === false
+    @test isbroadcastedcall(:(f.(a, b)), :g) === false
+
+end
+
 @testset "ensure_type tests" begin 
     import GraphPPL: ensure_type