Unify ChoiceMap and Selection via AddressTree

README.md

@@ -1,6 +1,6 @@
 # Gen.jl
 
-[![Build Status](https://travis-ci.org/probcomp/Gen.svg?branch=master)](https://travis-ci.org/probcomp/Gen.jl)
+[![Build Status](https://travis-ci.org/probcomp/Gen.jl.svg?branch=master)](https://travis-ci.org/probcomp/Gen.jl)
 [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://probcomp.github.io/Gen.jl/stable)
 [![](https://img.shields.io/badge/docs-dev-blue.svg)](https://probcomp.github.io/Gen.jl/dev)
 

docs/src/ref/choice_maps.md

@@ -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
+Value
+```
+
+### 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 Value 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 `Value`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
@@ -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.
 
@@ -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).

docs/src/ref/distributions.md

@@ -1,5 +1,11 @@
 # Probability Distributions
 
+In Gen, a probability distribution is a generative function which makes a single random choice
+and returns the value of this choice.  The choicemap for a probability distribution
+is always a [`Value`](@ref).  In addition to supporting the regular `GFI` methods,
+every distribution supports the methods [`random`](@ref) and [`logpdf`](@ref), described 
+in the [Distribution API](@ref custom_distributions).
+
 Gen provides a library of built-in probability distributions, and two ways of
 writing custom distributions, both of which are explained below:
 

docs/src/ref/extending.md

@@ -110,7 +110,7 @@ Gen's Distribution interface directly, as defined below.
 Probability distributions are singleton types whose supertype is `Distribution{T}`, where `T` indicates the data type of the random sample.
 
 ```julia
-abstract type Distribution{T} end
+abstract type Distribution{T} <: GenerativeFunction{T, DistributionTrace} end
 ```
 
 A new Distribution type must implement the following methods:
@@ -146,6 +146,9 @@ has_output_grad
 logpdf_grad
 ```
 
+Any custom distribution will automatically be a `GenerativeFunction` since `Distribution <: GenerativeFunction`;
+implementations of all GFI methods are automatically provided in terms of `random` and `logpdf`.
+
 ## Custom generative functions
 
 We recommend the following steps for implementing a new type of generative function, and also looking at the implementation for the [`DynamicDSLFunction`](@ref) type as an example.

examples/pmmh/model.jl

@@ -150,7 +150,7 @@ get_call_record(trace::CollapsedHMMTrace) = trace.vector.call
 has_choices(trace::CollapsedHMMTrace) = length(trace.vector.call.retval) > 0
 get_choices(trace::CollapsedHMMTrace) = CollapsedHMMChoiceMap(get_choices(trace.vector))
 
-struct CollapsedHMMChoiceMap <: ChoiceMap
+struct CollapsedHMMChoiceMap <: AddressTree{Value}
     y_assignment::VectorDistTraceChoiceMap
 end
 

src/Gen.jl

@@ -33,22 +33,23 @@ export load_generated_functions, @load_generated_functions
 # built-in extensions to the reverse mode AD
 include("backprop.jl")
 
-# addresses and address selections
-include("address.jl")
-
-# abstract and built-in concrete choice map data types
-include("choice_map.jl")
+# address and address trees
+# (including choicemaps and selections)
+include("address_tree/address_tree.jl")
 
 # a homogeneous trie data type (not for use as choice map)
 include("trie.jl")
 
+# built-in data types for arg-diff and ret-diff values
+include("diff.jl")
+
 # generative function interface
 include("gen_fn_interface.jl")
 
-# built-in data types for arg-diff and ret-diff values
-include("diff.jl")
+# distribution abstract type
+include("distribution.jl")
 
-# built-in probability disributions
+# built-in probability disributions; distribution dsl; combinators
 include("modeling_library/modeling_library.jl")
 
 # optimization of trainable parameters