Merge pull request #5 from biaslab/develop-var-options

feat(): where options support for datavar/constvar/datavar

Author: bvdmitri

src/GraphPPL.jl

@@ -104,17 +104,17 @@ argument_write_default_value(arg, default)          = Expr(:kw, arg, default)
 function write_argument_guard end
 
 """
-    write_randomvar_expression(backend, model, varexpr, arguments)
+    write_randomvar_expression(backend, model, varexpr, arguments, kwarguments)
 """
 function write_randomvar_expression end
 
 """
-    write_datavar_expression(backend, model, varexpr, type, arguments)
+    write_datavar_expression(backend, model, varexpr, type, arguments, kwarguments)
 """
 function write_datavar_expression end
 
 """
-    write_constvar_expression(backend, model, varexpr, arguments)
+    write_constvar_expression(backend, model, varexpr, arguments, kwarguments)
 """
 function write_constvar_expression end
 
@@ -138,6 +138,21 @@ function write_autovar_make_node_expression end
 """
 function write_node_options end
 
+"""
+    write_randomvar_options(backend, variable, options)
+"""
+function write_randomvar_options end
+
+"""
+    write_constvar_options(backend, variable, options)
+"""
+function write_constvar_options end
+
+"""
+    write_datavar_options(backend, variable, options)
+"""
+function write_datavar_options end
+
 include("backends/reactivemp.jl")
 
 __get_current_backend() = ReactiveMPBackend()
@@ -198,7 +213,7 @@ function generate_model_expression(backend, model_options, model_specification)
     end
 
     ms_args_const_init_block = map(ms_args_const_ids) do ms_arg_const_id
-        return write_constvar_expression(backend, model, first(ms_arg_const_id), [ last(ms_arg_const_id) ])
+        return write_constvar_expression(backend, model, first(ms_arg_const_id), [ last(ms_arg_const_id) ], [])
     end
 
     # Step 0: Check that all inputs are not AbstractVariables
@@ -222,6 +237,18 @@ function generate_model_expression(backend, model_options, model_specification)
 
             varexpr   =  @capture(varexpr, (nodeid_, varid_)) ? varexpr : :(($(gensym(:nnode)), $varexpr))
             return :($varexpr ~ $(fform)($((normalize_tilde_arguments(arguments))...); $(options...)))
+        elseif @capture(expression, varexpr_ = randomvar(arguments__) where { options__ })
+            return :($varexpr = randomvar($(arguments...); $(write_randomvar_options(backend, varexpr, options)...)))
+        elseif @capture(expression, varexpr_ = datavar(arguments__) where { options__ })
+            return :($varexpr = datavar($(arguments...); $(write_datavar_options(backend, varexpr, options)...)))
+        elseif @capture(expression, varexpr_ = constvar(arguments__) where { options__ })
+            return :($varexpr = constvar($(arguments...); $(write_constvar_options(backend, varexpr, options)...)))
+        elseif @capture(expression, varexpr_ = randomvar(arguments__))
+            return :($varexpr = randomvar($(arguments...); ))
+        elseif @capture(expression, varexpr_ = datavar(arguments__))
+            return :($varexpr = datavar($(arguments...); ))
+        elseif @capture(expression, varexpr_ = constvar(arguments__))
+            return :($varexpr = constvar($(arguments...); ))
         else
             return expression
         end
@@ -244,28 +271,28 @@ function generate_model_expression(backend, model_options, model_specification)
     # Step 2: Main pass
     ms_body = postwalk(ms_body) do expression
         # Step 2.1 Convert datavar calls
-        if @capture(expression, varexpr_ = datavar(arguments__)) 
+        if @capture(expression, varexpr_ = datavar(arguments__; kwarguments__)) 
             @assert varexpr ∉ varids "Invalid model specification: '$varexpr' id is duplicated"
             @assert length(arguments) >= 1 "datavar() call requires type specification as a first argument"
 
             push!(varids, varexpr)
 
-            type = arguments[1]
-            tail = arguments[2:end]
+            type_argument  = arguments[1]
+            tail_arguments = arguments[2:end]
 
-            return write_datavar_expression(backend, model, varexpr, type, tail)
+            return write_datavar_expression(backend, model, varexpr, type_argument, tail_arguments, kwarguments)
         # Step 2.2 Convert randomvar calls
-        elseif @capture(expression, varexpr_ = randomvar(arguments__))
+        elseif @capture(expression, varexpr_ = randomvar(arguments__; kwarguments__))
             @assert varexpr ∉ varids "Invalid model specification: '$varexpr' id is duplicated"
             push!(varids, varexpr)
 
-            return write_randomvar_expression(backend, model, varexpr, arguments)
+            return write_randomvar_expression(backend, model, varexpr, arguments, kwarguments)
         # Step 2.3 Conver constvar calls 
-        elseif @capture(expression, varexpr_ = constvar(arguments__))
+        elseif @capture(expression, varexpr_ = constvar(arguments__; kwarguments__))
             @assert varexpr ∉ varids "Invalid model specification: '$varexpr' id is duplicated"
             push!(varids, varexpr)
 
-            return write_constvar_expression(backend, model, varexpr, arguments)
+            return write_constvar_expression(backend, model, varexpr, arguments, kwarguments)
         # Step 2.2 Convert tilde expressions
         elseif @capture(expression, (nodeexpr_, varexpr_) ~ fform_(arguments__; kwarguments__))
             # println(expression)

src/backends/reactivemp.jl

@@ -8,16 +8,16 @@ function write_argument_guard(::ReactiveMPBackend, argument::Symbol)
     return :(@assert !($argument isa ReactiveMP.AbstractVariable) "It is not allowed to pass AbstractVariable objects to a model definition arguments. ConstVariables should be passed as their raw values.")
 end
 
-function write_randomvar_expression(::ReactiveMPBackend, model, varexp, arguments)
-    return :($varexp = ReactiveMP.randomvar($model, $(GraphPPL.fquote(varexp)), $(arguments...)))
+function write_randomvar_expression(::ReactiveMPBackend, model, varexp, arguments, kwarguments)
+    return :($varexp = ReactiveMP.randomvar($model, $(GraphPPL.fquote(varexp)), $(arguments...); $(kwarguments...)))
 end
 
-function write_datavar_expression(::ReactiveMPBackend, model, varexpr, type, arguments)
-    return :($varexpr = ReactiveMP.datavar($model, $(GraphPPL.fquote(varexpr)), ReactiveMP.PointMass{ GraphPPL.ensure_type($(type)) }, $(arguments...)))
+function write_datavar_expression(::ReactiveMPBackend, model, varexpr, type, arguments, kwarguments)
+    return :($varexpr = ReactiveMP.datavar($model, $(GraphPPL.fquote(varexpr)), ReactiveMP.PointMass{ GraphPPL.ensure_type($(type)) }, $(arguments...); $(kwarguments...)))
 end
 
-function write_constvar_expression(::ReactiveMPBackend, model, varexpr, arguments)
-    return :($varexpr = ReactiveMP.constvar($model, $(GraphPPL.fquote(varexpr)), $(arguments...)))
+function write_constvar_expression(::ReactiveMPBackend, model, varexpr, arguments, kwarguments)
+    return :($varexpr = ReactiveMP.constvar($model, $(GraphPPL.fquote(varexpr)), $(arguments...); $(kwarguments...)))
 end
 
 function write_as_variable(::ReactiveMPBackend, model, varexpr)
@@ -138,4 +138,27 @@ function write_fconstraint_option(form, variables, fconstraint)
     else
         error("Invalid factorisation constraint: $fconstraint")
     end
+end
+
+## 
+
+function write_randomvar_options(::ReactiveMPBackend, variable, options)
+    return map(options) do option
+        @capture(option, name_Symbol = value_) || error("Invalid variable options specification: $option. Should be in a form of 'name = value'")
+        return option
+    end
+end
+
+function write_constvar_options(::ReactiveMPBackend, variable, options)
+    return map(options) do option
+        @capture(option, name_Symbol = value_) || error("Invalid variable options specification: $option. Should be in a form of 'name = value'")
+        return option
+    end
+end
+
+function write_datavar_options(::ReactiveMPBackend, variable, options)
+    return map(options) do option
+        @capture(option, name_Symbol = value_) || error("Invalid variable options specification: $option. Should be in a form of 'name = value'")
+        return option
+    end
 end