Try #278:

Author: bors[bot]

Project.toml

@@ -5,6 +5,7 @@ version = "0.12.3"
 [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"

src/DynamicPPL.jl

@@ -6,6 +6,7 @@ using Distributions
 using Bijectors
 
 using AbstractMCMC: AbstractMCMC
+using BangBang: BangBang
 using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
@@ -67,6 +68,7 @@ export AbstractVarInfo,
     vectorize,
     # Model
     Model,
+    ContextualModel,
     getmissings,
     getargnames,
     generated_quantities,
@@ -81,6 +83,7 @@ export AbstractVarInfo,
     PriorContext,
     MiniBatchContext,
     PrefixContext,
+    ConditionContext,
     assume,
     dot_assume,
     observe,
@@ -99,6 +102,8 @@ export AbstractVarInfo,
     logprior,
     logjoint,
     pointwise_loglikelihoods,
+    condition,
+    decondition,
     # Convenience macros
     @addlogprob!,
     @submodel
@@ -129,5 +134,6 @@ include("prob_macro.jl")
 include("compat/ad.jl")
 include("loglikelihoods.jl")
 include("submodel_macro.jl")
+include("contextual_model.jl")
 
 end # module

src/compiler.jl

@@ -23,7 +23,8 @@ function isassumption(expr::Union{Symbol,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__)
+               $(DynamicPPL.inmissings)($vn, __model__) ||
+               $(DynamicPPL.contextual_isassumption)(__context__, $vn)
                 true
             else
                 # Evaluate the LHS
@@ -33,6 +34,9 @@ function isassumption(expr::Union{Symbol,Expr})
     end
 end
 
+contextual_isassumption(context::AbstractContext, vn) = false
+contextual_isassumption(context::ConditionContext, vn::VarName) = !(haskey(context, vn))
+
 # failsafe: a literal is never an assumption
 isassumption(expr) = :(false)
 
@@ -437,11 +441,14 @@ function build_output(modelinfo, linenumbernode)
     modeldef[:body] = MacroTools.@q begin
         $(linenumbernode)
         $evaluator = $(MacroTools.combinedef(evaluatordef))
-        return $(DynamicPPL.Model)(
-            $(QuoteNode(modeldef[:name])),
-            $evaluator,
+        return $(DynamicPPL.condition)(
+            $(DynamicPPL.Model)(
+                $(QuoteNode(modeldef[:name])),
+                $evaluator,
+                $allargs_namedtuple,
+                $defaults_namedtuple,
+            ),
             $allargs_namedtuple,
-            $defaults_namedtuple,
         )
     end
 

src/context_implementations.jl

@@ -133,6 +133,28 @@ function tilde_assume!(context, right, vn, inds, vi)
     return value
 end
 
+function tilde_assume!(context::ConditionContext, right, vn, inds, vi)
+    if haskey(context, vn)
+        # Extract value.
+        value = if inds isa Tuple{}
+            getfield(context.values, getsym(vn))
+        else
+            _getindex(getfield(context.values, getsym(vn)), inds)
+        end
+
+        # Should we even do this?
+        if haskey(vi, vn)
+            vi[vn] = value
+        end
+
+        tilde_observe!(context.context, right, value, vn, inds, vi)
+    else
+        value = tilde_assume!(context.context, right, vn, inds, vi)
+    end
+
+    return value
+end
+
 # observe
 """
     tilde_observe(context::SamplingContext, right, left, vname, vinds, vi)
@@ -217,6 +239,10 @@ function tilde_observe!(context, right, left, vi)
     return left
 end
 
+function tilde_observe!(context::ConditionContext, right, left, vi)
+    return tilde_observe!(context.context, right, left, vi)
+end
+
 function assume(rng, spl::Sampler, dist)
     return error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
 end
@@ -419,6 +445,28 @@ function dot_tilde_assume!(context, right, left, vn, inds, vi)
     return value
 end
 
+function dot_tilde_assume!(context::ConditionContext, right, left, vn, inds, vi)
+    if haskey(context, vn)
+        # Extract value.
+        value = if inds isa Tuple{}
+            getfield(context.values, sym)
+        else
+            _getindex(getfield(context.values, sym), inds)
+        end
+
+        # Should we even do this?
+        if haskey(vi, vn)
+            vi[vn] = value
+        end
+
+        dot_tilde_observe!(context.context, right, value, vn, inds, vi)
+    else
+        value = dot_tilde_assume!(context.context, right, left, vn, inds, vi)
+    end
+
+    return value
+end
+
 # `dot_assume`
 function dot_assume(
     dist::MultivariateDistribution, var::AbstractMatrix, vns::AbstractVector{<:VarName}, vi
@@ -637,6 +685,9 @@ function dot_tilde_observe!(context, right, left, vi)
     acclogp!(vi, logp)
     return left
 end
+function dot_tilde_observe!(context::ConditionContext, right, left, vi)
+    return dot_tilde_observe!(context.context, right, left, vi)
+end
 
 # Falls back to non-sampler definition.
 function dot_observe(::AbstractSampler, dist, value, vi)

src/contexts.jl

@@ -106,3 +106,73 @@ function prefix(::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
         VarName{Symbol(Prefix, PREFIX_SEPARATOR, Sym)}(vn.indexing)
     end
 end
+
+struct ConditionContext{Vars,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
+
+@generated function drop_missings(nt::NamedTuple{names,values}) where {names,values}
+    names_expr = Expr(:tuple)
+    values_expr = Expr(:tuple)
+
+    for (n, v) in zip(names, values.parameters)
+        if !(v <: Missing)
+            push!(names_expr.args, QuoteNode(n))
+            push!(values_expr.args, :(nt.$n))
+        end
+    end
+
+    return :(NamedTuple{$names_expr}($values_expr))
+end
+
+function ConditionContext(context::ConditionContext, child_context::AbstractContext)
+    return ConditionContext(context.values, child_context)
+end
+function ConditionContext(values::NamedTuple)
+    return ConditionContext(values, DefaultContext())
+end
+
+function ConditionContext(values::NamedTuple, context::AbstractContext)
+    values_wo_missing = drop_missings(values)
+    return ConditionContext{typeof(values_wo_missing)}(values_wo_missing, context)
+end
+
+# Try to avoid nested `ConditionContext`.
+function ConditionContext(values::NamedTuple{Vars}, context::ConditionContext) where {Vars}
+    # Note that this potentially overrides values from `context`, thus giving
+    # precedence to the outmost `ConditionContext`.
+    return ConditionContext(merge(context.values, values), context.context)
+end
+
+function Base.haskey(context::ConditionContext{vars}, vn::VarName{sym}) where {vars,sym}
+    # TODO: Add possibility of indexed variables, e.g. `x[1]`, etc.
+    return sym in vars
+end
+
+function Base.haskey(
+    context::ConditionContext{vars}, vn::AbstractArray{<:VarName{sym}}
+) where {vars,sym}
+    # TODO: Add possibility of indexed variables, e.g. `x[1]`, etc.
+    return sym in vars
+end
+
+# TODO: Can we maybe do this in a better way?
+# When no second argument is given, we remove _all_ conditioned variables.
+# TODO: Should we remove this and just return `context.context`?
+# That will work better if `Model` becomes like `ContextualModel`.
+decondition(context::ConditionContext) = ConditionContext(NamedTuple(), context.context)
+function decondition(context::ConditionContext, sym)
+    return ConditionContext(BangBang.delete!!(context.values, sym), context.context)
+end
+function decondition(context::ConditionContext, sym, syms...)
+    return decondition(
+        ConditionContext(BangBang.delete!!(context.values, sym), context.context), syms...
+    )
+end

src/contextual_model.jl

@@ -0,0 +1,29 @@
+struct ContextualModel{Ctx<:AbstractContext,M<:Model} <: AbstractModel
+    context::Ctx
+    model::M
+end
+
+function contextualize(model::AbstractModel, context::AbstractContext)
+    return ContextualModel(context, model)
+end
+
+# TODO: What do we do for other contexts? Could handle this in general if we had a
+# notion of wrapper-, primitive-context, etc.
+function (cmodel::ContextualModel{<:ConditionContext})(
+    varinfo::AbstractVarInfo, context::AbstractContext
+)
+    # Wrap `context` in the model-associated `ConditionContext`, but now using `context` as
+    # `ConditionContext` child.
+    return cmodel.model(varinfo, ConditionContext(cmodel.context.values, context))
+end
+
+condition(model::AbstractModel, values) = contextualize(model, ConditionContext(values))
+condition(model::AbstractModel; values...) = condition(model, (; values...))
+function condition(cmodel::ContextualModel{<:ConditionContext}, values)
+    return contextualize(cmodel.model, ConditionContext(values, cmodel.context))
+end
+
+decondition(model::AbstractModel, args...) = model
+function decondition(cmodel::ContextualModel{<:ConditionContext}, syms...)
+    return contextualize(cmodel.model, decondition(cmodel.context, syms...))
+end

src/model.jl

@@ -1,3 +1,5 @@
+abstract type AbstractModel <: AbstractProbabilisticProgram end
+
 """
     struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults}
         name::Symbol
@@ -32,8 +34,7 @@ julia> Model{(:y,)}(f, (x = 1.0, y = 2.0), (x = 42,)) # with special definition
 Model{typeof(f),(:x, :y),(:x,),(:y,),Tuple{Float64,Float64},Tuple{Int64}}(f, (x = 1.0, y = 2.0), (x = 42,))
 ```
 """
-struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults} <:
-       AbstractProbabilisticProgram
+struct Model{F,argnames,defaultnames,missings,Targs,Tdefaults} <: AbstractModel
     name::Symbol
     f::F
     args::NamedTuple{argnames,Targs}
@@ -82,7 +83,7 @@ Sample from the `model` using the `sampler` with random number generator `rng` a
 The method resets the log joint probability of `varinfo` and increases the evaluation
 number of `sampler`.
 """
-function (model::Model)(
+function (model::AbstractModel)(
     rng::Random.AbstractRNG,
     varinfo::AbstractVarInfo=VarInfo(),
     sampler::AbstractSampler=SampleFromPrior(),
@@ -91,7 +92,7 @@ function (model::Model)(
     return model(varinfo, SamplingContext(rng, sampler, context))
 end
 
-(model::Model)(context::AbstractContext) = model(VarInfo(), context)
+(model::AbstractModel)(context::AbstractContext) = model(VarInfo(), context)
 function (model::Model)(varinfo::AbstractVarInfo, context::AbstractContext)
     if Threads.nthreads() == 1
         return evaluate_threadunsafe(model, varinfo, context)
@@ -100,17 +101,17 @@ function (model::Model)(varinfo::AbstractVarInfo, context::AbstractContext)
     end
 end
 
-function (model::Model)(args...)
+function (model::AbstractModel)(args...)
     return model(Random.GLOBAL_RNG, args...)
 end
 
 # without VarInfo
-function (model::Model)(rng::Random.AbstractRNG, sampler::AbstractSampler, args...)
+function (model::AbstractModel)(rng::Random.AbstractRNG, sampler::AbstractSampler, args...)
     return model(rng, VarInfo(), sampler, args...)
 end
 
 # without VarInfo and without AbstractSampler
-function (model::Model)(rng::Random.AbstractRNG, context::AbstractContext)
+function (model::AbstractModel)(rng::Random.AbstractRNG, context::AbstractContext)
     return model(rng, VarInfo(), SampleFromPrior(), context)
 end
 

src/varinfo.jl

@@ -124,18 +124,18 @@ end
 
 function VarInfo(
     rng::Random.AbstractRNG,
-    model::Model,
+    model::AbstractModel,
     sampler::AbstractSampler=SampleFromPrior(),
     context::AbstractContext=DefaultContext(),
 )
     varinfo = VarInfo()
     model(rng, varinfo, sampler, context)
     return TypedVarInfo(varinfo)
 end
-VarInfo(model::Model, args...) = VarInfo(Random.GLOBAL_RNG, model, args...)
+VarInfo(model::AbstractModel, args...) = VarInfo(Random.GLOBAL_RNG, model, args...)
 
 # without AbstractSampler
-function VarInfo(rng::Random.AbstractRNG, model::Model, context::AbstractContext)
+function VarInfo(rng::Random.AbstractRNG, model::AbstractModel, context::AbstractContext)
     return VarInfo(rng, model, SampleFromPrior(), context)
 end
 

test/Project.toml

@@ -22,7 +22,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 AbstractMCMC = "2.1, 3.0"
 AbstractPPL = "0.1.3"
 Bijectors = "0.9.5"
-Distributions = "0.24, 0.25"
+Distributions = "< 0.25.11"
 DistributionsAD = "0.6.3"
 Documenter = "0.26.1, 0.27"
 ForwardDiff = "0.10.12"