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

Simplify ChoiceMap interface/architecture via ValueChoiceMaps #263

Draft
wants to merge 17 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all 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
83 changes: 80 additions & 3 deletions docs/src/ref/choice_maps.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,58 @@ ChoiceMap
Choice maps are constructed by users to express observations and/or constraints on the traces of generative functions.
Choice maps are also returned by certain Gen inference methods, and are used internally by various Gen inference methods.

A choicemap a tree, whose leaf nodes store a single value, and whose internal nodes provide addresses
for sub-choicemaps. Leaf nodes have type:
```@docs
ValueChoiceMap
```

### Example Usage Overview

Choicemaps store values nested in a tree where each node posesses an address for each subtree.
A leaf-node choicemap simply contains a value, and has its value looked up via:
```julia
value = choicemap[]
```
If a choicemap has a value choicemap at address `:a`, the value it stores is looked up via:
```julia
value = choicemap[:a]
```
A choicemap may also have a non-value choicemap stored at an address. For instance,
if a choicemap has another choicemap stored at address `:a`, and this internal choicemap
has a valuechoicemap stored at address `:b` and another at `:c`, we could perform the following lookups:
```julia
value1 = choicemap[:a => :b]
value2 = choicemap[:a => :c]
```
Nesting can be arbitrarily deep, and the keys can be arbitrary values; for instance
choicemaps can be constructed with values at the following nested addresses:
```julia
value = choicemap[:a => :b => :c => 4 => 1.63 => :e]
value = choicemap[:a => :b => :a => 2 => "alphabet" => :e]
```
To get a sub-choicemap, use `get_submap`:
```julia
value1 = choicemap[:a => :b]
submap = get_submap(choicemap, :a)
value1 == submap[:b] # is true

value_submap = get_submap(choicemap, :a => :b)
value_submap[] == value1 # is true
```
One can think of `ValueChoiceMap`s at storing being a choicemap which has a value at "nesting level zero",
while other choicemaps have values at "nesting level" one or higher.

### Interface

Choice maps provide the following methods:
```@docs
get_submap
get_submaps_shallow
has_value
get_value
get_submap
get_values_shallow
get_submaps_shallow
get_nonvalue_submaps_shallow
to_array
from_array
get_selected
Expand All @@ -23,7 +68,7 @@ Note that none of these methods mutate the choice map.

Choice maps also implement:

- `Base.isempty`, which tests of there are no random choices in the choice map
- `Base.isempty`, which returns `false` if the choicemap contains no value or submaps, and `true` otherwise.

- `Base.merge`, which takes two choice maps, and returns a new choice map containing all random choices in either choice map. It is an error if the choice maps both have values at the same address, or if one choice map has a value at an address that is the prefix of the address of a value in the other choice map.

Expand All @@ -50,3 +95,35 @@ choicemap
set_value!
set_submap!
```

### Implementing custom choicemap types

To implement a custom choicemap, one must implement
`get_submap` and `get_submaps_shallow`.
To avoid method ambiguity with the default
`get_submap(::ChoiceMap, ::Pair)`, one must implement both
```julia
get_submap(::CustomChoiceMap, addr)
```
and
```julia
get_submap(::CustomChoiceMap, addr::Pair)
```
To use the default implementation of `get_submap(_, ::Pair)`,
one may define
```julia
get_submap(c::CustomChoiceMap, addr::Pair) = _get_choicemap(c, addr)
```

Once `get_submap` and `get_submaps_shallow` are defined, default
implementations are provided for:
- `has_value`
- `get_value`
- `get_values_shallow`
- `get_nonvalue_submaps_shallow`
- `to_array`
- `get_selected`

If one wishes to support `from_array`, they must implement
`_from_array`, as described in the documentation for
[`from_array`](@ref).
2 changes: 1 addition & 1 deletion src/Gen.jl
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ include("backprop.jl")
include("address.jl")

# abstract and built-in concrete choice map data types
include("choice_map.jl")
include("choice_map/choice_map.jl")

# a homogeneous trie data type (not for use as choice map)
include("trie.jl")
Expand Down
Loading