Add fix and unfix (#488)

* aded FixedContext and everything that goes with it * initial work on making fix compat with compiler * added support for dot tilde * exported and added testing for fix and unfix * added equivalent support to condition * fixed some docstrings * added lots of documentation plus some doctests for fixing * added docs on fix * bump patch version * renamd getvalue and hasvalue for contexts to more descriptive get_conditioned_value and has_conditioned_value * formatting * Update src/model.jl * fixeed typo in docstring * fixed docs * Apply suggestions from code review Co-authored-by: David Widmann <[email protected]> * Update src/model.jl * Apply suggestions from code review * Update Project.toml * Update docs/src/api.md Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: Hong Ge <[email protected]> Co-authored-by: David Widmann <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
TuringLang · Jul 13, 2023 · 7d312cd · 7d312cd · yebai · Jul 13, 2023
1 parent e8172f0
commit 7d312cd
Show file tree

Hide file tree

Showing 7 changed files with 705 additions and 48 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.23.1"
+version = "0.23.2"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -82,6 +82,34 @@ Similarly, one can specify with [`AbstractPPL.decondition`](@ref) that certain,
 decondition
 ```
 
+## Fixing and unfixing
+
+We can also _fix_ a collection of variables in a [`Model`](@ref) to certain using [`fix`](@ref).
+
+This might seem quite similar to the aforementioned [`condition`](@ref) and its siblings,
+but they are indeed different operations:
+
+  - `condition`ed variables are considered to be _observations_, and are thus
+    included in the computation [`logjoint`](@ref) and [`loglikelihood`](@ref),
+    but not in [`logprior`](@ref).
+  - `fix`ed variables are considered to be _constant_, and are thus not included
+    in any log-probability computations.
+
+The differences are more clearly spelled out in the docstring of [`fix`](@ref) below.
+
+```@docs
+fix
+DynamicPPL.fixed
+```
+
+The difference between [`fix`](@ref) and [`condition`](@ref) is described in the docstring of [`fix`](@ref) above.
+
+Similarly, we can [`unfix`](@ref) variables, i.e. return them to their original meaning:
+
+```@docs
+unfix
+```
+
 ## Utilities
 
 It is possible to manually increase (or decrease) the accumulated log density from within a model function.
@@ -327,4 +355,3 @@ dot_tilde_assume
 tilde_observe
 dot_tilde_observe
 ```
-
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
@@ -119,6 +119,8 @@ export AbstractVarInfo,
     pointwise_loglikelihoods,
     condition,
     decondition,
+    fix,
+    unfix,
     # Convenience macros
     @addlogprob!,
     @submodel

diff --git a/src/compiler.jl b/src/compiler.jl
@@ -66,8 +66,8 @@ 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)
+    if hasconditioned(context, vn)
+        val = getconditioned(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
@@ -76,14 +76,48 @@ function contextual_isassumption(context::ConditionContext, vn)
         end
     end
 
-    # We might have nested contexts, e.g. `ContextionContext{.., <:PrefixContext{..., <:ConditionContext}}`
+    # We might have nested contexts, e.g. `ConditionContext{.., <: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
 
+isfixed(expr, vn) = false
+isfixed(::Union{Symbol,Expr}, vn) = :($(DynamicPPL.contextual_isfixed)(__context__, $vn))
+
+"""
+    contextual_isfixed(context, vn)
+
+Return `true` if `vn` is considered fixed by `context`.
+"""
+contextual_isfixed(::IsLeaf, context, vn) = false
+function contextual_isfixed(::IsParent, context, vn)
+    return contextual_isfixed(childcontext(context), vn)
+end
+function contextual_isfixed(context::AbstractContext, vn)
+    return contextual_isfixed(NodeTrait(context), context, vn)
+end
+function contextual_isfixed(context::PrefixContext, vn)
+    return contextual_isfixed(childcontext(context), prefix(context, vn))
+end
+function contextual_isfixed(context::FixedContext, vn)
+    if hasfixed(context, vn)
+        val = getfixed(context, vn)
+        # TODO: Do we even need the `>: Missing`, i.e. does it even help the compiler?
+        if eltype(val) >: Missing && val === missing
+            return false
+        else
+            return true
+        end
+    end
+
+    # We might have nested contexts, e.g. `FixedContext{.., <:PrefixContext{..., <:FixedContext}}`
+    # so we defer to `childcontext` if we haven't concluded that anything yet.
+    return contextual_isfixed(childcontext(context), vn)
+end
+
 # If we're working with, say, a `Symbol`, then we're not going to `view`.
 maybe_view(x) = x
 maybe_view(x::Expr) = :(@views($x))
@@ -341,12 +375,14 @@ function generate_tilde(left, right)
             $(AbstractPPL.drop_escape(varname(left))), $dist
         )
         $isassumption = $(DynamicPPL.isassumption(left, vn))
-        if $isassumption
+        if $(DynamicPPL.isfixed(left, vn))
+            $left = $(DynamicPPL.getfixed_nested)(__context__, $vn)
+        elseif $isassumption
             $(generate_tilde_assume(left, dist, vn))
         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)
+                $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
             end
 
             $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
@@ -400,12 +436,14 @@ function generate_dot_tilde(left, right)
             $(AbstractPPL.drop_escape(varname(left))), $right
         )
         $isassumption = $(DynamicPPL.isassumption(left, vn))
-        if $isassumption
+        if $(DynamicPPL.isfixed(left, vn))
+            $left .= $(DynamicPPL.getfixed_nested)(__context__, $vn)
+        elseif $isassumption
             $(generate_dot_tilde_assume(left, right, vn))
         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)
+                $left .= $(DynamicPPL.getconditioned_nested)(__context__, $vn)
             end
 
             $value, __varinfo__ = $(DynamicPPL.dot_tilde_observe!!)(