[Merged by Bors] - condition and decondition

Project.toml

@@ -1,10 +1,11 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.13.2"
+version = "0.14.0"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
+BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -15,6 +16,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 AbstractMCMC = "2, 3.0"
 AbstractPPL = "0.2"
+BangBang = "0.3"
 Bijectors = "0.5.2, 0.6, 0.7, 0.8, 0.9"
 ChainRulesCore = "0.9.7, 0.10"
 Distributions = "0.23.8, 0.24, 0.25"

src/DynamicPPL.jl

@@ -9,6 +9,7 @@ using AbstractMCMC: AbstractMCMC
 using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
+using BangBang: BangBang
 
 using Random: Random
 
@@ -81,6 +82,7 @@ export AbstractVarInfo,
     PriorContext,
     MiniBatchContext,
     PrefixContext,
+    ConditionContext,
     assume,
     dot_assume,
     observe,
@@ -99,6 +101,8 @@ export AbstractVarInfo,
     logprior,
     logjoint,
     pointwise_loglikelihoods,
+    condition,
+    decondition,
     # Convenience macros
     @addlogprob!,
     @submodel

src/compiler.jl

@@ -20,19 +20,66 @@ function isassumption(expr::Union{Symbol,Expr})
 
     return quote
         let $vn = $(varname(expr))
-            # This branch should compile nicely in all cases except for partial missing data
-            # For example, when `expr` is `:(x[i])` and `x isa Vector{Union{Missing, Float64}}`
-            if !$(DynamicPPL.inargnames)($vn, __model__) ||
-               $(DynamicPPL.inmissings)($vn, __model__)
-                true
+            if $(DynamicPPL.contextual_isassumption)(__context__, $vn)
+                # Considered an assumption by `__context__` which means either:
+                # 1. We hit the default implementation, e.g. using `DefaultContext`,
+                #    which in turn means that we haven't considered if it's one of
+                #    the model arguments, hence we need to check this.
+                # 2. We are working with a `ConditionContext` _and_ it's NOT in the model arguments,
+                #    i.e. we're trying to condition one of the latent variables.
+                #    In this case, the below will return `true` since the first branch
+                #    will be hit.
+                # 3. We are working with a `ConditionContext` _and_ it's in the model arguments,
+                #    i.e. we're trying to override the value. This is currently NOT supported.
+                #    TODO: Support by adding context to model, and use `model.args`
+                #    as the default conditioning. Then we no longer need to check `inargnames`
+                #    since it will all be handled by `contextual_isassumption`.
+                if !($(DynamicPPL.inargnames)($vn, __model__)) ||
+                   $(DynamicPPL.inmissings)($vn, __model__)
+                    true
+                else
+                    $(maybe_view(expr)) === missing
+                end
             else
-                # Evaluate the LHS
-                $(maybe_view(expr)) === missing
+                false
             end
         end
     end
 end
 
+"""
+    contextual_isassumption(context, vn)
+
+Return `true` if `vn` is considered an assumption by `context`.
+
+The default implementation for `AbstractContext` always returns `true`.
+"""
+contextual_isassumption(::IsLeaf, context, vn) = true
+function contextual_isassumption(::IsParent, context, vn)
+    return contextual_isassumption(childcontext(context), vn)
+end
+function contextual_isassumption(context::AbstractContext, vn)
+    return contextual_isassumption(NodeTrait(context), context, vn)
+end
+function contextual_isassumption(context::ConditionContext, vn)
+    if hasvalue(context, vn)
+        val = getvalue(context, vn)
+        # TODO: Do we even need the `>: Missing`, i.e. does it even help the compiler?
+        if eltype(val) >: Missing && val === missing
+            return true
+        else
+            return false
+        end
+    end
+
+    # We might have nested contexts, e.g. `ContextionContext{.., <:PrefixContext{..., <:ConditionContext}}`
+    # so we defer to `childcontext` if we haven't concluded that anything yet.
+    return contextual_isassumption(childcontext(context), vn)
+end
+function contextual_isassumption(context::PrefixContext, vn)
+    return contextual_isassumption(childcontext(context), prefix(context, vn))
+end
+
 # failsafe: a literal is never an assumption
 isassumption(expr) = :(false)
 
@@ -93,7 +140,7 @@ variables.
 
 # Example
 ```jldoctest; setup=:(using Distributions)
-julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(MvNormal(1, 1.0), randn(1, 2), @varname(x)); string(vns[end])
+julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(MvNormal([1.0, 1.0], [1.0 0.0; 0.0 1.0]), randn(2, 2), @varname(x)); string(vns[end])
 "x[:,2]"
 
 julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(Normal(), randn(1, 2), @varname(x[:])); string(vns[end])
@@ -351,6 +398,11 @@ function generate_tilde(left, right)
                 __varinfo__,
             )
         else
+            # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
+            if !$(DynamicPPL.inargnames)($vn, __model__)
+                $left = $(DynamicPPL.getvalue_nested)(__context__, $vn)
+            end
+
             $(DynamicPPL.tilde_observe!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),
@@ -395,6 +447,11 @@ function generate_dot_tilde(left, right)
                 __varinfo__,
             )
         else
+            # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
+            if !$(DynamicPPL.inargnames)($vn, __model__)
+                $left .= $(DynamicPPL.getvalue_nested)(__context__, $vn)
+            end
+
             $(DynamicPPL.dot_tilde_observe!)(
                 __context__,
                 $(DynamicPPL.check_tilde_rhs)($right),

src/context_implementations.jl

@@ -14,8 +14,17 @@ alg_str(spl::Sampler) = string(nameof(typeof(spl.alg)))
 require_gradient(spl::Sampler) = false
 require_particles(spl::Sampler) = false
 
-_getindex(x, inds::Tuple) = _getindex(view(x, first(inds)...), Base.tail(inds))
+_getindex(x, inds::Tuple) = _getindex(Base.maybeview(x, first(inds)...), Base.tail(inds))
 _getindex(x, inds::Tuple{}) = x
+_getvalue(x, vn::VarName{sym}) where {sym} = _getindex(getproperty(x, sym), vn.indexing)
+function _getvalue(x, vns::AbstractVector{<:VarName{sym}}) where {sym}
+    val = getproperty(x, sym)
+
+    # This should work with both cartesian and linear indexing.
+    return map(vns) do vn
+        _getindex(val, vn)
+    end
+end
 
 # assume
 """
@@ -177,13 +186,14 @@ tilde_observe(::PriorContext, sampler, right, left, vi) = 0
 function tilde_observe(context::MiniBatchContext, right, left, vi)
     return context.loglike_scalar * tilde_observe(context.context, right, left, vi)
 end
-function tilde_observe(context::MiniBatchContext, right, left, vname, vi)
-    return context.loglike_scalar * tilde_observe(context.context, right, left, vname, vi)
+function tilde_observe(context::MiniBatchContext, sampler, right, left, vi)
+    return context.loglike_scalar *
+           tilde_observe(context.context, sampler, right, left, vname, vi)
 end
 
 # `PrefixContext`
-function tilde_observe(context::PrefixContext, right, left, vname, vi)
-    return tilde_observe(context.context, right, left, prefix(context, vname), vi)
+function tilde_observe(context::PrefixContext, right, left, vi)
+    return tilde_observe(context.context, right, left, vi)
 end
 
 """

src/contexts.jl

@@ -251,3 +251,176 @@ function prefix(::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
         VarName{Symbol(Prefix, PREFIX_SEPARATOR, Sym)}(vn.indexing)
     end
 end
+
+struct ConditionContext{Names,Values,Ctx<:AbstractContext} <: AbstractContext
+    values::Values
+    context::Ctx
+
+    function ConditionContext{Values}(
+        values::Values, context::AbstractContext
+    ) where {names,Values<:NamedTuple{names}}
+        return new{names,typeof(values),typeof(context)}(values, context)
+    end
+end
+
+function ConditionContext(values::NamedTuple)
+    return ConditionContext(values, DefaultContext())
+end
+function ConditionContext(values::NamedTuple, context::AbstractContext)
+    return ConditionContext{typeof(values)}(values, context)
+end
+
+# Try to avoid nested `ConditionContext`.
+function ConditionContext(
+    values::NamedTuple{Names}, context::ConditionContext
+) where {Names}
+    # Note that this potentially overrides values from `context`, thus giving
+    # precedence to the outmost `ConditionContext`.
+    return ConditionContext(merge(context.values, values), childcontext(context))
+end
+
+function Base.show(io::IO, context::ConditionContext)
+    return print(io, "ConditionContext($(context.values), $(childcontext(context)))")
+end
+
+NodeTrait(context::ConditionContext) = IsParent()
+childcontext(context::ConditionContext) = context.context
+setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)
+
+"""
+    hasvalue(context, vn)
+
+Return `true` if `vn` is found in `context`.
+"""
+hasvalue(context, vn) = false
+
+function hasvalue(context::ConditionContext{vars}, vn::VarName{sym}) where {vars,sym}
+    return sym in vars
+end
+function hasvalue(
+    context::ConditionContext{vars}, vn::AbstractArray{<:VarName{sym}}
+) where {vars,sym}
+    return sym in vars
+end
+
+"""
+    getvalue(context, vn)
+
+Return value of `vn` in `context`.
+"""
+function getvalue(context::AbstractContext, vn)
+    return error("context $(context) does not contain value for $vn")
+end
+getvalue(context::ConditionContext, vn) = _getvalue(context.values, vn)
+
+"""
+    hasvalue_nested(context, vn)
+
+Return `true` if `vn` is found in `context` or any of its descendants.
+
+This is contrast to [`hasvalue`](@ref) which only checks for `vn` in `context`,
+not recursively checking if `vn` is in any of its descendants.
+"""
+function hasvalue_nested(context::AbstractContext, vn)
+    return hasvalue_nested(NodeTrait(hasvalue_nested, context), context, vn)
+end
+hasvalue_nested(::IsLeaf, context, vn) = hasvalue(context, vn)
+function hasvalue_nested(::IsParent, context, vn)
+    return hasvalue(context, vn) || hasvalue_nested(childcontext(context), vn)
+end
+function hasvalue_nested(context::PrefixContext, vn)
+    return hasvalue_nested(childcontext(context), prefix(context, vn))
+end
+
+"""
+    getvalue_nested(context, vn)
+
+Return the value of the parameter corresponding to `vn` from `context` or its descendants.
+
+This is contrast to [`getvalue`](@ref) which only returns the value `vn` in `context`,
+not recursively looking into its descendants.
+"""
+function getvalue_nested(context::AbstractContext, vn)
+    return getvalue_nested(NodeTrait(getvalue_nested, context), context, vn)
+end
+function getvalue_nested(::IsLeaf, context, vn)
+    return error("context $(context) does not contain value for $vn")
+end
+function getvalue_nested(context::PrefixContext, vn)
+    return getvalue_nested(childcontext(context), prefix(context, vn))
+end
+function getvalue_nested(::IsParent, context, vn)
+    return if hasvalue(context, vn)
+        getvalue(context, vn)
+    else
+        getvalue_nested(childcontext(context), vn)
+    end
+end
+
+"""
+    condition([context::AbstractContext,] values::NamedTuple)
+    condition([context::AbstractContext]; values...)
+
+Return `ConditionContext` with `values` and `context` if `values` is non-empty,
+otherwise return `context` which is [`DefaultContext`](@ref) by default.
+
+See also: [`decondition`](@ref)
+"""
+condition(; values...) = condition(DefaultContext(), NamedTuple(values))
+condition(values::NamedTuple) = condition(DefaultContext(), values)
+condition(context::AbstractContext, values::NamedTuple{()}) = context
+condition(context::AbstractContext, values::NamedTuple) = ConditionContext(values, context)
+condition(context::AbstractContext; values...) = condition(context, NamedTuple(values))
+
+"""
+    decondition(context::AbstractContext, syms...)
+
+Return `context` but with `syms` no longer conditioned on.
+
+Note that this recursively traverses contexts, deconditioning all along the way.
+
+See also: [`condition`](@ref)
+"""
+decondition(::IsLeaf, context, args...) = context
+function decondition(::IsParent, context, args...)
+    return setchildcontext(context, decondition(childcontext(context), args...))
+end
+decondition(context, args...) = decondition(NodeTrait(context), context, args...)
+function decondition(context::ConditionContext)
+    return decondition(childcontext(context))
+end
+function decondition(context::ConditionContext, sym)
+    return condition(
+        decondition(childcontext(context), sym), BangBang.delete!!(context.values, sym)
+    )
+end
+function decondition(context::ConditionContext, sym, syms...)
+    return decondition(
+        condition(
+            decondition(childcontext(context), syms...),
+            BangBang.delete!!(context.values, sym),
+        ),
+        syms...,
+    )
+end
+
+"""
+    conditioned(context::AbstractContext)
+
+Return `NamedTuple` of values that are conditioned on under context`.
+
+Note that this will recursively traverse the context stack and return
+a merged version of the condition values.
+"""
+function conditioned(context::AbstractContext)
+    return conditioned(NodeTrait(conditioned, context), context)
+end
+conditioned(::IsLeaf, context) = ()
+conditioned(::IsParent, context) = conditioned(childcontext(context))
+function conditioned(context::ConditionContext)
+    # Note the order of arguments to `merge`. The behavior of the rest of DPPL
+    # is that the outermost `context` takes precendence, hence when resolving
+    # the `conditioned` variables we need to ensure that `context.values` takes
+    # precedence over decendants of `context`.
+    return merge(context.values, conditioned(childcontext(context)))
+end