Merge branch 'master' into phg/varinfo-indices

Author: phipsgabler

.travis.yml

@@ -8,9 +8,7 @@ os:
   - osx
 julia:
   - 1.0
-  - 1.1
-  - 1.2
-  - 1.3
+  - 1
   - nightly
 matrix:
   allow_failures:
@@ -19,6 +17,6 @@ matrix:
 notifications:
   email: false
 after_success:
-  - if [[ $TRAVIS_JULIA_VERSION = 1.3 ]] && [[ $TRAVIS_OS_NAME = linux ]]; then
+  - if [[ $TRAVIS_JULIA_VERSION = 1 ]] && [[ $TRAVIS_OS_NAME = linux ]]; then
       julia -e 'using Pkg; Pkg.add("Coverage"); using Coverage; Codecov.submit(process_folder())';
     fi

Project.toml

@@ -10,7 +10,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 
 [compat]
-AbstractMCMC = "0.4, 0.5"
+AbstractMCMC = "0.4, 0.5, 1.0"
 Bijectors = "0.5.2, 0.6"
 Distributions = "0.22, 0.23"
 MacroTools = "0.5.1"

src/compiler.jl

@@ -27,45 +27,44 @@ function wrong_dist_errormsg(l)
 end
 
 """
-    @preprocess(data_vars, missing_vars, ex)
+    @isassumption(model, expr)
 
-Let `ex` be `x[1]`. This macro returns `@varname x[1]` in any of the following cases:
+Let `expr` be `x[1]`. `vn` is an assumption in the following cases:
     1. `x` was not among the input data to the model,
     2. `x` was among the input data to the model but with a value `missing`, or
     3. `x` was among the input data to the model with a value other than missing, 
-    but `x[1] === missing`.
-Otherwise, the value of `x[1]` is returned.
+       but `x[1] === missing`.
+When `expr` is not an expression or symbol (i.e., a literal), this expands to `false`.
 """
-macro preprocess(data_vars, missing_vars, ex)
-    ex
-end
-macro preprocess(model, ex::Union{Symbol, Expr})
-    sym = gensym(:sym)
-    lhs = gensym(:lhs)
-    return esc(quote
-        # Extract symbol
-        $sym = Val($(vsym(ex)))
+macro isassumption(model, expr::Union{Symbol, Expr})
+    # Note: never put a return in this... don't forget it's a macro!
+    vn = gensym(:vn)
+    
+    return quote
+        $vn = @varname($expr)
+        
         # This branch should compile nicely in all cases except for partial missing data
-        # For example, when `ex` is `x[i]` and `x isa Vector{Union{Missing, Float64}}`
-        if !DynamicPPL.inargnames($sym, $model) || DynamicPPL.inmissings($sym, $model)
-            $(varname(ex)), $(vinds(ex))
+        # For example, when `expr` is `x[i]` and `x isa Vector{Union{Missing, Float64}}`
+        if !DynamicPPL.inargnames($vn, $model) || DynamicPPL.inmissings($vn, $model)
+            true
         else
-            if DynamicPPL.inargnames($sym, $model)
+            if DynamicPPL.inargnames($vn, $model)
                 # Evaluate the lhs
-                $lhs = $ex
-                if $lhs === missing
-                    $(varname(ex)), $(vinds(ex))
-                else
-                    $lhs
-                end
+                $expr === missing
             else
                 throw("This point should not be reached. Please report this error.")
             end
         end
-    end)
+    end |> esc
+end
+
+macro isassumption(model, expr)
+    # failsafe: a literal is never an assumption
+    false
 end
 
 
+
 #################
 # Main Compiler #
 #################
@@ -246,40 +245,40 @@ function replace_sampler!(model_info)
     return model_info
 end
 
-# The next function is defined that way because .~ gives a parsing error in Julia 1.0
 """
-\"""
     replace_tilde!(model_info)
 
-Replaces `~` expressions with observation or assumption expressions, updating `model_info`.
-\"""
+Replace `~` and `.~` expressions with observation or assumption expressions, updating `model_info`.
+"""
 function replace_tilde!(model_info)
-    ex = model_info[:main_body]
-    ex = MacroTools.postwalk(ex) do x 
-        if @capture(x, @M_ L_ ~ R_) && M == Symbol("@__dot__")
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
+    # Apply the `@.` macro first.
+    expr = model_info[:main_body]
+    dottedexpr = MacroTools.postwalk(apply_dotted, expr)
+
+    # Check for tilde operators.
+    tildeexpr = MacroTools.postwalk(dottedexpr) do x
+        # Check dot tilde first.
+        dotargs = getargs_dottilde(x)
+        if dotargs !== nothing
+            L, R = dotargs
+            return generate_dot_tilde(L, R, model_info)
         end
-    end
-    $(VERSION >= v"1.1" ? "ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ .~ R_)
-            generate_dot_tilde(L, R, model_info)
-        else
-            x
-        end
-    end" : "")
-    ex = MacroTools.postwalk(ex) do x
-        if @capture(x, L_ ~ R_)
-            generate_tilde(L, R, model_info)
-        else
-            x
+
+        # Check tilde.
+        args = getargs_tilde(x)
+        if args !== nothing
+            L, R = args
+            return generate_tilde(L, R, model_info)
         end
+
+        return x
     end
-    model_info[:main_body] = ex
+
+    # Update the function body.
+    model_info[:main_body] = tildeexpr
+
     return model_info
 end
-""" |> Meta.parse |> eval
 
 # """ Unbreak code highlighting in Emacs julia-mode
 
@@ -300,32 +299,36 @@ function generate_tilde(left, right, model_info)
     lp = gensym(:lp)
     vn = gensym(:vn)
     inds = gensym(:inds)
-    preprocessed = gensym(:preprocessed)
+    isassumption = gensym(:isassumption)
     assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))
+    
     if left isa Symbol || left isa Expr
         ex = quote
             $temp_right = $right
             $assert_ex
-            $preprocessed = DynamicPPL.@preprocess($model, $left)
-            if $preprocessed isa Tuple
-                $vn, $inds = $preprocessed
-                $out = DynamicPPL.tilde($ctx, $sampler, $temp_right, $vn, $inds, $vi)
+            
+            $vn, $inds = $(varname(left)), $(vinds(left))
+            $isassumption = DynamicPPL.@isassumption($model, $left)
+            if $isassumption 
+                $out = DynamicPPL.tilde_assume($ctx, $sampler, $temp_right, $vn, $inds, $vi)
                 $left = $out[1]
                 DynamicPPL.acclogp!($vi, $out[2])
             else
                 DynamicPPL.acclogp!(
                     $vi,
-                    DynamicPPL.tilde($ctx, $sampler, $temp_right, $preprocessed, $vi),
+                    DynamicPPL.tilde_observe($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi),
                 )
             end
         end
     else
+        # we have a literal, which is automatically an observation
         ex = quote
             $temp_right = $right
             $assert_ex
+            
             DynamicPPL.acclogp!(
                 $vi,
-                DynamicPPL.tilde($ctx, $sampler, $temp_right, $left, $vi),
+                DynamicPPL.tilde_observe($ctx, $sampler, $temp_right, $left, $vi),
             )
         end
     end
@@ -335,48 +338,51 @@ end
 """
     generate_dot_tilde(left, right, model_info)
 
-This function returns the expression that replaces `left .~ right` in the model body. If `preprocessed isa VarName`, then a `dot_assume` block will be run. Otherwise, a `dot_observe` block will be run.
+This function returns the expression that replaces `left .~ right` in the model body. If
+`preprocessed isa VarName`, then a `dot_assume` block will be run. Otherwise, a `dot_observe` block
+will be run.
 """
 function generate_dot_tilde(left, right, model_info)
     model = model_info[:main_body_names][:model]
     vi = model_info[:main_body_names][:vi]
     ctx = model_info[:main_body_names][:ctx]
     sampler = model_info[:main_body_names][:sampler]
     out = gensym(:out)
-    temp_left = gensym(:temp_left)
     temp_right = gensym(:temp_right)
-    preprocessed = gensym(:preprocessed)
+    isassumption = gensym(:isassumption)
     lp = gensym(:lp)
     vn = gensym(:vn)
     inds = gensym(:inds)
     assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))
+    
     if left isa Symbol || left isa Expr
         ex = quote
             $temp_right = $right
             $assert_ex
-            $preprocessed = DynamicPPL.@preprocess($model, $left)
-            if $preprocessed isa Tuple
-                $vn, $inds = $preprocessed
-                $temp_left = $left
-                $out = DynamicPPL.dot_tilde($ctx, $sampler, $temp_right, $temp_left, $vn, $inds, $vi)
+
+            $vn, $inds = $(varname(left)), $(vinds(left))
+            $isassumption = DynamicPPL.@isassumption($model, $left)
+            
+            if $isassumption
+                $out = DynamicPPL.dot_tilde_assume($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi)
                 $left .= $out[1]
                 DynamicPPL.acclogp!($vi, $out[2])
             else
-                $temp_left = $preprocessed
                 DynamicPPL.acclogp!(
                     $vi,
-                    DynamicPPL.dot_tilde($ctx, $sampler, $temp_right, $temp_left, $vi),
+                    DynamicPPL.dot_tilde_observe($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi),
                 )
             end
         end
     else
+        # we have a literal, which is automatically an observation
         ex = quote
-            $temp_left = $left
             $temp_right = $right
             $assert_ex
+            
             DynamicPPL.acclogp!(
                 $vi,
-                DynamicPPL.dot_tilde($ctx, $sampler, $temp_right, $temp_left, $vi),
+                DynamicPPL.dot_tilde_observe($ctx, $sampler, $temp_right, $left, $vi),
             )
         end
     end
@@ -416,7 +422,7 @@ function build_output(model_info)
     model_gen = model_info[:name]
     # Main body of the model
     main_body = model_info[:main_body]
-    
+
     unwrap_data_expr = Expr(:block)
     for var in arg_syms
         temp_var = gensym(:temp_var)

src/context_implementations.jl

@@ -35,6 +35,19 @@ function tilde(ctx::MiniBatchContext, sampler, right, left::VarName, inds, vi)
     return tilde(ctx.ctx, sampler, right, left, inds, vi)
 end
 
+"""
+    tilde_assume(ctx, sampler, right, vn, inds, vi)
+
+This method is applied in the generated code for assumed variables, e.g., `x ~ Normal()` where
+`x` does not occur in the model inputs.
+
+Falls back to `tilde(ctx, sampler, right, vn, inds, vi)`.
+"""
+function tilde_assume(ctx, sampler, right, vn, inds, vi)
+    return tilde(ctx, sampler, right, vn, inds, vi)
+end
+
+
 function _tilde(sampler, right, vn::VarName, vi)
     return assume(sampler, right, vn, vi)
 end
@@ -56,6 +69,30 @@ function tilde(ctx::MiniBatchContext, sampler, right, left, vi)
     return ctx.loglike_scalar * tilde(ctx.ctx, sampler, right, left, vi)
 end
 
+"""
+    tilde_observe(ctx, sampler, right, left, vname, vinds, vi)
+
+This method is applied in the generated code for observed variables, e.g., `x ~ Normal()` where
+`x` does occur in the model inputs.
+
+Falls back to `tilde(ctx, sampler, right, left, vi)` ignoring the information about variable
+name and indices; if needed, these can be accessed through this function, though.
+"""
+function tilde_observe(ctx, sampler, right, left, vname, vinds, vi)
+    return tilde(ctx, sampler, right, left, vi)
+end
+
+"""
+    tilde_observe(ctx, sampler, right, left, vi)
+
+This method is applied in the generated code for observed constants, e.g., `1.0 ~ Normal()`.
+Falls back to `tilde(ctx, sampler, right, left, vi)`.
+"""
+function tilde_observe(ctx, sampler, right, left, vi)
+    return tilde(ctx, sampler, right, left, vi)
+end
+
+
 _tilde(sampler, right, left, vi) = observe(sampler, right, left, vi)
 
 function assume(spl::Sampler, dist)
@@ -151,6 +188,19 @@ function dot_tilde(
     return _dot_tilde(sampler, dist, left, vns, vi)
 end
 
+"""
+    dot_tilde_assume(ctx, sampler, right, left, vn, inds, vi)
+
+This method is applied in the generated code for assumed vectorized variables, e.g., `x .~
+MvNormal()` where `x` does not occur in the model inputs.
+
+Falls back to `dot_tilde(ctx, sampler, right, left, vn, inds, vi)`.
+"""
+function dot_tilde_assume(ctx, sampler, right, left, vn, inds, vi)
+    return dot_tilde(ctx, sampler, right, left, vn, inds, vi)
+end
+
+
 function get_vns_and_dist(dist::NamedDist, var, vn::VarName)
     return get_vns_and_dist(dist.dist, var, dist.name)
 end
@@ -314,6 +364,30 @@ function dot_tilde(ctx::MiniBatchContext, sampler, right, left, vi)
     return ctx.loglike_scalar * dot_tilde(ctx.ctx, sampler, right, left, left, vi)
 end
 
+"""
+    dot_tilde_observe(ctx, sampler, right, left, vname, vinds, vi)
+
+This method is applied in the generated code for vectorized observed variables, e.g., `x .~
+MvNormal()` where `x` does occur the model inputs.
+
+Falls back to `dot_tilde(ctx, sampler, right, left, vi)` ignoring the information about variable
+name and indices; if needed, these can be accessed through this function, though.
+"""
+function dot_tilde_observe(ctx, sampler, right, left, vn, inds, vi)
+    return dot_tilde(ctx, sampler, right, left, vi)
+end
+
+"""
+    dot_tilde_observe(ctx, sampler, right, left, vi)
+
+This method is applied in the generated code for vectorized observed constants, e.g., `[1.0] .~
+MvNormal()`.  Falls back to `dot_tilde(ctx, sampler, right, left, vi)`.
+"""
+function dot_tilde_observe(ctx, sampler, right, left, vi)
+    return dot_tilde(ctx, sampler, right, left, vi)
+end
+
+
 function _dot_tilde(sampler, right, left::AbstractArray, vi)
     return dot_observe(sampler, right, left, vi)
 end