Merge pull request #18 from biaslab/dev-node-specification-aliases

Feature: aliases and keyword distribution specification in the model macro

Author: bvdmitri

src/backends/reactivemp.jl

@@ -455,4 +455,87 @@ end
 
 function write_meta_specification_entry(::ReactiveMPBackend, F, N, meta) 
     return :(ReactiveMP.MetaSpecificationEntry(Val($F), Val($N), $meta))
+end
+
+# Aliases
+
+ReactiveMPNodeAliases = (
+    (
+        (expression) -> @capture(expression, a_ || b_) ? :(ReactiveMP.OR($a, $b)) : expression,     
+        "`a || b`: alias for `OR(a, b)` node (operator precedence between `||`, `&&`, `->` and `!` is the same as in Julia)."
+    ),
+    (
+        (expression) -> @capture(expression, a_ && b_) ? :(ReactiveMP.AND($a, $b)) : expression,    
+        "`a && b`: alias for `AND(a, b)` node (operator precedence `||`, `&&`, `->` and `!` is the same as in Julia)."
+    ), 
+    (
+        (expression) -> @capture(expression, a_ -> b_) ? :(ReactiveMP.IMPLY($a, $b)) : expression,  
+        "`a -> b`: alias for `IMPLY(a, b)` node (operator precedence `||`, `&&`, `->` and `!` is the same as in Julia)."
+    ),
+    (
+        (expression) -> @capture(expression, (¬a_) | (!a_)) ? :(ReactiveMP.NOT($a)) : expression,            
+        "`¬a` and `!a`: alias for `NOT(a)` node (Unicode `\\neg`, operator precedence `||`, `&&`, `->` and `!` is the same as in Julia)."
+    ),
+    (
+        (expression) -> @capture(expression, +(args__)) ? fold_linear_operator_call(expression) : expression, 
+        "`a + b + c`: alias for `(a + b) + c`"
+    ),
+    (
+        (expression) -> @capture(expression, *(args__)) ? fold_linear_operator_call(expression) : expression, 
+        "`a * b * c`: alias for `(a * b) * c`"
+    ),
+    (
+        (expression) -> @capture(expression, (Normal | Gaussian)((μ)|(m)|(mean) = mean_, (σ²)|(τ⁻¹)|(v)|(var)|(variance) = var_)) ? :(NormalMeanVariance($mean, $var)) : expression, 
+        "`Normal(μ|m|mean = ..., σ²|τ⁻¹|v|var|variance = ...)` alias for `NormalMeanVariance(..., ...)` node. `Gaussian` could be used instead `Normal` too."
+    ),
+    (
+        (expression) -> @capture(expression, (Normal | Gaussian)((μ)|(m)|(mean) = mean_, (τ)|(γ)|(σ⁻²)|(w)|(p)|(prec)|(precision) = prec_)) ? :(NormalMeanPrecision($mean, $prec)) : expression, 
+        "`Normal(μ|m|mean = ..., τ|γ|σ⁻²|w|p|prec|precision = ...)` alias for `NormalMeanVariance(..., ...)` node. `Gaussian` could be used instead `Normal` too."
+    ),
+    (
+        (expression) -> @capture(expression, (MvNormal | MvGaussian)((μ)|(m)|(mean) = mean_, (Σ)|(V)|(Λ⁻¹)|(cov)|(covariance) = cov_)) ? :(MvNormalMeanCovariance($mean, $cov)) : expression, 
+        "`MvNormal(μ|m|mean = ..., Σ|V|Λ⁻¹|cov|covariance = ...)` alias for `MvNormalMeanCovariance(..., ...)` node. `MvGaussian` could be used instead `MvNormal` too."
+    ),
+    (
+        (expression) -> @capture(expression, (MvNormal | MvGaussian)((μ)|(m)|(mean) = mean_, (Λ)|(W)|(Σ⁻¹)|(prec)|(precision) = prec_)) ? :(MvNormalMeanPrecision($mean, $prec)) : expression, 
+        "`MvNormal(μ|m|mean = ..., Λ|W|Σ⁻¹|prec|precision = ...)` alias for `MvNormalMeanPrecision(..., ...)` node. `MvGaussian` could be used instead `MvNormal` too."
+    ),
+    ((expression) -> @capture(expression, (Normal | Gaussian)(args__)) ? :(error("Please use a specific version of the `Normal` (`Gaussian`) distribution (e.g. `NormalMeanVariance` or aliased version `Normal(mean = ..., variance|precision = ...)`).")) : expression, missing),
+    ((expression) -> @capture(expression, (MvNormal | MvGaussian)(args__)) ? :(error("Please use a specific version of the `MvNormal` (`MvGaussian`) distribution (e.g. `MvNormalMeanCovariance` or aliased version `MvNormal(mean = ..., covariance|precision = ...)`).")) : expression, missing),
+    (
+        (expression) -> @capture(expression, Gamma((α)|(a)|(shape) = shape_, (θ)|(β⁻¹)|(scale) = scale_)) ? :(GammaShapeScale($shape, $scale)) : expression,
+        "`Gamma(α|a|shape = ..., θ|β⁻¹|scale = ...)` alias for `GammaShapeScale(..., ...) node.`"
+    ),
+    (
+        (expression) -> @capture(expression, Gamma((α)|(a)|(shape) = shape_, (β)|(θ⁻¹)|(rate) = rate_)) ? :(GammaShapeRate($shape, $rate)) : expression,
+        "`Gamma(α|a|shape = ..., β|θ⁻¹|rate = ...)` alias for `GammaShapeRate(..., ...) node.`"
+    ),
+)
+
+function show_tilderhs_alias(::ReactiveMPBackend, io = stdout)
+    foreach(skipmissing(map(last, ReactiveMPNodeAliases))) do alias 
+        println(io, "- ", alias)
+    end
+end
+
+function apply_alias_transformation(notanexpression, alias)
+    # We always short-circuit on non-expression
+    return (notanexpression, true)
+end
+
+function apply_alias_transformation(expression::Expr, alias)
+    _expression = first(alias)(expression)
+    # Returns potentially modified expression and a Boolean flag, 
+    # which indicates if expression actually has been modified
+    return (_expression, _expression !== expression)
+end
+
+function write_inject_tilderhs_aliases(::ReactiveMPBackend, model, tilderhs)
+    return postwalk(tilderhs) do expression
+        # We short-circuit if `mflag` is true
+        _expression, _ = foldl(ReactiveMPNodeAliases; init = (expression, false)) do (expression, mflag), alias
+            return mflag ? (expression, true) : apply_alias_transformation(expression, alias)
+        end
+        return _expression
+    end
 end

src/model.jl

@@ -211,6 +211,31 @@ function write_default_model_constraints end
 """
 function write_default_model_meta end
 
+"""
+    write_inject_tilderhs_aliases(backend, model, tilderhs)
+"""
+function write_inject_tilderhs_aliases end
+
+"""
+    show_tilderhs_alias(backend, io)
+"""
+function show_tilderhs_alias end
+
+"""
+
+```julia
+@model [ model_options ] function model_name(model_arguments...; model_keyword_arguments...)
+    # model description
+end
+```
+
+`@model` macro generates a function that returns an equivalent graph-representation of the given probabilistic model description.
+
+## Supported alias in the model specification
+$(begin io = IOBuffer(); show_tilderhs_alias(__get_current_backend(), io); String(take!(io)) end)
+"""
+macro model end
+
 macro model(model_specification)
     return esc(:(@model [] $model_specification))
 end
@@ -278,7 +303,18 @@ function generate_model_expression(backend, model_options, model_specification)
     # Doing so can lead to undefined behaviour
     ms_args_checks = map((ms_arg) -> write_argument_guard(backend, ms_arg), ms_args_guard_ids)
 
-    # Step 1: Probabilistic arguments normalisation
+    # Step 1: Inject node's aliases 
+    ms_body = postwalk(ms_body) do expression 
+        if @capture(expression,  (lhs_ ~ rhs_ where { options__ }) | (lhs_ .~ rhs_ where { options__ }))
+            return :($lhs ~ $(write_inject_tilderhs_aliases(backend, model, rhs)) where { $options... })
+        elseif @capture(expression,  (lhs_ ~ rhs_) | (lhs_ .~ rhs_))
+            return :($lhs ~ $(write_inject_tilderhs_aliases(backend, model, rhs)))
+        else 
+            return expression
+        end
+    end
+
+    # Step 2: Probabilistic arguments normalisation
     ms_body = prewalk(ms_body) do expression
         if @capture(expression, 
             (varexpr_ ~ fform_(arguments__) where { options__ }) | (varexpr_ ~ fform_(arguments__)) |
@@ -329,9 +365,9 @@ function generate_model_expression(backend, model_options, model_specification)
         return expression
     end
 
-    # Step 2: Main pass
+    # Step 3: Main pass
     ms_body = postwalk(ms_body) do expression
-        # Step 2.1 Convert datavar calls
+        # Step 3.1 Convert datavar calls
         if @capture(expression, varexpr_ = datavar(arguments__; options__)) 
             @assert length(arguments) >= 1 "The expression `$expression` is incorrect. datavar(::Type, [ dims... ]) requires `Type` as a first argument."
 
@@ -340,14 +376,14 @@ function generate_model_expression(backend, model_options, model_specification)
             dvoptions      = write_datavar_options(backend, varexpr, type_argument, options)
 
             return write_datavar_expression(backend, model, varexpr, dvoptions, type_argument, tail_arguments)
-        # Step 2.2 Convert randomvar calls
+        # Step 3.2 Convert randomvar calls
         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 Convert constvar calls 
+        # Step 3.3 Convert constvar calls 
         elseif @capture(expression, varexpr_ = constvar(arguments__))
             return write_constvar_expression(backend, model, varexpr, arguments)
-        # Step 2.2 Convert tilde expressions
+        # Step 3.2 Convert tilde expressions
         elseif @capture(expression, ((nodeexpr_, varexpr_) ~ fform_(arguments__; kwarguments__)) | ((nodeexpr_, varexpr_) .~ fform_(arguments__; kwarguments__)))
 
             varexpr, short_id, full_id = parse_varexpr(varexpr)
@@ -382,7 +418,7 @@ function generate_model_expression(backend, model_options, model_specification)
         end
     end
 
-    # Step 3: Final pass
+    # Step 4: Final pass
     final_pass_exceptions = (x) -> @capture(x, (some_ -> body_) | (function some_(args__) body_ end) | (some_(args__) = body_))
     final_pass_target     = (x) -> @capture(x, return ret_)
 

src/utils.jl

@@ -82,4 +82,20 @@ getref(expr) = isref(expr) ? (view(expr.args, 2:lastindex(expr.args))...,) : ()
 Checks if `x` is of type `Type` 
 """
 ensure_type(x::Type) = true
-ensure_type(x)       = false
+ensure_type(x)       = false
+
+fold_linear_operator_call(any) = any
+
+fold_linear_operator_call_first_arg(::typeof(foldl), args) = args[begin + 1]
+fold_linear_operator_call_tail_args(::typeof(foldl), args) = args[begin+2:end]
+
+fold_linear_operator_call_first_arg(::typeof(foldr), args) = args[end]
+fold_linear_operator_call_tail_args(::typeof(foldr), args) = args[begin+1:end-1]
+
+function fold_linear_operator_call(expr::Expr, fold = foldl)
+    if @capture(expr, op_(args__, )) && length(args) > 2
+        return fold((res, el) -> Expr(:call, op, res, el), fold_linear_operator_call_tail_args(fold, expr.args); init = fold_linear_operator_call_first_arg(fold, expr.args))
+    else
+        return expr 
+    end
+end

test/utils.jl

@@ -2,6 +2,7 @@ module UtilsTests
 
 using Test
 using GraphPPL
+using MacroTools
 
 @testset "issymbol tests" begin
     import GraphPPL: issymbol
@@ -85,4 +86,20 @@ end
     @test ensure_type(1.0) === false
 end
 
+@testset "fold_linear_operator_call" begin 
+    import GraphPPL: fold_linear_operator_call
+
+    @test @capture(fold_linear_operator_call(:(+a)), +a)
+    @test @capture(fold_linear_operator_call(:(a + b)), a + b)
+    @test @capture(fold_linear_operator_call(:(a + b + c)), (a + b) + c)
+    @test @capture(fold_linear_operator_call(:(a + b + c + d)), ((a + b) + c) + d)
+    @test @capture(fold_linear_operator_call(:(a + b + c + d), foldr), (a + (b + (c + d))))
+
+    @test @capture(fold_linear_operator_call(:(a * b)), a * b)
+    @test @capture(fold_linear_operator_call(:(a * b * c)), (a * b) * c)
+    @test @capture(fold_linear_operator_call(:(a * b * c * d)), ((a * b) * c) * d)
+    @test @capture(fold_linear_operator_call(:(a * b * c * d), foldr), (a * (b * (c * d))))
+
+end
+
 end