Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add fix and unfix #488

Merged
merged 21 commits into from
Jul 13, 2023
Merged
Show file tree
Hide file tree
Changes from 20 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
28 changes: 27 additions & 1 deletion docs/src/api.md
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,33 @@ 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.
yebai marked this conversation as resolved.
Show resolved Hide resolved
torfjelde marked this conversation as resolved.
Show resolved Hide resolved

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
torfjelde marked this conversation as resolved.
Show resolved Hide resolved
```

## Utilities

It is possible to manually increase (or decrease) the accumulated log density from within a model function.
Expand Down Expand Up @@ -321,4 +348,3 @@ dot_tilde_assume
tilde_observe
dot_tilde_observe
```

2 changes: 2 additions & 0 deletions src/DynamicPPL.jl
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,8 @@ export AbstractVarInfo,
pointwise_loglikelihoods,
condition,
decondition,
fix,
unfix,
# Convenience macros
@addlogprob!,
@submodel
Expand Down
52 changes: 45 additions & 7 deletions src/compiler.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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))
Expand Down Expand Up @@ -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!!)(
Expand Down Expand Up @@ -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!!)(
Expand Down
Loading