Merge branch 'master' into step

Author: devmotion

src/compiler.jl

@@ -208,7 +208,6 @@ function generate_mainbody!(found, expr::Expr, args)
     return Expr(expr.head, map(x -> generate_mainbody!(found, x, args), expr.args)...)
 end
 
-# """ Unbreak code highlighting in Emacs julia-mode
 
 
 """
@@ -218,7 +217,7 @@ Generate an `observe` expression for data variables and `assume` expression for
 variables.
 """
 function generate_tilde(left, right, args)
-    @gensym tmpright tmpleft
+    @gensym tmpright
     top = [:($tmpright = $right),
            :($tmpright isa Union{$Distribution,AbstractVector{<:$Distribution}}
              || throw(ArgumentError($DISTMSG)))]
@@ -227,48 +226,32 @@ function generate_tilde(left, right, args)
         @gensym out vn inds
         push!(top, :($vn = $(varname(left))), :($inds = $(vinds(left))))
 
-        assumption = [
-            :($out = $(DynamicPPL.tilde_assume)(_context, _sampler, $tmpright, $vn, $inds,
-                                                _varinfo)),
-            :($(DynamicPPL.acclogp!)(_varinfo, $out[2])),
-            :($left = $out[1])
-        ]
-
         # It can only be an observation if the LHS is an argument of the model
         if vsym(left) in args
             @gensym isassumption
             return quote
                 $(top...)
                 $isassumption = $(DynamicPPL.isassumption(left))
                 if $isassumption
-                    $(assumption...)
+                    $left = $(DynamicPPL.tilde_assume)(
+                        _context, _sampler, $tmpright, $vn, $inds, _varinfo)
                 else
-                    $tmpleft = $left
-                    $(DynamicPPL.acclogp!)(
-                        _varinfo,
-                        $(DynamicPPL.tilde_observe)(_context, _sampler, $tmpright, $tmpleft,
-                                                    $vn, $inds, _varinfo)
-                    )
-                    $tmpleft
+                    $(DynamicPPL.tilde_observe)(
+                        _context, _sampler, $tmpright, $left, $vn, $inds, _varinfo)
                 end
             end
         end
 
         return quote
             $(top...)
-            $(assumption...)
+            $left = $(DynamicPPL.tilde_assume)(_context, _sampler, $tmpright, $vn, $inds, _varinfo)
         end
     end
 
     # If the LHS is a literal, it is always an observation
     return quote
         $(top...)
-        $tmpleft = $left
-        $(DynamicPPL.acclogp!)(
-            _varinfo,
-            $(DynamicPPL.tilde_observe)(_context, _sampler, $tmpright, $tmpleft, _varinfo)
-        )
-        $tmpleft
+        $(DynamicPPL.tilde_observe)(_context, _sampler, $tmpright, $left, _varinfo)
     end
 end
 
@@ -278,7 +261,7 @@ end
 Generate the expression that replaces `left .~ right` in the model body.
 """
 function generate_dot_tilde(left, right, args)
-    @gensym tmpright tmpleft
+    @gensym tmpright
     top = [:($tmpright = $right),
            :($tmpright isa Union{$Distribution,AbstractVector{<:$Distribution}}
              || throw(ArgumentError($DISTMSG)))]
@@ -287,49 +270,33 @@ function generate_dot_tilde(left, right, args)
         @gensym out vn inds
         push!(top, :($vn = $(varname(left))), :($inds = $(vinds(left))))
 
-        assumption = [
-            :($out = $(DynamicPPL.dot_tilde_assume)(_context, _sampler, $tmpright, $left,
-                                                    $vn, $inds, _varinfo)),
-            :($(DynamicPPL.acclogp!)(_varinfo, $out[2])),
-            :($left .= $out[1])
-        ]
-
         # It can only be an observation if the LHS is an argument of the model
         if vsym(left) in args
             @gensym isassumption
             return quote
                 $(top...)
                 $isassumption = $(DynamicPPL.isassumption(left))
                 if $isassumption
-                    $(assumption...)
+                    $left .= $(DynamicPPL.dot_tilde_assume)(
+                        _context, _sampler, $tmpright, $left, $vn, $inds, _varinfo)
                 else
-                    $tmpleft = $left
-                    $(DynamicPPL.acclogp!)(
-                        _varinfo,
-                        $(DynamicPPL.dot_tilde_observe)(_context, _sampler, $tmpright,
-                                                        $tmpleft, $vn, $inds, _varinfo)
-                    )
-                    $tmpleft
+                    $(DynamicPPL.dot_tilde_observe)(
+                        _context, _sampler, $tmpright, $left, $vn, $inds, _varinfo)
                 end
             end
         end
 
         return quote
             $(top...)
-            $(assumption...)
+            $left .= $(DynamicPPL.dot_tilde_assume)(
+                _context, _sampler, $tmpright, $left, $vn, $inds, _varinfo)
         end
     end
 
     # If the LHS is a literal, it is always an observation
     return quote
         $(top...)
-        $tmpleft = $left
-        $(DynamicPPL.acclogp!)(
-            _varinfo,
-            $(DynamicPPL.dot_tilde_observe)(_context, _sampler, $tmpright, $tmpleft,
-                                            _varinfo)
-        )
-        $tmpleft
+        $(DynamicPPL.dot_tilde_observe)(_context, _sampler, $tmpright, $left, _varinfo)
     end
 end
 

src/context_implementations.jl

@@ -38,13 +38,15 @@ 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.
+Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
+accumulate the log probability, and return the sampled value.
 
 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)
+    value, logp = tilde(ctx, sampler, right, vn, inds, vi)
+    acclogp!(vi, logp)
+    return value
 end
 
 
@@ -72,24 +74,30 @@ 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.
+Handle observed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
+accumulate the log probability, and return the observed value.
 
-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.
+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)
+    logp = tilde(ctx, sampler, right, left, vi)
+    acclogp!(vi, logp)
+    return left
 end
 
 """
     tilde_observe(ctx, sampler, right, left, vi)
 
-This method is applied in the generated code for observed constants, e.g., `1.0 ~ Normal()`.
+Handle observed constants, e.g., `1.0 ~ Normal()`, accumulate the log probability, and return the
+observed value.
+
 Falls back to `tilde(ctx, sampler, right, left, vi)`.
 """
 function tilde_observe(ctx, sampler, right, left, vi)
-    return tilde(ctx, sampler, right, left, vi)
+    logp = tilde(ctx, sampler, right, left, vi)
+    acclogp!(vi, logp)
+    return left
 end
 
 
@@ -103,24 +111,48 @@ function observe(spl::Sampler, weight)
     error("DynamicPPL.observe: unmanaged inference algorithm: $(typeof(spl))")
 end
 
+# If parameters exist, they are used and not overwritten.
 function assume(
-    spl::Union{SampleFromPrior, SampleFromUniform},
+    spl::SampleFromPrior,
     dist::Distribution,
     vn::VarName,
     vi::VarInfo,
 )
     if haskey(vi, vn)
         if is_flagged(vi, vn, "del")
             unset_flag!(vi, vn, "del")
-            r = spl isa SampleFromUniform ? init(dist) : rand(dist)
+            r = rand(dist)
             vi[vn] = vectorize(dist, r)
+            settrans!(vi, false, vn)
             setorder!(vi, vn, get_num_produce(vi))
         else
-        r = vi[vn]
+            r = vi[vn]
         end
     else
-        r = isa(spl, SampleFromUniform) ? init(dist) : rand(dist)
+        r = rand(dist)
         push!(vi, vn, r, dist, spl)
+        settrans!(vi, false, vn)
+    end
+    return r, Bijectors.logpdf_with_trans(dist, r, istrans(vi, vn))
+end
+
+# Always overwrites the parameters with new ones.
+function assume(
+    spl::SampleFromUniform,
+    dist::Distribution,
+    vn::VarName,
+    vi::VarInfo,
+)
+    if haskey(vi, vn)
+        unset_flag!(vi, vn, "del")
+        r = init(dist)
+        vi[vn] = vectorize(dist, r)
+        settrans!(vi, true, vn)
+        setorder!(vi, vn, get_num_produce(vi))
+    else
+        r = init(dist)
+        push!(vi, vn, r, dist, spl)
+        settrans!(vi, true, vn)
     end
     # NOTE: The importance weight is not correctly computed here because
     #       r is genereated from some uniform distribution which is different from the prior
@@ -191,13 +223,15 @@ 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.
+Handle broadcasted assumed variables, e.g., `x .~ MvNormal()` (where `x` does not occur in the
+model inputs), accumulate the log probability, and return the sampled value.
 
 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)
+    value, logp = dot_tilde(ctx, sampler, right, left, vn, inds, vi)
+    acclogp!(vi, logp)
+    return value
 end
 
 
@@ -367,24 +401,30 @@ 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.
+Handle broadcasted observed values, e.g., `x .~ MvNormal()` (where `x` does occur the model inputs),
+accumulate the log probability, and return the observed value.
 
 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)
+    logp = dot_tilde(ctx, sampler, right, left, vi)
+    acclogp!(vi, logp)
+    return left
 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)`.
+Handle broadcasted observed constants, e.g., `[1.0] .~ MvNormal()`, accumulate the log
+probability, and return the observed value.
+
+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)
+    logp = dot_tilde(ctx, sampler, right, left, vi)
+    acclogp!(vi, logp)
+    return left
 end
 
 

src/utils.jl

@@ -117,8 +117,8 @@ end
 
 # ROBUST INITIALISATIONS
 # Uniform rand with range 2; ref: https://mc-stan.org/docs/2_19/reference-manual/initialization.html
-randrealuni() = Real(2rand())
-randrealuni(args...) = map(Real, 2rand(args...))
+randrealuni() = 4 * rand() - 2
+randrealuni(args...) = 4 .* rand(args...) .- 2
 
 const Transformable = Union{TransformDistribution, SimplexDistribution, PDMatDistribution}
 

src/varinfo.jl

@@ -420,6 +420,7 @@ end
 
 # Get all vns of variables belonging to spl
 _getvns(vi::AbstractVarInfo, spl::Sampler) = _getvns(vi, spl.selector, Val(getspace(spl)))
+_getvns(vi::AbstractVarInfo, spl::Union{SampleFromPrior, SampleFromUniform}) = _getvns(vi, Selector(), Val(()))
 _getvns(vi::UntypedVarInfo, s::Selector, space) = view(vi.metadata.vns, _getidcs(vi, s, space))
 function _getvns(vi::TypedVarInfo, s::Selector, space)
     return _getvns(vi.metadata, _getidcs(vi, s, space))
@@ -820,6 +821,7 @@ Return the current value(s) of the random variables sampled by `spl` in `vi`.
 The value(s) may or may not be transformed to Euclidean space.
 """
 getindex(vi::AbstractVarInfo, spl::SampleFromPrior) = copy(getall(vi))
+getindex(vi::AbstractVarInfo, spl::SampleFromUniform) = copy(getall(vi))
 getindex(vi::UntypedVarInfo, spl::Sampler) = copy(getval(vi, _getranges(vi, spl)))
 function getindex(vi::TypedVarInfo, spl::Sampler)
     # Gets the ranges as a NamedTuple

src/varname.jl

@@ -139,6 +139,10 @@ _issubrange(i::Colon, j::ConcreteIndex) = true
 
 A macro that returns an instance of `VarName` given the symbol or expression of a Julia variable, 
 e.g. `@varname x[1,2][1+5][45][3]` returns `VarName{:x}(((1, 2), (6,), (45,), (3,)))`.
+
+!!! compat "Julia 1.5"
+    Using `begin` in an indexing expression to refer to the first index requires at least
+    Julia 1.5.
 """
 macro varname(expr::Union{Expr, Symbol})
     return esc(varname(expr))
@@ -177,8 +181,12 @@ end
 """
     @vinds(expr)
 
-Returns a tuple of tuples of the indices in `expr`. For example, `@vinds x[1, :][2]` returns 
+Returns a tuple of tuples of the indices in `expr`. For example, `@vinds x[1, :][2]` returns
 `((1, Colon()), (2,))`.
+
+!!! compat "Julia 1.5"
+    Using `begin` in an indexing expression to refer to the first index requires at least
+    Julia 1.5.
 """
 macro vinds(expr::Union{Expr, Symbol})
     return esc(vinds(expr))
@@ -188,7 +196,11 @@ vinds(expr::Symbol) = Expr(:tuple)
 function vinds(expr::Expr)
     if Meta.isexpr(expr, :ref)
         ex = copy(expr)
-        Base.replace_ref_end!(ex)
+        @static if VERSION < v"1.5.0-DEV.666"
+            Base.replace_ref_end!(ex)
+        else
+            Base.replace_ref_begin_end!(ex)
+        end
         last = Expr(:tuple, ex.args[2:end]...)
         init = vinds(ex.args[1]).args
         return Expr(:tuple, init..., last)
@@ -197,7 +209,6 @@ function vinds(expr::Expr)
     end
 end
 
-
 @generated function inargnames(::VarName{s}, ::Model{_F, argnames}) where {s, argnames, _F}
     return s in argnames
 end

test/compiler.jl

@@ -543,6 +543,14 @@ end
         var_expr = :(x[2:3,2:3][[1,2],[1,2]])
         @test vsym(var_expr) == :x
         @test vinds(var_expr) == :(((2:3, 2:3), ([1, 2], [1, 2])))
+
+        var_expr = :(x[end])
+        @test vsym(var_expr) == :x
+        @test vinds(var_expr) == :((($lastindex(x),),))
+
+        var_expr = :(x[1, end])
+        @test vsym(var_expr) == :x
+        @test vinds(var_expr) == :(((1, $lastindex(x, 2)),))
     end
     @testset "user-defined variable name" begin
         @model f1() = begin