[WIP] Add pure logdensity to Model

src/DynamicPPL.jl

@@ -120,12 +120,14 @@ include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
 include("varinfo.jl")
+include("simple_varinfo.jl")
 include("threadsafe.jl")
 include("context_implementations.jl")
 include("compiler.jl")
 include("prob_macro.jl")
 include("compat/ad.jl")
 include("loglikelihoods.jl")
+include("bijectors.jl")
 include("submodel_macro.jl")
 
 end # module

src/bijectors.jl

@@ -0,0 +1,32 @@
+@generated function _bijector(md::NamedTuple{names}; tuplify=false) where {names}
+    expr = Expr(:tuple)
+    for n in names
+        e = quote
+            if length(md.$n.dists) == 1 &&
+               md.$n.dists[1] isa $(Distributions.UnivariateDistribution)
+                $(Bijectors).bijector(md.$n.dists[1])
+            elseif tuplify
+                $(Bijectors.Stacked)(
+                    map($(Bijectors).bijector, tuple(md.$n.dists...)), md.$n.ranges
+                )
+            else
+                $(Bijectors.Stacked)(map($(Bijectors).bijector, md.$n.dists), md.$n.ranges)
+            end
+        end
+        push!(expr.args, e)
+    end
+
+    return quote
+        bs = NamedTuple{$names}($expr)
+        return $(Bijectors).NamedBijector(bs)
+    end
+end
+
+"""
+    bijector(varinfo::VarInfo; tuplify=false)
+
+Returns a `NamedBijector` which can transform different variants of `varinfo`.
+
+If `tuplify` is true, then a type-stable bijector will be returned.
+"""
+Bijectors.bijector(vi::TypedVarInfo; kwargs...) = _bijector(vi.metadata; kwargs...)

src/compiler.jl

@@ -84,11 +84,17 @@ end
 
 function model(mod, linenumbernode, expr, warn)
     modelinfo = build_model_info(expr)
+    modelinfo_logjoint = deepcopy(modelinfo)
 
     # Generate main body
     modelinfo[:body] = generate_mainbody(mod, modelinfo[:modeldef][:body], warn)
 
-    return build_output(modelinfo, linenumbernode)
+    # Generate logjoint
+    modelinfo_logjoint[:body] = generate_mainbody_logdensity(
+        mod, modelinfo_logjoint[:modeldef][:body], warn
+    )
+
+    return build_output(modelinfo, modelinfo_logjoint, linenumbernode)
 end
 
 """
@@ -339,14 +345,7 @@ hasmissing(T::Type{<:AbstractArray{TA}}) where {TA<:AbstractArray} = hasmissing(
 hasmissing(T::Type{<:AbstractArray{>:Missing}}) = true
 hasmissing(T::Type) = false
 
-"""
-    build_output(modelinfo, linenumbernode)
-
-Builds the output expression.
-"""
-function build_output(modelinfo, linenumbernode)
-    ## Build the anonymous evaluator from the user-provided model definition.
-
+function build_evaluator(modelinfo)
     # Remove the name.
     evaluatordef = deepcopy(modelinfo[:modeldef])
     delete!(evaluatordef, :name)
@@ -369,6 +368,55 @@ function build_output(modelinfo, linenumbernode)
     # Replace the user-provided function body with the version created by DynamicPPL.
     evaluatordef[:body] = modelinfo[:body]
 
+    return evaluatordef
+end
+
+function build_logjoint(modelinfo)
+    # Remove the name.
+    def = deepcopy(modelinfo[:modeldef])
+    def[:name] = :logjoint
+
+    # Add the internal arguments to the user-specified arguments (positional + keywords).
+    @gensym T
+    def[:args] = vcat(
+        [
+            :(__model__::$(DynamicPPL.Model)),
+            :(__sampler__::$(DynamicPPL.AbstractSampler)),
+            :(__context__::$(DynamicPPL.AbstractContext)),
+            :(__variables__),
+        ],
+        modelinfo[:allargs_exprs],
+        [Expr(:kw, :(::Type{$T}), :Float64), ]
+    )
+
+    # Add the type-parameter.
+    def[:whereparams] = (def[:whereparams]..., T)
+
+    # Delete the keyword arguments.
+    def[:kwargs] = []
+
+    # Replace the user-provided function body with the version created by DynamicPPL.
+    def[:body] = quote
+        __lp__ = zero($T)
+        $(modelinfo[:body])
+        return __lp__
+    end
+
+    return def
+end
+
+"""
+    build_output(modelinfo, linenumbernode)
+
+Builds the output expression.
+"""
+function build_output(modelinfo, modelinfo_logjoint, linenumbernode)
+    ## Build logjoint.
+    logjointdef = build_logjoint(modelinfo_logjoint)
+
+    ## Build the anonymous evaluator from the user-provided model definition.
+    evaluatordef = build_evaluator(modelinfo)
+
     ## Build the model function.
 
     # Extract the named tuple expression of all arguments and the default values.
@@ -378,18 +426,20 @@ function build_output(modelinfo, linenumbernode)
     # Update the function body of the user-specified model.
     # We use a name for the anonymous evaluator that does not conflict with other variables.
     modeldef = modelinfo[:modeldef]
-    @gensym evaluator
+    @gensym evaluator logjoint
     # We use `MacroTools.@q begin ... end` instead of regular `quote ... end` to ensure
     # that no new `LineNumberNode`s are added apart from the reference `linenumbernode`
     # to the call site
     modeldef[:body] = MacroTools.@q begin
         $(linenumbernode)
         $evaluator = $(MacroTools.combinedef(evaluatordef))
+        $logjoint = $(MacroTools.combinedef(logjointdef))
         return $(DynamicPPL.Model)(
             $(QuoteNode(modeldef[:name])),
             $evaluator,
             $allargs_namedtuple,
             $defaults_namedtuple,
+            $logjoint,
         )
     end
 
@@ -447,3 +497,150 @@ end
 
 floatof(::Type{T}) where {T<:Real} = typeof(one(T) / one(T))
 floatof(::Type) = Real # fallback if type inference failed
+
+#####################
+###  `logdensity` ###
+#####################
+function generate_mainbody_logdensity(mod, expr, warn)
+    return generate_mainbody_logdensity!(mod, Symbol[], expr, warn)
+end
+
+generate_mainbody_logdensity!(mod, found, x, warn) = x
+function generate_mainbody_logdensity!(mod, found, sym::Symbol, warn)
+    if sym in DEPRECATED_INTERNALNAMES
+        newsym = Symbol(:_, sym, :__)
+        Base.depwarn(
+            "internal variable `$sym` is deprecated, use `$newsym` instead.",
+            :generate_mainbody_logdensity!,
+        )
+        return generate_mainbody_logdensity!(mod, found, newsym, warn)
+    end
+
+    if warn && sym in INTERNALNAMES && sym ∉ found
+        @warn "you are using the internal variable `$sym`"
+        push!(found, sym)
+    end
+
+    return sym
+end
+function generate_mainbody_logdensity!(mod, found, expr::Expr, warn)
+    # Do not touch interpolated expressions
+    expr.head === :$ && return expr.args[1]
+
+    # If it's a macro, we expand it
+    if Meta.isexpr(expr, :macrocall)
+        return generate_mainbody_logdensity!(
+            mod, found, macroexpand(mod, expr; recursive=true), warn
+        )
+    end
+
+    # If it's a return, we instead return `__lp__`.
+    if Meta.isexpr(expr, :return)
+        returnbody = Expr(
+            :block,
+            map(x -> generate_mainbody_logdensity!(mod, found, x, warn), expr.args)...,
+        )
+        return :($(returnbody); return __lp__)
+    end
+
+    # Modify dotted tilde operators.
+    args_dottilde = getargs_dottilde(expr)
+    if args_dottilde !== nothing
+        L, R = args_dottilde
+        left = generate_mainbody_logdensity!(mod, found, L, warn)
+        return Base.remove_linenums!(
+            generate_dot_tilde_logdensity(
+                left, generate_mainbody_logdensity!(mod, found, R, warn)
+            ),
+        )
+    end
+
+    # Modify tilde operators.
+    args_tilde = getargs_tilde(expr)
+    if args_tilde !== nothing
+        L, R = args_tilde
+        left = generate_mainbody_logdensity!(mod, found, L, warn)
+        return Base.remove_linenums!(
+            generate_tilde_logdensity(
+                left, generate_mainbody_logdensity!(mod, found, R, warn)
+            ),
+        )
+    end
+
+    return Expr(
+        expr.head,
+        map(x -> generate_mainbody_logdensity!(mod, found, x, warn), expr.args)...,
+    )
+end
+
+function generate_tilde_logdensity(left, right)
+    # If the LHS is a literal, it is always an observation
+    if !(left isa Symbol || left isa Expr)
+        return quote
+            __lp__ += $(DynamicPPL.tilde_observe)(
+                __context__,
+                __sampler__,
+                $(DynamicPPL.check_tilde_rhs)($right),
+                $left,
+                __varinfo__,
+            )
+        end
+    end
+
+    @gensym vn inds isassumption
+
+    # If it's not present in args of the model, we need to extract it from `__variables__`.
+    return quote
+        $vn = $(varname(left))
+        $inds = $(vinds(left))
+        $isassumption = $(DynamicPPL.isassumption(left))
+        if $isassumption
+            $(vsym(left)) = __variables__.$(vsym(left))
+        end
+        __lp__ += $(DynamicPPL.tilde_observe)(
+            __context__,
+            __sampler__,
+            $(DynamicPPL.check_tilde_rhs)($right),
+            $left,
+            $vn,
+            $inds,
+            nothing,
+        )
+    end
+end
+
+function generate_dot_tilde_logdensity(left, right)
+    # If the LHS is a literal, it is always an observation
+    if !(left isa Symbol || left isa Expr)
+        return quote
+            __lp__ += $(DynamicPPL.dot_tilde_observe)(
+                __context__,
+                __sampler__,
+                $(DynamicPPL.check_tilde_rhs)($right),
+                $left,
+                __varinfo__,
+            )
+        end
+    end
+
+    # Otherwise it is determined by the model or its value,
+    # if the LHS represents an observation
+    @gensym vn inds isassumption
+    return quote
+        $vn = $(varname(left))
+        $inds = $(vinds(left))
+        $isassumption = $(DynamicPPL.isassumption(left)) || $left === missing
+        if $isassumption
+            $(vsym(left)) = __variables__.$(vsym(left))
+        end
+        __lp__ += $(DynamicPPL.dot_tilde_observe)(
+            __context__,
+            __sampler__,
+            $(DynamicPPL.check_tilde_rhs)($right),
+            $left,
+            $vn,
+            $inds,
+            nothing,
+        )
+    end
+end