From c7310d5caee2c4ed8fcc004e407eae0af59393e5 Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Thu, 28 Mar 2024 00:06:44 -0500 Subject: [PATCH 01/19] Documentation revision start --- docs/.gitignore | 1 - docs/Project.toml | 3 +- docs/build_docs_locally.sh | 6 - docs/make.jl | 42 +- docs/pages.jl | 44 + docs/src/{ref => api}/choice_maps.md | 0 docs/src/{ref => api}/combinators.md | 0 docs/src/{ref => api}/distributions.md | 0 docs/src/{ref => api}/extending.md | 0 docs/src/{ref => api}/gfi.md | 0 docs/src/{ref => api}/importance.md | 0 .../internals/language_implementation.md | 0 .../internals/parameter_optimization.md | 0 docs/src/{ref => api}/learning.md | 0 docs/src/{ref => api}/map.md | 0 docs/src/{ref => api}/mcmc.md | 0 docs/src/{ref => api}/modeling.md | 0 .../{ref => api}/parameter_optimization.md | 0 docs/src/{ref => api}/pf.md | 0 docs/src/{ref => api}/selections.md | 0 docs/src/{ref => api}/trace_translators.md | 0 docs/src/{ref => api}/vi.md | 0 docs/src/assets/header.css | 138 ++ docs/src/assets/header.js | 96 ++ docs/src/assets/theme.css | 27 + .../src/getting_started/gaussian_processes.md | 0 .../getting_started/hidden_markov_model.md | 0 .../infinite_mixture_models.md | 0 .../linear_regression.md} | 41 +- .../getting_started/variational_inference.md | 0 docs/src/index.md | 72 +- docs/src/ref/images/rjmcmc.png | Bin 26056 -> 0 bytes docs/src/tutorials/combinators.md | 140 ++ docs/src/tutorials/languages/static_dsl.md | 19 + docs/src/tutorials/modeling_in_gen.md | 1388 +++++++++++++++++ docs/tex/mcmc.tex | 208 --- docs/tex/references.bib | 39 - 37 files changed, 1938 insertions(+), 326 deletions(-) delete mode 100644 docs/.gitignore delete mode 100755 docs/build_docs_locally.sh create mode 100644 docs/pages.jl rename docs/src/{ref => api}/choice_maps.md (100%) rename docs/src/{ref => api}/combinators.md (100%) rename docs/src/{ref => api}/distributions.md (100%) rename docs/src/{ref => api}/extending.md (100%) rename docs/src/{ref => api}/gfi.md (100%) rename docs/src/{ref => api}/importance.md (100%) rename docs/src/{ref => api}/internals/language_implementation.md (100%) rename docs/src/{ref => api}/internals/parameter_optimization.md (100%) rename docs/src/{ref => api}/learning.md (100%) rename docs/src/{ref => api}/map.md (100%) rename docs/src/{ref => api}/mcmc.md (100%) rename docs/src/{ref => api}/modeling.md (100%) rename docs/src/{ref => api}/parameter_optimization.md (100%) rename docs/src/{ref => api}/pf.md (100%) rename docs/src/{ref => api}/selections.md (100%) rename docs/src/{ref => api}/trace_translators.md (100%) rename docs/src/{ref => api}/vi.md (100%) create mode 100644 docs/src/assets/header.css create mode 100644 docs/src/assets/header.js create mode 100644 docs/src/assets/theme.css create mode 100644 docs/src/getting_started/gaussian_processes.md create mode 100644 docs/src/getting_started/hidden_markov_model.md create mode 100644 docs/src/getting_started/infinite_mixture_models.md rename docs/src/{getting_started.md => getting_started/linear_regression.md} (55%) create mode 100644 docs/src/getting_started/variational_inference.md delete mode 100644 docs/src/ref/images/rjmcmc.png create mode 100644 docs/src/tutorials/combinators.md create mode 100644 docs/src/tutorials/languages/static_dsl.md create mode 100644 docs/src/tutorials/modeling_in_gen.md delete mode 100644 docs/tex/mcmc.tex delete mode 100644 docs/tex/references.bib diff --git a/docs/.gitignore b/docs/.gitignore deleted file mode 100644 index ba39cc531..000000000 --- a/docs/.gitignore +++ /dev/null @@ -1 +0,0 @@ -Manifest.toml diff --git a/docs/Project.toml b/docs/Project.toml index 01db6a719..6dfe81a83 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,6 +1,7 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Gen = "ea4f424c-a589-11e8-07c0-fd5c91b9da4a" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" [compat] -Documenter = "0.27" +Documenter = "1" diff --git a/docs/build_docs_locally.sh b/docs/build_docs_locally.sh deleted file mode 100755 index 40d6d9002..000000000 --- a/docs/build_docs_locally.sh +++ /dev/null @@ -1,6 +0,0 @@ -#!/bin/sh - -# run this script from the Gen/ directory, it will generate HTML documentation under docs/build - -julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' -julia --project=docs/ docs/make.jl diff --git a/docs/make.jl b/docs/make.jl index 9c641d022..6807a0c4c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,37 +1,19 @@ +# Run: julia --project make.jl using Documenter, Gen +include("pages.jl") makedocs( - sitename = "Gen", modules = [Gen], - pages = [ - "Home" => "index.md", - "Getting Started" => "getting_started.md", - "Tutorials" => "tutorials.md", - "Modeling Languages and APIs" => [ - "Generative Functions" => "ref/gfi.md", - "Probability Distributions" => "ref/distributions.md", - "Built-in Modeling Language" => "ref/modeling.md", - "Generative Function Combinators" => "ref/combinators.md", - "Choice Maps" => "ref/choice_maps.md", - "Selections" => "ref/selections.md", - "Optimizing Trainable Parameters" => "ref/parameter_optimization.md", - "Trace Translators" => "ref/trace_translators.md", - "Extending Gen" => "ref/extending.md" - ], - "Standard Inference Library" => [ - "Importance Sampling" => "ref/importance.md", - "MAP Optimization" => "ref/map.md", - "Markov chain Monte Carlo" => "ref/mcmc.md", - "MAP Optimization" => "ref/map.md", - "Particle Filtering" => "ref/pf.md", - "Variational Inference" => "ref/vi.md", - "Learning Generative Functions" => "ref/learning.md" - ], - "Internals" => [ - "Optimizing Trainable Parameters" => "ref/internals/parameter_optimization.md", - "Modeling Language Implementation" => "ref/internals/language_implementation.md" - ] - ] + # linkcheck = true, + doctest = false, + clean = true, + warnonly = true, + # warnonly = [:missing_docs], + format = Documenter.HTML(; + assets = String["assets/header.js", "assets/header.css", "assets/theme.css"] + ), + sitename = "Gen.jl", + pages = pages, ) deploydocs( diff --git a/docs/pages.jl b/docs/pages.jl new file mode 100644 index 000000000..0227a90f3 --- /dev/null +++ b/docs/pages.jl @@ -0,0 +1,44 @@ +pages = [ + "Home" => "index.md", + "Getting Started" => [ + "Example 1: Linear Regression" => "getting_started/linear_regression.md", + "Example 2: Hidden Markov Models" => "getting_started/linear_regression.md", + ], + "Tutorials" => [ + "Basics" => [ + "tutorials/modeling_in_gen.md", + "tutorials/combinators.md", + ], + "Advanced" => [ + "tutorials/modeling_in_gen.md", + # "tutorials/data_driven_proposals.md", + ], + "Modeling Languages" => [ + "tutorials/languages/static_dsl.md", + ], + ], + "Modeling Languages and APIs" => [ + "Generative Functions" => "api/gfi.md", + "Probability Distributions" => "api/distributions.md", + "Built-in Modeling Language" => "api/modeling.md", + "Generative Function Combinators" => "api/combinators.md", + "Choice Maps" => "api/choice_maps.md", + "Selections" => "api/selections.md", + "Optimizing Trainable Parameters" => "api/parameter_optimization.md", + "Trace Translators" => "api/trace_translators.md", + "Extending Gen" => "api/extending.md" + ], + "Standard Inference Library" => [ + "Importance Sampling" => "api/importance.md", + "MAP Optimization" => "api/map.md", + "Markov chain Monte Carlo" => "api/mcmc.md", + "MAP Optimization" => "api/map.md", + "Particle Filtering" => "api/pf.md", + "Variational Inference" => "api/vi.md", + "Learning Generative Functions" => "api/learning.md" + ], + "Internals" => [ + "Optimizing Trainable Parameters" => "api/internals/parameter_optimization.md", + "Modeling Language Implementation" => "api/internals/language_implementation.md" + ] +] diff --git a/docs/src/ref/choice_maps.md b/docs/src/api/choice_maps.md similarity index 100% rename from docs/src/ref/choice_maps.md rename to docs/src/api/choice_maps.md diff --git a/docs/src/ref/combinators.md b/docs/src/api/combinators.md similarity index 100% rename from docs/src/ref/combinators.md rename to docs/src/api/combinators.md diff --git a/docs/src/ref/distributions.md b/docs/src/api/distributions.md similarity index 100% rename from docs/src/ref/distributions.md rename to docs/src/api/distributions.md diff --git a/docs/src/ref/extending.md b/docs/src/api/extending.md similarity index 100% rename from docs/src/ref/extending.md rename to docs/src/api/extending.md diff --git a/docs/src/ref/gfi.md b/docs/src/api/gfi.md similarity index 100% rename from docs/src/ref/gfi.md rename to docs/src/api/gfi.md diff --git a/docs/src/ref/importance.md b/docs/src/api/importance.md similarity index 100% rename from docs/src/ref/importance.md rename to docs/src/api/importance.md diff --git a/docs/src/ref/internals/language_implementation.md b/docs/src/api/internals/language_implementation.md similarity index 100% rename from docs/src/ref/internals/language_implementation.md rename to docs/src/api/internals/language_implementation.md diff --git a/docs/src/ref/internals/parameter_optimization.md b/docs/src/api/internals/parameter_optimization.md similarity index 100% rename from docs/src/ref/internals/parameter_optimization.md rename to docs/src/api/internals/parameter_optimization.md diff --git a/docs/src/ref/learning.md b/docs/src/api/learning.md similarity index 100% rename from docs/src/ref/learning.md rename to docs/src/api/learning.md diff --git a/docs/src/ref/map.md b/docs/src/api/map.md similarity index 100% rename from docs/src/ref/map.md rename to docs/src/api/map.md diff --git a/docs/src/ref/mcmc.md b/docs/src/api/mcmc.md similarity index 100% rename from docs/src/ref/mcmc.md rename to docs/src/api/mcmc.md diff --git a/docs/src/ref/modeling.md b/docs/src/api/modeling.md similarity index 100% rename from docs/src/ref/modeling.md rename to docs/src/api/modeling.md diff --git a/docs/src/ref/parameter_optimization.md b/docs/src/api/parameter_optimization.md similarity index 100% rename from docs/src/ref/parameter_optimization.md rename to docs/src/api/parameter_optimization.md diff --git a/docs/src/ref/pf.md b/docs/src/api/pf.md similarity index 100% rename from docs/src/ref/pf.md rename to docs/src/api/pf.md diff --git a/docs/src/ref/selections.md b/docs/src/api/selections.md similarity index 100% rename from docs/src/ref/selections.md rename to docs/src/api/selections.md diff --git a/docs/src/ref/trace_translators.md b/docs/src/api/trace_translators.md similarity index 100% rename from docs/src/ref/trace_translators.md rename to docs/src/api/trace_translators.md diff --git a/docs/src/ref/vi.md b/docs/src/api/vi.md similarity index 100% rename from docs/src/ref/vi.md rename to docs/src/api/vi.md diff --git a/docs/src/assets/header.css b/docs/src/assets/header.css new file mode 100644 index 000000000..155b10a3b --- /dev/null +++ b/docs/src/assets/header.css @@ -0,0 +1,138 @@ + +@media all and (max-width: 560px) { + header.navigation { + position: fixed !important; + left:0; + top: 0; + width: 100%; + } + + header.navigation div.container { + margin-left: 0rem; + } + + header.navigation div.container nav.navbar { + min-height: 1rem !important; + } + + header.navigation div.container nav.navbar ul.navbar-nav { + min-height: 1rem !important; + margin-left: 0.5rem !important; + } + + header.navigation div.container nav.navbar ul.navbar-nav li.small-item { + visibility: visible !important; + display: block !important; + margin: 0.5rem; + } + + header.navigation div.container nav.navbar ul.navbar-nav li.nav-item { + visibility: hidden; + display: none; + } + + header.navigation div.container nav.navbar ul.navbar-nav li.nav-item a { + visibility: hidden; + display: none; + } + + html:not(.theme--documenter-dark) body #documenter .docs-main { + margin-top: 2rem !important; + } +} + +@media all and (max-width: 1055px) and (min-width: 561px){ + header.navigation { + position: fixed !important; + left:0; + top: 0; + width: 100%; + } + + header.navigation div.container { + margin-left: 0rem; + } + + header.navigation div.container nav.navbar ul.navbar-nav { + width: 80% !important; + } +} + +@media all and (min-width: 1056px) { + header.navigation { + position: fixed !important; + left:0; + top: 0; + width: 100%; + } + + header.navigation div.container { + margin-left: 18rem; + } + +} + +html.theme--documenter-dark header.navigation { + background-color: #1f2424 !important; +} + +html.theme--documenter-dark header.navigation div.container { + border-bottom: 1px solid #5e6d6f; +} + +html.theme--documenter-dark header.navigation div.container nav.navbar { + background-color: #1f2424 !important; +} + +html.theme--documenter-dark header.navigation div.container nav.navbar ul.navbar-nav li.nav-item a.nav-link { + color: white; + transition: color 100ms; +} + +html.theme--documenter-dark header.navigation div.container nav.navbar ul.navbar-nav li.nav-item a.nav-link:hover { + color: #0aa8a7 +} + +html header.navigation { + background-color: white !important; +} + +html header.navigation div.container { + border-bottom: 1px solid #dbdbdb; +} + +html header.navigation div.container nav.navbar ul.navbar-nav li.nav-item a.nav-link { + color: #222; + transition: color 100ms; +} + +html header.navigation div.container nav.navbar ul.navbar-nav li.nav-item a.nav-link:hover { + color: #0aa8a7 +} + +header.navigation { + z-index: 3; +} + +header.navigation div.container nav.navbar ul.navbar-nav { + margin-left: 4rem; + min-height: 3.25rem; + width: 70%; + display: flex; + align-self: auto; + flex-direction: row; + justify-content: space-around; +} + +header.navigation div.container nav.navbar ul.navbar-nav li.nav-item { + align-self: stretch; + align-content: space-around; + justify-content: center; + display: flex; + flex-direction: column; +} + +header.navigation div.container nav.navbar ul.navbar-nav li.small-item { + visibility: hidden; + display: none; +} \ No newline at end of file diff --git a/docs/src/assets/header.js b/docs/src/assets/header.js new file mode 100644 index 000000000..b25445a32 --- /dev/null +++ b/docs/src/assets/header.js @@ -0,0 +1,96 @@ + +// We add a simple `onload` hook to inject the custom header for our `HTML`-generated pages +window.onload = function() { + //
+ const header = document.createElement('header') + header.classList.add("navigation") + header.appendChild((() => { + const container = document.createElement('div') + container.classList.add('container') + container.appendChild((() => { + const nav = document.createElement('nav') + nav.classList.add("navbar") + + nav.appendChild((() => { + const ul = document.createElement("ul") + ul.classList.add("navbar-nav") + + ul.appendChild((() => { + const smalllink = document.createElement('li') + smalllink.classList.add('small-item') + smalllink.appendChild((() => { + const a = document.createElement('a') + a.classList.add("nav-link") + a.href = 'https://gen.dev' + a.innerHTML = 'Gen.jl' + a.title = 'Gen.jl' + return a + })()) + return smalllink + })()) + + const items = [ + { title: "Home", link: "https://gen.dev" }, + { title: "Get Started", link: "https://github.com" }, + { title: "Documentation", link: "https://www.gen.dev/docs/stable/" }, + { title: "Examples", link: "https://probcomp.github.io/Gen.jl/" }, + { title: "Papers", link: "http://probcomp.csail.mit.edu" }, + { title: "Contact", link: "http://localhost:1313/" }, + { title: "GitHub", link: "https://github.com/probcomp/", icon: ["fab", "fa-github"] }, + ] + + items.forEach((item) => { + ul.appendChild(((item) => { + const li = document.createElement("li") + li.classList.add("nav-item") + li.appendChild((() => { + const a = document.createElement("a") + + if (item.icon !== undefined) { + a.appendChild((() => { + const i = document.createElement("i") + i.classList.add(...(item.icon)) + return i + })()) + } + + a.classList.add("nav-link") + a.href = item.link + a.title = item.title + + a.appendChild((() => { + const span = document.createElement("span") + span.innerHTML = ` ${item.title}` + return span + })()) + + return a + })()) + return li + })(item)) + }) + return ul + })()) + return nav + })()) + return container + })()) + + const documenterTarget = document.querySelector('#documenter'); + + documenterTarget.parentNode.insertBefore(header, documenterTarget); + + // Edit broken links in the examples, see issue #70 + const editOnGithubLinkTarget = document.querySelector('.docs-edit-link'); + + if (editOnGithubLinkTarget) { + const link = editOnGithubLinkTarget.href; + if (link.includes('docs/src/examples')) { + const fixedLink = link.replace('docs/src/', '').replace('.md', '.ipynb'); + editOnGithubLinkTarget.href = fixedLink; + console.log('Fixed link for the example: ', fixedLink) + } + } + +} + diff --git a/docs/src/assets/theme.css b/docs/src/assets/theme.css new file mode 100644 index 000000000..67017e105 --- /dev/null +++ b/docs/src/assets/theme.css @@ -0,0 +1,27 @@ +html body #documenter .docs-main { + margin-top: 4rem; +} + +html body #documenter .docs-sidebar.visible { + margin-top: 0px !important; +} + +html div.light-biglogo { + display: block; +} + +html.theme--documenter-dark div.light-biglogo { + display: none; +} + +html div.dark-biglogo { + display: none; +} + +html.theme--documenter-dark div.dark-biglogo { + display: block; +} + +#documenter .docs-main article.content img { + width: 100%; +} \ No newline at end of file diff --git a/docs/src/getting_started/gaussian_processes.md b/docs/src/getting_started/gaussian_processes.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/getting_started/hidden_markov_model.md b/docs/src/getting_started/hidden_markov_model.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/getting_started/infinite_mixture_models.md b/docs/src/getting_started/infinite_mixture_models.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/getting_started.md b/docs/src/getting_started/linear_regression.md similarity index 55% rename from docs/src/getting_started.md rename to docs/src/getting_started/linear_regression.md index 01680544b..5a936127c 100644 --- a/docs/src/getting_started.md +++ b/docs/src/getting_started/linear_regression.md @@ -1,34 +1,29 @@ -# Getting Started +# [Getting Started with Gen.jl: Linear Regression](@id gen_linear_regression) -## Installation -First, obtain Julia 1.3 or later, available [here](https://julialang.org/downloads/). +Let's write a short Gen program that does Bayesian linear regression - that is, given a set of observation points in the xy-plane, we want to find a line that fits them "well". What does "well" mean in Bayesian linear regression? Well one way of interpreting this is by proposing a line that makes the data highly likely - as in a high probability of occuring. -The Gen package can be installed with the Julia package manager. From the Julia REPL, type `]` to enter the Pkg REPL mode and then run: -``` -pkg> add Gen -``` -To test the installation locally, you can run the tests with: -```julia -using Pkg; Pkg.test("Gen") -``` +First, we need to define a _generative model_ that describes how we believe the points were generated. -## Example - -Let's write a short Gen program that does Bayesian linear regression: given a set of points in the (x, y) plane, we want to find a line that fits them well. +```math +\mu \sim N(0,2)\\ +b \sim N(0,10)\\ +\epsilon_i \sim N(0,1)\\ +y_i | x_i \sim \mu x_i + b + \epsilon_i +``` -There are three main components to a typical Gen program. +This model first randomly samples a slope ``\mu`` and an intercept ``b`` from normal distributions to define the line ``y=mx+b``. Next each x-coordinate is evaluated and perturbed with a little noise. Now let's write this as a probabilistic program. -First, we define a _generative model_: a Julia function, extended with some extra syntax, that, conceptually, simulates a fake dataset. The model below samples `slope` and `intercept` parameters, and then for each of the x-coordinates that it accepts as input, samples a corresponding y-coordinate. We name the random choices we make with `@trace`, so we can refer to them in our inference program. +The description of the line is a mathematical one, but we can write it using normal code constructs. The _generative model_ is a Julia function with a _tilde_ (~) operator for sampling. Observe that the function below looks almost the same as the generative model. -```julia +```@example linear_regression using Gen @gen function my_model(xs::Vector{Float64}) - slope = @trace(normal(0, 2), :slope) - intercept = @trace(normal(0, 10), :intercept) + slope ~ normal(0, 2) + intercept ~ normal(0, 10) for (i, x) in enumerate(xs) - @trace(normal(slope * x + intercept, 1), "y-$i") + {"y-$i"} ~ normal(slope * x + intercept, 1) end end ``` @@ -38,7 +33,7 @@ Inference programs are regular Julia code, and make use of Gen's standard infere The inference program below takes in a data set, and runs an iterative MCMC algorithm to fit `slope` and `intercept` parameters: -```julia +```@example linear_regression function my_inference_program(xs::Vector{Float64}, ys::Vector{Float64}, num_iters::Int) # Create a set of constraints fixing the # y coordinates to the observed y values @@ -53,7 +48,7 @@ function my_inference_program(xs::Vector{Float64}, ys::Vector{Float64}, num_iter # Iteratively update the slope then the intercept, # using Gen's metropolis_hastings operator. - for iter=1:num_iters + for _=1:num_iters (trace, _) = metropolis_hastings(trace, select(:slope)) (trace, _) = metropolis_hastings(trace, select(:intercept)) end @@ -67,7 +62,7 @@ end Finally, we run the inference program on some data, and get the results: -```julia +```@example linear_regression xs = [1., 2., 3., 4., 5., 6., 7., 8., 9., 10.] ys = [8.23, 5.87, 3.99, 2.59, 0.23, -0.66, -3.53, -6.91, -7.24, -9.90] (slope, intercept) = my_inference_program(xs, ys, 1000) diff --git a/docs/src/getting_started/variational_inference.md b/docs/src/getting_started/variational_inference.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/index.md b/docs/src/index.md index 161192b99..bdb9af6fd 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -2,27 +2,63 @@ *A general-purpose probabilistic programming system with programmable inference, embedded in Julia* -```@contents -Pages = [ - "getting_started.md", - "tutorials.md", - "guide.md", -] -Depth = 2 +- What does Gen provide. +!!! note + `Gen.jl` is still under active development. If you find a a bug or wish to share ideas for improvement, feel free to visit the Github site or contact at us here. + +## Installation + +The Gen package can be installed with the Julia package manager. From the Julia REPL, type `]` to enter the Pkg REPL mode and then run: +``` +pkg> add Gen ``` +!!! note + Alternatively, + + ```julia + julia> import Pkg; Pkg.add("Gen") + ``` + +To test the installation locally, you can run the tests with: +```julia +using Pkg; Pkg.test("Gen") +``` + +## Getting Started +To see a overview of the package, check out the [examples](getting_started/linear_regression.md). For a deep-dive on how to do inference with Gen.jl, check out the [tutorials](tutorials/modeling_in_gen.md). -Reference ```@contents Pages = [ - "ref/modeling.md", - "ref/combinators.md", - "ref/assignments.md", - "ref/selections.md", - "ref/parameter_optimization.md", - "ref/inference.md", - "ref/gfi.md", - "ref/distributions.md" - "ref/extending.md", -] + "getting_started/linear_regression.md", + "getting_started/hidden_markov_model.md", + "getting_started/gaussian_processes.md" + ] Depth = 2 ``` + +## Contributing +See the [Developer's Guide](https://gen.dev) on how to contribute to the Gen ecosystem. + +## Supporting and Citing +- this repo is academic research +- please star repo +- please cite +```bibtex +@inproceedings{Cusumano-Towner:2019:GGP:3314221.3314642, + author = {Cusumano-Towner, Marco F. and Saad, Feras A. and Lew, Alexander K. and Mansinghka, Vikash K.}, + title = {Gen: A General-purpose Probabilistic Programming System with Programmable Inference}, + booktitle = {Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation}, + series = {PLDI 2019}, + year = {2019}, + isbn = {978-1-4503-6712-7}, + location = {Phoenix, AZ, USA}, + pages = {221--236}, + numpages = {16}, + url = {http://doi.acm.org/10.1145/3314221.3314642}, + doi = {10.1145/3314221.3314642}, + acmid = {3314642}, + publisher = {ACM}, + address = {New York, NY, USA}, + keywords = {Markov chain Monte Carlo, Probabilistic programming, sequential Monte Carlo, variational inference}, +} +``` \ No newline at end of file diff --git a/docs/src/ref/images/rjmcmc.png b/docs/src/ref/images/rjmcmc.png deleted file mode 100644 index ce1cdb8d1f26632547aa6c545e73bacb643504e9..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 26056 zcmb?@Rajh6vt{G%?t$RJ-Ccu6pb747K?99ja1E9aAduj$!JQD?J-7uA?z8#t+_^9F zFb~7m{Q;-t^x0Lls%q_89igeNh=op$4uL?hl$GROK_Ku@2n3D|6&W0%5g$qczYtwz zl;5C&9~i3nNAN$Ivy!eW1cGDobcZXD%(npth1}$!Zm*pz-8@a+TRk zpLtlkceQbHvcL4fi(<1@ z^7zO))nV?ru!kv=2fnZmJZ%*LJgt4gTc9!PEMbrW?QQbx8voY5yL47(uUyT^M8|bf`C6{l!|f?qj14N`rphQbl?DL1kcNgr{~Cgi^gYrdUiZ> z?*Ff+ZDiC+DQ_s<5z7`HT@)=2wjj8nH49axh#>OpLLvz}V-+H#P6re-8{0 zw2+0M+S=J^>FZ;LnP|Rw6Y}rhdnj89cS~JO4dL6jZ@Cxh*49klXlwV{-^10`)*6#q zcAHff6`^BxAs#rXsbNjpx_56zMMi#WEj7Pm?jR3OXq_A*E`{KiOTp1gK%s2Ji>@Q$ z1!ZN}?*zNJr0c(Z`-Z3IW-mZN8WuQa{pG9`%U|L+_?vcdu%D)B%R^Na&#b+`L^(-stL*)f3>SSeS5;(n%QU>8UV9zb0gWvT%G2 zgjce!0H;!rmzV!_qnm=%=h9i)+Dd(UaUkpBa?&s@Ip$jpWn*S#4eB??_tX6Q_b=pG z!3*B$y>f6xFD1~Flfl`#;Amn{qf5hF%{ z;QSKPjeWk&*3On2}LYCVb>t@7_g@kG}%1&w)J~U85+Id#+iQtz9hpo8zU5 ziq`Ab1T|ReyN8PnE@rdDB~+pw-KE{Zax@IkB(8Dkx4OC!%q+xi6ac=El zKIJ(mccbnMud^Ze1Z=8|6KP2=J`aTf*8`5rU2ljK?YaDSZ!c$NuWzxNX79Y*utmh? zul}#Iw_v$SwWEMX(qdyt&Y?}>3J3_8F4Q>`6ct6L#mCgxKdq0$LRe#?NZOF)liv_W zA)TL}cX}ZHMxm=MdaJDsmJ=*7HLyU%yHeVd7GU}B-@gmF|0Q%2S!yt)8mJoGJD2Wj z9kZsa^cl6RmM+x}U)+gY-H=9`rk58f08u0B_bW^QE^ae`s@NQzP|7qYbk06@mPf#C zDG=T{xZ~N~km*k+#v~j4BfPzHyV`k|`guHj2etulUPC1>a&qz@^({7Hl~L}ZnLRB7 z1FE|39(fk*#28P3Xk>(wsm?U%6O->oR}q^j>Iqx|odFpO3rm{YFn=Paqy&>~AWV%^ z@wK5Lb&bEfyZiHA!xRgbzmtXWdd&keO6JMNYOnHwUzX2=6SB`tPn%BsP&1P-&Sf_Y zg#YMY0WKPkASP!Io6z6W-94u1oYlDbHbr&MN;jn!^|`Wa@s#Q5){wQmy_~#!=&?9X zjl@uvq>+$tBw>`oyJ>TtK#MtLq6mr@1?$@R_@wE?#Ki2wi^N3y)kgxB;MoS3=g|te z&alJTO0!AQo+|rU3T15YfyR|krg~#2LTO$m_0zj4$D9Ym)OjJHKf4Z+H#PunJi6-! zUeOma!~fNHV%Ke5UDK?^PbN|884Eh>$X$f0v9U2kZv}nDxwqQBL{&f=>0Po7fFW8Xt?UDN)Z491E||)B|pdiHxAiNbuwWy1FdU z6F&erXlVsvQ3#;+zjknNm^A}Ff-nLwqSKrQ%>qAgok25+v|t@Q^&`?tM&Oj8?ucOu z5f{p?+2@UQbsoRJ#3%g%pRh7A5*>s+mUi)5QXgwNloIV?B3mqb0X?N`Y)LZQoC{izFmo#J<+;1)eEY(fVsl9i!zfXCo zNgN`V1m3N84`WaQx8b*w_0T4WJphkdhK8}%r`x#z$<`vel?IT*gjhH^6~Xmrs^x<9 zRnI4=qfG=cx^wPrKU)FU)6+9*TWdz{>eLKhIul5)%hTgxK3G&-ESzEydh6Ztb)oq~ zwR2l#Wdg+m6P6fQLcFThq0fS@YyFo@AOE(KZuFR6D}sEfZG-*uxLsQO}9zCCE^ zI(x3V;oEn^KUe2K0q|CKI)7qU0T#FFKeM_Z0?S+lD1i-`nVKRxyHeWYc=2NW%v&7% z+w=C#8)xwD5TQ)G4qxwG2QQZVzy_ZL&tyE-((tjr; zA<1oQ%-EX8$So*91xWQ~K^m3T1@{wiiSuh8Oda40s_U=x^;MH}64PdXx^Cx%87sXF z>`=*9c`wf$ftzTZ*xx*B%++pN`|SzK8)LtHqjPh2hs)NTRNk*(qev?D`O(;pdQli_ zicd?gtIFCzE>ivns)_qNq#M77UPh_vqh5lKl@$~4KM=ZHTe!3#^yp6yu&EUx*@>|N zN4UX>Q3(F`@83rQVPHD-{MQXfr`N`{?-DXIvBFH&s#-$guU3SxJ({Sludi_<);vvk z!5jQ?){0RaBUQyPLH43l*3p^ySpMO-14S6(mP@hpW305WkcE8pc%+6y@ zhBjNYlanTgs!aUal>YMTS88($i=3vW&%?yZ;~!4XK3|nDJ%-+$!J1D(1rr;0J}f*q zc>SAT>XNwpOI7}jtD0+AkB>j@sIXxuBz{c(bT)S3ZHW25i!Teo&odP(Y)%*}(8=4F zvKQuJtR#GpTMi8&0e)rFI)klrjz|7eUy9=ymAgV)F=m%rX;L=~@#`8l?B&b}XnQb^4I>I2N7XGXJNO^;#{onL$8UFHhA z>X&V7iSTV4NzO(8+x>;v78}pScZr1fNN29+UOji{hPBNZHE7eKR?E`2QleEvb6$rx z+s}wp2wIdwa1*Q61a~mgD=vrI!f6if@r0a5nfWT3_PrL*g!~$A!>?8kY1RrgvR`+{9K%BrvIpiuQ(99R9Q)Kxo!}wO?py2)zK0F1mr#&BB-)NollAdr zl%lmTv=c=qT#+Sb&k?*})Qen4$amiNV|#6!D)HK1-<|kx8wKo z?J1`vQg-R<6mbNxTT@?)c;d^9MlEc4RiTdrMIFRp<_3cnd60yTT<7&7Pvrk}r9i?jr@; zkRdt<4gb%NSf2-wA<$nO-HmtdCY)^XFTdkP^{jhq$SB#%J*QT4=SgSdIQ{x6=o4gw z3Wm)|=|{VEDn3lz0ONq%#a(1Z9YxRl=*di(`nL6(9A>rcc|2q0-!tD_JIm$Sm1d~t z|AQKmRqf&PR;C6X!!l&Lkoz8!jTi&8VKdwEzzUSe;dQ51Vhy54 zjvvXm-;cj{1|_9STm8UZXD=B=Aeoer9gX_=m!KsJ#j+P9KYwkKRKaJ>OUkFPy|7R0 zsbaGn%J?*C7J`@f8^G-kkp!g>Py)y@CiR^8GIzWT@Y9eKK^9Ta|DpNz_w4W`GclR61aFCgyCUq|{m9L@AXJUMp z{OiE_gd=R&9}Xln9$Y7hIjV=EC&cp;atAvmZt_nY*6UryJ@7A&8fiw7C3t_x(c%ui zcXdrjNf9XApj%GDi5o${{xXU*H#P&u+5AgqQW!!m^rmbabSsr$+&K#u@jw4OKbB zx;M9QaQk$%GZe$x(Gesp6V*ph{4%uv==q$@IjIOUWYPL2^_Hs?JK(77hJ*YaxpVe9 zr07C>Rb%E64PR@ce7iW`P+rfU;(fBP%>CPT(;?HRD=s*jSbyes^GBnv*#O7r zfh)y2WyR)x;>z+~@T7;qh;m|k?c6)mI3+VveCskwE^t&@>h`B!YtPG0=^V}W17}%! zB3hNN4D`Z~Y}^;|ELpQ3yA88Q?tSlZmE|>+2%WtstzBya6L^$Tf0#LL>2qNsG!4YM zUbs5Da5MG7$qGVCR`*e9eVg-_k&DTVPavnMJYbE9uS9 z$08bNGk6y0;{CS?3!uU87&Y%5{c zf-AvWILH4YQeBosFHK@?)qn8>YRwOjq=vJb8s)8`7>e77Mi}apb8Y{L3O5RgJ6-*E zH_(@<&yXt>w0Ch53gtE@>VNK2npB9a~l%Rxh5AskkfC(v(T#O8VbVv1nybvj2OZ!8>;z}8= zyIcLtvfFFwAFG7jBW9f8HFL@c6gzw&s_F5Jdw2mkWb<L~k&%hl)I#TtzY(lF_s9wI2O9hSh+Tr{H(LF?u4;t)poGtc{A=6E2g)5GNGyi`+r{R=55Lp~ux`J6#9Qqc zpzcn}5{m$cb`p6_wE7cdS?K3_0o#QPB{~IAV_@gc0NZc|M7Wz_=O-32M`Zk{I}lhu z$z)nNN22khP(eTbwXGd}bF(}})2EnU@+)3t<>F>^@t{4T&X6rbd6|jlUiB1l03Fic zZXq!}II&Xv*o0F~W8WcL&k?Rf2r-^Bb*k9ud#^9<*XXgY zb7aw9&B(trcz#;e=k8iAD{F@!VgG;&@CQUiy|t5aVYN-D94%oz${_&{9(Fy6Myo6S zJ%`gTS599TvEf2aR1mqFp4J^9NcY-uRE#tGe_1 z*fk58ztmG{b-trs`_iwf?IG@&`3Jf0!mv0vGvGeS1bH+OU|BU*SU#*YdoBCy5;Rvm z200tCB(3b}494AFFoOC3qpc}}M(ZK4@qWkiLuFBXkHd@84lYc$P+J&mK60sV;bgPc zHFV~@!8-)r1|qGK+Jn)>^;E8bVwdXjGB#C8kT#%zupw!8Z%CHCLR2hmh_$lieW^3u z``e_WZ@Dv}vLuaKVuJr*IpWxY?|iPzA2S$!?(DmtpiQ56EV_Od;JB4i*!t(Sy6;Il zDpi;qA)9CkA|T1ZgE`NX?bi^b5AMK&saXQs!Gp>E-jUg&qdY-HE-k%KsW55AI<3wX zPrwsrrKFaI4w?CDahNKd%x+aWBwz-c4~3mXsm#*iUn#;-<^~GdrV#c>izh4d^~<*? zBc5F=lA`x>>;oYPyT9e{oC$nl>rvPrlT`fSyGoH)J+cm4uN<8dSB0?eIW#7nB(eyx zQ*prUxSxCFB8-io@UwdZ=sZa=kwHf*TdJK4NKU{v14{9~*w)AsqX_!(*f2(jln|o%Z8%KVL7Y#Q|56Yt>Ff$SJa~j@i(% z^Vc(H=}^ip4#8I)N6Zl0kY-yWk4n0}dfKmMj>;@nZ-RkQZGpH%YH#Y_9a{n6o$lih z$>yrX#D`|T3Npzy3oBsTS&0?P)EY#;#R^ovb+uYooZT`SJXvtetR8#@FOF%fu`vfZ!x&ug^&e2wq z!<^jZ(brCn!Mxtd@=dCx0mjt(G7uRYxAZoYEn4y2iOs|>v}tAkpxS4tZJlpKpHm@` zGh~z?cv+Ss-Ky*5B)>ck=iCbs!n337GkN^!$o91m-NqeQmTAE|RJ$%tLwEJ;>RIZ? z<6j3aDqP$HD;CSv53U9jty~s!Q5nQ^=d@72hVyKpz*|u&4s2}H-qbar` zL8y$In|_q1b36a3dC`Z1UnTyZB}UxKrCkb!(_#A(i1`z-R9qc4J{h|9lYIYz397^X zyJNgr)tr@n$HmL}SO&*Xa+a?70!37!hho_%Za?~X9iryG`yTDEjl<14n5eosTRrJp z@p6pD@O)6CcK>9x%>*MPz;I-QPZRJZVt^PG*&-!ddSImRptiyCxubO0?TFBp;O-xI zm2Li~J1z?Z)7|zXEE;1en4L?|=u3G1;dm(9-nol(DXM!6Z11cQX^vKQN`I|w$EpT# zVDS%%F;|BvDe}`NjiUM5W7bR1`QG17EL?PkFxF~@abK{Url(Gvc+m2b*Hzu-=&cXS zA3@I*PbVtHuJp3(%?S(Bv!Bhgcl9I3TRFrV)h{XSD%aO1?v8p|C$#nJVT>Dzqwn}0 z=#yl-$X+1p|A`Ob&=QL68Y32UOdBu-DMyX*@mdKVsyJohDMSCB_BD(A57UjRFzfET z2`bLiAU4xoy1@Ln#PvX zR|IEefgEQaZW6K*Q1A%I#gf+>p2-$(+F3wTuo^WOf7JE`{&S#d-#jryI!d0P(0ChZ+IKJd2-|2J0B-Y;(dJTPaRUa93!JfzGazjauk70 zic|k>c@ys)2227w?m#Tip=rq38%Ddg0DXn~Z??y^s-3I-XBpg6ZqN#SjKy~9fDahA zs6Ji`x9Z?(m#?lqp8M%2c58I+Qvv#DF{{ z^JT9bID7I(GsnGcMS;`^d8l4;Ofm81%I{F#?4ka48qd!s3-z%k+hYPI20Rg{?<-O3 z`~;)J7o1Ofa$Iexh5DGXC9s92ZAtJaEp^jf5l`Tw8;0JBNl&8uANJMl28hAHoTy4F9IAwDzmw?q5I8>Fea(ti zu|2||j=v^W(Vk^2XG~BvXz^5b->P)VtH@8TS2&kdbP^rjC$fGG@S$%)NLu|^TCWDf zy>6{3BqS_)zR&OS&rE#6uzGTe_E9fv zu3?ln`i42b-mqaCSh9_q3v10O;QR9O5;a)e!^F;G$@B~*0swgZt4iK+^*(M%l?YrI zFadZfv}nlZDO#}yT}UG+0axefWgo}2Mc7_1aKHPaI7dBctw_;x734aQt=7YcK2(S< zUUgjDbAhODk~oNV)2Cq^(_+bk#wP_5fQqfApK-DX+sbj<%vNjIILLkK8|B^Z3x;0< zuS~(o^mR|ML5=z3VBL}u;yB?8;rX8FifR=A?rksMwX<4Wt!4Z=t$9z0lHV6Dk&l8i z>5CkT(3<@@nDP_l$YCvN#HTldD`otlR~P&kGJbAtu0!zphRHFYXt)n%iE_DRxPF4s zlMPX(%9b1!`=E2+c&m@$9Bxc%LFeg`CrOd%*>Vc3cxFFd!Yu?~sMq80^3XebH^(c> z8IvL-*)j~jyM=>X>qg?4B{3R=Nu;6}($H~iUtEpuoq;0%`D@yqfCUVYY!(72GR}5X z4~W#u+a5G%iTXa&F}IrPnatxikj*YGOqSYB-M`hot}gl7MpgTOTV zuU;P}(aOzD1W3R+hLzO*qp9JR##eX*2g6!slpfcqhVwPp-H;phdAMG_L6G}|{0Fa7 z%@-V)WH2{!{xvnIP8$aGk@1V57phV4={K!yCU_T=N9}>9OING0_gDQxDkuAYUJEVF z#Wk}bJv-=N?Bk(>+MazrsV#UtLZG(H#OFLKnUyK6V7w`ct7 zdrA05Y%+2$bwp+}i{hm9Cjp4Cw$^pJWDYF#WFA6%+F`8vqWE(rsDTLqTYg^@uek?q z!GF&JWOf*Yz|C-w=|70hu(b!5Q+66%@}tH1&YR1+UPAs=b;hl3oq3%A-|P4se${fGl@cT)*0RZO)v0a2T@=E^=r2%?%E^n- zg9(=YB?^1}%S@W}`9`(vfRHwX?da4t5WHR*tn-qlb(FwmX?@Jn(|F_r`zp+Bn+Gz@ zo)!r6us8d2aLMF1&ShYh;l^T*z%hHXl=i6yO&k|J0#)}VOY$pjF$rv4o9CYE7R>ZJ zt{;p&z35Wyp(f^5|Fq zVPNck>qq%lj4;`F$a3O%vo;j3E&DQ+16X{Sq&N7d?U*q}N1`rk&K7FG=tj9wMpRq+ z@Dr5Xou%wxH3}2u4KR3J9GVa?HMs>i{rV)!2b02$T^~Kfx%kUR)_7x(Z2X-$T`r{W zuo+SC(T@{kz_qRAHVgLm0S#+8QEE(+`(`3MzX#u!C}_BQV=j#Zwq`%Wub9zLR0x2qC$t>{ zjZMY}a-kN9eW$9`m|*s+zl|L^P8m%o$8j2flQ>^RP+`gSv9pI zwRp$PX}DE|nGu1Ivr>}s@xVwaDjDa#52tSs4*PWcOxc&DQ2##^BI+xZ#G*CAyVC6Z zfy8XOY15zGe#-5j?^8Q>ZAz}bR=N(lqVum}BWJkQVxhM_fV=rZeI>+U{sGkZWfBj9 zuG++sUxFFc!ltS7AAehw9}CCy9-IfR3pMoNatg;1ysO{=29%%rx4$V9KM zw7mez63sir(urM}uS{k-$w`6Bum@AIn=S5!7KxAYn5}+2u$=1_C#*O8 z!b$iZ(1+~ZGMqY*60p4325T{9PU$Y0117AnGPj9586^I#yHODiKN0(yBg-l2r=Uy6)zZ4 z4*%YQqxkuNw(p-NW#4MO6BR)t8{EoV3-M%fDdFM5je$(*@!P2P&rBh4ab151sD#{d zP}JiD3ugUpTdNCh7T}&Be!I1i9F$lRG~pw|?$2LGRxFtX%H@CWgR{$0VTR#MI>e?x zk(j@mD#_7#!6eR_nrTbr9mKJnSKlRDge5yl8=-LwQ++0D1|1oId2yu2YV_WiC|B_(&V{)mUpKQc?3 zc)Xj1zf=NEwMjuLh4x@4r;&J&Z+_{{j=QNs80Cn+GcwA^rPTNMm6WK{XGBq&hWY`( z%i&g%#xfP>+f9xUrMkHqEM3YfR7(~c=x}fz_qc}Lu~w!Kh8l(>&1JJc?Oc4SKIvPs z1U(>V5J^ss#{LZ8%=UA&$2sXo8nt*JipJF-Frs#l z2`EjWw~Nfbm2s305lVILMhQ>K^^|uJ?c@cP?ftsF<#!(heQ**iw@ifOPs}gclG1@ zm1){4rS{-(XXIR^S{|L2QqR>*-e?`ywQoHKx zinBGM)OTG5M`($HDRC~~X+RJrMI;OKud@aJ5Dk3>lehnFxmdLb!>EAPrW#4Zb?RQwK- z4=9M=CZ%s=E{);;mIn~h^!oymJWqL?3lPSIhp*TD?n~I@XDw9{K6aG8IF*gO7PyL1 zpcr<`@ll#^5<3$SpOySA#82>kpZ{UGM^M3kX>8Ju(t5scS>HF8GfRrP9Os_npHPaGWs7&~PV@|G4W8Opx_)@#6~`Uljkx z4gFCO^(%}ulHoVfti-DsJ+MOM?AXEW^`Ks3N>C^5TL}*g@!Ljp%bu`c9 zM3o58VyX9*Oev3dSgRg=6Gu5Nv)LYF?~h@E9LtRam(8P)d?vSV(`eQ85R$vscx;Wd z-g_~1UF8ssQsf84k>#&#{9}52-=y|Fy<#OU+ef2^9@6R%A#+qIrNCh*+Y?!jb19aF z%=Dq5y`c|2*LjYkS`;67^1DT61=K6@=lJKky6E=h`cVtoTV3nxO;v#r4AA=0!EnIB zl}{k(wNaoKqUuTdYmKWje*rv~-Gg^!Ukuanc%Ky<1h6!GyF6?o+j=u!A6IJPx$gfh zD~7b;K`1Xh3pvgnTi_A0$#G5xXdDlv70U^1cwCfXmQ4Mt2peoq3O&_3S1)paQD-=_ zLiIR?`zAbhvy&r3NqOeY7s%d6Xd+4`Ouq)fgHWCR6d)Hjf5P*@lnwt;Wd0>JM6GIB?%JIyE!nggF_6&6VKuL|9K^fYH|u zO#e))hN`qm(+P8sH-*FvK{*tnonOMY`;lL>nuiD+d`ewQb0qIE^j$KOXmM9rTE}Wx^R(JS!o_S+I`!-z;tpoOkiR(WAD1+`I$rd zepP$%R0aX^VE-?0kK}a1g3>3`h-&BYhpT(cdUT5sPmrSI8h@*iyMv zAF_EAaYmQ9y^Ox1<>eoWeeU0QMP`beej6{~u- zE>O$x?T(|lv%|pzwenml61B852;*Rm!e9L#F@<7-^;vrb?;^I=(pf)}`pI3YF(5r{ zcC_9}(yWH!3>#1Ug$JBx=SiIu3GL&x1E64Q<>^a%uOEsZg)Ome*G!16%Ut7Ld7mH+xyeA#wV@rylsndjf16dJeWc z;Qt8C2IUKK1pcXt@^?N_8(W&q=R~==o6uy2uG|>GMakqRkX0K#Hu6})SF2nh?&-_A z5tR77aJCh!*in!wg7BOMbPfa%IKewf!rQ|n_6Oz<)P;P+m^JgIGdX12qMYhdFzxDy z-&Ne1H4BLuI#Z%VDDk~-Ul;OQ5JQl%l#pZ?bC$q)6oUx)Kzw-H?CA|jR!Av7k;ghp zC7^vY6P)K+MY1gd6nj|L0oH=SX~JV_mHAgjo5nG^n8UBX525@ho+ap+lZGq=>)|u( zy@jz7VzAryV1Hz>!IT)V$3c%}lCg>s_3(gwlT$sGhr9XCd-A7Wz4`29P3=-}SyHlZgjzOFINDpA5b=*qodqU@->*}H4AuR8BR<4jqoHxJLo%zSmA!a*e9dPgSWQO*jDWuJS`v z>PrhnrS+fF1NVJ!dqjOe*W=2vlZVTw z3P_n8^(4`9YFQyH%gY#3%*GjQKD%VD9e_;j1_1i z0KQovFj#^s8h6Y@xJUA;*tF_qn7@O4{&JUO^*osUiCJ_M?U`94wjB?`->F&?zS=0k zqIx*2sGZ-BFUCP29(j|qBq7xBcPc8$s*$Y_N_c4!31}a%*KFeH)$T$m6x)VQp#+-l z;EsT*sw%wHN+=r_DJB>#t$?x>6cpexIz>lJ%R)%G>Z%LQ3M-!T%36C3+KS^hUB=V} za;)_82Y7QF*VaR{Hm0MSIJ2VW#VslAYS;HhP7VBcc{i^t(zL#|)tk_IP}%hXN(+z? zUjL#*2rcthd&-R>PwU_cF0@ZDI2yTcPX8?4mv_9&_Vh=gEI`0CLTHB!ND4ljqp?GO zN;N*B?82UddZg50`H|#dnIRZq+H~u^@n$Bh-uyoJ&JXKD8(cjJ4`jEh!}1dJ!|5+s zV67{b&Ujr?ti&Cg=N+3TSi}uDIM~!oqweU(AFj=9tkpi{GRXw>%=kzZZY7npdPsc0 zUM)9@XT7Fq=X4SsVCo^`%^##xQYK644KN`s`C8sloW*ap%wne$G_`c35d3JGj-~R$ zsCev6tyc>Az$vfGjtmazBToNz^Z;nF&O;`_@5k%Nl%PkjmVsM_wmFh!V>0}B} z>eP5AKFKeMjQxNm)RKB6Gik^!#=o*gP%_A+GZL~lpYZYyEoCMg&NExe6&f@(gWZ5) zN74QgRbL+WJKi{r3JGjihSjq5JL6=~deGH5#&5Xw70o@d4SjmS<(Hwk^(RTn=hyWQv9fe%vW1D>Fa?dAooA{*9N)n;zsp2u+cfl|=w! z?Dh!NY*^3}qz8ewne(DokC)Um+#T>&mRnV8qm2R1U@Il}LQ2kZq*)yaR z`_P$CVxQ{J_8y9B#FIQNrHaNI=secu3>(#pj(#8dLoKmE~VPAI5SX|Nu4?)k5t7A7kmj}k(qWTNm-<@uW9}a6&DK$Gd>px5;^=5 zgY^|~?Y}w0j68U^Nl_eV{rrd`~)uI(*{evZ$xj(68R2~CnDhA2>v zQ!{(vj`}sZ#&PZuNol)NznmyXuvNFnr@>RkZ2_Ok15!;pB)!XM@3`KOWp1<>Y4ZVb zU<|<=JN<}5bA@9d&Femy4-0nyt~|*oi%w#X3gG{cf_)(;UBX7-SDm`GZ->mdd zFwQZ$ek_&A*i+Prrq3pee+VX>p5}JwEZ?JPH|hsD85x9N)WMacXEjugCPs`w;{czD zUOf@Y;*T$`<1i&(QSAV;!uQBoyULe}x7C zfsNDx6$OOTd1Z&3rTm<+%Cz%Om1|!CwuE zlIKKGX$8ZETzD=v$Sf=@9YFOFL|qMpMW5H|?dO7SkJ^kHmM127#gj^2`&{{N8=(zC zQT6>ek572-h>Q*LQBqYj0vXpLJP5+xQ9+mww&jlJQZt~ifZ5%u5&$%stm2^-~W zdK0St$kTo={=LT+&N8osA@dz65)9BOeh+vvh`QwGe2LITZL(hoz4^Nby7AoER$oza z3JRdbdCD_DAkt+okoy;9xo4W`tL)O>$(^t27in(iYm`rXHE6g-ce4iAS5iJzl6A#( zf0beFt@czb@QJCH>&>GvTH$OY3;xJmScI-u?t8KdjX(^`q-h$kt)j(H~Nk%6Y+2PR_XR*%VBE31Bukg2fw@L!~VK0V!Yf7 zBn||-jFcjBvXCx_jzqf=yEW`%mSO@sB(|+zQu3>`+RS5oTMXloDS(b^%z@OUFiC%3 ztihBM4K8Yef2ysdewYFd^N0|jQiC8z^$0r5fBYwUgx(tD5}E#z+vbBGvEk!`cg-0h zPCnQPiV-A%t_9VAIR@YiGkD!n&%oF-!*z`HxSu)%!ri=c8{K=_vybwaTM7Ye?i#sk%NttizC>F3QG`F=gxfuYLk$ ziKH0iB;Oy0R!#V7W$;^N2y92kOCOSyE9B@_>X-r~4LfkMjvLx=#a#s?dv$8C z*xhi#O?nsUsZ%T^+J`9gb(*o%3KkLWtca9;z$7HD_Z%*bYIjlvo*y#76DFaO<@WdH z(f@d*MEf!b!x>)>PQc)ULBIg@HL+6JxNG1*YtpPkOuf|QQiNb%9jt+)GvJH-omlv| zTdRd?M(?Aznaj$MK6bz=frO9B*$mDx7*@!}W4`jlIYv5XN4dH^P`t`$&*>+k4N-NC zILiAH%W=;r+4GW~jwOpruF;nL6MHkZ*ScbR`sFn(3Sq8Wf}Cv!I@EQ{*Z zDj?SE&I=l<)B7O4KAg4q6GijGkuTTbnjK=4(;a?IsO#v}q53nmuDA&#)pQ$oRNB`@ zK$I4(@Enre&_MdGRA4pu+3TQBs5aye3NjWXa2a&1*aKAT?*c*aSL#jxQPl7%^!XHnD(}eSUxKcT7HEQnO!USnW1;_O8DNI?GU-g`v(#P7c zaRbG}4sQ}?I}6EI1b^sUEuP=Cc|ltS2fS>U`Cy)- zvb|S^5qw7JS7lSlHl_!p-^xG)@oVyR4P2W4 z)QYcdM)a9AFd;XqZ9o1PK5_y@IK~h)zVb=8Z3> z|Lev2pvDCCe8?@^K-`XFB{T%0h)R2sSi&6v6zyvP>q3}b$oT*Myo|sfBvmlrWM!3c zcjpISgC5lr^!Kl=H6Qt#w{PKAM$p|w{QiEEf*?FY6nI&Nnl8f3*h9?#4e9O(u!5vM zY>tOCUn+X~hVa?Ur==|Y-|!vicmE5Z4Sv}@jm&*G-LiHPprsG{-|vLm@k8x11ah`3 zobcbX001o@|Ctn{F0-Bla^RetCt>`_f)Y^5bAK|!e>#64m?s7&POiA5E|yC2&i?6} z$LNvN9|8p8`-?wa-%}`|Mu=) zhFRAP-yIKM1_H6H7T%EnQuar8k<8;}!*U$&SX>!XlcGL)+yAHU6QB=EtUTty+c}a_ zS{rlu6*aRQ2#jPZ#rg>za&vo2(9FcyDut%-=q~mP($blOE zNddhnDk=)<4WRV{f_0$WCG!x@pFqo>U|Bpm1O(98*_o=PD;ljbv$MXU*|B7#a!7pt zM@zAcH$g;YbwEq~ksn7kI8`AR&Gt8q?cQ{$*VSscEHx)HbK2k*zD#1;lB1%w_H^E~ zG`JXitg(-ZpuL8sgihIFc@J7v%iBBBcw2(n&x9wvlcUXIfOygW@xIeHYXqfuDsq2+ zpF2ixY+@ofCi4Q>J@A5cxZ<>R2^# zuZ^15fxfW>h`OJ?FJOnec#6C(V#%Dxq$>=|jFtJ*z-IiwrgWTIo&p^&FE2wIJzQa{ zKjgryh}40G$>#@Y{2+IP#P+1rt$}!ncdg$)*J5GjE&J@V_p_g~ z-?R7BykG?#3NrNy0~f{YCxTt^xJ&!h2LIrSp9KMT19h7UncC_L#VR~Zp^m%5`d#Qd zUGN9Hr=(V#kChtsg;`Myuln;}?>R1qE}jV+ij*kv5ZZTez4okcrf^b z!6JXOTkT>9)xW8yj4n`F1+pNZ$p-iX^o$6&xvfm_THc%aW()}mV0W*+kx?Hgxpo9U zhWh&YAH{D701LtZY%Brwz@Yqk7R>L4p^_3Ag&%3l9~Y^T@^azV&!mL8R}~A;xdz4Y zq%GxwHzx>H(h(`i>83zup#E8m_TJa$A3-tuIjHV;1W5+YS$6@EDiObs$NL8YtsZ)tao#Q?&{3r-D``VqxQL;$YFsX{V&9*jxhjFGZu4~ zpG(Ljf3aG`cuWMtkLy51%=*&B-o2#oOUAXzgzZM{1GP(wUcFW9Fo3ykvkmCc* zIS6I`(Z}r8oI3M&IiR1^c^`it`X_&bl{`#jbjyR!?8CCo{IT?qORx%2N4DyjFaXfrB3Ak*m^A^P z-0@lmy3sSjRn*kf6rfbFtr)TjZj~H)``kT1r)^RHlg{;8~ zWOrk_uZYlm)s^|GG8UwJtRAk6`icT0qZR#O-kYV-o@%4BdMvG5(8=v*rc;29APxz8 za09E9*wZSljbLRNC+x*%7q<`4E9@YYKF|LxeLjtqx~Ba1avHpkRZly0PpQL|vo;zh z>RfLUd}ahs940fmr8**sDQzJ##0_BIk`^@F{n3ZRI`ozB^ z9B%{!u)LoRL#I1p5z zzhCMD;Z7U$Mp^il6tmRFav4B>c8yhgRw6G#R8*AEu;f6dw_<5|nax#JAdZr#xh*Lv z>Cw2Hg6i+xot=z~jC&IOs1QQ^I?rI33aD*1yi#MWK4c^M!E=E%HZCq4? z!QEY)?5={2IxNzG5MMGsdcM#KfAhBu(fb<#|5t~h=q&-TZ2dPz{rSVXFNK!lT>Csr zDULfvW~=}f5;;8!!KNw5Gi@Drz$HVmg*PUD41cC zf1@00zCEBhHO|knZTbO6D}OwP6v4ptp}L)+REJ*mdYP@_H@{+u&;@~74#*u!3vvp4A{b9^OL zU5HMAF*B3y)2klP)hP#={Vw^lL2=Zj;?mMt=H`PyJL$8XJo#T+^SvI&tF=JP7b9_- zWA?Rin`a{?z8s~Edj|SgD+^OoVovnNYs`zzhNi0W zDt-C#Wmn?fgL5GBkHvZYzc&x9;KQ&>hk#(bEj+uIaX}DY;o8=>$AS%m!;Jc}hL5$h zv>efB>Wuw8?|ericDpjU5qs~ilnDEx%FcN=XJa`0eyiUw#eG@??sDL16{!QGLeYRQ z5(0^8#VQ79-n!34MXKj8L5Xl{ZKpE`q}RaSzH`UynC&KcQR|$A2J9o!lunQH;097d zL&H@xQg2=Sj4n`t1q%K8n8>c0iYqyHw5hWvk{A!lkbf*H!LMBDp%BDte)%lRPh2>{ zk(jdTD2+U{Tb%q{vPz;=zshwJi7B5-tqKngmdZl6P%6;jy1OsgR2vNRS`2#xuOA$+ zS&cnnTBu$F=h#lI z>Hh~@rC``*@;LlYw@v*Bo4oX@oQaBxiiW8!P6A<1F4*kA4f}Zi;|szq<6RA*qjokc zub0iCxDm2)3aWe9>vmp}a4QbyP+0^Ve%1nXR^f(JR8(9L`U$iH=H_Cc0!LrlQC)ZU z10qfivmJbYJ=SmdbBWg_#IVh9gFQDX6C_IhHgfm%H;zPhy9(%Uc6N5W2ecFaHYf;K zjj5cEi_6@2T5WklwE<1FL0=6=TU#3oo7AydH#NeZ&aTx@3C_bws>ey~)nap%x@z3d z;#KM9XA}3S@SyXRu}tga)jDuhyeJamh+!fZguxD3%8roVN%c^X4}^#d%k{fTrPn*R zCk?9G?2U$PM7Oi|jY)7zOP?ec;?=3au(rs=kgSVU9vXX>dVxw%x4#>c=%Gp@z%76% z3oW10$OZAMnnCqtP1D7$odkK(w}^YCZfBLZNyy5^$=Sy~uFyib-9<`bI*sf1+407K zeq=vuvj9|XlQ4Zhp=#USxe2EGHiDpQbc@r`O`+1Sa;Bc|N#iGvpHLyv_K?JBkxx6J zK>BmjmaW8<%kh(+4K-18%D`s~_i_8j829!Pj&QDad!y-~jcca6#w!W664ounRhEbAOkm5Xp zULLt(u8tD7h3|>nLV>&No?YXBRh1NW?&7RloTOvNs(bqmArP|rNK)Lj9c1@R9<=Bh znV|e0Qj}a!ymJBuqogDkj1)!+BM-Tdfo-PS3=iP&A$L&nn4BFcF=_){=oWhx-2ek2{Fn1EVPmfQvSGa&4YQDOte6l=MpT^q!FAYvRe_4n z;xOlXY6KM3Vg~x^KqDg~HxD*ZE9>hcvS)v^eAx(*ud}ObbEegQip>;dZhnFA%Gi#(fMrsc0q*(i*)xxg8R8O?W_x3)kYMu1^PiC1Ryj0c+QCHj>zR-KE61@!jU;X5C`g!<~ROQ+-?t#3O+IVo)1VKP|pHc?7T- zdfK4PjbDmaIXRD=#xIjSsQ1c#^ZbJ+`T5aEvv#qb5e-7^?8fMtmoGJ95!F($*dd8< zYn1M+rkY$n6h6DdsF1iFeL5oGT^xL(GQ3`_-{`1T(r8d@xyGLwR<(M4G|S1E$9;Xr zmdAZ|EY6;h%oZZ5Y5AnOi~-`Q)Gc5md!Y?l)Q=hV?}(aMT~F^?g$maD<@6jf!QrPL z&0B255ik`ec}^9l+L6+ziAE{-#g+zjTJN&pV6QB5Mo~9pR#ujVZ%8Et9xX%BYurrE z=8A9Nzsaq=NzVm|x~ir(=qBw*UbHQY9L4Lm-05$1y1OGnE(nG3z?Q?K+e83(LE~GE3dDo;tI&RU0YN0 zR#(7cs#I4%H4lq0Hg>aM4f*6-B-(8|J3AY3w!~hQ-0nAD4$S%q1>38wIe?V{e*NTZ zY01f7zyp+ZE7`U6^?JFUmpAvoEllJ{iFhKjH6>n zaBwPXY;v|sO9i`a2r(XQPgbOsW7gC;K-@bPq1UWq$5=-6ytxxheAm zVR7-G?(S}=!S^vhIp}~E66LA?%%Mv`8qU<(`cgUKTbu>=$FoXYw$Hcv%l>uovk5OF~{{Y{>BV5NCZCvx&`L;jl6>2b#(VDe@n7Qyv43&2Hi$1e8GUz36=#c0q zBC4{8UoU@7NcPw~p1ot>-O-;Cw0n4sLU1O_b+A_wp>R=)4pLL61&60Ss(43!18JSi zTdm<1bO2BZz{+ncC~h$vnG13z(h1*gh3&DGQ5!>H<}EOSk%Oy)nDP&+9k{nf(1BD^ zQ1vV2$j}1f-V%;X?TY|z{8IfaUQo46LW?@Bw+#Z|@J0|Ser8Yz)+?gZzO|LSDkh)Z zVb8li2nspqC8VTCzK`C^A}=LjBKYB)4v@&2#~oi~P9bP}d-*!@skbXz`u63&$A`Tn z8`?)!v+1<4eeydpz1%>Js8}5Sc?-sdF7R=%ps{xO&_auS>Ed#MDH#hfp zkv??uZVh{wi40r8a6y&xQAA{9d3E&@_gkD0ni2w9|HaPx*s40VmA5VAf-{2Vo|wwo zJxf3XoSY!Dn$yko^|r^U>0#mFu8d`24<}iGTi}}}06FnEJ6t3L3mLDSujms1&VbAk zgGB;V#^YL}?AysnqZbYiDf=qQ%GpfmAWVGg@89~)9YqJCPQiB(8jw!|&|lK8((N~I z&f`Awf^TuYe~L}hJi_vpv&)fQkbYW!){~YB&4E&G4C(XUrpijv4A+*5lAHJ=tE*pH z**^PKYTuQX4mBu#3+$4p#pL`*|70UbzI|myaLVV{?$4_5$(nm}8(44j9j`ZR(|+(? zle~eY{mUTr7KD%=uay<56o3Fw5_I$4px~fS1Y2#0;<%T? zqEB~MN#a)=v&{d|7h)-E0Ffx#)kvlG)16py$Y91r!S4 z!WP--NDElX3^$aBpxB{;nQ26I`F^sPOlYlJLf)pguo zbcq`i84;0Od)n|SIG8-Jg4ou|>eKx0nG{E;DpSPV6nfmDsJG;Y>RP8O1F`S^Z3YMp zr3I;8eJjJL|1g_(*;tG#EiQvb#3n81KSQbh;y27js^?wsVb1$ru z73M4s5RW5#a(V!ly`sKHs&tqb!x8OBa?;;V3zhd->2_k>uIx8qLU~b8P{`IkcW@}R zPfqJcfk1n2`>2+-YE&Tc;PoIaUwY~3=_}9}5Z5?$F#YHGdUp(;E7m0@rsmP}0Qv>@ zH^E1?X!Hd3H1atyeyY(M@$A{%^)54E^E<5LBTx>GVWX{uHadA1nRHEB$W*g0r+$f{ zqm12C0}g}b_`web#k|F`;k{PhW|x+}4$2tpBD?MDM^bmT5Q;;D#dTd*82$EI{kmIK zaRb8P{&MSbHk5dBBpb-2uCDIVoB`+Q13~kaal_K+iHvM_qv5TT^D)NCxHv|2P0i$U zz`K?`&O3yB%5TV^mmjBp2^DGz&dIr+Z?`}8_gIi?up1Y?d)I%!_iTIN z=%AfX?)BlJCsE7I(uOyYz9+LBc9B*h@9-m~^;0z?#@;lQm6ZoQ`c1!&8URo5y0D`c z6U?6|!WHn+o0p1;Dp)xS(3l8ghy9V?!-e|gm6iT^={{qFA^>BBMY1PuSi*Y_H`?fg zB_&4(_v@L{aNHTk$?b{(aQ`+CjFWb)t_`%Zad502zQ5SU2nTeMIX(vDkVNqW(0}#lr zO@>XhwrV-z${!a4VB%4lc}f_=wSCxw^73-7AAhRA?3p<_J_1Z?8v}=N5nUvLv4y`w z%T0G(U)8J^`&6OfYQ#(g=V1^&2V8|KK)%+LHwIEzSV(dYD;A=hwXo`h5i~1L%v#jV z?pYjB+x6nPU>nU@J9`bslS2~!uA?3pIHwm^zZ|@-BxOHb=(9=loI;g@zreUinr5h3 zxmAs`rzd7jve~|^(+u=Y$JZp$E5KopM}kb_a(MzzD=j-+~aE5kh3&bstF^0g=zKId%4XZ1onMY7YPk zbK9QVXA29UIn9VGQOo(+zC2L(oReg>&JB~A+9&L~qq*UM1={?H<8<06b~b@HH7qJh z^+KFl&Z@J3&iuA5V}RF%GT!_Oj@M^69t%w;eT;`n)}w7VskF}JiaZRh;Nj^cDo(9KaH;sikBY0)zKP_fGL6v%Qfp2;Qm#f)zWhfodx>Lo{n*e@*gbYC_6ivhYufq zcja(jn-Dfh=(7a4h29Oc8X+I9vCh=ClJadU9lxE>o0Q%QAW%?H@ETJnG0~kQ>b1#w z3Zs$DDerYIkg8j^-A}dAJw9R6Eg?v3@~?fBC?WW-fM<^17%gNYR`hdpE;&0pyGAO( zAhEp`M4Ycafk(>kI$&!Pn7bP$0m`dm89&5~o;_2POD~mvP6lQy8e~eE@7-R8v?OQu zxUFRHn2Qk+k&zHSBS1Wdgy3d_*R&U|f)ohIk`S^!|K)y!7zu|*6ZilG@_>+n`TM5c zy@`7FfrduVYI)6qF1njXM8oz_N&{O~R`z}8U~G(SVR13jnhh);mP@{`b&7<1^+iPT zUt}pG`G^oquH8Ot$htoZhWj+q_`&B-?B`73RFQ)<9V#aY(l7KzC zU+0}H_1$xH98Z`>dWD+9)N(-Ioe#b9x8>zc_7PHP%CIPj+Lv|0;1oJ=3MDv&{|e>8 zlNo1Bswc5JpIJG}lo$S=Z6C3e4sb^-nS!{uxZQyMhE4KY3`P>T$+v+4l#LCgt*z}> z%gV%*l(vb9R8L{;nUy-5PCOG6(>EA9uvl7lT>-}J<*h)4R&=Bz^g)hJ>z6Mi3kwU& zP`;0rMJh!qK<_yjz>CfASf4#JGEx$VQ@E1^wlY)d>FY!Uw^-lk(;s*5ly4|c*#4AT0mx)KC=eO=YFx3{;ZcKiVk11wDw ztmi>4<|Dwj+CX|(Om{Fs(6R~Y_w(n^4Rmt2ixG2@y7EOXIL45P2o3;P;bPRlVqgg# zJf;GYcYx>2lipc_pCsk+TODS?QQhWn0IY*nL+Dp>>C7Q;g`$!YTw_A80ZP_PmXdN0 zd80&|e`Y?xg#6_m^G}eblKRm3;|GHi#&CS?1R?kqY*cB|#Ae{7{3128Mc1A&hIZt3 zxf+SruO!6Be`_O8Nl8H%8p88+@6>NDE-peMyUH6I!{M;kU}Imh{IT0sg2W^Ga%ee} ztovm%v)^1HNKoYH5cM~O)s>)7C{o~DEh_R+#t1~q@87>M*wC(7gpm;#&u7o=+*Mnh z{3~EiUoQ4@sPMh66YkUe*4HP>*sGIoe(RwOR0ZJO&Q8_Sm##(x`!sTIH1+gk1of|;@L$a`Kw-M6yF$g=$`9@N*uux$?4Ggmm3GcLI$xQ?U@-B%s~BO8 zJm8SPm-O`XGOXF|B&lOf-glak;<^08!)d_!-75L5hm8HLuHt94-?~=LKJvhDS$c@b zq3las@Mh*W6OZ_>fVlK+?U$}VAowtW#T?kjYi6<^Gfe5?zle#8_d$Y+{f{9K&+y&g zH*bj_ngXz2*}4^{wlL!Zha027k-xYU00zOjfBOGVFz8=?@c#g)|G}H*6#dH6`vQ}) Sr{Ggph#n|u!b=p8Z~hn3j3Yn* diff --git a/docs/src/tutorials/combinators.md b/docs/src/tutorials/combinators.md new file mode 100644 index 000000000..3f31498fe --- /dev/null +++ b/docs/src/tutorials/combinators.md @@ -0,0 +1,140 @@ +# Generative Combinators + +Generative function combinators are Julia functions that take one or more generative functions as input and return a new generative function. Generative function combinators are used to express patterns of repeated computation that appear frequently in generative models. Some generative function combinators are similar to higher order functions from functional programming languages. + +## Map combinator + +In the schematic below, the kernel is denoted ``\mathcal{G}_{\mathrm{k}}``. +```@raw html +
+ schematic of map combinator +
+``` + +For example, consider the following generative function, which makes one random choice at address `r^2`: +```@example map_combinator +using Gen +@gen function foo(x, y, z) + r ~ normal(x^2 + y^2 + z^2, 1.0) + return r +end +``` +We apply the map combinator to produce a new generative function `bar`: +```@example map_combinator +bar = Map(foo) +``` +We can then obtain a trace of `bar`: +```@example map_combinator +trace, _ = generate(bar, ([0.0, 0.5], [0.5, 1.0], [1.0, -1.0])) +trace +``` +This causes `foo` to be invoked twice, once with arguments `(0.0, 0.5, 1.0)` in address namespace `1` and once with arguments `(0.5, 1.0, -1.0)` in address namespace `2`. + +```@example map_combinator +get_choices(trace) +``` +If the resulting trace has random choices: +then the return value is: + +```@example map_combinator +get_retval(trace) +``` + +## Unfold combinator + +In the schematic below, the kernel is denoted ``\mathcal{G}_{\mathrm{k}}``. +The initial state is denoted ``y_0``, the number of applications is ``n``, and the remaining arguments to the kernel not including the state, are ``z``. +```@raw html +
+ schematic of unfold combinator +
+``` + +For example, consider the following kernel, with state type `Bool`, which makes one random choice at address `:z`: +```@example unfold_combinator +using Gen +@gen function foo(t::Int, y_prev::Bool, z1::Float64, z2::Float64) + y = @trace(bernoulli(y_prev ? z1 : z2), :y) + return y +end +``` +We apply the map combinator to produce a new generative function `bar`: +```@example unfold_combinator +bar = Unfold(foo) +``` +We can then obtain a trace of `bar`: +```@example unfold_combinator +trace, _ = generate(bar, (5, false, 0.05, 0.95)) +trace +``` +This causes `foo` to be invoked five times. +The resulting trace may contain the following random choices: +```@example unfold_combinator +get_choices(trace) +``` +then the return value is: +```@example unfold_combinator +get_retval(trace) +``` + +## Switch combinator + +```@raw html +
+ schematic of switch combinator +
+``` + +Consider the following constructions: + +```@example switch_combinator +using Gen +@gen function bang((grad)(x::Float64), (grad)(y::Float64)) + std::Float64 = 3.0 + z = @trace(normal(x + y, std), :z) + return z +end + +@gen function fuzz((grad)(x::Float64), (grad)(y::Float64)) + std::Float64 = 3.0 + z = @trace(normal(x + 2 * y, std), :z) + return z +end + +sc = Switch(bang, fuzz) +``` + +This creates a new generative function `sc`. We can then obtain the trace of `sc`: + +```@example switch_combinator +trace, _ = simulate(sc, (2, 5.0, 3.0)) +trace +``` + +The resulting trace contains the subtrace from the branch with index `2` - in this case, a call to `fuzz`: + +```@example switch_combinator +get_choices(trace) +``` +``` +│ +└── :z : 13.552870875213735 +``` + +## Recurse combinator + +TODO: document me + +```@raw html +
+ schematic of recurse combinatokr +
+``` + +TODO: document me + +```@raw html +
+ schematic of recurse combinatokr +
+``` \ No newline at end of file diff --git a/docs/src/tutorials/languages/static_dsl.md b/docs/src/tutorials/languages/static_dsl.md new file mode 100644 index 000000000..03d9d9273 --- /dev/null +++ b/docs/src/tutorials/languages/static_dsl.md @@ -0,0 +1,19 @@ +# Reduce Runtime with the Static Modeling Language + +!!! note + note + +!!! tip + tip + +!!! warning + warning + +!!! danger + danger + +!!! details + details + +!!! note "optional title" + what note \ No newline at end of file diff --git a/docs/src/tutorials/modeling_in_gen.md b/docs/src/tutorials/modeling_in_gen.md new file mode 100644 index 000000000..07cef4e94 --- /dev/null +++ b/docs/src/tutorials/modeling_in_gen.md @@ -0,0 +1,1388 @@ +# Introduction to Modeling in Gen + +Welcome! In this tutorial, you'll get your feet wet with Gen, a +multi-paradigm platform for probabilistic modeling and inference. By +"multi-paradigm," we mean that Gen supports many different approaches to +modeling and inference: + +- Unsupervised learning and posterior inference in generative models using + Monte Carlo, variational, EM, and stochastic gradient techniques. + +- Supervised learning of conditional inference models (e.g. supervised + classification and regression). + +- Hybrid approaches including amortized inference / inference compilation, + variational autoencoders, and semi-supervised learning. + +Don't worry if you haven't seen some of these approaches before. One goal of +these tutorials will be to introduce you to a subset of them, +from a unified probabilistic programming perspective. + +#### In this Tutorial + +Approaching a problem from a probabilistic perspective requires both *modeling* +and *inference*: + +- **Modeling**: You first need to frame the problem — and any assumptions you + bring to the table — as a probabilistic model. A huge variety of problems can + be viewed from a modeling & inference lens, if you set them up properly. + **This notebook is about how to think of problems in this light, and how to use Gen** + **to formally specify your assumptions and the tasks you wish to solve.** + +- **Inference**: You then need to do the hard part: inference, that is, + solving the problem. In this notebook, we'll use a particularly simple + *generic* inference algorithm: importance sampling with the prior as our + proposal distributions. With enough computation, the algorithm + can in theory solve any modeling and inference problem, but in practice, for most problems of + interest, it is too slow to achieve accurate results in a reasonable amount of time. + **Future tutorials introduce some of Gen's** + **programmable inference features**, + which let you tailor the inference algorithm for use with more complex models (Gen will still automate the math!). + +Throughout this +tutorial, +we will emphasize key degrees of modeling flexibility +afforded by the probabilistic programming approach, for example: + +- Using a stochastic branching and function abstraction to express + uncertainty about which of multiple models is appropriate. + +- Representing models with an unbounded number of parameters (a 'Bayesian + non-parametric' model) using loops and recursion. + +We'll also introduce a technique for validating a model and +inference algorithm by predicting new data from inferred parameters, and +comparing this data to the observed data set. + +However, this tutorial does not exhaustively cover all features of Gen's modeling language. +For example, Gen's modeling combinators and its static modeling language +enable improved performance, but are not covered here. + +## Gen and Julia + +Gen is a package for the Julia language. The package can be loaded with: + + +```@example tutorial_1 +using Gen +``` + +Gen programs typically consist of a combination of (i) probabilistic models written in modeling languages and (ii) inference programs written in regular Julia code. Gen provides a built-in modeling language that is itself based on Julia. This notebook uses the [Plots.jl](https://github.com/JuliaPlots/Plots.jl) Julia package for plotting. + +```@example tutorial_1 +using Plots +``` + +## Generative Functions + +Probabilistic models are represented in Gen as *generative functions*. +Generative functions are used to represent a variety of different types of +probabilistic computations including generative models, inference models, +custom proposal distributions, and variational approximations (see the [Gen +documentation](https://www.gen.dev/docs/stable/ref/gfi/) or the +[paper](https://dl.acm.org/doi/10.1145/3314221.3314642)). In this +tutorial, +we focus on implementing _generative models_. A generative model represents +a data-generating process; as such, it encodes any assumptions we have about +our data and our problem domain. + + +The simplest way to construct a generative function is by using the [built-in +modeling DSL](https://www.gen.dev/docs/stable/ref/modeling/). Generative +functions written in the built-in modeling DSL are based on Julia function +definition syntax, but are prefixed with the `@gen` macro: + +```julia +@gen function function_name_here(input_arguments) + # Function body... +end +``` + +The function represents the data-generating process we are modeling. +Conceptually, every time we run the function, it should generate a new +"synthetic dataset" in line with our assumptions. Along the way, it will make +random choices; each random choice it makes can be thought of as adding +a random variable to a probabilistic model. + +Within the function body, most Julia code is permitted, but random choices use +special syntax that annotates them with an _address_: + +```julia +{addr} ~ distribution(parameters) +``` + +A simple example of such an invocation is a normal distribution parametrized +with mean 0 and standard deviation 1: + +```julia +my_variable = {:my_variable_address} ~ normal(0, 1) +``` + +Every random choice must be given an _address_, which can be +an arbitrary value—but we often use a symbol. +(`:my_variable_address` is a symbol in the Julia language.) +Think of the address as the name of a particular random choice, which +is distinct from the name of the variable. For example, consider +the following code: + +```julia +x = {:initial_x} ~ normal(0, 1) +if x < 0 + x = x + ({:addition_to_x} ~ normal(2, 1)) +end +``` + +This code manipulates a single variable, `x`, but may make up to two random +choices: `:initial_x` and `:addition_to_x`. + + +Note that we can only use `~` to give addresses to _random choices_. +The following will _not_ work because the code is trying to trace the +expression `sin(x)` which is an invocation of an ordinary Julia function, not +a distribution. +```Julia +# INVALID: +my_variable = {:not_a_random_choice} ~ sin(x) +``` +(We will see a bit later that it is _also_ possible to use `~` +to sample from helper _generative functions_, not just primitive +distributions like `normal`. But for now, think of `~` as being +for making random choices.) + +### Example 1: Linear regression + +Suppose we have a dataset of points $(x, y)$ in the plane, and we'd like +to infer a likely slope and intercept that explains their (linear) relationship. +To approach this problem from a probabilistic perspective, we first need to +develop a model. The model answers the question: how might this dataset have +come to be? It also encodes our assumptions, e.g., our assumption that our +data is explained by a linear relationship between $x$ and $y$. + +The generative function below represents a probabilistic model of a linear +relationship in the x-y plane. Given a set of $x$ coordinates, it randomly +chooses a line in the plane and generates corresponding $y$ coordinates so +that each $(x, y)$ is near the line. We might think of this function as +modeling house prices as a function of square footage, or the measured volume +of a gas as a function of its measured temperature. + + +```@example tutorial_1 +@gen function line_model(xs::Vector{Float64}) + # We begin by sampling a slope and intercept for the line. + # Before we have seen the data, we don't know the values of + # these parameters, so we treat them as random choices. The + # distributions they are drawn from represent our prior beliefs + # about the parameters: in this case, that neither the slope nor the + # intercept will be more than a couple points away from 0. + slope = ({:slope} ~ normal(0, 1)) + intercept = ({:intercept} ~ normal(0, 2)) + + # We define a function to compute y for a given x + function y(x) + return slope * x + intercept + end + + # Given the slope and intercept, we can sample y coordinates + # for each of the x coordinates in our input vector. + for (i, x) in enumerate(xs) + # Note that we name each random choice in this loop + # slightly differently: the first time through, + # the name (:y, 1) will be used, then (:y, 2) for + # the second point, and so on. + ({(:y, i)} ~ normal(y(x), 0.1)) + end + + # Most of the time, we don't care about the return + # value of a model, only the random choices it makes. + # It can sometimems be useful to return something + # meaningful, however; here, we return the function `y`. + return y +end; +``` + +The generative function takes as an argument a vector of x-coordinates. We create one below: + + +```@example tutorial_1 +xs = [-5., -4., -3., -2., -1., 0., 1., 2., 3., 4., 5.]; +``` + +Given this vector, the generative function samples a random choice +representing the slope of a line from a normal distribution with mean 0 and +standard deviation 1, and a random choice representing the intercept of a +line from a normal distribution with mean 0 and standard deviation 2. In +Bayesian statistics terms, these distributions are the *prior distributions* +of the slope and intercept respectively. Then, the function samples values +for the y-coordinates corresponding to each of the provided x-coordinates. + +This generative function returns a function `y` encoding the slope +and intercept. +We can run the model like we run a regular Julia function: + + +```@example tutorial_1 +y = line_model(xs) +``` + + + + + y (generic function with 1 method) + + + +This gives us the return value of the model, but we +may be more interested in _the values of the random choices_ that +`line_model` makes. **Crucially, each random choice is annotated with a +unique *address*.** A random choice is assigned an address using the `{addr} ~ ...` +keyword. Addresses can be any Julia value. In this program, there are two +types of addresses used -- Julia symbols and tuples of symbols and integers. +Note that within the `for` loop, the same line of code is executed multiple +times, but each time, the random choice it makes is given a distinct address. + +Although the random choices are not included in the return value, they *are* +included in the *execution trace* of the generative function. We can run the +generative function and obtain its trace using the [` +simulate`](https://www.gen.dev/docs/stable/ref/gfi/#Gen.simulate) method +from the Gen API: + + +```@example tutorial_1 +trace = Gen.simulate(line_model, (xs,)); +``` + +This method takes the function to be executed, and a tuple of arguments to the function, and returns a trace and a second value that we will not be using in this tutorial. When we print the trace, we see that it is a complex data structure. + + +```@example tutorial_1 +println(trace) +``` + + +A trace of a generative function contains various information about an execution of the function. For example, it contains the arguments on which the function was run, which are available with the API method `get_args`: + + +```@example tutorial_1 +Gen.get_args(trace) +``` + + +The trace also contains the value of the random choices, stored in a map from address to value called a *choice map*. This map is available through the API method [`get_choices`](): + + +```@example tutorial_1 +Gen.get_choices(trace) +``` + + + +We can pull out individual values from this map using Julia's subscripting syntax `[...]`: + + +```@example tutorial_1 +choices = Gen.get_choices(trace) +choices[:slope] +``` + +We can also read the value of a random choice directly from the trace, without having to use `get_choices` first: + + +```@example tutorial_1 +trace[:slope] +``` + +The return value is also recorded in the trace, and is accessible with the `get_retval` API method: + + +```@example tutorial_1 +Gen.get_retval(trace) +``` + +Or we can access the return value directly from the trace via the syntactic sugar `trace[]`: + + +```@example tutorial_1 +trace[] +``` + + +In order to understand the probabilistic behavior of a generative function, it is helpful to be able to visualize its traces. Below, we define a function that uses PyPlot to render a trace of the generative function above. The rendering shows the x-y data points and the line that is represented by the slope and intercept choices. + + +```@example tutorial_1 +function render_trace(trace; show_data=true) + + # Pull out xs from the trace + xs, = get_args(trace) + + xmin = minimum(xs) + xmax = maximum(xs) + + # Pull out the return value, useful for plotting + y = get_retval(trace) + + # Draw the line + test_xs = collect(range(-5, stop=5, length=1000)) + fig = plot(test_xs, map(y, test_xs), color="black", alpha=0.5, label=nothing, + xlim=(xmin, xmax), ylim=(xmin, xmax)) + + if show_data + ys = [trace[(:y, i)] for i=1:length(xs)] + + # Plot the data set + scatter!(xs, ys, c="black", label=nothing) + end + + return fig +end; +``` + + +```@example tutorial_1 +render_trace(trace) +``` + + +Because a generative function is stochastic, we need to visualize many runs in order to understand its behavior. The cell below renders a grid of traces. + + +```@example tutorial_1 +function grid(renderer::Function, traces) + Plots.plot(map(renderer, traces)...) +end; +``` + + +```@example tutorial_1 +traces = [Gen.simulate(line_model, (xs,)) for _=1:12] +grid(render_trace, traces) +``` + + + + + + + + + +------------------------------------------------- +### Exercise + +Write a generative function that uses the same address twice. Run it to see what happens. + +------------------------------- +### Exercise + +Write a model that generates a sine wave with random phase, period and +amplitude, and then generates y-coordinates from a given vector of +x-coordinates by adding noise to the value of the wave at each x-coordinate. +Use a `gamma(1, 1)` prior distribution for the period, and a `gamma(1, 1)` +prior distribution on the amplitude (see +[`Gen.gamma`](https://www.gen.dev/docs/stable/ref/distributions/#Gen.gamma)). +Sampling from a Gamma distribution will ensure to give us postive real values. +Use a uniform distribution between 0 and $2\pi$ for the phase (see +[`Gen.uniform`](https://www.gen.dev/docs/stable/ref/distributions/#Gen.uniform)). + +The sine wave should implement: + +$ y(x) = a \sin(2\pi \frac{1}{p} x + \varphi)$, + +where $a$ is the amplitude, $p$ is the period and $\varphi$ is the phase. In +Julia the constant $\pi$ can be expressed as either `pi` or `π` (unicode). + +When calling `trace = Gen.simulate(sine_model, (xs,))`, the following choices should appear: +- amplitude: `trace[:amplitude]` +- period: `trace[:period]` +- phase: `trace[:phase]` + +We have provided some starter code for the sine wave model: + + +```@example tutorial_1 +@gen function sine_model(xs::Vector{Float64}) + + # < your code here, for sampling a phase, period, and amplitude > + + function y(x) + return 1 # < Edit this line to compute y for a given x > + end + + for (i, x) in enumerate(xs) + {(:y, i)} ~ normal(y(x), 0.1) + end + + return y # We return the y function so it can be used for plotting, below. +end; +``` + + +```@example tutorial_1 +traces = [Gen.simulate(sine_model, (xs,)) for _=1:12]; +``` + + +```@example tutorial_1 +grid(render_trace, traces) +``` + + + + + + + + + +
+ + +## Posterior Inference + +Of course, we don't really care about generating lots of pictures of lines +(or sine waves). We'd really like to begin with an actual dataset of observed +$(x, y)$ points, and infer the corresponding slope and intercept (or phase, +period, and amplitude). This task is called _posterior inference_. + +We now will provide a data set of y-coordinates and try to draw inferences +about the process that generated the data. We begin with the following data +set: + + +```@example tutorial_1 +ys = [6.75003, 6.1568, 4.26414, 1.84894, 3.09686, 1.94026, 1.36411, -0.83959, -0.976, -1.93363, -2.91303]; +``` + + +```@example tutorial_1 +scatter(xs, ys, color="black", label=nothing, title="Observed data (linear)", xlabel="X", ylabel="Y") +``` + + + + + + + + + +We will assume that the line model was responsible for generating the data, +and infer values of the slope and intercept that explain the data. + +To do this, we write a simple *inference program* that takes the model we are +assuming generated our data, the data set, and the amount of computation to +perform, and returns a trace of the function that is approximately sampled +from the _posterior distribution_ on traces of the function, given the +observed data. That is, the inference program will try to find a trace that +well explains the dataset we created above. We can inspect that trace to find +estimates of the slope and intercept of a line that fits the data. + +Functions like `importance_resampling` expect us to provide a _model_ and +also an _choice map_ representing our data set and relating it to the model. +A choice map maps random choice addresses from the model to values from our +data set. Here, we want to tie model addresses like `(:y, 4)` to data set +values like `ys[4]`: + + +```@example tutorial_1 +function do_inference(model, xs, ys, amount_of_computation) + + # Create a choice map that maps model addresses (:y, i) + # to observed values ys[i]. We leave :slope and :intercept + # unconstrained, because we want them to be inferred. + observations = Gen.choicemap() + for (i, y) in enumerate(ys) + observations[(:y, i)] = y + end + + # Call importance_resampling to obtain a likely trace consistent + # with our observations. + (trace, _) = Gen.importance_resampling(model, (xs,), observations, amount_of_computation); + return trace +end; +``` + +We can run the inference program to obtain a trace, and then visualize the result: + + +```@example tutorial_1 +trace = do_inference(line_model, xs, ys, 100) +render_trace(trace) +``` + + + + + + + + + +We see that `importance_resampling` found a reasonable slope and intercept to explain the data. We can also visualize many samples in a grid: + + +```@example tutorial_1 +traces = [do_inference(line_model, xs, ys, 100) for _=1:10]; +grid(render_trace, traces) +``` + + + + + + + + + +We can see here that there is some uncertainty: with our limited data, we can't be 100% sure exactly where the line is. We can get a better sense for the variability in the posterior distribution by visualizing all the traces in one plot, rather than in a grid. Each trace is going to have the same observed data points, so we only plot those once, based on the values in the first trace: + + +```@example tutorial_1 +function overlay(renderer, traces; same_data=true, args...) + fig = renderer(traces[1], show_data=true, args...) + + xs, = get_args(traces[1]) + xmin = minimum(xs) + xmax = maximum(xs) + + for i=2:length(traces) + y = get_retval(traces[i]) + test_xs = collect(range(-5, stop=5, length=1000)) + fig = plot!(test_xs, map(y, test_xs), color="black", alpha=0.5, label=nothing, + xlim=(xmin, xmax), ylim=(xmin, xmax)) + end + return fig +end; +``` + + +```@example tutorial_1 +traces = [do_inference(line_model, xs, ys, 100) for _=1:10]; +overlay(render_trace, traces) +``` + + + + + + + + + +-------------- + +### Exercise + +The results above were obtained for `amount_of_computation = 100`. Run the algorithm with this value set to `1`, `10`, and `1000`, etc. Which value seems like a good tradeoff between accuracy and running time? Discuss. + +------------------ + +### Exercise +Consider the following data set. + + +```@example tutorial_1 +ys_sine = [2.89, 2.22, -0.612, -0.522, -2.65, -0.133, 2.70, 2.77, 0.425, -2.11, -2.76]; +``` + + +```@example tutorial_1 +scatter(xs, ys_sine, color="black", label=nothing) +``` + + + + + + + + + +Write an inference program that generates traces of `sine_model` that explain this data set. Visualize the resulting distribution of traces. Temporarily change the prior distribution on the period to be `gamma(1, 1)` (by changing and re-running the cell that defines `sine_model` from a previous exercise). Can you explain the difference in inference results when using `gamma(1, 1)` vs `gamma(5, 1)` prior on the period? How much computation did you need to get good results? + +## Predicting new data + +What if we'd want to predict `ys` given `xs`? + +Using the API method +[`generate`](https://www.gen.dev/docs/dev/ref/gfi/#Gen.generate), we +can generate a trace of a generative function in which the values of certain +random choices are constrained to given values. The constraints are a choice +map that maps the addresses of the constrained random choices to their +desired values. + +For example: + + + +```@example tutorial_1 +constraints = Gen.choicemap() +constraints[:slope] = 0. +constraints[:intercept] = 0. +(trace, _) = Gen.generate(line_model, (xs,), constraints) +render_trace(trace) +``` + + + + + + + + + +Note that the random choices corresponding to the y-coordinates are still made randomly. Run the cell above a few times to verify this. + + +We will use the ability to run constrained executions of a generative +function to predict the value of the y-coordinates at new x-coordinates by +running new executions of the model generative function in which the random +choices corresponding to the parameters have been constrained to their +inferred values. We have provided a function below (`predict_new_data`) that +takes a trace, and a vector of new x-coordinates, and returns a vector of +predicted y-coordinates corresponding to the x-coordinates in `new_xs`. We +have designed this function to work with multiple models, so the set of +parameter addresses is an argument (`param_addrs`): + + + +```@example tutorial_1 +function predict_new_data(model, trace, new_xs::Vector{Float64}, param_addrs) + + # Copy parameter values from the inferred trace (`trace`) + # into a fresh set of constraints. + constraints = Gen.choicemap() + for addr in param_addrs + constraints[addr] = trace[addr] + end + + # Run the model with new x coordinates, and with parameters + # fixed to be the inferred values. + (new_trace, _) = Gen.generate(model, (new_xs,), constraints) + + # Pull out the y-values and return them. + ys = [new_trace[(:y, i)] for i=1:length(new_xs)] + return ys +end; +``` + +To illustrate, we call the function above given the previous trace (which +constrained slope and intercept to be zero). + + +```@example tutorial_1 +predict_new_data(line_model, trace, [1., 2., 3.], [:slope, :intercept]) +``` + + + + + 3-element Vector{Float64}: + 0.02694273341948937 + -0.05578345643577398 + 0.0031286541463625053 + + + +The cell below defines a function that first performs inference on an +observed data set `(xs, ys)`, and then runs `predict_new_data` to generate +predicted y-coordinates. It repeats this process `num_traces` times, and +returns a vector of the resulting y-coordinate vectors. + + +```@example tutorial_1 +function infer_and_predict(model, xs, ys, new_xs, param_addrs, num_traces, amount_of_computation) + pred_ys = [] + for i=1:num_traces + trace = do_inference(model, xs, ys, amount_of_computation) + push!(pred_ys, predict_new_data(model, trace, new_xs, param_addrs)) + end + pred_ys +end; +``` + +To illustrate, we generate predictions at `[1., 2., 3.]` given one (approximate) posterior +trace. + + +```@example tutorial_1 +pred_ys = infer_and_predict(line_model, xs, ys, [1., 2., 3.], [:slope, :intercept], 1, 1000) +``` + + + + + 1-element Vector{Any}: + [1.0032271405440183, -0.13506083352207435, -0.9583356345851033] + + + +Finally, we define a cell that plots the observed data set `(xs, ys)` as red dots, and the predicted data as small black dots. + + +```@example tutorial_1 +function plot_predictions(xs, ys, new_xs, pred_ys; title="predictions") + fig = scatter(xs, ys, color="red", label="observed data", title=title) + for (i, pred_ys_single) in enumerate(pred_ys) + scatter!(new_xs, pred_ys_single, color="black", alpha=0.1, label=i == 1 ? "predictions" : nothing) + end + return fig +end; +``` + +Recall the original dataset for the line model. The x-coordinates span the interval -5 to 5. + + +```@example tutorial_1 +scatter(xs, ys, color="red", label="observed data") +``` + + + + + + + + + +We will use the inferred values of the parameters to predict y-coordinates for x-coordinates in the interval 5 to 10 from which data was not observed. We will also predict new data within the interval -5 to 5, and we will compare this data to the original observed data. Predicting new data from inferred parameters, and comparing this new data to the observed data is the core idea behind *posterior predictive checking*. This tutorial does not intend to give a rigorous overview behind techniques for checking the quality of a model, but intends to give high-level intuition. + + +```@example tutorial_1 +new_xs = collect(range(-5, stop=10, length=100)); +``` + +We generate and plot the predicted data: + + +```@example tutorial_1 +pred_ys = infer_and_predict(line_model, xs, ys, new_xs, [:slope, :intercept], 20, 1000) +plot_predictions(xs, ys, new_xs, pred_ys) +``` + + + + + + + + + +The results look reasonable, both within the interval of observed data and in the extrapolated predictions on the right. + +Now consider the same experiment run with the following data set, which has significantly more noise. + + +```@example tutorial_1 +ys_noisy = [5.092, 4.781, 2.46815, 1.23047, 0.903318, 1.11819, 2.10808, 1.09198, 0.0203789, -2.05068, 2.66031]; +``` + + +```@example tutorial_1 +pred_ys = infer_and_predict(line_model, xs, ys_noisy, new_xs, [:slope, :intercept], 20, 1000) +plot_predictions(xs, ys_noisy, new_xs, pred_ys) +``` + + + + + + + + + +It looks like the generated data is less noisy than the observed data in the regime where data was observed, and it looks like the forecasted data is too overconfident. This is a sign that our model is mis-specified. In our case, this is because we have assumed that the noise has value 0.1. However, the actual noise in the data appears to be much larger. We can correct this by making the noise a random choice as well and inferring its value along with the other parameters. + +We first write a new version of the line model that samples a random choice for the noise from a `gamma(1, 1)` prior distribution. + + +```@example tutorial_1 +@gen function line_model_fancy(xs::Vector{Float64}) + slope = ({:slope} ~ normal(0, 1)) + intercept = ({:intercept} ~ normal(0, 2)) + + function y(x) + return slope * x + intercept + end + + noise = ({:noise} ~ gamma(1, 1)) + for (i, x) in enumerate(xs) + {(:y, i)} ~ normal(slope * x + intercept, noise) + end + return y +end; +``` + +Then, we compare the predictions using inference of the unmodified and modified models on the `ys` data set: + + +```@example tutorial_1 +pred_ys = infer_and_predict(line_model, xs, ys, new_xs, [:slope, :intercept], 20, 1000) +fixed_noise_plot = plot_predictions(xs, ys, new_xs, pred_ys; title="fixed noise") + +pred_ys = infer_and_predict(line_model_fancy, xs, ys, new_xs, [:slope, :intercept, :noise], 20, 10000) +inferred_noise_plot = plot_predictions(xs, ys, new_xs, pred_ys; title="inferred noise") + +plot(fixed_noise_plot, inferred_noise_plot) +``` + + + + + + + + + +Notice that there is more uncertainty in the predictions made using the modified model. + +We also compare the predictions using inference of the unmodified and modified models on the `ys_noisy` data set: + + +```@example tutorial_1 +pred_ys = infer_and_predict(line_model, xs, ys_noisy, new_xs, [:slope, :intercept], 20, 1000) +fixed_noise_plot = plot_predictions(xs, ys_noisy, new_xs, pred_ys; title="fixed noise") + + +pred_ys = infer_and_predict(line_model_fancy, xs, ys_noisy, new_xs, [:slope, :intercept, :noise], 20, 10000) +inferred_noise_plot = plot_predictions(xs, ys_noisy, new_xs, pred_ys; title="inferred noise") + +plot(fixed_noise_plot, inferred_noise_plot) +``` + + + + + + + + + +Notice that while the unmodified model was very overconfident, the modified model has an appropriate level of uncertainty, while still capturing the general negative trend. + +------------------------- +### Exercise + +Write a modified version of the sine model that makes noise into a random choice. Compare the predicted data with the observed data using `infer_and_predict` and `plot_predictions` for the unmodified and modified models, and for the `ys_sine` and `ys_noisy` data sets. Discuss the results. Experiment with the amount of inference computation used. The amount of inference computation will need to be higher for the model with the noise as a random choice. + +We have provided you with starter code: + + +```@example tutorial_1 +@gen function sine_model_fancy(xs::Vector{Float64}) + + # < your code here > + + for (i, x) in enumerate(xs) + {(:y, i)} ~ normal(0., 0.1) # < edit this line > + end + return nothing +end; +``` + + +```@example tutorial_1 +# Modify the line below to experiment with the amount_of_computation parameter +pred_ys = infer_and_predict(sine_model, xs, ys_sine, new_xs, [], 20, 1) +fixed_noise_plot = plot_predictions(xs, ys_sine, new_xs, pred_ys; title="Fixed noise level") + +# Modify the line below to experiment with the amount_of_computation parameter +pred_ys = infer_and_predict(sine_model_fancy, xs, ys_sine, new_xs, [], 20, 1) +inferred_noise_plot = plot_predictions(xs, ys_sine, new_xs, pred_ys; title="Inferred noise level") + +Plots.plot(fixed_noise_plot, inferred_noise_plot) +``` + + +```@example tutorial_1 +# Modify the line below to experiment with the amount_of_computation parameter +pred_ys = infer_and_predict(sine_model, xs, ys_noisy, new_xs, [], 20, 1) +fixed_noise_plot = plot_predictions(xs, ys_noisy, new_xs, pred_ys; title="Fixed noise level") + +# Modify the line below to experiment with the amount_of_computation parameter +pred_ys = infer_and_predict(sine_model_fancy, xs, ys_noisy, new_xs, [], 20, 1) +inferred_noise_plot = plot_predictions(xs, ys_noisy, new_xs, pred_ys; title="Inferred noise level") + +Plots.plot(fixed_noise_plot, inferred_noise_plot) +``` + +
+ + +## Calling other generative functions + +In addition to making random choices, generative functions can invoke other +generative functions. To illustrate this, we will write a probabilistic model +that combines the line model and the sine model. This model is able to +explain data using either model, and which model is chosen will depend on the +data. This is called *model selection*. + +A generative function can invoke another generative function in three ways: + +- **(NOT RECOMMENDED)** using regular Julia function call syntax: `f(x)` + +- using the `~` snytax with an address for the call: `{addr} ~ f(x)` + +- using the `~` syntax with a wildcard address: `{*} ~ f(x)` + +When invoking using regular function call syntax, the random choices made by +the callee function are not traced at all, and Gen cannot reason about them during inference. +When invoking using `~` with a _wildcard_ address (`{*} ~ f(x)`), the random choices of the +callee function are imported directly into the caller's trace. So, for example, if `f` makes a choice +called `:f_choice`, then the caller's trace will have a choice called `:f_choice` too. +Note that a downside of this is that if `f` is called _twice_ by the same caller, then the two +choices called `:f_choice` will clash, leading to an error. +In this case, it is best to provide an address (`{addr} ~ f(x)`): `f`'s random choices will +be placed under the _namespace_ `addr`. + + +```@example tutorial_1 +@gen function foo() + {:y} ~ normal(0, 1) +end + +@gen function bar() + {:x} ~ bernoulli(0.5) + # Call `foo` with a wildcard address. + # Its choices (:y) will appear directly + # within the trace of `bar`. + {*} ~ foo() +end + +@gen function bar_using_namespace() + {:x} ~ bernoulli(0.5) + # Call `foo` with the address `:z`. + # The internal choice `:y` of `foo` + # will appear in our trace at the + # hierarchical address `:z => :y`. + {:z} ~ foo() +end; +``` + +We first show the addresses sampled by `bar`: + + +```@example tutorial_1 +trace = Gen.simulate(bar, ()) +Gen.get_choices(trace) +``` + + + + + │ + ├── :y : 0.6475752978557953 + │ + └── :x : true + + + + +And the addresses sampled by `bar_using_namespace`: + + +```@example tutorial_1 +trace = Gen.simulate(bar_using_namespace, ()) +Gen.get_choices(trace) +``` + + + + + │ + ├── :x : true + │ + └── :z + │ + └── :y : 0.3801328264226968 + + + + +Using `{addr} ~ f()`, with a namespace, can help avoid address collisions for complex models. + +A hierarchical address is represented as a Julia `Pair`, where the first element of the pair is the first element of the address and the second element of the pair is the rest of the address: + + +```@example tutorial_1 +trace[Pair(:z, :y)] +``` + + + + + 0.3801328264226968 + + + +Julia uses the `=>` operator as a shorthand for the `Pair` constructor, so we can access choices at hierarchical addresses like: + + +```@example tutorial_1 +trace[:z => :y] +``` + + + + + 0.3801328264226968 + + + +If we have a hierarchical address with more than two elements, we can construct the address by chaining the `=>` operator: + + +```@example tutorial_1 +@gen function baz() + {:a} ~ bar_using_namespace() +end + +trace = simulate(baz, ()) + +trace[:a => :z => :y] +``` + + + + + 0.5622991768999647 + + + +Note that the `=>` operator associated right, so this is equivalent to: + + +```@example tutorial_1 +trace[Pair(:a, Pair(:z, :y))] +``` + + + + + 0.5622991768999647 + + + +Now, we write a generative function that combines the line and sine models. It makes a Bernoulli random choice (e.g. a coin flip that returns true or false) that determines which of the two models will generate the data. + + +```@example tutorial_1 +@gen function combined_model(xs::Vector{Float64}) + if ({:is_line} ~ bernoulli(0.5)) + # Call line_model_fancy on xs, and import + # its random choices directly into our trace. + return ({*} ~ line_model_fancy(xs)) + else + # Call sine_model_fancy on xs, and import + # its random choices directly into our trace + return ({*} ~ sine_model_fancy(xs)) + end +end; +``` + +We visualize some traces, and see that sometimes it samples linear data and other times sinusoidal data. + + +```@example tutorial_1 +traces = [Gen.simulate(combined_model, (xs,)) for _=1:12]; +grid(render_trace, traces) +``` + + +We run inference using this combined model on the `ys` data set and the `ys_sine` data set. + + +```@example tutorial_1 +traces = [do_inference(combined_model, xs, ys, 10000) for _=1:10]; +linear_dataset_plot = overlay(render_trace, traces) +traces = [do_inference(combined_model, xs, ys_sine, 10000) for _=1:10]; +sine_dataset_plot = overlay(render_trace, traces) +Plots.plot(linear_dataset_plot, sine_dataset_plot) +``` + + +The results should show that the line model was inferred for the `ys` data set, and the sine wave model was inferred for the `ys_sine` data set. + +------- + +### Exercise + +Construct a data set for which it is ambiguous whether the line or sine wave model is best. Visualize the inferred traces using `render_combined` to illustrate the ambiguity. Write a program that takes the data set and returns an estimate of the posterior probability that the data was generated by the sine wave model, and run it on your data set. + +Hint: To estimate the posterior probability that the data was generated by the sine wave model, run the inference program many times to compute a large number of traces, and then compute the fraction of those traces in which `:is_line` is false. + +
+ +## Modeling with an unbounded number of parameters + +Gen's built-in modeling language can be used to express models that use an +unbounded number of parameters. This section walks you through development of +a model of data that does not a-priori specify an upper bound on the +complexity of the model, but instead infers the complexity of the model as +well as the parameters. This is a simple example of a *Bayesian +nonparametric* model. + +We will consider two data sets: + + +```@example tutorial_1 +xs_dense = collect(range(-5, stop=5, length=50)) +ys_simple = fill(1., length(xs_dense)) .+ randn(length(xs_dense)) * 0.1 +ys_complex = [Int(floor(abs(x/3))) % 2 == 0 ? 2 : 0 for x in xs_dense] .+ randn(length(xs_dense)) * 0.1; +``` + + +```@example tutorial_1 +simple_plot = scatter(xs_dense, ys_simple, color="black", label=nothing, title="ys-simple", ylim=(-1, 3)) +complex_plot = scatter(xs_dense, ys_complex, color="black", label=nothing, title="ys-complex", ylim=(-1, 3)) +Plots.plot(simple_plot, complex_plot) +``` + + + + + + + + + +The data set on the left appears to be best explained as a contant function +with some noise. The data set on the right appears to include two +changepoints, with a constant function in between the changepoints. We want a +model that does not a-priori choose the number of changepoints in the data. +To do this, we will recursively partition the interval into regions. We +define a Julia data structure that represents a binary tree of intervals; +each leaf node represents a region in which the function is constant. + + +```@example tutorial_1 +struct Interval + l::Float64 + u::Float64 +end +``` + + +```@example tutorial_1 +abstract type Node end + +struct InternalNode <: Node + left::Node + right::Node + interval::Interval +end + +struct LeafNode <: Node + value::Float64 + interval::Interval +end +``` + +We now write a generative function that randomly creates such a tree. Note the use of recursion in this function to create arbitrarily large trees representing arbitrarily many changepoints. Also note that we assign the address namespaces `:left` and `:right` to the calls made for the two recursive calls to `generate_segments`. + + +```@example tutorial_1 +@gen function generate_segments(l::Float64, u::Float64) + interval = Interval(l, u) + if ({:isleaf} ~ bernoulli(0.7)) + value = ({:value} ~ normal(0, 1)) + return LeafNode(value, interval) + else + frac = ({:frac} ~ beta(2, 2)) + mid = l + (u - l) * frac + # Call generate_segments recursively! + # Because we will call it twice -- one for the left + # child and one for the right child -- we use + # addresses to distinguish the calls. + left = ({:left} ~ generate_segments(l, mid)) + right = ({:right} ~ generate_segments(mid, u)) + return InternalNode(left, right, interval) + end +end; +``` + +We also define some helper functions to visualize traces of the `generate_segments` function. + + +```@example tutorial_1 +function render_node!(node::LeafNode) + plot!([node.interval.l, node.interval.u], [node.value, node.value], label=nothing, linewidth=5) +end + +function render_node!(node::InternalNode) + render_node!(node.left) + render_node!(node.right) +end; +``` + + +```@example tutorial_1 +function render_segments_trace(trace; xlim=(0,1)) + node = get_retval(trace) + fig = plot(xlim=xlim, ylim=(-3, 3)) + render_node!(node) + return fig +end; +``` + +We generate 12 traces from this function and visualize them below. We plot the piecewise constant function that was sampled by each run of the generative function. Different constant segments are shown in different colors. Run the cell a few times to get a better sense of the distribution on functions that is represented by the generative function. + + +```@example tutorial_1 +traces = [Gen.simulate(generate_segments, (0., 1.)) for i=1:12] +grid(render_segments_trace, traces) +``` + + + + + + + + + +Because we only sub-divide an interval with 30% probability, most of these sampled traces have only one segment. + +Now that we have a generative function that generates a random piecewise-constant function, we write a model that adds noise to the resulting constant functions to generate a data set of y-coordinates. The noise level will be a random choice. + + +```@example tutorial_1 +# get_value_at searches a binary tree for +# the leaf node containing some value. +function get_value_at(x::Float64, node::LeafNode) + @assert x >= node.interval.l && x <= node.interval.u + return node.value +end + +function get_value_at(x::Float64, node::InternalNode) + @assert x >= node.interval.l && x <= node.interval.u + if x <= node.left.interval.u + get_value_at(x, node.left) + else + get_value_at(x, node.right) + end +end + +# Our full model +@gen function changepoint_model(xs::Vector{Float64}) + node = ({:tree} ~ generate_segments(minimum(xs), maximum(xs))) + noise = ({:noise} ~ gamma(0.5, 0.5)) + for (i, x) in enumerate(xs) + {(:y, i)} ~ normal(get_value_at(x, node), noise) + end + return node +end; +``` + +We write a visualization for `changepoint_model` below: + + +```@example tutorial_1 +function render_changepoint_model_trace(trace; show_data=true) + xs = Gen.get_args(trace)[1] + node = Gen.get_retval(trace) + fig = render_segments_trace(trace; xlim=(minimum(xs), maximum(xs))) + render_node!(node) + if show_data + ys = [trace[(:y, i)] for i=1:length(xs)] + scatter!(xs, ys, c="gray", label=nothing, alpha=0.3, markersize=3) + end +end; +``` + +Finally, we generate some simulated data sets and visualize them on top of the underlying piecewise constant function from which they were generated: + + +```@example tutorial_1 +traces = [Gen.simulate(changepoint_model, (xs_dense,)) for i=1:12] +grid(render_changepoint_model_trace, traces) +``` + + + + + + + + + +Notice that the amount of variability around the piecewise constant mean function differs from trace to trace. + +Now we perform inference for the simple data set: + + +```@example tutorial_1 +traces = [do_inference(changepoint_model, xs_dense, ys_simple, 10000) for _=1:12]; +grid(render_changepoint_model_trace, traces) +``` + + + + + + + + + +We see that we inferred that the mean function that explains the data is a constant with very high probability. + +For inference about the complex data set, we use more computation. You can experiment with different amounts of computation to see how the quality of the inferences degrade with less computation. Note that we are using a very simple generic inference algorithm in this tutorial, which really isn't suited for this more complex task. In later tutorials, we will learn how to write more efficient algorithms, so that accurate results can be obtained with significantly less computation. We will also see ways of annotating the model for better performance, no matter the inference algorithm. + +!!! warning + The following cell may run for 2-3 minutes. + + +```@example tutorial_1 +traces = [do_inference(changepoint_model, xs_dense, ys_complex, 100000) for _=1:12]; +grid(render_changepoint_model_trace, traces) +``` + + + +The results show that more segments are inferred for the more complex data set. + +------ +### Exercise +Write a function that takes a data set of x- and y-coordinates and plots the histogram of the probability distribution on the number of changepoints. +Show the results for the `ys_simple` and `ys_complex` data sets. + +Hint: The return value of `changepoint_model` is the tree of `Node` values. Walk this tree. + +------- + +### Exercise +Write a new version of `changepoint_model` that uses `{*} ~ ...` without an address to make the recursive calls. + +Hint: You will need to guarantee that all addresses are unique. How can you label each node in a binary tree using an integer? \ No newline at end of file diff --git a/docs/tex/mcmc.tex b/docs/tex/mcmc.tex deleted file mode 100644 index e381f9b62..000000000 --- a/docs/tex/mcmc.tex +++ /dev/null @@ -1,208 +0,0 @@ -\documentclass{article} -\usepackage{geometry} -\usepackage[utf8]{inputenc} -\usepackage{amssymb} -\usepackage{amsmath} -\usepackage{amsthm} -\usepackage{url} - -% bibliography -\PassOptionsToPackage{numbers, compress}{natbib} -\usepackage{natbib} - -\newcommand{\reals}{\mathbb{R}} -\newcommand{\normal}{\mathcal{N}} -\newcommand{\pathx}{x} -\newcommand{\pathy}{y} -\newcommand{\measx}{\tilde{x}} -\newcommand{\measy}{\tilde{y}} - -\newtheorem{prop}{Proposition} - -\title{Efficient Automated Calculation of the Jacobian Determinant for Reversible Jump MCMC} -\author{Marco Cusumano-Towner} -\date{} - -\begin{document} - -\maketitle - -\section{Introduction} -Reversible jump MCMC~\citep{green1995reversible} (RJMCMC) is a general framework for constructing Markov chain Monte Carlo kernels that satisfy detailed balance, while permitting transdimensional and/or deterministic moves. -Some generic transdimensional MCMC kernels that are implemented by various probabilistic programming systems, including `single-site resimulation Metropolis-Hastings' ~\citep{goodman2008church,wingate2011lightweight} can be justified as instances of RJMCMC. -However, it is possible to construct much more efficient MCMC kernels that are specialized to a model using the RJMCMC framework. -The Gen~\citep{cusumano2019gen} probabilistic programming system allows users to program their own custom RJMCMC kernels that are specialized to the model, by specifying an auxiliary generative function called a \emph{proposal} and an \emph{involution} (a function that is its own inverse) that maps tuples of the form $(t, u)$ to tuples of the form $(t', u')$ where $t$ and $t'$ are traces of the model generative function, and $u$ and $u'$ are traces of the proposal generative function. -In general, both $t$ and $u$ (and by necessity $t'$ and $u'$) may contain the values of continuous random choices, in which case the involution must be constructed from a family of differentiable bijections on the values of these continuous random choices. -When applying a particular RJMCMC kernel, the determinant of the Jacobian of one of these bijections is required for the acceptance probability calculation. -This document describes how to implement this Jacobian determinant computation efficiently using automatic differentiation, while exploiting the sparsity of the Jacobian. -The approach uses a domain-specific language (DSL) for specifying families of bijections that makes the sparsity structure of the Jacobian explicit. - -\section{Example} - -We consider a variant of the changepoint model that was used to introduce RJMCMC in~\citep{green1995reversible}. -Briefly, we put a prior on a class of piecewise constant functions on an interval $[0, x_{max}]$ where the number of changepoints is random and Poisson distributed. -Let $n \in \{1, 2, \ldots, \}$ denote the number of segments, so that the number of changepoints is $n-1 \in \{0, 1, 2, \ldots\}$. -Let $x_i$ denote changepoint $i$, for $i=1,\ldots,n-1$. -Then, segment $1$ has range $[0, x_1)$, segment $d$ has range $[x_{n-1}, x_{max}]$ and segment $i$ has range $[x_{i-1}, x_i)$ for $i=2,\ldots,n-1$. -Let $r_i \in (0, \infty)$ denote the rate for segment $i$ for $i=1,\ldots,n$. -%There is some likelihood function $l_{\mathcal{D}}(n, \mathbf{x}, \mathbf{r})$ where $\mathcal{D}$ denotes some observed data. - -Suppose we are performing a birth-death move as described in~\citep{green1995reversible}, in which we either insert a new changepoint or remove an existing changepoint, and adjust the rates. -Let $(n, \mathbf{r}, \mathbf{x})$ denote the the previous hypothesis (corresponding to the previous \emph{trace} of the generative function representing the model; we can omit the observations $\mathcal{D}$ from the trace for to simplify notation because they are held fixed). -The proposal generates the following random choices, dependent on the given trace. -If $n = 1$, then only a birth move is possible, because there are no changepoints to remove. -If $n > 1$, then either a birth or death move is possible and we randomly pick sample $b \in \{0, 1\}$, where $b = 1$ indicates a birth move, from a Bernooulli distribution. -If a birth move is chosen, then the proposal samples an existing segment $i \in \{1, \ldots, n\}$ that will be split into two by the new changepoint from a discrete distribution. -The proposal then samples a value $u_1 \in (x_{i-1}, x_i)$ for the changepoint from a uniform continuous distribution (where $x_{i-1}$ and $x_i$ are the bounds of segment $i$ in the previous trace); this variable represents the locaion of the newly proposed changepoint. -The proposal also samples a value $u_2 \in (0, 1)$ that will be used to determine the new rates $r_{i}'$ and $r_{i+1}'$ from the previous rate $r_i$ and $x_i$ and $u_1$ and $x_{i+1}$ (as described in~\citep{green1995reversible} the new rates are determined such that their harmonic mean, weighted by the size of the segment, is equal to the previous rate). -If a death move is chosen, then the proposal samples an existing changepoint $i \in \{1, \ldots, n-1\}$ to remove. - -Suppose the current trace contains $n > 1$ segments, and we are proposing a birth move that is a transition from $n$ segments to $n+1$ segments. -Then, $\mathbf{r} \in \mathbb{R}^n$ and $\mathbf{r}' \in \mathbb{R}^{n+1}$, and $\mathbf{x} \in \mathbb{R}^{n-1}$ and $\mathbf{x'} \in \mathbb{R}^n$. -Specificaly, suppose that $n = 3$ and that $i = 2$ (so we are splitting the middle of three segments). -Then, the differentiable bijection from $(\mathbf{r}, \mathbf{x}, u_1, u_2) \in \mathbb{R}^{7}$ to $(\mathbf{r}', \mathbf{x}') \in \mathbb{R}^{7}$ is: -\[ -\begin{array}{rl} -x_1' &= x_1 \\ -x_2' &= u_1 \\ -x_3' &= x_2 \\ -r_1' &= r_1 \\ -r_2' &= f_1(r_2, x_1, u_1, x_2, u_2) \\ -r_3' &= f_2(r_2, x_1, u_1, x_2, u_2)\\ -r_4' &= r_3 -\end{array} -\] -for some functions $f_1$ and $f_2$ whose details will not concern us here. -Then, the Jacobian (where columns are inputs and rows are outputs) is as follows, where rows and columns are labeled: -\[ -\begin{array}{c|ccccccc} -& x_1 & x_2 & r_1 & r_2 & r_3 & u_1 & u_2 \\ -\hline -x_1' & 1 & 0 & 0 & 0 & 0 & 0 & 0 \\ -x_2' & 0 & 0 & 0 & 0 & 0 & 1 & 0 \\ -x_3' & 0 & 1 & 0 & 0 & 0 & 0 & 0 \\ -r_1' & 0 & 0 & 1 & 0 & 0 & 0 & 0 \\ -r_2' & {\scriptstyle \partial f_1 / \partial x_1} - & {\scriptstyle \partial f_1 / \partial x_2} - & 0 - & {\scriptstyle \partial f_1 / \partial r_2} - & 0 - & {\scriptstyle \partial f_1 / \partial u_1} - & {\scriptstyle \partial f_1 / \partial u_2} \\ -r_3' & {\scriptstyle \partial f_2 / \partial x_1} - & {\scriptstyle \partial f_2 / \partial x_2} - & 0 - & {\scriptstyle \partial f_2 / \partial r_2} - & 0 - & {\scriptstyle \partial f_2 / \partial u_1} - & {\scriptstyle \partial f_2 / \partial u_2} \\ -r_4' & 0 & 0 & 0 & 0 & 1 & 0 & 0 -\end{array} -\] -Note that the following rows have a single 1 and all other entries 0: $x_1', x_2', x_3', r_1', r_4'$. -Therefore, by cofactor expansion, the absolute value of the determinant of this matrix is equivalent to the absolute value of the determinant of the following submatrix of the Jacobian: -\[ -\begin{array}{c|ccccccc} -& r_2 & u_2 \\ -\hline -r_2' & {\scriptstyle \partial f_1 / \partial r_2} & {\scriptstyle \partial f_1 / \partial u_2} \\ -r_3' & {\scriptstyle \partial f_2 / \partial r_2} & {\scriptstyle \partial f_2 / \partial u_2} \\ -\end{array} -\] -To compute the Jacobian determinant, it turns out that we only need to compute $4$ of the $49$ elements of the Jacobian. -Also, somewhat surprisingly, we do not need to compute all elements that are not $0$ or $1$. - -\section{Automatically Exploiting the Sparse Jacobian Structure} - -This section describes how to automatically perform the optimization recognized above in a general setting for arbitrary generative models and RJMCMC kernels, while using automatic differentiation. -Let $t \in \mathbb{R}^K$ denote the vector of values of continuous random choices in the previous model trace. -Let $u \in \mathbb{R}^L$ denote the vector of values of continuous random choices in the proposal trace. -Similarly, let $t' \in \mathbb{R}^{K'}$ and $u' \in \mathbb{R}^{L'}$ denote the vector of values of continuous random choices in the new model trace, and new proposal trace. -We assume that $K + L = K' + L'$ (this is the \emph{dimension matching} requirement), and that we are given a bijection $h : \mathbb{R}^{K + L} \to \mathbb{R}^{K + L}$ that is differentiable and whose Jacobian has nonzero determinant everwhere.\footnote{It may be possible to generalize and allow functions that are not differentiable at a measure zero set of points.} -This section discusses how to efficiently compute the absolute value of the determinant of the Jacobian of $h$, which we denote by $J \in \mathbb{R}^{(K + L) \times (K + L)}$. - -$J$ is often sparse due to the fact that often a large fraction of the entries in $(t', u')$ are obtained by \emph{copying} the value of some entry in $(t, u)$. -We denote the subvector of entries of $t'$ that are produced by copying as $t'_c$ and similarly for $u'$ (subvector of copies is denoted $u'_c$). -We denote the subvectors containing other entries of $t'$ and $u'$ by $t'_w$ and $u'_w$ respectively (where $w$ stands for `write'). -Let $t_c$ and $u_c$ denote the subvectors of $t$ and $u$ respectively such that each entry is copied to at least one entry in $(t', u')$. -We denote the subvectors containing the remaining entries of $t$ and $u$ as $t_r$ and $u_r$ (where $r$ stands for `read'). -The Jacobian of $h$ (up to reordering of rows and columns, which does not change the absolute value of the determinant) then has the following block structure, where white space indicates a zero submatrix (and where sets of rows and columns are labeled on right): -\[ -J = \left[ -\begin{array}{cccc} -J_{11} & J_{12} & & \\ -J_{21} & J_{22} & & \\ -J_{31} & J_{32} & J_{33} & J_{34}\\ -J_{41} & J_{42} & J_{43} & J_{44} -\end{array} -\right] -\;\;\;\;\; -\begin{array}{c|cccc} -& t_c & u_c & t_r & u_r\\ -\hline -t'_c & J_{11} & J_{12} & & \\ -u'_c & J_{21} & J_{22} & & \\ -t'_w & J_{31} & J_{32} & J_{33} & J_{34}\\ -u'_w & J_{41} & J_{42} & J_{43} & J_{44} -\end{array} -\] - -We will show that the absolute value of the determinant of this matrix is the same as the absolute value of the determinant of the following submatrix (which we will show is square), which much smaller than $J$ in the common case when a reversible jump move only modifies a subset of the random choices in a model: -\[ -\left[ -\begin{array}{cc} -J_{33} & J_{34}\\ -J_{43} & J_{44} -\end{array} -\right] -\] - -\begin{prop} -The submatrix $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ is square. -\end{prop} -\begin{proof} -Each row of $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ contains exactly one $1$ and all other entries are $0$, by the definition of $t_c'$ and $u_c'$. -Each column of $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ contains at least one $1$, by the definition of $t_c$ and $u_c$. -Therefore, the number of rows must be greater than or equal to the number of columns. -Since we assume that $J$ has nonzero determinant, all its rows must be linearly independent. -Together with the block structure of $J$ shown above, this implies that rows of $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ are linearly independent, and in particular that no two rows are identical. -Since each row of $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ contains exactly one $1$, the number of columns in $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ must be greater than or equal to the number of rows. -\end{proof} - -\begin{prop} -The submatrix $[J_{33} \; J_{34}; J_{43} \; J_{44}]$ is square. -\end{prop} -\begin{proof} -$J$ is square and $[J_{11} \; J_{12}; J_{21} \; J_{22}]$ is square. -\end{proof} - -\begin{prop} -$|\det(J)| = |\det([J_{33} \; J_{34}; J_{43} \; J_{44}])|$. -\end{prop} -\begin{proof} -Each row of $J$ corresponding to an entry of $t_c'$ or an entry of $u_c'$ has a single $1$ entry and other entries $0$ by definition of $t_c'$ and $u_c'$. -Then, the result follows by consecutive co-factor expansion of the determinant of $J$, deleting all rows in $t_c'$ and all columns in $u_c'$. -\end{proof} - -Automating the computation of $|\det(J)|$ by computing instead the determinant of the submatrix allows us to avoid evaluating the entire $J$, which scales quadratically in the maximum number of random choices in the model, and allows us to avoid computing the full determinant, which scales cubically (albeit likely smaller smaller lower-order factors than computing the Jacobian itself). -We only need to evaluate the submatrix $[J_{33} \; J_{34}; J_{43} \; J_{44}]$. -However, this requires knowing which entries in $t$ and $u$ correspond to $(t_r, u_r)$, and the which entries in $t'$ and $u'$ correspond to $(t'_w, u'_w)$. -This is made possible using a domain-specific language (DSL) for defining the bijection $h$ with separate language constructs for: -\begin{enumerate} -\item Reading the values of continuous random choices from $t$ and $u$. -\item Writing the values of continuous random choices to $t'$ and $u'$. -\item Copying the values of continuous random choices from $t$ and $u$ to $t'$ and $u'$. -If the user instead express a copy using a read operation followed by a write operation, but this will result in unecessary extra computation. -\item Tagging read operations for of continuous random choices read from $t$ as \emph{retained}, which asserts that the random choice will not be deleted by the reversible jump move. -If the same address is also not written to in $t'$, and the user assertion is true, then the random choice will be copied implicitly from $t$, which means it can be treated as part of $t_c$, reducing the size of the Jacobian submatrix that needs to be computed. -This is optional -- not tagging a read operation as retained when it is results in unecessary computation. -\end{enumerate} -Automatic differentiation is used to compute the entries of $[J_{33} \; J_{34}; J_{43} \; J_{44}]$. -Finally, note that specification of the involution on discrete random choices may or may not be separated from the specification of the family of bijections $h$ on the continuous random choices. -In the current Gen implementation, a single unified DSL is used to specify the entire involution and the approach described in this document is used for Jacobian determinant correction computation. - -\bibliographystyle{unsrt} -\bibliography{references} - - -\end{document} diff --git a/docs/tex/references.bib b/docs/tex/references.bib deleted file mode 100644 index f7bd2a17d..000000000 --- a/docs/tex/references.bib +++ /dev/null @@ -1,39 +0,0 @@ -@article{green1995reversible, - title={Reversible jump {{M}}arkov chain {{M}}onte {{C}}arlo computation and {{B}}ayesian model determination}, - author={Green, Peter J}, - journal={Biometrika}, - volume={82}, - number={4}, - pages={711--732}, - year={1995}, - publisher={Oxford University Press} -} - -@inproceedings{goodman2008church, -title = {Church: a language for generative models}, -author = {Goodman, Noah and Mansinghka, Vikash and Roy, Daniel M. and Bonawitz, Keith and Tenenbaum, Joshua B.}, -booktitle = {Proceedings of the 24th Annual Conference on Uncertainty in Artificial Intelligence}, -series = {UAI 2008}, -pages = {220--229}, -year = {2008}, -publisher = {AUAI Press}, -} - -@inproceedings{wingate2011lightweight, - title={Lightweight Implementations of Probabilistic Programming Languages Via Transformational Compilation.}, - author={Wingate, David and Stuhlm{\"u}ller, Andreas and Goodman, Noah D}, - booktitle={AISTATS}, - pages={770--778}, - year={2011} -} - -@inproceedings{cusumano2019gen, - title={Gen: a general-purpose probabilistic programming system with programmable inference}, - author={Cusumano-Towner, Marco F and Saad, Feras A and Lew, Alexander K and Mansinghka, Vikash K}, - booktitle={Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation}, - pages={221--236}, - year={2019}, - organization={ACM} -} - - From 8043ab34996d9a72700cb9f4855fda19aaeb321b Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Wed, 10 Apr 2024 11:28:10 -0400 Subject: [PATCH 02/19] Doc changes --- docs/pages.jl | 13 +- docs/src/api/gfi.md | 348 +------------- docs/src/api/mcmc.md | 371 --------------- docs/src/api/trace_translators.md | 438 +----------------- docs/src/api/vi.md | 13 +- docs/src/getting_started/mcmc_cipher.md | 0 docs/src/misc/changelog.md | 0 docs/src/misc/contributing.md | 0 docs/src/misc/references.md | 0 docs/src/tutorials/gfi.md | 355 ++++++++++++++ .../{api => tutorials/languages}/extending.md | 0 docs/src/tutorials/mcmc.md | 369 +++++++++++++++ docs/src/tutorials/trace_translators.md | 435 +++++++++++++++++ docs/src/tutorials/vi.md | 20 + 14 files changed, 1195 insertions(+), 1167 deletions(-) create mode 100644 docs/src/getting_started/mcmc_cipher.md create mode 100644 docs/src/misc/changelog.md create mode 100644 docs/src/misc/contributing.md create mode 100644 docs/src/misc/references.md create mode 100644 docs/src/tutorials/gfi.md rename docs/src/{api => tutorials/languages}/extending.md (100%) create mode 100644 docs/src/tutorials/mcmc.md create mode 100644 docs/src/tutorials/trace_translators.md create mode 100644 docs/src/tutorials/vi.md diff --git a/docs/pages.jl b/docs/pages.jl index 0227a90f3..615c37d75 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -8,13 +8,18 @@ pages = [ "Basics" => [ "tutorials/modeling_in_gen.md", "tutorials/combinators.md", + "tutorials/gfi.md", + "tutorials/mcmc.md", + "tutorials/vi.md" ], "Advanced" => [ "tutorials/modeling_in_gen.md", + "tutorials/trace_translators.md" # "tutorials/data_driven_proposals.md", ], "Modeling Languages" => [ "tutorials/languages/static_dsl.md", + "tutorials/languages/extending.md" ], ], "Modeling Languages and APIs" => [ @@ -26,7 +31,6 @@ pages = [ "Selections" => "api/selections.md", "Optimizing Trainable Parameters" => "api/parameter_optimization.md", "Trace Translators" => "api/trace_translators.md", - "Extending Gen" => "api/extending.md" ], "Standard Inference Library" => [ "Importance Sampling" => "api/importance.md", @@ -40,5 +44,10 @@ pages = [ "Internals" => [ "Optimizing Trainable Parameters" => "api/internals/parameter_optimization.md", "Modeling Language Implementation" => "api/internals/language_implementation.md" - ] + ], + "Miscellanea" => [ + "Changelog" => "misc/changelog.md", + "Contributing" => "misc/contributing.md", + "References" => "misc/references.md" + ] ] diff --git a/docs/src/api/gfi.md b/docs/src/api/gfi.md index c6b03fc7b..805efb97b 100644 --- a/docs/src/api/gfi.md +++ b/docs/src/api/gfi.md @@ -1,356 +1,14 @@ -# Generative Functions +## Generative Functions -One of the core abstractions in Gen is the **generative function**. -Generative functions are used to represent a variety of different types of probabilistic computations including generative models, inference models, custom proposal distributions, and variational approximations. - -## Introduction - -Generative functions are represented by the following abstact type: ```@docs GenerativeFunction ``` -There are various kinds of generative functions, which are represented by concrete subtypes of [`GenerativeFunction`](@ref). -For example, the [Built-in Modeling Language](@ref) allows generative functions to be constructed using Julia function definition syntax: -```julia -@gen function foo(a, b=0) - if @trace(bernoulli(0.5), :z) - return a + b + 1 - else - return a + b - end -end -``` -Users can also extend Gen by implementing their own [Custom generative function types](@ref), which can be new modeling languages, or just specialized optimized implementations of a fragment of a specific model. - -Generative functions behave like Julia functions in some respects. -For example, we can call a generative function `foo` on arguments and get an output value using regular Julia call syntax: -```julia-repl -julia> foo(2, 4) -7 -``` -However, generative functions are distinct from Julia functions because they support additional behaviors, described in the remainder of this section. - - -## Mathematical concepts - -Generative functions represent computations that accept some arguments, may use randomness internally, return an output, and cannot mutate externally observable state. -We represent the randomness used during an execution of a generative function as a **choice map** from unique **addresses** to values of random choices, denoted ``t : A \to V`` where ``A`` is a finite (but not a priori bounded) address set and ``V`` is a set of possible values that random choices can take. -In this section, we assume that random choices are discrete to simplify notation. -We say that two choice maps ``t`` and ``s`` **agree** if they assign the same value for any address that is in both of their domains. - -Generative functions may also use **non-addressable randomness**, which is not included in the map ``t``. -We denote non-addressable randomness by ``r``. -Untraced randomness is useful for example, when calling black box Julia code that implements a randomized algorithm. - -The observable behavior of every generative function is defined by the following mathematical objects: - -### Input type -The set of valid argument tuples to the function, denoted ``X``. - -### Probability distribution family -A family of probability distributions ``p(t, r; x)`` on maps ``t`` from random choice addresses to their values, and non-addressable randomness ``r``, indexed by arguments ``x``, for all ``x \in X``. -Note that the distribution must be normalized: -```math -\sum_{t, r} p(t, r; x) = 1 \;\; \mbox{for all} \;\; x \in X -``` -This corresponds to a requirement that the function terminate with probabability 1 for all valid arguments. -We use ``p(t; x)`` to denote the marginal distribution on the map ``t``: -```math -p(t; x) := \sum_{r} p(t, r; x) -``` -And we denote the conditional distribution on non-addressable randomness ``r``, given the map ``t``, as: -```math -p(r | t; x) := p(t, r; x) / p(t; x) -``` - -### Return value function -A (deterministic) function ``f`` that maps the tuple ``(x, t)`` of the arguments and the choice map to the return value of the function (which we denote by ``y``). -Note that the return value cannot depend on the non-addressable randomness. - -### Auxiliary state -Generative functions may expose additional **auxiliary state** associated with an execution, besides the choice map and the return value. -This auxiliary state is a function ``z = h(x, t, r)`` of the arguments, choice map, and non-addressable randomness. -Like the choice map, the auxiliary state is indexed by addresses. -We require that the addresses of auxiliary state are disjoint from the addresses in the choice map. -Note that when a generative function is called within a model, the auxiliary state is not available to the caller. -It is typically used by inference programs, for logging and for caching the results of deterministic computations that would otherwise need to be reconstructed. - -### Internal proposal distribution family -A family of probability distributions ``q(t; x, u)`` on maps ``t`` from random choice addresses to their values, indexed by tuples ``(x, u)`` where ``u`` is a map from random choice addresses to values, and where ``x`` are the arguments to the function. -It must satisfy the following conditions: -```math -\sum_{t} q(t; x, u) = 1 \;\; \mbox{for all} \;\; x \in X, u -``` -```math -p(t; x) > 0 \mbox{ if and only if } q(t; x, u) > 0 \mbox{ for all } u \mbox{ where } u \mbox{ and } t \mbox{ agree } -``` -```math -q(t; x, u) > 0 \mbox{ implies that } u \mbox{ and } t \mbox{ agree }. -``` -There is also a family of probability distributions ``q(r; x, t)`` on non-addressable randomness, that satisfies: -```math -q(r; x, t) > 0 \mbox{ if and only if } p(r | t, x) > 0 -``` - - -## Traces - -An **execution trace** (or just *trace*) is a record of an execution of a generative function. -Traces are the primary data structures manipulated by Gen inference programs. -There are various methods for producing, updating, and inspecting traces. -Traces contain: - -- the arguments to the generative function - -- the choice map - -- the return value - -- auxiliary state - -- other implementation-specific state that is not exposed to the caller or user of the generative function, but is used internally to facilitate e.g. incremental updates to executions and automatic differentiation - -- any necessary record of the non-addressable randomness - - -Different concrete types of generative functions use different data structures and different Julia types for their traces, but traces are subtypes of [`Trace`](@ref). -```@docs -Trace -``` -The concrete trace type that a generative function uses is the second type parameter of the [`GenerativeFunction`](@ref) abstract type. -For example, the trace type of [`DynamicDSLFunction`](@ref) is `DynamicDSLTrace`. - -A generative function can be executed to produce a trace of the execution using [`simulate`](@ref): -```julia -trace = simulate(foo, (a, b)) -``` -A traced execution that satisfies constraints on the choice map can be generated using [`generate`](@ref): -```julia -trace, weight = generate(foo, (a, b), choicemap((:z, false))) -``` - -There are various methods for inspecting traces, including: - -- [`get_args`](@ref) (returns the arguments to the function) - -- [`get_retval`](@ref) (returns the return value of the function) - -- [`get_choices`](@ref) (returns the choice map) - -- [`get_score`](@ref) (returns the log probability that the random choices took the values they did) - -- [`get_gen_fn`](@ref) (returns a reference to the generative function) - -You can also access the values in the choice map and the auxiliary state of the trace by passing the address to [`Base.getindex`](@ref). -For example, to retrieve the value of random choice at address `:z`: -```julia -z = trace[:z] -``` - -When a generative function has default values specified for trailing arguments, those arguments can be left out when calling [`simulate`](@ref), [`generate`](@ref), and other functions provided by the generative function interface. The default values will automatically be filled in: -```julia -julia> trace = simulate(foo, (2,)); -julia> get_args(trace) -(2, 0) -``` - -## Updating traces - -It is often important to incrementally modify the trace of a generative function (e.g. within MCMC, numerical optimization, sequential Monte Carlo, etc.). -In Gen, traces are **functional data structures**, meaning they can be treated as immutable values. -There are several methods that take a trace of a generative function as input and return a new trace of the generative function based on adjustments to the execution history of the function. -We will illustrate these methods using the following generative function: +add example ```julia -@gen function bar() - val = @trace(bernoulli(0.3), :a) - if @trace(bernoulli(0.4), :b) - val = @trace(bernoulli(0.6), :c) && val - else - val = @trace(bernoulli(0.1), :d) && val - end - val = @trace(bernoulli(0.7), :e) && val - return val +@gen function model() end ``` -Suppose we have a trace (`trace`) of `bar` with initial choices: -``` -│ -├── :a : false -│ -├── :b : true -│ -├── :c : false -│ -└── :e : true -``` -Note that address `:d` is not present because the branch in which `:d` is sampled was not taken because random choice `:b` had value `true`. - -### Update -The [`update`](@ref) method takes a trace and generates an adjusted trace that is consistent with given changes to the arguments to the function, and changes to the values of random choices made. - -**Example.** -Suppose we run [`update`](@ref) on the example `trace`, with the following constraints: -``` -│ -├── :b : false -│ -└── :d : true -``` -```julia -constraints = choicemap((:b, false), (:d, true)) -(new_trace, w, _, discard) = update(trace, (), (), constraints) -``` -Then `get_choices(new_trace)` will be: -``` -│ -├── :a : false -│ -├── :b : false -│ -├── :d : true -│ -└── :e : true -``` -and `discard` will be: -``` -│ -├── :b : true -│ -└── :c : false -``` -Note that the discard contains both the previous values of addresses that were overwritten, and the values for addresses that were in the previous trace but are no longer in the new trace. -The weight (`w`) is computed as: -```math -p(t; x) = 0.7 × 0.4 × 0.4 × 0.7 = 0.0784\\ -p(t'; x') = 0.7 × 0.6 × 0.1 × 0.7 = 0.0294\\ -w = \log p(t'; x')/p(t; x) = \log 0.0294/0.0784 = \log 0.375 -``` - -**Example.** -Suppose we run [`update`](@ref) on the example `trace`, with the following constraints, which *do not* contain a value for `:d`: -``` -│ -└── :b : false -``` -```julia -constraints = choicemap((:b, false)) -(new_trace, w, _, discard) = update(trace, (), (), constraints) -``` -Then `get_choices(new_trace)` will be: -``` -│ -├── :a : false -│ -├── :b : false -│ -├── :d : true -│ -└── :e : true -``` -with probability 0.1, or: -``` -│ -├── :a : false -│ -├── :b : false -│ -├── :d : false -│ -└── :e : true -``` -with probability 0.9. -Also, `discard` will be: -``` -│ -├── :b : true -│ -└── :c : false -``` -If the former case occurs and `:d` is assigned to `true`, then the weight (`w`) is computed as: -```math -p(t; x) = 0.7 × 0.4 × 0.4 × 0.7 = 0.0784\\ -p(t'; x') = 0.7 × 0.6 × 0.1 × 0.7 = 0.0294\\ -q(t'; x', t + u) = 0.1\\ -w = \log p(t'; x')/(p(t; x) q(t'; x', t + u)) = \log 0.0294/(0.0784 \cdot 0.1) = \log (3.75) -``` - - -### Regenerate -The [`regenerate`](@ref) method takes a trace and generates an adjusted trace that is consistent with a change to the arguments to the function, and also generates new values for selected random choices. - -**Example.** -Suppose we run [`regenerate`](@ref) on the example `trace`, with selection `:a` and `:b`: -```julia -(new_trace, w, _) = regenerate(trace, (), (), select(:a, :b)) -``` -Then, a new value for `:a` will be sampled from `bernoulli(0.3)`, and a new value for `:b` will be sampled from `bernoulli(0.4)`. -If the new value for `:b` is `true`, then the previous value for `:c` (`false`) will be retained. -If the new value for `:b` is `false`, then a new value for `:d` will be sampled from `bernoulli(0.7)`. -The previous value for `:c` will always be retained. -Suppose the new value for `:a` is `true`, and the new value for `:b` is `true`. -Then `get_choices(new_trace)` will be: -``` -│ -├── :a : true -│ -├── :b : true -│ -├── :c : false -│ -└── :e : true -``` -The weight (`w`) is ``\log 1 = 0``. - - -### Argdiffs - -In addition to the input trace, and other arguments that indicate how to adjust the trace, each of these methods also accepts an **args** argument and an **argdiffs** argument, both of which are tuples. -The args argument contains the new arguments to the generative function, which may differ from the previous arguments to the generative function (which can be retrieved by applying [`get_args`](@ref) to the previous trace). -In many cases, the adjustment to the execution specified by the other arguments to these methods is 'small' and only affects certain parts of the computation. -Therefore, it is often possible to generate the new trace and the appropriate log probability ratios required for these methods without revisiting every state of the computation of the generative function. - -To enable this, the argdiffs argument provides additional information about the *difference* between each of the previous arguments to the generative function, and its new argument value. -This argdiff information permits the implementation of the update method to avoid inspecting the entire argument data structure to identify which parts were updated. -Note that the correctness of the argdiff is in general not verified by Gen---passing incorrect argdiff information may result in incorrect behavior. - -The trace update methods for all generative functions above should accept at least the following types of argdiffs: -```@docs -NoChange -UnknownChange -``` -Generative functions may also be able to process more specialized diff data types for each of their arguments, that allow more precise information about the different to be supplied. - -### Retdiffs - -To enable generative functions that invoke other functions to efficiently make use of incremental computation, the trace update methods of generative functions also return a **retdiff** value, which provides information about the difference in the return value of the previous trace an the return value of the new trace. - -## Differentiable programming - - -The trace of a generative function may support computation of gradients of its log probability with respect to some subset of (i) its arguments, (ii) values of random choice, and (iii) any of its **trainable parameters** (see below). - -To compute gradients with respect to the arguments as well as certain selected random choices, use: - -- [`choice_gradients`](@ref) - -To compute gradients with respect to the arguments, and to increment a stateful gradient accumulator for the trainable parameters of the generative function, use: - -- [`accumulate_param_gradients!`](@ref) - -A generative function statically reports whether or not it is able to compute gradients with respect to each of its arguments, through the function [`has_argument_grads`](@ref). - -### Trainable parameters -The **trainable parameters** of a generative function are (unlike arguments and random choices) *state* of the generative function itself, and are not contained in the trace. -Generative functions that have trainable parameters maintain *gradient accumulators* for these parameters, which get incremented by the gradient induced by the given trace by a call to [`accumulate_param_gradients!`](@ref). -Users then use these accumulated gradients to update to the values of the trainable parameters. - -### Return value gradient -The set of elements (either arguments, random choices, or trainable parameters) for which gradients are available is called the **gradient source set**. -If the return value of the function is conditionally dependent on any element in the gradient source set given the arguments and values of all other random choices, for all possible traces of the function, then the generative function requires a *return value gradient* to compute gradients with respect to elements of the gradient source set. -This static property of the generative function is reported by [`accepts_output_grad`](@ref). - -## Generative function interface The complete set of methods in the generative function interface (GFI) is: ```@docs diff --git a/docs/src/api/mcmc.md b/docs/src/api/mcmc.md index 65cb8b426..e5cb61a11 100644 --- a/docs/src/api/mcmc.md +++ b/docs/src/api/mcmc.md @@ -1,375 +1,4 @@ # Markov chain Monte Carlo (MCMC) - -Markov chain Monte Carlo (MCMC) is an approach to inference which involves initializing a hypothesis and then repeatedly sampling a new hypotheses given the previous hypothesis by making a change to the previous hypothesis. -The function that samples the new hypothesis given the previous hypothesis is called the **MCMC kernel** (or `kernel' for short). -If we design the kernel appropriately, then the distribution of the hypotheses will converge to the conditional (i.e. posterior) distribution as we increase the number of times we apply the kernel. - -Gen includes primitives for constructing MCMC kernels and composing them into MCMC algorithms. -Although Gen encourages you to write MCMC algorithms that converge to the conditional distribution, Gen does not enforce this requirement. -You may use Gen's MCMC primitives in other ways, including for stochastic optimization. - -For background on MCMC see [1]. - -[1] Andrieu, Christophe, et al. "An introduction to MCMC for machine learning." Machine learning 50.1-2 (2003): 5-43. [Link](https://www.cs.ubc.ca/~arnaud/andrieu_defreitas_doucet_jordan_intromontecarlomachinelearning.pdf). - -## MCMC in Gen -Suppose we are doing inference in the following toy model: -```julia -@gen function model() - x = @trace(bernoulli(0.5), :x) # a latent variable - @trace(normal(x ? -1. : 1., 1.), :y) # the variable that will be observed -end -``` - -To do MCMC, we first need to obtain an initial trace of the model. -Recall that a trace encodes both the observed data and hypothesized values of latent variables. -We can obtain an initial trace that encodes the observed data, and contains a randomly initialized hypothesis, using [`generate`](@ref), e.g.: -```julia -observations = choicemap((:y, 1.23)) -trace, = generate(model, (), observations) -``` - -Then, an MCMC algorithm is Gen is implemented simply by writing Julia `for` loop, which repeatedly applies a kernel, which is a regular Julia function: -```julia -for i=1:100 - trace = kernel(trace) -end -``` - -## Built-in Stationary Kernels -However, we don't expect to be able to use any function for `kernel` and expect to converge to the conditional distribution. -To converge to the conditional distribution, the kernels must satisfy some properties. -One of these properties is that the kernel is **stationary** with respect to the conditional distribution. -Gen's inference library contains a number of functions for constructing stationary kernels: - -- [`metropolis_hastings`](@ref) with alias [`mh`](@ref), which has three variants with differing tradeoffs between ease-of-use and efficiency. The simplest variant simply requires you to select the set of random choices to be updated, without specifying how. The middle variant allows you to use custom proposals that encode problem-specific heuristics, or custom proposals based on neural networks that are trained via amortized inference. The most sophisticated variant allows you to specify any kernel in the [reversible jump MCMC](https://people.maths.bris.ac.uk/~mapjg/papers/RJMCMCBka.pdf) framework. - -- [`mala`](@ref), which performs a Metropolis Adjusted Langevin algorithm update on a set of selected random choices. - -- [`hmc`](@ref), which performs a Hamiltonian Monte Carlo update on a set of selected random choices. - -- [`elliptical_slice`](@ref), which performs an elliptical slice sampling update on a selected multivariate normal random choice. - -For example, here is an MCMC inference algorithm that uses [`mh`](@ref): -```julia -function do_inference(y, num_iters) - trace, = generate(model, (), choicemap((:y, y))) - xs = Float64[] - for i=1:num_iters - trace, = mh(trace, select(:x)) - push!(xs, trace[:x]) - end - xs -end -``` - -Note that each of the kernel functions listed above stationary with respect to the joint distribution on traces of the model, but may not be stationary with respect to the intended conditional distribution, which is determined by the set of addresses that consititute the observed data. -If a kernel modifies the values of any of the observed data, then the kernel is not stationary with respect to the conditional distribution. -Therefore, you should **ensure that your MCMC kernels never propose to the addresses of the observations**. - -Note that stationarity with respect to the conditional distribution alone is not sufficient for a kernel to converge to the posterior with infinite iterations. -Other requirements include that the chain is **irreducible** (it is possible to get from any state to any other state in a finite number of steps), and **aperiodicity**, which is a more complex requirement that is satisfied when kernels have some probability of staying in the same state, which most of the primitive kernels above satisfy. -We refer interested readers to [1] for additional details on MCMC convergence. - -## Enabling Dynamic Checks - -Gen does not statically guarantee that kernels (either ones built-in or composed with the [Composite Kernel DSL](@ref)) are stationary. -However, you can enable dynamic checks that will detect common bugs that break stationarity. -To enable the dynamic checks we pass a keyword argument beyond those of the kernel itself: -```julia -new_trace = k(trace, 2, check=true) -``` -Note that these checks aim to detect when a kernel is not stationary with respect to the model's **joint** distribution. -To add an additional dynamic check for violation of stationarity with respect to the *conditional* distribution (conditioned on observations), we pass in an additional keyword argument containing a choice map with the observations: -```julia -new_trace = k(traced, 2, check=true, observations=choicemap((:y, 1.2))) -``` -If `check` is set to `false`, then the observation check is not performed. - -## Composite Kernel DSL - -You can freely compose the primitive kernels listed above into more complex kernels. -Common types of composition including e.g. cycling through multiple kernels, randomly choosing a kernel to apply, and choosing which kernel to apply based on the current state. -However, not all such compositions of stationary kernels will result in kernels that are themselves stationary. - -Gen's **Composite Kernel DSL** is an embedded inference DSL that allows for more safe composition of MCMC kernels, by formalizing properties of the compositions that are sufficient for stationarity, encouraging compositions with these properties, and dynamically checking for violation of these properties. -Although the DSL does not *guarantee* stationarity of the composite kernels, its dynamic checks do catch common cases of non-stationary kernels. -The dynamic checks can be enabled and disabled as needed (e.g. enabled during testing and prototyping and disabled during deployment for higher performance). - -The DSL consists of a macro -- [`@kern`](@ref) for composing stationary kernels from primitive stationary kernels and composite stationary kernels, and two additional macros: ---- [`@pkern`](@ref) for declaring Julia functions to be custom primitive stationary kernels, and [`@rkern`](@ref) for declaring the reversal of a custom primitive kernel (these two macros are advanced features not necessary for standard MCMC algorithms). - -### Composing Stationary Kernels - -The [`@kern`](@ref) macro defines a composite MCMC kernel in a restricted DSL that is based on Julia's own function definition syntax. - -Suppose we are doing inference in the following model: -```julia -@gen function model() - n = @trace(geometric(0.5), :n) - total = 0. - for i=1:n - total += @trace(normal(0, 1), (:x, i)) - end - @trace(normal(total, 1.), :y) - total -end -``` - -Here is an example composite kernel for MCMC in this model: -```julia -@kern function my_kernel(trace) - - # cycle through the x's and do a random walk update on each one - for i in 1:trace[:n] - trace ~ mh(trace, random_walk_proposal, (i,)) - end - - # repeatedly pick a random x and do a random walk update on it - if trace[:n] > 0 - for rep in 1:10 - let i ~ uniform_discrete(1, trace[:n]) - trace ~ mh(trace, random_walk_proposal, (i,)) - end - end - end - - # remove the last x, or add a new one, a random number of times - let n_add_remove_reps ~ uniform_discrete(0, max_n_add_remove) - for rep in 1:n_add_remove_reps - trace ~ mh(trace, add_remove_proposal, (), add_remove_involution) - end - end -end -``` - -In the DSL, the first argument (`trace` in this case) represents the trace on which the kernel is acting. -the kernel may have additional arguments. -The code inside the body can read from the trace (e.g. `trace[:n]` reads the value of the random choice `:n`). -Finally, the return value of the composite kernel is automatically set to the trace. -NOTE: It is not permitted to assign to the trace variable, except with `~` expressions. -Also note that stationary kernels, when treated as Julia functions, return a tuple, where the first element is the trace and the remaining arguments are metadata. -When applying these kernels with `~` syntax within the DSL, it is not necessary to unpack the tuple (the metadata is ignored automatically). - -The language constructs supported by this DSL are: - -**Applying a stationary kernel.** -To apply a kernel, the syntax `trace ~ k(trace, args..)` is used. -Note that the `check` and `observations` keyword arguments (see [Enabling Dynamic Checks](@ref)) should not be used here; they will be added automatically. - -**For loops.** -The range of the for loop may be a deterministic function of the trace (as in `trace[:n]` above). -The range must be *invariant* under all possible executions of the body of the for loop. -For example, the random walk based kernel embedded in the for loop in our example above cannot modify the value of the random choice `:n` in the trace. - -**If-end expressions** -The predicate condition may be a deterministic function of the trace, but it also must be invariant (i.e. remain true) under all possible executions of the body. - -**Deterministic let expressions.** -We can use `let x = value .. end` to bind values to a variable, but the expression on the right-hand-side must be deterministic function of its free variables, its value must be invariant under all possible executions of the body. - -**Stochastic let expressions.** -We can use `let x ~ dist(args...) .. end` to sample a stochastic value and bind to a variable, but the expression on the right-hand-side must be the application of a Gen [`Distribution`](@ref) to arguments, and the distribution and its arguments must be invariant under all possible executions of the body. - - -### Declaring primitive kernels for use in composite kernels - -Note that all calls to built-in kernels like [`mh`](@ref) should be stationary, but that users are also free to declare their own arbitrary code as stationary. -The [`@pkern`](@ref) macro declares a Julia function as a stationary MCMC kernel, for use with the MCMC Kernel DSL. -The following custom primitive kernel permutes the random variables using random permutation generated from outside of Gen: -```julia -@pkern function permute_move(trace; check=false, observations=EmptyChoiceMap()) - perm = Random.randperm(trace[:n]) - constraints = choicemap() - for (i, j) in enumerate(perm) - constraints[(:x, i)] = trace[(:x, j)] - constraints[(:x, j)] = trace[(:x, i)] - end - trace, = update(trace, (), (), constraints) - metadata = nothing - trace, metadata -end -``` -The first argument to the function should be the trace, and the function must have keyword arguments `check` and `observations` (see [Enabling Dynamic Checks](@ref)). -The return value should be a tuple where the first element is the new trace (and any remaining elements are optional metadata). - -**Primitive kernels are Julia functions.** -Note that although we will be invoking these kernels within [`@kern`](@ref) functions, these kernels can still be called like a regular Julia function. -```julia -new_trace = permute_move(trace, 2) -``` -Indeed, they are just regular Julia functions, but with some extra information attached so that the composite kernel DSL knows they have been declared as stationary kernels. - - -## Involutive MCMC - -Gen's most flexible variant of [`metropolis_hastings`](@ref), called **Involutive MCMC**, allows users to specify any MCMC kernel in the reversible jump MCMC (RJMCMC) framework [2]. -Involution MCMC allows you to express a broad class of custom MCMC kernels that are not expressible using the other, simpler variants of Metropolis-Hastings supported by Gen. -These kernels are particularly useful for inferring the structure (e.g. control flow) of a model. - -[2] Green, Peter J. "Reversible jump Markov chain Monte Carlo computation and Bayesian model determination." Biometrika 82.4 (1995): 711-732. [Link](https://academic.oup.com/biomet/article-abstract/82/4/711/252058) - -An involutive MCMC kernel in Gen takes as input a previous trace of the model (whose choice map we will denote by ``t``), and performs three phases to obtain a new trace of the model: - -- First, it traces the execution of a **proposal**, which is an auxiliary generative function that takes the previous trace of the model as its first argument. Mathematically, we will denote the choice map associated with the trace of the proposal by ``u``. The proposal can of course be defined using the [Built-In Modeling Languages](@ref), just like the model itself. However, unlike many other uses of proposals in Gen, these proposals *can make random choices at addresses that the model does not*. - -- Next, it takes the tuple ``(t, u)`` and passes it into an **involution** (denoted mathematically by ``h``), which is a function that returns a new tuple ``(t', u')``, where ``t'`` is the choice map for a new proposed trace of the model, and ``u'`` are random choices for a new trace of the proposal. The defining property of the involution is that *it is invertible*, and *it is its own inverse*; i.e. ``(t, u) = h(h(t, u))``. Intuitively, ``u'`` is a description of a way that the proposal could be reversed, taking ``t'`` to ``t``. - -- Finally, it computes an acceptance probability, which involves computing certain derivatives associated with the involution, and stochastically accepts or rejects the proposed model trace according to this probability. The involution is typically defined using the [Trace Transform DSL](@ref), in which case the acceptance probability calculation is fully automated. - -### Example -Consider the following generative model of two pieces of observed data, at addresses `:y1` and `:y2`. -```julia -@gen function model() - if ({:z} ~ bernoulli(0.5)) - m1 = ({:m1} ~ gamma(1, 1)) - m2 = ({:m2} ~ gamma(1, 1)) - else - m = ({:m} ~ gamma(1, 1)) - (m1, m2) = (m, m) - end - {:y1} ~ normal(m1, 0.1) - {:y2} ~ normal(m2, 0.1) -end -``` -Because this model has stochastic control flow, it represents two distinct structural hypotheses about how the observed data could have been generated: -If `:z` is `true` then we enter the first branch, and we hypothesize that the two data points were generated from separate means, sampled at addresses `:m1` and `:m2`. -If `:z` is `false` then we enter the second branch, and we hypohesize that there is a single mean that explains both data points, sampled at address `:m`. - -We want to construct an MCMC kernel that is able to transition between these two distinct structural hypotheses. -We could construct such a kernel with the simpler 'selection' variant of Metropolis-Hastings, by selecting the address `:z`, e.g.: -```julia -select_mh_structure_kernel(trace) = mh(trace, select(:z))[1] -``` -Sometimes, this kernel would propose to change the value of `:z`. -We could interleave this kernel with another kernel that does inference over the mean random choices, without changing the structure, e.g.: -```julia -@gen function fixed_structure_proposal(trace) - if trace[:z] - {:m1} ~ normal(trace[:m1], 0.1) - {:m2} ~ normal(trace[:m2], 0.1) - else - {:m} ~ normal(trace[:m], 0.1) - end -end - -fixed_structure_kernel(trace) = mh(trace, fixed_structure_proposal, ())[1] -``` -Combining these together, and applying to particular data and with a specific initial hypotheses: -```julia -(y1, y2) = (1.0, 1.3) -trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.2))) -for iter=1:100 - trace = select_mh_structure_kernel(trace) - trace = fixed_structure_kernel(trace) -end -``` -However, this algorithm will not be very efficient, because the internal proposal used by the selection variant of MH is not specialized to the model. -In particular, when switching from the model with a single mean to the model with two means, the values of the new addresses `:m1` and `:m2` will be proposed from the prior distribution. -This is wasteful, since if we have inferred an accurate value for `:m`, we expect the values for `:m1` and `:m2` to be near this value. -The same is true when proposing a structure change in the opposite direction. -That means it will take many more steps to get an accurate estimate of the posterior probability distribution on the two structures. - -We would like to use inferred values for `:m1` and `:m2` to inform our proposal for the value of `:m`. -For example, we could take the geometric mean: -```julia -m = sqrt(m1 * m2) -``` -However, there are many combinations of `m1` and `m2` that have the same geometric mean. -In other words, the geometric mean is not *invertible*. -However, if we return the additional degree of freedom alongside the geometric mean (`u`), then we do have an invertible function: -```julia -function merge_means(m1, m2) - m = sqrt(m1 * m2) - u = m1 / (m1 + m2) - (m, u) -end -``` -The inverse function is: -```julia -function split_mean(m, u) - m1 = m * sqrt((u / (1 - u))) - m2 = m * sqrt(((1 - u) / u)) - (m1, m2) -end -``` -We use these two functions to construct an involution, and we use this involution with [`metropolis_hastings`](@ref) to construct an MCMC kernel that we call a 'split/merge' kernel, because it either splits a parameter value, or merges two parameter values. -The proposal is responsible for generating the extra degree of freedom when splitting: -```julia -@gen function split_merge_proposal(trace) - if trace[:z] - # currently two segments, switch to one - else - # currently one segment, switch to two - {:u} ~ uniform_continuous(0, 1) - end -end -``` -Finally, we write the involution itself, using the [Trace Transform DSL](@ref): -```julia -@transform split_merge_involution (model_in, aux_in) to (model_out, aux_out) begin - if @read(model_in[:z], :discrete) - - # currently two means, switch to one - @write(model_out[:z], false, :discrete) - m1 = @read(model_in[:m1], :continuous) - m2 = @read(model_in[:m2], :continuous) - (m, u) = merge_means(m1, m2) - @write(model_out[:m], m, :continuous) - @write(aux_out[:u], u, :continuous) - else - - # currently one mean, switch to two - @write(model_out[:z], true, :discrete) - m = @read(model_in[:m], :continuous) - u = @read(aux_in[:u], :continuous) - (m1, m2) = split_mean(m, u) - @write(model_out[:m1], m1, :continuous) - @write(model_out[:m2], m2, :continuous) - end -end -``` -The body of this function reads values from ``(t, u)`` at specific addresses and writes values to ``(t', u')`` at specific addresses, where ``t`` and ``t'`` are called 'model' choice maps, and ``u`` and ``u'`` are called 'proposal' choice maps. -Note that the inputs and outputs of this function are **not** represented in the same way as arguments or return values of regular Julia functions --- they are implicit and can only be read from and written to, respectively, using a set of special macros (listed below). -You should convince yourself that this function is invertible and its own inverse. - -Finally, we compose a structure-changing MCMC kernel using this involution: -```julia -split_merge_kernel(trace) = mh(trace, split_merge_proposal, (), split_merge_involution)[1] -``` -We then compose this move with the fixed structure move, and run it on the observed data: -```julia -(y1, y2) = (1.0, 1.3) -trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.))) -for iter=1:100 - trace = split_merge_kernel(trace) - trace = fixed_structure_kernel(trace) -end -``` -We can then compare the results to the results from the Markov chain that used the selection-based structure-changing kernel: - -![rjmcmc plot](images/rjmcmc.png) - -We see that if we initialize the Markov chains from the same state with a single mean (`:z` is `false`) then the selection-based kernel fails to accept any moves to the two-mean structure within 100 iterations, whereas the split-merge kernel transitions back and forth many times, -If we repeated the selection-based kernel for enough iterations, it would eventually transition back and forth at the same rate as the split-merge. -The split-merge kernel gives a much more efficient inference algorithm for estimating the posterior probability on the two structures. - - -## Reverse Kernels -The **reversal** of a stationary MCMC kernel with distribution ``k_1(t'; t)``, for model with distribution ``p(t; x)``, is another MCMC kernel with distribution: -```math -k_2(t; t') := \frac{p(t; x)}{p(t'; x)} k_1(t'; t) -``` -For custom primitive kernels declared with [`@pkern`](@ref), users can declare the reversal kernel with the [`@rkern`](@ref) macro: -```julia -@rkern k1 : k2 -``` -This also assigns `k1` as the reversal of `k2`. -The composite kernel DSL automatically generates the reversal kernel for composite kernels, and built-in stationary kernels like [`mh`](@ref). -The reversal of a kernel (primitive or composite) can be obtained with [`reversal`](@ref). - - -## API ```@docs metropolis_hastings mh diff --git a/docs/src/api/trace_translators.md b/docs/src/api/trace_translators.md index cec0cf963..30209823e 100644 --- a/docs/src/api/trace_translators.md +++ b/docs/src/api/trace_translators.md @@ -1,440 +1,4 @@ -# Trace Translators - -While [Generative Functions](@ref) define probability distributions on traces, **Trace Translators** convert from one space of traces to another space of traces. -Trace translators are building blocks of inference programs that utilize multiple model representations, like [Involutive MCMC](@ref). - -Trace translators are significantly more general than [Bijectors](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Bijector). -Trace translators can (i) convert between spaces of traces that include mixed numeric discrete random choices, as well as stochastic control flow, and (ii) convert between spaces for which there is no one-to-one correspondence (e.g. between models of different dimensionality, or between discrete and continuous models). -Bijectors are limited to deterministic transformations between real-valued vectors of constant dimension. - -## Deterministic Trace Translators - -Inference programs manipulate traces, but they also keep track of probabilities and probability densities associated with these traces. -Suppose we have two generative functions `p1` and `p2`. -Given a trace `t2` of `p2` we can easily compute the probability (or probability density) that the trace would have been generated by `p2` using `get_score(t2)`. -But suppose we want to construct the trace of `p2` first sampling a trace `t1` of `p1` and then applying a deterministic transformation to that trace to obtain `t2`. -How can we compute the probability that a trace `t2` would have been produced by this process? -This probability is needed if, for example, `p2` defines a probabilistic model and want to use `p1` as a proposal distribution within importance sampling. -If we produce `t2` via an arbitrary deterministic transformation of the random choices in `t1`, then computing the necessary probability is difficult. - -If we restrict ourselves to deterministic transformations that are *bijections* (one-to-one correspondences) from the set of traces of `p1` to the set of traces of `p2`, then the problem is much simplified. -If the transformation is a bijection this means that (i) each trace of `p1` gets mapped to a different trace of `p2`, and (ii) for every trace of `p2` there is some trace of `p1` that maps to it. -Bijective transformations between traces are useful components of inference programs because the probability that a given trace `t2` of `p2` would have been produced by first sampling from `p1` and then applying the transform can be computed simply as the probability that `p1` would produce the (unique) trace `t1` that gets mapped to the given trace by the transform. -Conceptually, bijective trace transforms *preserve probability*. -When trace transforms operate on traces with continuous random choices, computing probability densities of the transformed traces requires computing a Jacobian associated with the continuous part of the transformation. - -Gen provides a DSL for expressing bijections between spaces of traces, called the **Trace Transform DSL**. -We introduce this DSL via an example. -Below are two generative functions. -The first samples polar coordinates and the second uses cartesian coordinates. -```julia -@gen function p1() - r ~ inv_gamma(1, 1) - theta ~ uniform(-pi/2, pi/2) -end -``` -```julia -@gen function p2() - x ~ normal(0, 1) - y ~ normal(0, 1) -end -``` - -### Defining a trace transform with the Trace Transform DSL - -The following trace transform DSL program defines a transformation (called `f`) that transforms traces of `p1` into traces of `p2`: -```julia -@transform f (t1) to (t2) begin - r = @read(t1[:r], :continuous) - theta = @read(t1[:theta], :continuous) - @write(t2[:x], r * cos(theta), :continuous) - @write(t2[:y], r * sin(theta), :continuous) -end -``` -This transform reads values of random choices in the input trace (`t1`) at specific addresses (indicated by the syntax `t1[addr]`) using `@read` and writes values to the output trace (`t2`) using `@write`. -Each read and write operation is labeled with whether the random choice is discrete or continuous. -The section [Trace Transform DSL](@ref) defines the DSL in more detail. - -It is usually a good idea to write the inverse of the bijection. -The inverse can provide a dynamic check that the transform truly is a bijection. -The inverse of the above transformation is: -```julia -@transform finv (t2) to (t1) begin - x = @read(t2[:x], :continuous) - y = @read(t2[:y], :continuous) - r = sqrt(x^2 + y^2) - @write(t1[:r], sqrt(x^2 + y^2), :continuous) - @write(t1[:theta], atan(y, x), :continuous) -end -``` -We can inform Gen that two transforms are inverses of one another using [`pair_bijections!`](@ref): -```julia -pair_bijections!(f, finv) -``` - -### Wrapping a trace transform in a trace translator - -Note that the transform DSL code does not specify what the two generative functions are, or what the arguments to these generative functions are. -This information will be required for computing probabilities and probability densities of traces. -We provide this information by constructing a **Trace Translator** that wraps the transform along with this transformation: -```julia -translator = DeterministicTraceTranslator(p2, (), choicemap(), f) -``` -We then can then apply the translator to a trace of `p1` using function call syntax. -The translator returns a trace of `p2` and a log-weight that we can use to compute the probability (density) of the resulting trace: -```julia -t2, log_weight = translator(t1) -``` -Specifically, the log probability (density) that the trace `t2` was produced by first sampling `t1 = simulate(p1, ())` and then applying the trace translator, is: -```julia -get_score(t1) + log_weight -``` -Let's unpack the previous few code blocks in more detail. -First, note that we did not pass in the source generative function (`p1`) or its arguments (`()`) when we constructed the trace translator. -This is because this information will be attached to the input trace `t1` itself. -We *did* need to pass in the target generative function (`p2`) and its arguments (`()`) however, because this information is not included in `t1`. - -In this case, because continuous random choices are involved, the probabilities are probability densities, and the trace translator used automatic differentiation of the code in the transform `f` to compute a change-of-variables Jacobian that is necessary to compute the correct probability density of the new trace `t2`. - -### Observations - -Typically, there are observations associated with one or both of the generative functions involved, and we have values for these in a choice map, so we do not want the trace translator to be responsible for populating the values of these observed random choices. -For example, suppose we want to condition `p2` on an observed random choice `z`: -```julia -@gen function p2() - x ~ normal(0, 1) - y ~ normal(0, 1) - z ~ normal(x + y, 0.1) -end -observations = choicemap() -observations[:z] = 2.3 -``` -When `p2` has observations, these can be passed in as an additional argument to the trace translator constructor: -```julia -translator = DeterministicTraceTranslator(p2, (), observations, f) -``` - -### Discrete random choices and stochastic control flow - -Trace transforms and trace translators interoperate seamlessly with discrete random choices and stochastic control flow. -Both the input and the output traces can contain a mix of discrete and continuous choices, and arbitrary stochastic control flow. -Consider the following simple example, where we use two different discrete representations to represent probability distributions the integers 0-7: -```julia -@gen function p1() - bit1 ~ bernoulli(0.5) - bit2 ~ bernoulli(0.5) - bit3 ~ bernoulli(0.5) -end -``` -```julia -@gen function p2() - n ~ categorical([0.1, 0.1, 0.1, 0.2, 0.2, 0.15, 0.15]) -end -``` -We define the forward and inverse transforms: -```julia -@transform f (t1) to (t2) begin - n = ( - @read(t1[:bit1], :discrete) * 1 + - @read(t1[:bit2], :discrete) * 2 + - @read(t1[:bit3], :discrete) * 4) - @write(t2[:n], n, :discrete) -end -``` -```julia -@transform finv (t2) to (t1) begin - bits = digits(@read(t2[:n], :discrete), base=2) - @write(t1[:bit1], bits[1], :discrete) - @write(t1[:bit2], bits[2], :discrete) - @write(t1[:bit3], bits[3], :discrete) -end -``` -Here is an example that includes discrete random choices, stochastic control flow, and continuous random choices. -```julia -@gen function p1() - if ({:branch} ~ bernoulli(0.5)) - x ~ normal(0, 1) - else - other ~ categorical([0.3, 0.7]) - end -end -``` -```julia -@gen function p2() - k ~ uniform_discrete(1, 4) - if k <= 2 - y ~ gamma(1, 1) - end -end -``` -Note that transformations between spaces of traces need not be intuitive (although they probably should)! -Try to convince yourself that the functions below are indeed a pair of bijections between the traces of these two generative functions. -```julia -@transform f (t1) to (t2) begin - if @read(t1[:branch], :discrete) - x = @read(t1[:x], :continuous) - if x > 0 - @write(t2[:k], 2, :discrete) - else - @write(t2[:k], 1, :discrete) - end - @write(t2[:y], abs(x), :continuous) - else - other = @read(t1[:other], :discrete) - @write(t2[:k], (other == 1) ? 3 : 4, :discrete) - end -end -``` -```julia -@transform finv (t2) to (t1) begin - k = @read(t2[:k], :discrete) - if k <= 2 - y = @read(t2[:y], :continuous) - @write(t2[:x], (k == 1) ? -y : y, :continuous) - else - @write(t1[:other], (k == 3) ? 1 : 2, :discrete) - end -end -``` - -## General Trace Translators - -Note that for two arbitrary generative functions `p1` and `p2` there may not exist any one-to-one correspondence between their spaces of traces. -For example, consider a generative function `p1` that samples points within the unit square ``[0, 1]^2`` -```julia -@gen function p1() - x ~ uniform(0, 1) - y ~ uniform(0, 1) -end -``` -and another generative function `p2` that samples one of 100 possible discrete values, each value representing one cell of the unit square: -```julia -@gen function p2() - i ~ uniform_discrete(1, 10) # interval [(i-1)/10, i/10] - j ~ uniform_discrete(1, 10) # interval [(j-1)/10, j/10] -end -``` -There is no one-to-one correspondence between the spaces of traces of these two generative functions: -The first is an uncountably infinite set, and the other is a finite set with 100 elements in it. - -However, there is an intuitive notion of correspondence that we would like to be able to encode. -Each discrete cell ``(i, j)`` corresponds to a subset of the unit square ``[(i - 1)/10, i/10] \times [(j-1)/10, j/10]``. - -We can express this correspondence (and any correspondence between two arbitrary generative functions) by introducing two auxiliary generative functions `q1` and `q2`. -The first function `q1` will take a trace of `p1` as input, and the second function `q2` will take a trace of `p2` as input. -Then, instead of a transfomation between traces of `p1` and traces of `p2` our trace transform will transform between (i) the space of pairs of traces of `p1` and `q1` and (ii) the space of pairs of traces of `p2` and `q2`. -We construct `q1` and `q2` so that the two spaces have the same size, and a one-to-one correspondence is possible. - -For our example above, we construct `q2` to sample the coordinate (``[0, 0.1]^2``) relative to the cell. -We construct `q1` to be empty--there is already a mapping from each trace of `p1` to each trace of `p2` that simply identifies what cell ``(i, j)`` a given point in ``[0, 1]^2`` is in, so no extra random choices are needed. -```julia -@gen function q1(p1_trace) -end - -@gen function q2(p2_trace) - dx ~ uniform(0.0, 0.1) - dy ~ uniform(0.0, 0.1) -end -``` -### Trace transforms between pairs of traces - -To handle general trace translators that require auxiliary probability distributions, the trace trace DSL supports defining transformations between *pairs* of traces. -For example, the following defines a trace transform that maps from pairs of traces of `p1` and `q1` to pairs of traces of `p2` and `q2`: - -```julia -@transform f (p1_trace, q1_trace) to (p2_trace, q2_trace) begin - x = @read(p1_trace[:x], :continuous) - y = @read(p1_trace[:y], :continuous) - i = ceil(x * 10) - j = ceil(y * 10) - @write(p2_trace[:i], i, :discrete) - @write(p2_trace[:j], j, :discrete) - @write(q2_trace[:dx], x - (i-1)/10, :continuous) - @write(q2_trace[:dy], y - (j-1)/10, :continuous) -end -``` -and the inverse transform: -```julia -@transform f_inv (p2_trace, q2_trace) to (p1_trace, q1_trace) begin - i = @read(p2_trace[:i], :discrete) - j = @read(p2_trace[:j], :discrete) - dx = @read(q2_trace[:dx], :continuous) - dy = @read(q2_trace[:dy], :continuous) - x = (i-1)/10 + dx - y = (j-1)/10 + dy - @write(p1_trace[:x], x, :continuous) - @write(p1_trace[:y], y, :continuous) -end -``` -which we associate as inverses: -```julia -pair_bijections!(f, f_inv) -``` - -### Constructing a general trace translator - -We now wrap the transform above into a general trace translator, by providing the three probabilistic programs `p2`, `q1`, `q2` that it uses (a reference to `p1` will included in the input trace), and the arguments to these functions. -```julia -translator = GeneralTraceTranslator( - p_new=p2, - p_new_args=(), - new_observations=choicemap(), - q_forward=q1, - q_forward_args=(), - q_backward=q2, - q_backward_args=(), - f=f) -``` -Then, we can apply the trace translator to a trace (`t1`) of `p1` and get a trace (`t2`) of `p2` and a log-weight: -```julia -t2, log_weight = translator(t1) -``` - - -## Symmetric Trace Translators - -When the previous and new generative functions (e.g. `p1` and `p2` in the previous example) are the same, and their arguments are the same, and `q_forward` and `q_backward` (and their arguments) are also identical, we call this the trace translator a **Symmetric Trace Translator**. -Symmetric trace translators are important because they form the basis of [Involutive MCMC](@ref). -Instead of translating a trace of one generative function to the trace of another generative function, they translate a trace of a generative function to another trace of the *same* generative function. - -Symmetric trace translators have the interesting property that the function `f` is an **involution**, or a function that is its own inverse. -To indicate that a trace transform is an involution, use [`is_involution!`](@ref). - -Because symmetric trace translators translate within the same generative function, their implementation uses [`update`](@ref) to incrementally modify the trace from the previous to the new trace. -This has two benefits when the previous and new traces have random choices that aren't modified between them: -(i) the incremental modification may be more efficient than writing the new trace entirely from scratch, and -(ii) the transform DSL program does not need to specify a value for addresses whose value is not changed from the previous trace. - -## Simple Extending Trace Translators - -Simple extending trace translators extend an existing trace with new random -choices sampled from a proposal distribution, as well as any new observations. -The arguments of the trace may also be updated. This type of trace translation -is the basic operation used in [Particle Filtering](@ref). For example, -we might have a model that sequentially samples new latent variables `(:z, t)` -and observations `(:x, t)` for each timestep `t`: - -```julia -@gen function model(T::Int) - for t in 1:T - z = {(:z, t)} ~ normal(0, 1) - x = {(:x, t)} ~ normal(z, 1) - end -end -``` - -Each time we observe a new `(:x ,t)`, we might want to propose `(:z, t)` so that -it is close in value: - -```julia -@gen function proposal(trace::Trace, x::Real) - t = get_args(trace)[1] + 1 - {(:z, t)} ~ normal(x, 1) -end -``` - -Suppose we initially generated a trace up to timestep `t=1`, e.g. by calling -`t1 = simulate(model, (1,))`. Now we observe `(:x, 2)` to be `5.0`. By -constructing a simple extending trace translator, we can simultaneously -update the trace `t1` with new arguments, introduce the new observation -at `(:x, 2)`, and propose a likely value for `(:z, 2)`: - -```julia -translator = SimpleExtendingTraceTranslator( - p_new_args=(2,), p_argdiffs=(UnknownChange(),), - new_observations=choicemap((:x, 2) => 5.0), - q_forward=proposal, q_forward_args=(5.0,)) -t2, log_weight = translator(t1) -``` - -Similar functionality can be achieved through a combination of [`propose`](@ref) -on the proposal and [`update`](@ref) on the original trace, but using a trace -translator provides a nice layer of abstraction. - -## Trace Transform DSL - -The **Trace Transform DSL** is a differentiable programming language for writing deterministic transformations of traces. -Programs written in this DSL are called *transforms*. -Transforms read the value of random choices from input trace(s) and write values of random choices to output trace(s). -These programs are not typically executed directly by users, but are instead wrapped into one of the several forms of trace translators listed above ([`GeneralTraceTranslator`](@ref), and [`SymmetricTraceTranslator`](@ref)). - -A transform is identified with the [`@transform`](@ref) macro, and uses one of the following four syntactic forms for the signature (the name of the transform, and the names of the input and output traces are all user-defined varibles; the only keywords are `@transform`, `to`, `begin`, and `end`): - -*A transform from one trace to another, without extra parameters* -```julia -@transform f t1 to t2 begin - ... -end -``` -*A transform from one trace to another, with extra parameters* -```julia -@transform f(x, y, ..) t1 to t2 begin - ... -end -``` -*A transform from pairs of traces to pairs of traces, without extra parameters* -```julia -@transform f (t1, t2) to (t3, t4) begin - ... -end -``` -*A transform from one trace to another, with extra parameters* -```julia -@transform f(x, y, ..) (t1, t2) to (t3, t4) begin - ... -end -``` -The extra parameters are optional, and can be used to pass arguments to a transform function that is invoked, from another transform function, using the [`@tcall`](@ref) macro. -For example: -```julia -@transform g(x) t1 to t2 begin - ... -end -@transform f t1 to t2 begin - x = .. - @tcall(g(x)) -end -``` -The callee transform function (`g` above) reads and writes to the same input and output traces as the caller transform function (`f` above). Note that the input and output traces can be assigned different names in the caller and the callee. - -The body of a transform reads the values of random choices at addresses in the input trace(s), performs computation using regular Julia code (provided this code can be differentiated with [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)) and writes valeus of random choices at addresses in the output trace(s). -In the body [`@read`](@ref) expressions read a value from a particular address of an input trace: -``` -val = @read(, ) -``` -where `` is an expression of the form `[]` where `` must be the name of an input trace in the transform's signature. -The `` is either `:continuous` or `:discrete`, and indicates whether the random choice is discrete or continuous (in measure-theoretic terms, whether it uses the counting measure, or a Lebesgue-measure a Euclidean space of some dimension). -Similarly, [`@write`](@ref) expressions write a value to a particular address in an output trace: -``` -@write(, , ) -``` -Sometimes trace transforms need to directly copy the value from one address in an input trace to one address in an output trace. -In these cases, it is recommended to use the explicit [`@copy`](@ref) expression: -``` -@copy(, ) -``` -where `` and `` are of the form `[]` as above. -Note you can also copy collections of multiple random choices under an address namespace in an input trace to an address namespace in an output trace. -For example, -``` -@copy(trace1[:foo], trace2[:bar]) -``` -will copy every random choice in `trace1` with an address of the form `:foo => ` to `trace2` at address `:bar => `. - -It is also possible to read the *return value* from an input trace using the following syntax, but this value must be discrete (in the local neighborhood of traces, the return value must be constant as a function of all continuous random choices in input traces): -``` -val = @read([], :discrete) -``` -This feature is useful when the generative function precomputes a quantity as part of its return value, and we would like to reuse this value, instead of having to recompute it as part of the transform. -The `discrete' requirement is needed because the transform DSL does not currently backpropagate through the return value (this feature could be added in the future). - -Tips for defining valid transforms: - -- If you find yourself copying the same continuous source address to multiple locations, it probably means your transform is not valid (the Jacobian matrix will have rows that are identical, and so the Jacobian determinant will be zero). - -- You can gain some confidence that your transform is valid by enabling dynamic checks (`check=true`) in the trace translator that uses it. - -## API +## Trace Translators ```@docs @transform diff --git a/docs/src/api/vi.md b/docs/src/api/vi.md index 3b638449d..4d55fb43c 100644 --- a/docs/src/api/vi.md +++ b/docs/src/api/vi.md @@ -1,18 +1,7 @@ -# Variational Inference - -Variational inference involves optimizing the parameters of a variational family to maximize a lower bound on the marginal likelihood called the ELBO. -In Gen, variational families are represented as generative functions, and variational inference typically involves optimizing the trainable parameters of generative functions. - -## Black box variational inference +## Variational inference There are two procedures in the inference library for performing black box variational inference. Each of these procedures can also train the model using stochastic gradient descent, as in a variational autoencoder. ```@docs black_box_vi! black_box_vimco! ``` - -## Reparametrization trick - -To use the reparametrization trick to reduce the variance of gradient estimators, users currently need to write two versions of their variational family, one that is reparametrized and one that is not. -Gen does not currently include inference library support for this. -We plan add add automated support for reparametrization and other variance reduction techniques in the future. diff --git a/docs/src/getting_started/mcmc_cipher.md b/docs/src/getting_started/mcmc_cipher.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/misc/changelog.md b/docs/src/misc/changelog.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/misc/contributing.md b/docs/src/misc/contributing.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/misc/references.md b/docs/src/misc/references.md new file mode 100644 index 000000000..e69de29bb diff --git a/docs/src/tutorials/gfi.md b/docs/src/tutorials/gfi.md new file mode 100644 index 000000000..28a2d30f5 --- /dev/null +++ b/docs/src/tutorials/gfi.md @@ -0,0 +1,355 @@ +# [Generative Function Interface](@id gfi) + +One of the core abstractions in Gen is the **generative function**. +Generative functions are used to represent a variety of different types of probabilistic computations including generative models, inference models, custom proposal distributions, and variational approximations. + +## Introduction + +Generative functions are represented by the following abstact type: + +There are various kinds of generative functions, which are represented by concrete subtypes of [`GenerativeFunction`](@ref). +For example, the [Built-in Modeling Language](@ref) allows generative functions to be constructed using Julia function definition syntax: +```@example gfi_example +using Gen # hide +@gen function foo(a, b=0) + if @trace(bernoulli(0.5), :z) + return a + b + 1 + else + return a + b + end +end +``` +Users can also extend Gen by implementing their own [Custom generative function types](@ref), which can be new modeling languages, or just specialized optimized implementations of a fragment of a specific model. + +Generative functions behave like Julia functions in some respects. +For example, we can call a generative function `foo` on arguments and get an output value using regular Julia call syntax: +```@example gfi_example +foo(2, 4) +``` + +However, generative functions are distinct from Julia functions because they support additional behaviors, described in the remainder of this section. + + +## Mathematical concepts + +Generative functions represent computations that accept some arguments, may use randomness internally, return an output, and cannot mutate externally observable state. +We represent the randomness used during an execution of a generative function as a **choice map** from unique **addresses** to values of random choices, denoted ``t : A \to V`` where ``A`` is a finite (but not a priori bounded) address set and ``V`` is a set of possible values that random choices can take. +In this section, we assume that random choices are discrete to simplify notation. +We say that two choice maps ``t`` and ``s`` **agree** if they assign the same value for any address that is in both of their domains. + +Generative functions may also use **non-addressable randomness**, which is not included in the map ``t``. +We denote non-addressable randomness by ``r``. +Untraced randomness is useful for example, when calling black box Julia code that implements a randomized algorithm. + +The observable behavior of every generative function is defined by the following mathematical objects: + +### Input type +The set of valid argument tuples to the function, denoted ``X``. + +### Probability distribution family +A family of probability distributions ``p(t, r; x)`` on maps ``t`` from random choice addresses to their values, and non-addressable randomness ``r``, indexed by arguments ``x``, for all ``x \in X``. +Note that the distribution must be normalized: +```math +\sum_{t, r} p(t, r; x) = 1 \;\; \mbox{for all} \;\; x \in X +``` +This corresponds to a requirement that the function terminate with probabability 1 for all valid arguments. +We use ``p(t; x)`` to denote the marginal distribution on the map ``t``: +```math +p(t; x) := \sum_{r} p(t, r; x) +``` +And we denote the conditional distribution on non-addressable randomness ``r``, given the map ``t``, as: +```math +p(r | t; x) := p(t, r; x) / p(t; x) +``` + +### Return value function +A (deterministic) function ``f`` that maps the tuple ``(x, t)`` of the arguments and the choice map to the return value of the function (which we denote by ``y``). +Note that the return value cannot depend on the non-addressable randomness. + +### Auxiliary state +Generative functions may expose additional **auxiliary state** associated with an execution, besides the choice map and the return value. +This auxiliary state is a function ``z = h(x, t, r)`` of the arguments, choice map, and non-addressable randomness. +Like the choice map, the auxiliary state is indexed by addresses. +We require that the addresses of auxiliary state are disjoint from the addresses in the choice map. +Note that when a generative function is called within a model, the auxiliary state is not available to the caller. +It is typically used by inference programs, for logging and for caching the results of deterministic computations that would otherwise need to be reconstructed. + +### Internal proposal distribution family +A family of probability distributions ``q(t; x, u)`` on maps ``t`` from random choice addresses to their values, indexed by tuples ``(x, u)`` where ``u`` is a map from random choice addresses to values, and where ``x`` are the arguments to the function. +It must satisfy the following conditions: +```math +\sum_{t} q(t; x, u) = 1 \;\; \mbox{for all} \;\; x \in X, u +``` +```math +p(t; x) > 0 \mbox{ if and only if } q(t; x, u) > 0 \mbox{ for all } u \mbox{ where } u \mbox{ and } t \mbox{ agree } +``` +```math +q(t; x, u) > 0 \mbox{ implies that } u \mbox{ and } t \mbox{ agree }. +``` +There is also a family of probability distributions ``q(r; x, t)`` on non-addressable randomness, that satisfies: +```math +q(r; x, t) > 0 \mbox{ if and only if } p(r | t, x) > 0 +``` + + +## Traces + +An **execution trace** (or just *trace*) is a record of an execution of a generative function. +Traces are the primary data structures manipulated by Gen inference programs. +There are various methods for producing, updating, and inspecting traces. +Traces contain: + +- the arguments to the generative function + +- the choice map + +- the return value + +- auxiliary state + +- other implementation-specific state that is not exposed to the caller or user of the generative function, but is used internally to facilitate e.g. incremental updates to executions and automatic differentiation + +- any necessary record of the non-addressable randomness + + +Different concrete types of generative functions use different data structures and different Julia types for their traces, but traces are subtypes of [`Trace`](@ref). +```@docs +Trace +``` +The concrete trace type that a generative function uses is the second type parameter of the [`GenerativeFunction`](@ref) abstract type. +For example, the trace type of [`DynamicDSLFunction`](@ref) is `DynamicDSLTrace`. + +A generative function can be executed to produce a trace of the execution using [`simulate`](@ref): +```julia +trace = simulate(foo, (a, b)) +``` +A traced execution that satisfies constraints on the choice map can be generated using [`generate`](@ref): +```julia +trace, weight = generate(foo, (a, b), choicemap((:z, false))) +``` + +There are various methods for inspecting traces, including: + +- [`get_args`](@ref) (returns the arguments to the function) + +- [`get_retval`](@ref) (returns the return value of the function) + +- [`get_choices`](@ref) (returns the choice map) + +- [`get_score`](@ref) (returns the log probability that the random choices took the values they did) + +- [`get_gen_fn`](@ref) (returns a reference to the generative function) + +You can also access the values in the choice map and the auxiliary state of the trace by passing the address to [`Base.getindex`](@ref). +For example, to retrieve the value of random choice at address `:z`: +```julia +z = trace[:z] +``` + +When a generative function has default values specified for trailing arguments, those arguments can be left out when calling [`simulate`](@ref), [`generate`](@ref), and other functions provided by the generative function interface. The default values will automatically be filled in: +```julia +julia> trace = simulate(foo, (2,)); +julia> get_args(trace) +(2, 0) +``` + +## Updating traces + +It is often important to incrementally modify the trace of a generative function (e.g. within MCMC, numerical optimization, sequential Monte Carlo, etc.). +In Gen, traces are **functional data structures**, meaning they can be treated as immutable values. +There are several methods that take a trace of a generative function as input and return a new trace of the generative function based on adjustments to the execution history of the function. +We will illustrate these methods using the following generative function: +```@example gfi_example +@gen function bar() + val = @trace(bernoulli(0.3), :a) + if @trace(bernoulli(0.4), :b) + val = @trace(bernoulli(0.6), :c) && val + else + val = @trace(bernoulli(0.1), :d) && val + end + val = @trace(bernoulli(0.7), :e) && val + return val +end +``` +Suppose we have a trace (`trace`) of `bar` with initial choices: +```@example gfi_example +a = 1 +``` +``` +│ +├── :a : false +│ +├── :b : true +│ +├── :c : false +│ +└── :e : true +``` +Note that address `:d` is not present because the branch in which `:d` is sampled was not taken because random choice `:b` had value `true`. + +### Update +The [`update`](@ref) method takes a trace and generates an adjusted trace that is consistent with given changes to the arguments to the function, and changes to the values of random choices made. + +**Example.** +Suppose we run [`update`](@ref) on the example `trace`, with the following constraints: +``` +│ +├── :b : false +│ +└── :d : true +``` +```julia +constraints = choicemap((:b, false), (:d, true)) +(new_trace, w, _, discard) = update(trace, (), (), constraints) +``` +Then `get_choices(new_trace)` will be: +``` +│ +├── :a : false +│ +├── :b : false +│ +├── :d : true +│ +└── :e : true +``` +and `discard` will be: +``` +│ +├── :b : true +│ +└── :c : false +``` +Note that the discard contains both the previous values of addresses that were overwritten, and the values for addresses that were in the previous trace but are no longer in the new trace. +The weight (`w`) is computed as: +```math +p(t; x) = 0.7 × 0.4 × 0.4 × 0.7 = 0.0784\\ +p(t'; x') = 0.7 × 0.6 × 0.1 × 0.7 = 0.0294\\ +w = \log p(t'; x')/p(t; x) = \log 0.0294/0.0784 = \log 0.375 +``` + +**Example.** +Suppose we run [`update`](@ref) on the example `trace`, with the following constraints, which *do not* contain a value for `:d`: +``` +│ +└── :b : false +``` +```@example gfi_example +constraints = choicemap((:b, false)) +(new_trace, w, _, discard) = update(trace, (), (), constraints) +``` +Then `get_choices(new_trace)` will be: +``` +│ +├── :a : false +│ +├── :b : false +│ +├── :d : true +│ +└── :e : true +``` +with probability 0.1, or: +``` +│ +├── :a : false +│ +├── :b : false +│ +├── :d : false +│ +└── :e : true +``` +with probability 0.9. +Also, `discard` will be: +``` +│ +├── :b : true +│ +└── :c : false +``` +If the former case occurs and `:d` is assigned to `true`, then the weight (`w`) is computed as: +```math +p(t; x) = 0.7 × 0.4 × 0.4 × 0.7 = 0.0784\\ +p(t'; x') = 0.7 × 0.6 × 0.1 × 0.7 = 0.0294\\ +q(t'; x', t + u) = 0.1\\ +w = \log p(t'; x')/(p(t; x) q(t'; x', t + u)) = \log 0.0294/(0.0784 \cdot 0.1) = \log (3.75) +``` + + +### Regenerate +The [`regenerate`](@ref) method takes a trace and generates an adjusted trace that is consistent with a change to the arguments to the function, and also generates new values for selected random choices. + +**Example.** +Suppose we run [`regenerate`](@ref) on the example `trace`, with selection `:a` and `:b`: +```@example gfi_example +(new_trace, w, _) = regenerate(trace, (), (), select(:a, :b)) +``` +Then, a new value for `:a` will be sampled from `bernoulli(0.3)`, and a new value for `:b` will be sampled from `bernoulli(0.4)`. +If the new value for `:b` is `true`, then the previous value for `:c` (`false`) will be retained. +If the new value for `:b` is `false`, then a new value for `:d` will be sampled from `bernoulli(0.7)`. +The previous value for `:c` will always be retained. +Suppose the new value for `:a` is `true`, and the new value for `:b` is `true`. +Then `get_choices(new_trace)` will be: +```@example gfi_example +get_choices(new_trace) +``` + +The weight (`w`) is ``\log 1 = 0``. + + +### Argdiffs + +In addition to the input trace, and other arguments that indicate how to adjust the trace, each of these methods also accepts an **args** argument and an **argdiffs** argument, both of which are tuples. +The args argument contains the new arguments to the generative function, which may differ from the previous arguments to the generative function (which can be retrieved by applying [`get_args`](@ref) to the previous trace). +In many cases, the adjustment to the execution specified by the other arguments to these methods is 'small' and only affects certain parts of the computation. +Therefore, it is often possible to generate the new trace and the appropriate log probability ratios required for these methods without revisiting every state of the computation of the generative function. + +To enable this, the argdiffs argument provides additional information about the *difference* between each of the previous arguments to the generative function, and its new argument value. +This argdiff information permits the implementation of the update method to avoid inspecting the entire argument data structure to identify which parts were updated. +Note that the correctness of the argdiff is in general not verified by Gen---passing incorrect argdiff information may result in incorrect behavior. + +The trace update methods for all generative functions above should accept at least the following types of argdiffs: +```@docs +NoChange +UnknownChange +``` +Generative functions may also be able to process more specialized diff data types for each of their arguments, that allow more precise information about the different to be supplied. + +### Retdiffs + +To enable generative functions that invoke other functions to efficiently make use of incremental computation, the trace update methods of generative functions also return a **retdiff** value, which provides information about the difference in the return value of the previous trace an the return value of the new trace. + +## Differentiable programming + + +The trace of a generative function may support computation of gradients of its log probability with respect to some subset of (i) its arguments, (ii) values of random choice, and (iii) any of its **trainable parameters** (see below). + +To compute gradients with respect to the arguments as well as certain selected random choices, use: + +- [`choice_gradients`](@ref) + +To compute gradients with respect to the arguments, and to increment a stateful gradient accumulator for the trainable parameters of the generative function, use: + +- [`accumulate_param_gradients!`](@ref) + +A generative function statically reports whether or not it is able to compute gradients with respect to each of its arguments, through the function [`has_argument_grads`](@ref). + +### Trainable parameters +The **trainable parameters** of a generative function are (unlike arguments and random choices) *state* of the generative function itself, and are not contained in the trace. +Generative functions that have trainable parameters maintain *gradient accumulators* for these parameters, which get incremented by the gradient induced by the given trace by a call to [`accumulate_param_gradients!`](@ref). +Users then use these accumulated gradients to update to the values of the trainable parameters. + +### Return value gradient +The set of elements (either arguments, random choices, or trainable parameters) for which gradients are available is called the **gradient source set**. +If the return value of the function is conditionally dependent on any element in the gradient source set given the arguments and values of all other random choices, for all possible traces of the function, then the generative function requires a *return value gradient* to compute gradients with respect to elements of the gradient source set. +This static property of the generative function is reported by [`accepts_output_grad`](@ref). diff --git a/docs/src/api/extending.md b/docs/src/tutorials/languages/extending.md similarity index 100% rename from docs/src/api/extending.md rename to docs/src/tutorials/languages/extending.md diff --git a/docs/src/tutorials/mcmc.md b/docs/src/tutorials/mcmc.md new file mode 100644 index 000000000..b0917a152 --- /dev/null +++ b/docs/src/tutorials/mcmc.md @@ -0,0 +1,369 @@ +# Markov chain Monte Carlo (MCMC) + +Markov chain Monte Carlo (MCMC) is an approach to inference which involves initializing a hypothesis and then repeatedly sampling a new hypotheses given the previous hypothesis by making a change to the previous hypothesis. +The function that samples the new hypothesis given the previous hypothesis is called the **MCMC kernel** (or `kernel' for short). +If we design the kernel appropriately, then the distribution of the hypotheses will converge to the conditional (i.e. posterior) distribution as we increase the number of times we apply the kernel. + +Gen includes primitives for constructing MCMC kernels and composing them into MCMC algorithms. +Although Gen encourages you to write MCMC algorithms that converge to the conditional distribution, Gen does not enforce this requirement. +You may use Gen's MCMC primitives in other ways, including for stochastic optimization. + +For background on MCMC see [1]. + +[1] Andrieu, Christophe, et al. "An introduction to MCMC for machine learning." Machine learning 50.1-2 (2003): 5-43. [Link](https://www.cs.ubc.ca/~arnaud/andrieu_defreitas_doucet_jordan_intromontecarlomachinelearning.pdf). + +## MCMC in Gen +Suppose we are doing inference in the following toy model: +```julia +@gen function model() + x = @trace(bernoulli(0.5), :x) # a latent variable + @trace(normal(x ? -1. : 1., 1.), :y) # the variable that will be observed +end +``` + +To do MCMC, we first need to obtain an initial trace of the model. +Recall that a trace encodes both the observed data and hypothesized values of latent variables. +We can obtain an initial trace that encodes the observed data, and contains a randomly initialized hypothesis, using [`generate`](@ref), e.g.: +```julia +observations = choicemap((:y, 1.23)) +trace, = generate(model, (), observations) +``` + +Then, an MCMC algorithm is Gen is implemented simply by writing Julia `for` loop, which repeatedly applies a kernel, which is a regular Julia function: +```julia +for i=1:100 + trace = kernel(trace) +end +``` + +## Built-in Stationary Kernels +However, we don't expect to be able to use any function for `kernel` and expect to converge to the conditional distribution. +To converge to the conditional distribution, the kernels must satisfy some properties. +One of these properties is that the kernel is **stationary** with respect to the conditional distribution. +Gen's inference library contains a number of functions for constructing stationary kernels: + +- [`metropolis_hastings`](@ref) with alias [`mh`](@ref), which has three variants with differing tradeoffs between ease-of-use and efficiency. The simplest variant simply requires you to select the set of random choices to be updated, without specifying how. The middle variant allows you to use custom proposals that encode problem-specific heuristics, or custom proposals based on neural networks that are trained via amortized inference. The most sophisticated variant allows you to specify any kernel in the [reversible jump MCMC](https://people.maths.bris.ac.uk/~mapjg/papers/RJMCMCBka.pdf) framework. + +- [`mala`](@ref), which performs a Metropolis Adjusted Langevin algorithm update on a set of selected random choices. + +- [`hmc`](@ref), which performs a Hamiltonian Monte Carlo update on a set of selected random choices. + +- [`elliptical_slice`](@ref), which performs an elliptical slice sampling update on a selected multivariate normal random choice. + +For example, here is an MCMC inference algorithm that uses [`mh`](@ref): +```julia +function do_inference(y, num_iters) + trace, = generate(model, (), choicemap((:y, y))) + xs = Float64[] + for i=1:num_iters + trace, = mh(trace, select(:x)) + push!(xs, trace[:x]) + end + xs +end +``` + +Note that each of the kernel functions listed above stationary with respect to the joint distribution on traces of the model, but may not be stationary with respect to the intended conditional distribution, which is determined by the set of addresses that consititute the observed data. +If a kernel modifies the values of any of the observed data, then the kernel is not stationary with respect to the conditional distribution. +Therefore, you should **ensure that your MCMC kernels never propose to the addresses of the observations**. + +Note that stationarity with respect to the conditional distribution alone is not sufficient for a kernel to converge to the posterior with infinite iterations. +Other requirements include that the chain is **irreducible** (it is possible to get from any state to any other state in a finite number of steps), and **aperiodicity**, which is a more complex requirement that is satisfied when kernels have some probability of staying in the same state, which most of the primitive kernels above satisfy. +We refer interested readers to [1] for additional details on MCMC convergence. + +## Enabling Dynamic Checks + +Gen does not statically guarantee that kernels (either ones built-in or composed with the [Composite Kernel DSL](@ref)) are stationary. +However, you can enable dynamic checks that will detect common bugs that break stationarity. +To enable the dynamic checks we pass a keyword argument beyond those of the kernel itself: +```julia +new_trace = k(trace, 2, check=true) +``` +Note that these checks aim to detect when a kernel is not stationary with respect to the model's **joint** distribution. +To add an additional dynamic check for violation of stationarity with respect to the *conditional* distribution (conditioned on observations), we pass in an additional keyword argument containing a choice map with the observations: +```julia +new_trace = k(traced, 2, check=true, observations=choicemap((:y, 1.2))) +``` +If `check` is set to `false`, then the observation check is not performed. + +## Composite Kernel DSL + +You can freely compose the primitive kernels listed above into more complex kernels. +Common types of composition including e.g. cycling through multiple kernels, randomly choosing a kernel to apply, and choosing which kernel to apply based on the current state. +However, not all such compositions of stationary kernels will result in kernels that are themselves stationary. + +Gen's **Composite Kernel DSL** is an embedded inference DSL that allows for more safe composition of MCMC kernels, by formalizing properties of the compositions that are sufficient for stationarity, encouraging compositions with these properties, and dynamically checking for violation of these properties. +Although the DSL does not *guarantee* stationarity of the composite kernels, its dynamic checks do catch common cases of non-stationary kernels. +The dynamic checks can be enabled and disabled as needed (e.g. enabled during testing and prototyping and disabled during deployment for higher performance). + +The DSL consists of a macro -- [`@kern`](@ref) for composing stationary kernels from primitive stationary kernels and composite stationary kernels, and two additional macros: +--- [`@pkern`](@ref) for declaring Julia functions to be custom primitive stationary kernels, and [`@rkern`](@ref) for declaring the reversal of a custom primitive kernel (these two macros are advanced features not necessary for standard MCMC algorithms). + +### Composing Stationary Kernels + +The [`@kern`](@ref) macro defines a composite MCMC kernel in a restricted DSL that is based on Julia's own function definition syntax. + +Suppose we are doing inference in the following model: +```julia +@gen function model() + n = @trace(geometric(0.5), :n) + total = 0. + for i=1:n + total += @trace(normal(0, 1), (:x, i)) + end + @trace(normal(total, 1.), :y) + total +end +``` + +Here is an example composite kernel for MCMC in this model: +```julia +@kern function my_kernel(trace) + + # cycle through the x's and do a random walk update on each one + for i in 1:trace[:n] + trace ~ mh(trace, random_walk_proposal, (i,)) + end + + # repeatedly pick a random x and do a random walk update on it + if trace[:n] > 0 + for rep in 1:10 + let i ~ uniform_discrete(1, trace[:n]) + trace ~ mh(trace, random_walk_proposal, (i,)) + end + end + end + + # remove the last x, or add a new one, a random number of times + let n_add_remove_reps ~ uniform_discrete(0, max_n_add_remove) + for rep in 1:n_add_remove_reps + trace ~ mh(trace, add_remove_proposal, (), add_remove_involution) + end + end +end +``` + +In the DSL, the first argument (`trace` in this case) represents the trace on which the kernel is acting. +the kernel may have additional arguments. +The code inside the body can read from the trace (e.g. `trace[:n]` reads the value of the random choice `:n`). +Finally, the return value of the composite kernel is automatically set to the trace. +NOTE: It is not permitted to assign to the trace variable, except with `~` expressions. +Also note that stationary kernels, when treated as Julia functions, return a tuple, where the first element is the trace and the remaining arguments are metadata. +When applying these kernels with `~` syntax within the DSL, it is not necessary to unpack the tuple (the metadata is ignored automatically). + +The language constructs supported by this DSL are: + +**Applying a stationary kernel.** +To apply a kernel, the syntax `trace ~ k(trace, args..)` is used. +Note that the `check` and `observations` keyword arguments (see [Enabling Dynamic Checks](@ref)) should not be used here; they will be added automatically. + +**For loops.** +The range of the for loop may be a deterministic function of the trace (as in `trace[:n]` above). +The range must be *invariant* under all possible executions of the body of the for loop. +For example, the random walk based kernel embedded in the for loop in our example above cannot modify the value of the random choice `:n` in the trace. + +**If-end expressions** +The predicate condition may be a deterministic function of the trace, but it also must be invariant (i.e. remain true) under all possible executions of the body. + +**Deterministic let expressions.** +We can use `let x = value .. end` to bind values to a variable, but the expression on the right-hand-side must be deterministic function of its free variables, its value must be invariant under all possible executions of the body. + +**Stochastic let expressions.** +We can use `let x ~ dist(args...) .. end` to sample a stochastic value and bind to a variable, but the expression on the right-hand-side must be the application of a Gen [`Distribution`](@ref) to arguments, and the distribution and its arguments must be invariant under all possible executions of the body. + + +### Declaring primitive kernels for use in composite kernels + +Note that all calls to built-in kernels like [`mh`](@ref) should be stationary, but that users are also free to declare their own arbitrary code as stationary. +The [`@pkern`](@ref) macro declares a Julia function as a stationary MCMC kernel, for use with the MCMC Kernel DSL. +The following custom primitive kernel permutes the random variables using random permutation generated from outside of Gen: +```julia +@pkern function permute_move(trace; check=false, observations=EmptyChoiceMap()) + perm = Random.randperm(trace[:n]) + constraints = choicemap() + for (i, j) in enumerate(perm) + constraints[(:x, i)] = trace[(:x, j)] + constraints[(:x, j)] = trace[(:x, i)] + end + trace, = update(trace, (), (), constraints) + metadata = nothing + trace, metadata +end +``` +The first argument to the function should be the trace, and the function must have keyword arguments `check` and `observations` (see [Enabling Dynamic Checks](@ref)). +The return value should be a tuple where the first element is the new trace (and any remaining elements are optional metadata). + +**Primitive kernels are Julia functions.** +Note that although we will be invoking these kernels within [`@kern`](@ref) functions, these kernels can still be called like a regular Julia function. +```julia +new_trace = permute_move(trace, 2) +``` +Indeed, they are just regular Julia functions, but with some extra information attached so that the composite kernel DSL knows they have been declared as stationary kernels. + + +## Involutive MCMC + +Gen's most flexible variant of [`metropolis_hastings`](@ref), called **Involutive MCMC**, allows users to specify any MCMC kernel in the reversible jump MCMC (RJMCMC) framework [2]. +Involution MCMC allows you to express a broad class of custom MCMC kernels that are not expressible using the other, simpler variants of Metropolis-Hastings supported by Gen. +These kernels are particularly useful for inferring the structure (e.g. control flow) of a model. + +[2] Green, Peter J. "Reversible jump Markov chain Monte Carlo computation and Bayesian model determination." Biometrika 82.4 (1995): 711-732. [Link](https://academic.oup.com/biomet/article-abstract/82/4/711/252058) + +An involutive MCMC kernel in Gen takes as input a previous trace of the model (whose choice map we will denote by ``t``), and performs three phases to obtain a new trace of the model: + +- First, it traces the execution of a **proposal**, which is an auxiliary generative function that takes the previous trace of the model as its first argument. Mathematically, we will denote the choice map associated with the trace of the proposal by ``u``. The proposal can of course be defined using the [Built-In Modeling Languages](@ref), just like the model itself. However, unlike many other uses of proposals in Gen, these proposals *can make random choices at addresses that the model does not*. + +- Next, it takes the tuple ``(t, u)`` and passes it into an **involution** (denoted mathematically by ``h``), which is a function that returns a new tuple ``(t', u')``, where ``t'`` is the choice map for a new proposed trace of the model, and ``u'`` are random choices for a new trace of the proposal. The defining property of the involution is that *it is invertible*, and *it is its own inverse*; i.e. ``(t, u) = h(h(t, u))``. Intuitively, ``u'`` is a description of a way that the proposal could be reversed, taking ``t'`` to ``t``. + +- Finally, it computes an acceptance probability, which involves computing certain derivatives associated with the involution, and stochastically accepts or rejects the proposed model trace according to this probability. The involution is typically defined using the [Trace Transform DSL](@ref), in which case the acceptance probability calculation is fully automated. + +### Example +Consider the following generative model of two pieces of observed data, at addresses `:y1` and `:y2`. +```julia +@gen function model() + if ({:z} ~ bernoulli(0.5)) + m1 = ({:m1} ~ gamma(1, 1)) + m2 = ({:m2} ~ gamma(1, 1)) + else + m = ({:m} ~ gamma(1, 1)) + (m1, m2) = (m, m) + end + {:y1} ~ normal(m1, 0.1) + {:y2} ~ normal(m2, 0.1) +end +``` +Because this model has stochastic control flow, it represents two distinct structural hypotheses about how the observed data could have been generated: +If `:z` is `true` then we enter the first branch, and we hypothesize that the two data points were generated from separate means, sampled at addresses `:m1` and `:m2`. +If `:z` is `false` then we enter the second branch, and we hypohesize that there is a single mean that explains both data points, sampled at address `:m`. + +We want to construct an MCMC kernel that is able to transition between these two distinct structural hypotheses. +We could construct such a kernel with the simpler 'selection' variant of Metropolis-Hastings, by selecting the address `:z`, e.g.: +```julia +select_mh_structure_kernel(trace) = mh(trace, select(:z))[1] +``` +Sometimes, this kernel would propose to change the value of `:z`. +We could interleave this kernel with another kernel that does inference over the mean random choices, without changing the structure, e.g.: +```julia +@gen function fixed_structure_proposal(trace) + if trace[:z] + {:m1} ~ normal(trace[:m1], 0.1) + {:m2} ~ normal(trace[:m2], 0.1) + else + {:m} ~ normal(trace[:m], 0.1) + end +end + +fixed_structure_kernel(trace) = mh(trace, fixed_structure_proposal, ())[1] +``` +Combining these together, and applying to particular data and with a specific initial hypotheses: +```julia +(y1, y2) = (1.0, 1.3) +trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.2))) +for iter=1:100 + trace = select_mh_structure_kernel(trace) + trace = fixed_structure_kernel(trace) +end +``` +However, this algorithm will not be very efficient, because the internal proposal used by the selection variant of MH is not specialized to the model. +In particular, when switching from the model with a single mean to the model with two means, the values of the new addresses `:m1` and `:m2` will be proposed from the prior distribution. +This is wasteful, since if we have inferred an accurate value for `:m`, we expect the values for `:m1` and `:m2` to be near this value. +The same is true when proposing a structure change in the opposite direction. +That means it will take many more steps to get an accurate estimate of the posterior probability distribution on the two structures. + +We would like to use inferred values for `:m1` and `:m2` to inform our proposal for the value of `:m`. +For example, we could take the geometric mean: +```julia +m = sqrt(m1 * m2) +``` +However, there are many combinations of `m1` and `m2` that have the same geometric mean. +In other words, the geometric mean is not *invertible*. +However, if we return the additional degree of freedom alongside the geometric mean (`u`), then we do have an invertible function: +```julia +function merge_means(m1, m2) + m = sqrt(m1 * m2) + u = m1 / (m1 + m2) + (m, u) +end +``` +The inverse function is: +```julia +function split_mean(m, u) + m1 = m * sqrt((u / (1 - u))) + m2 = m * sqrt(((1 - u) / u)) + (m1, m2) +end +``` +We use these two functions to construct an involution, and we use this involution with [`metropolis_hastings`](@ref) to construct an MCMC kernel that we call a 'split/merge' kernel, because it either splits a parameter value, or merges two parameter values. +The proposal is responsible for generating the extra degree of freedom when splitting: +```julia +@gen function split_merge_proposal(trace) + if trace[:z] + # currently two segments, switch to one + else + # currently one segment, switch to two + {:u} ~ uniform_continuous(0, 1) + end +end +``` +Finally, we write the involution itself, using the [Trace Transform DSL](@ref): +```julia +@transform split_merge_involution (model_in, aux_in) to (model_out, aux_out) begin + if @read(model_in[:z], :discrete) + + # currently two means, switch to one + @write(model_out[:z], false, :discrete) + m1 = @read(model_in[:m1], :continuous) + m2 = @read(model_in[:m2], :continuous) + (m, u) = merge_means(m1, m2) + @write(model_out[:m], m, :continuous) + @write(aux_out[:u], u, :continuous) + else + + # currently one mean, switch to two + @write(model_out[:z], true, :discrete) + m = @read(model_in[:m], :continuous) + u = @read(aux_in[:u], :continuous) + (m1, m2) = split_mean(m, u) + @write(model_out[:m1], m1, :continuous) + @write(model_out[:m2], m2, :continuous) + end +end +``` +The body of this function reads values from ``(t, u)`` at specific addresses and writes values to ``(t', u')`` at specific addresses, where ``t`` and ``t'`` are called 'model' choice maps, and ``u`` and ``u'`` are called 'proposal' choice maps. +Note that the inputs and outputs of this function are **not** represented in the same way as arguments or return values of regular Julia functions --- they are implicit and can only be read from and written to, respectively, using a set of special macros (listed below). +You should convince yourself that this function is invertible and its own inverse. + +Finally, we compose a structure-changing MCMC kernel using this involution: +```julia +split_merge_kernel(trace) = mh(trace, split_merge_proposal, (), split_merge_involution)[1] +``` +We then compose this move with the fixed structure move, and run it on the observed data: +```julia +(y1, y2) = (1.0, 1.3) +trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.))) +for iter=1:100 + trace = split_merge_kernel(trace) + trace = fixed_structure_kernel(trace) +end +``` +We can then compare the results to the results from the Markov chain that used the selection-based structure-changing kernel: + +![rjmcmc plot](images/rjmcmc.png) + +We see that if we initialize the Markov chains from the same state with a single mean (`:z` is `false`) then the selection-based kernel fails to accept any moves to the two-mean structure within 100 iterations, whereas the split-merge kernel transitions back and forth many times, +If we repeated the selection-based kernel for enough iterations, it would eventually transition back and forth at the same rate as the split-merge. +The split-merge kernel gives a much more efficient inference algorithm for estimating the posterior probability on the two structures. + + +## Reverse Kernels +The **reversal** of a stationary MCMC kernel with distribution ``k_1(t'; t)``, for model with distribution ``p(t; x)``, is another MCMC kernel with distribution: +```math +k_2(t; t') := \frac{p(t; x)}{p(t'; x)} k_1(t'; t) +``` +For custom primitive kernels declared with [`@pkern`](@ref), users can declare the reversal kernel with the [`@rkern`](@ref) macro: +```julia +@rkern k1 : k2 +``` +This also assigns `k1` as the reversal of `k2`. +The composite kernel DSL automatically generates the reversal kernel for composite kernels, and built-in stationary kernels like [`mh`](@ref). +The reversal of a kernel (primitive or composite) can be obtained with [`reversal`](@ref). \ No newline at end of file diff --git a/docs/src/tutorials/trace_translators.md b/docs/src/tutorials/trace_translators.md new file mode 100644 index 000000000..8bf381559 --- /dev/null +++ b/docs/src/tutorials/trace_translators.md @@ -0,0 +1,435 @@ +# Trace Translators + +While [Generative Functions](@ref) define probability distributions on traces, **Trace Translators** convert from one space of traces to another space of traces. +Trace translators are building blocks of inference programs that utilize multiple model representations, like [Involutive MCMC](@ref). + +Trace translators are significantly more general than [Bijectors](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/Bijector). +Trace translators can (i) convert between spaces of traces that include mixed numeric discrete random choices, as well as stochastic control flow, and (ii) convert between spaces for which there is no one-to-one correspondence (e.g. between models of different dimensionality, or between discrete and continuous models). +Bijectors are limited to deterministic transformations between real-valued vectors of constant dimension. + +## Deterministic Trace Translators + +Inference programs manipulate traces, but they also keep track of probabilities and probability densities associated with these traces. +Suppose we have two generative functions `p1` and `p2`. +Given a trace `t2` of `p2` we can easily compute the probability (or probability density) that the trace would have been generated by `p2` using `get_score(t2)`. +But suppose we want to construct the trace of `p2` first sampling a trace `t1` of `p1` and then applying a deterministic transformation to that trace to obtain `t2`. +How can we compute the probability that a trace `t2` would have been produced by this process? +This probability is needed if, for example, `p2` defines a probabilistic model and want to use `p1` as a proposal distribution within importance sampling. +If we produce `t2` via an arbitrary deterministic transformation of the random choices in `t1`, then computing the necessary probability is difficult. + +If we restrict ourselves to deterministic transformations that are *bijections* (one-to-one correspondences) from the set of traces of `p1` to the set of traces of `p2`, then the problem is much simplified. +If the transformation is a bijection this means that (i) each trace of `p1` gets mapped to a different trace of `p2`, and (ii) for every trace of `p2` there is some trace of `p1` that maps to it. +Bijective transformations between traces are useful components of inference programs because the probability that a given trace `t2` of `p2` would have been produced by first sampling from `p1` and then applying the transform can be computed simply as the probability that `p1` would produce the (unique) trace `t1` that gets mapped to the given trace by the transform. +Conceptually, bijective trace transforms *preserve probability*. +When trace transforms operate on traces with continuous random choices, computing probability densities of the transformed traces requires computing a Jacobian associated with the continuous part of the transformation. + +Gen provides a DSL for expressing bijections between spaces of traces, called the **Trace Transform DSL**. +We introduce this DSL via an example. +Below are two generative functions. +The first samples polar coordinates and the second uses cartesian coordinates. +```julia +@gen function p1() + r ~ inv_gamma(1, 1) + theta ~ uniform(-pi/2, pi/2) +end +``` +```julia +@gen function p2() + x ~ normal(0, 1) + y ~ normal(0, 1) +end +``` + +### Defining a trace transform with the Trace Transform DSL + +The following trace transform DSL program defines a transformation (called `f`) that transforms traces of `p1` into traces of `p2`: +```julia +@transform f (t1) to (t2) begin + r = @read(t1[:r], :continuous) + theta = @read(t1[:theta], :continuous) + @write(t2[:x], r * cos(theta), :continuous) + @write(t2[:y], r * sin(theta), :continuous) +end +``` +This transform reads values of random choices in the input trace (`t1`) at specific addresses (indicated by the syntax `t1[addr]`) using `@read` and writes values to the output trace (`t2`) using `@write`. +Each read and write operation is labeled with whether the random choice is discrete or continuous. +The section [Trace Transform DSL](@ref) defines the DSL in more detail. + +It is usually a good idea to write the inverse of the bijection. +The inverse can provide a dynamic check that the transform truly is a bijection. +The inverse of the above transformation is: +```julia +@transform finv (t2) to (t1) begin + x = @read(t2[:x], :continuous) + y = @read(t2[:y], :continuous) + r = sqrt(x^2 + y^2) + @write(t1[:r], sqrt(x^2 + y^2), :continuous) + @write(t1[:theta], atan(y, x), :continuous) +end +``` +We can inform Gen that two transforms are inverses of one another using [`pair_bijections!`](@ref): +```julia +pair_bijections!(f, finv) +``` + +### Wrapping a trace transform in a trace translator + +Note that the transform DSL code does not specify what the two generative functions are, or what the arguments to these generative functions are. +This information will be required for computing probabilities and probability densities of traces. +We provide this information by constructing a **Trace Translator** that wraps the transform along with this transformation: +```julia +translator = DeterministicTraceTranslator(p2, (), choicemap(), f) +``` +We then can then apply the translator to a trace of `p1` using function call syntax. +The translator returns a trace of `p2` and a log-weight that we can use to compute the probability (density) of the resulting trace: +```julia +t2, log_weight = translator(t1) +``` +Specifically, the log probability (density) that the trace `t2` was produced by first sampling `t1 = simulate(p1, ())` and then applying the trace translator, is: +```julia +get_score(t1) + log_weight +``` +Let's unpack the previous few code blocks in more detail. +First, note that we did not pass in the source generative function (`p1`) or its arguments (`()`) when we constructed the trace translator. +This is because this information will be attached to the input trace `t1` itself. +We *did* need to pass in the target generative function (`p2`) and its arguments (`()`) however, because this information is not included in `t1`. + +In this case, because continuous random choices are involved, the probabilities are probability densities, and the trace translator used automatic differentiation of the code in the transform `f` to compute a change-of-variables Jacobian that is necessary to compute the correct probability density of the new trace `t2`. + +### Observations + +Typically, there are observations associated with one or both of the generative functions involved, and we have values for these in a choice map, so we do not want the trace translator to be responsible for populating the values of these observed random choices. +For example, suppose we want to condition `p2` on an observed random choice `z`: +```julia +@gen function p2() + x ~ normal(0, 1) + y ~ normal(0, 1) + z ~ normal(x + y, 0.1) +end +observations = choicemap() +observations[:z] = 2.3 +``` +When `p2` has observations, these can be passed in as an additional argument to the trace translator constructor: +```julia +translator = DeterministicTraceTranslator(p2, (), observations, f) +``` + +### Discrete random choices and stochastic control flow + +Trace transforms and trace translators interoperate seamlessly with discrete random choices and stochastic control flow. +Both the input and the output traces can contain a mix of discrete and continuous choices, and arbitrary stochastic control flow. +Consider the following simple example, where we use two different discrete representations to represent probability distributions the integers 0-7: +```julia +@gen function p1() + bit1 ~ bernoulli(0.5) + bit2 ~ bernoulli(0.5) + bit3 ~ bernoulli(0.5) +end +``` +```julia +@gen function p2() + n ~ categorical([0.1, 0.1, 0.1, 0.2, 0.2, 0.15, 0.15]) +end +``` +We define the forward and inverse transforms: +```julia +@transform f (t1) to (t2) begin + n = ( + @read(t1[:bit1], :discrete) * 1 + + @read(t1[:bit2], :discrete) * 2 + + @read(t1[:bit3], :discrete) * 4) + @write(t2[:n], n, :discrete) +end +``` +```julia +@transform finv (t2) to (t1) begin + bits = digits(@read(t2[:n], :discrete), base=2) + @write(t1[:bit1], bits[1], :discrete) + @write(t1[:bit2], bits[2], :discrete) + @write(t1[:bit3], bits[3], :discrete) +end +``` +Here is an example that includes discrete random choices, stochastic control flow, and continuous random choices. +```julia +@gen function p1() + if ({:branch} ~ bernoulli(0.5)) + x ~ normal(0, 1) + else + other ~ categorical([0.3, 0.7]) + end +end +``` +```julia +@gen function p2() + k ~ uniform_discrete(1, 4) + if k <= 2 + y ~ gamma(1, 1) + end +end +``` +Note that transformations between spaces of traces need not be intuitive (although they probably should)! +Try to convince yourself that the functions below are indeed a pair of bijections between the traces of these two generative functions. +```julia +@transform f (t1) to (t2) begin + if @read(t1[:branch], :discrete) + x = @read(t1[:x], :continuous) + if x > 0 + @write(t2[:k], 2, :discrete) + else + @write(t2[:k], 1, :discrete) + end + @write(t2[:y], abs(x), :continuous) + else + other = @read(t1[:other], :discrete) + @write(t2[:k], (other == 1) ? 3 : 4, :discrete) + end +end +``` +```julia +@transform finv (t2) to (t1) begin + k = @read(t2[:k], :discrete) + if k <= 2 + y = @read(t2[:y], :continuous) + @write(t2[:x], (k == 1) ? -y : y, :continuous) + else + @write(t1[:other], (k == 3) ? 1 : 2, :discrete) + end +end +``` + +## General Trace Translators + +Note that for two arbitrary generative functions `p1` and `p2` there may not exist any one-to-one correspondence between their spaces of traces. +For example, consider a generative function `p1` that samples points within the unit square ``[0, 1]^2`` +```julia +@gen function p1() + x ~ uniform(0, 1) + y ~ uniform(0, 1) +end +``` +and another generative function `p2` that samples one of 100 possible discrete values, each value representing one cell of the unit square: +```julia +@gen function p2() + i ~ uniform_discrete(1, 10) # interval [(i-1)/10, i/10] + j ~ uniform_discrete(1, 10) # interval [(j-1)/10, j/10] +end +``` +There is no one-to-one correspondence between the spaces of traces of these two generative functions: +The first is an uncountably infinite set, and the other is a finite set with 100 elements in it. + +However, there is an intuitive notion of correspondence that we would like to be able to encode. +Each discrete cell ``(i, j)`` corresponds to a subset of the unit square ``[(i - 1)/10, i/10] \times [(j-1)/10, j/10]``. + +We can express this correspondence (and any correspondence between two arbitrary generative functions) by introducing two auxiliary generative functions `q1` and `q2`. +The first function `q1` will take a trace of `p1` as input, and the second function `q2` will take a trace of `p2` as input. +Then, instead of a transfomation between traces of `p1` and traces of `p2` our trace transform will transform between (i) the space of pairs of traces of `p1` and `q1` and (ii) the space of pairs of traces of `p2` and `q2`. +We construct `q1` and `q2` so that the two spaces have the same size, and a one-to-one correspondence is possible. + +For our example above, we construct `q2` to sample the coordinate (``[0, 0.1]^2``) relative to the cell. +We construct `q1` to be empty--there is already a mapping from each trace of `p1` to each trace of `p2` that simply identifies what cell ``(i, j)`` a given point in ``[0, 1]^2`` is in, so no extra random choices are needed. +```julia +@gen function q1(p1_trace) +end + +@gen function q2(p2_trace) + dx ~ uniform(0.0, 0.1) + dy ~ uniform(0.0, 0.1) +end +``` +### Trace transforms between pairs of traces + +To handle general trace translators that require auxiliary probability distributions, the trace trace DSL supports defining transformations between *pairs* of traces. +For example, the following defines a trace transform that maps from pairs of traces of `p1` and `q1` to pairs of traces of `p2` and `q2`: + +```julia +@transform f (p1_trace, q1_trace) to (p2_trace, q2_trace) begin + x = @read(p1_trace[:x], :continuous) + y = @read(p1_trace[:y], :continuous) + i = ceil(x * 10) + j = ceil(y * 10) + @write(p2_trace[:i], i, :discrete) + @write(p2_trace[:j], j, :discrete) + @write(q2_trace[:dx], x - (i-1)/10, :continuous) + @write(q2_trace[:dy], y - (j-1)/10, :continuous) +end +``` +and the inverse transform: +```julia +@transform f_inv (p2_trace, q2_trace) to (p1_trace, q1_trace) begin + i = @read(p2_trace[:i], :discrete) + j = @read(p2_trace[:j], :discrete) + dx = @read(q2_trace[:dx], :continuous) + dy = @read(q2_trace[:dy], :continuous) + x = (i-1)/10 + dx + y = (j-1)/10 + dy + @write(p1_trace[:x], x, :continuous) + @write(p1_trace[:y], y, :continuous) +end +``` +which we associate as inverses: +```julia +pair_bijections!(f, f_inv) +``` + +### Constructing a general trace translator + +We now wrap the transform above into a general trace translator, by providing the three probabilistic programs `p2`, `q1`, `q2` that it uses (a reference to `p1` will included in the input trace), and the arguments to these functions. +```julia +translator = GeneralTraceTranslator( + p_new=p2, + p_new_args=(), + new_observations=choicemap(), + q_forward=q1, + q_forward_args=(), + q_backward=q2, + q_backward_args=(), + f=f) +``` +Then, we can apply the trace translator to a trace (`t1`) of `p1` and get a trace (`t2`) of `p2` and a log-weight: +```julia +t2, log_weight = translator(t1) +``` + + +## Symmetric Trace Translators + +When the previous and new generative functions (e.g. `p1` and `p2` in the previous example) are the same, and their arguments are the same, and `q_forward` and `q_backward` (and their arguments) are also identical, we call this the trace translator a **Symmetric Trace Translator**. +Symmetric trace translators are important because they form the basis of [Involutive MCMC](@ref). +Instead of translating a trace of one generative function to the trace of another generative function, they translate a trace of a generative function to another trace of the *same* generative function. + +Symmetric trace translators have the interesting property that the function `f` is an **involution**, or a function that is its own inverse. +To indicate that a trace transform is an involution, use [`is_involution!`](@ref). + +Because symmetric trace translators translate within the same generative function, their implementation uses [`update`](@ref) to incrementally modify the trace from the previous to the new trace. +This has two benefits when the previous and new traces have random choices that aren't modified between them: +(i) the incremental modification may be more efficient than writing the new trace entirely from scratch, and +(ii) the transform DSL program does not need to specify a value for addresses whose value is not changed from the previous trace. + +## Simple Extending Trace Translators + +Simple extending trace translators extend an existing trace with new random +choices sampled from a proposal distribution, as well as any new observations. +The arguments of the trace may also be updated. This type of trace translation +is the basic operation used in [Particle Filtering](@ref). For example, +we might have a model that sequentially samples new latent variables `(:z, t)` +and observations `(:x, t)` for each timestep `t`: + +```julia +@gen function model(T::Int) + for t in 1:T + z = {(:z, t)} ~ normal(0, 1) + x = {(:x, t)} ~ normal(z, 1) + end +end +``` + +Each time we observe a new `(:x ,t)`, we might want to propose `(:z, t)` so that +it is close in value: + +```julia +@gen function proposal(trace::Trace, x::Real) + t = get_args(trace)[1] + 1 + {(:z, t)} ~ normal(x, 1) +end +``` + +Suppose we initially generated a trace up to timestep `t=1`, e.g. by calling +`t1 = simulate(model, (1,))`. Now we observe `(:x, 2)` to be `5.0`. By +constructing a simple extending trace translator, we can simultaneously +update the trace `t1` with new arguments, introduce the new observation +at `(:x, 2)`, and propose a likely value for `(:z, 2)`: + +```julia +translator = SimpleExtendingTraceTranslator( + p_new_args=(2,), p_argdiffs=(UnknownChange(),), + new_observations=choicemap((:x, 2) => 5.0), + q_forward=proposal, q_forward_args=(5.0,)) +t2, log_weight = translator(t1) +``` + +Similar functionality can be achieved through a combination of [`propose`](@ref) +on the proposal and [`update`](@ref) on the original trace, but using a trace +translator provides a nice layer of abstraction. + +## Trace Transform DSL + +The **Trace Transform DSL** is a differentiable programming language for writing deterministic transformations of traces. +Programs written in this DSL are called *transforms*. +Transforms read the value of random choices from input trace(s) and write values of random choices to output trace(s). +These programs are not typically executed directly by users, but are instead wrapped into one of the several forms of trace translators listed above ([`GeneralTraceTranslator`](@ref), and [`SymmetricTraceTranslator`](@ref)). + +A transform is identified with the [`@transform`](@ref) macro, and uses one of the following four syntactic forms for the signature (the name of the transform, and the names of the input and output traces are all user-defined varibles; the only keywords are `@transform`, `to`, `begin`, and `end`): + +*A transform from one trace to another, without extra parameters* +```julia +@transform f t1 to t2 begin + ... +end +``` +*A transform from one trace to another, with extra parameters* +```julia +@transform f(x, y, ..) t1 to t2 begin + ... +end +``` +*A transform from pairs of traces to pairs of traces, without extra parameters* +```julia +@transform f (t1, t2) to (t3, t4) begin + ... +end +``` +*A transform from one trace to another, with extra parameters* +```julia +@transform f(x, y, ..) (t1, t2) to (t3, t4) begin + ... +end +``` +The extra parameters are optional, and can be used to pass arguments to a transform function that is invoked, from another transform function, using the [`@tcall`](@ref) macro. +For example: +```julia +@transform g(x) t1 to t2 begin + ... +end +@transform f t1 to t2 begin + x = .. + @tcall(g(x)) +end +``` +The callee transform function (`g` above) reads and writes to the same input and output traces as the caller transform function (`f` above). Note that the input and output traces can be assigned different names in the caller and the callee. + +The body of a transform reads the values of random choices at addresses in the input trace(s), performs computation using regular Julia code (provided this code can be differentiated with [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)) and writes valeus of random choices at addresses in the output trace(s). +In the body [`@read`](@ref) expressions read a value from a particular address of an input trace: +``` +val = @read(, ) +``` +where `` is an expression of the form `[]` where `` must be the name of an input trace in the transform's signature. +The `` is either `:continuous` or `:discrete`, and indicates whether the random choice is discrete or continuous (in measure-theoretic terms, whether it uses the counting measure, or a Lebesgue-measure a Euclidean space of some dimension). +Similarly, [`@write`](@ref) expressions write a value to a particular address in an output trace: +``` +@write(, , ) +``` +Sometimes trace transforms need to directly copy the value from one address in an input trace to one address in an output trace. +In these cases, it is recommended to use the explicit [`@copy`](@ref) expression: +``` +@copy(, ) +``` +where `` and `` are of the form `[]` as above. +Note you can also copy collections of multiple random choices under an address namespace in an input trace to an address namespace in an output trace. +For example, +``` +@copy(trace1[:foo], trace2[:bar]) +``` +will copy every random choice in `trace1` with an address of the form `:foo => ` to `trace2` at address `:bar => `. + +It is also possible to read the *return value* from an input trace using the following syntax, but this value must be discrete (in the local neighborhood of traces, the return value must be constant as a function of all continuous random choices in input traces): +``` +val = @read([], :discrete) +``` +This feature is useful when the generative function precomputes a quantity as part of its return value, and we would like to reuse this value, instead of having to recompute it as part of the transform. +The `discrete' requirement is needed because the transform DSL does not currently backpropagate through the return value (this feature could be added in the future). + +Tips for defining valid transforms: + +- If you find yourself copying the same continuous source address to multiple locations, it probably means your transform is not valid (the Jacobian matrix will have rows that are identical, and so the Jacobian determinant will be zero). + +- You can gain some confidence that your transform is valid by enabling dynamic checks (`check=true`) in the trace translator that uses it. diff --git a/docs/src/tutorials/vi.md b/docs/src/tutorials/vi.md new file mode 100644 index 000000000..00653b9d1 --- /dev/null +++ b/docs/src/tutorials/vi.md @@ -0,0 +1,20 @@ +# Variational Inference (@id vi) + +Variational inference involves optimizing the parameters of a variational family to maximize a lower bound on the marginal likelihood called the ELBO. +In Gen, variational families are represented as generative functions, and variational inference typically involves optimizing the trainable parameters of generative functions. + +Let's create a generative model representing a joint set of coins. + +Now let's fit a mean field approximation to the coins: + +```@example vi_tutorial +@gen function independent_coins() + +end +``` + +## Reparametrization trick + +To use the reparametrization trick to reduce the variance of gradient estimators, users currently need to write two versions of their variational family, one that is reparametrized and one that is not. +Gen does not currently include inference library support for this. +We plan add add automated support for reparametrization and other variance reduction techniques in the future. \ No newline at end of file From 172e9014cf4b5a81b9c6a67f088fcc2cd996957f Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Tue, 16 Apr 2024 10:43:32 -0700 Subject: [PATCH 03/19] swap in --- docs/README | 16 +++++++++++++ docs/pages.jl | 2 +- .../getting_started/hidden_markov_model.md | 23 +++++++++++++++++++ 3 files changed, 40 insertions(+), 1 deletion(-) create mode 100644 docs/README diff --git a/docs/README b/docs/README new file mode 100644 index 000000000..16e69157d --- /dev/null +++ b/docs/README @@ -0,0 +1,16 @@ +# Website Docs +- `pages.jl` to find skeleton of website. +- `make.jl` to build the website index. + +# How to add tutorials +Currently you must write the tutorial directly in the docs rather than a source file (e.g. Quarto). See `getting_started` or `tutorials` for examples. + +Code snippets must use the triple backtick with a label to run. The environment carries over so long as the labels match. Example: + +```@example tutorial_1 +x = rand() +``` + +```@example tutorial_1 +print(x) +``` \ No newline at end of file diff --git a/docs/pages.jl b/docs/pages.jl index 615c37d75..1c2460fd3 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -2,7 +2,7 @@ pages = [ "Home" => "index.md", "Getting Started" => [ "Example 1: Linear Regression" => "getting_started/linear_regression.md", - "Example 2: Hidden Markov Models" => "getting_started/linear_regression.md", + "Example 2: Hidden Markov Models" => "getting_started/hidden_markov_model.md", ], "Tutorials" => [ "Basics" => [ diff --git a/docs/src/getting_started/hidden_markov_model.md b/docs/src/getting_started/hidden_markov_model.md index e69de29bb..13a31504d 100644 --- a/docs/src/getting_started/hidden_markov_model.md +++ b/docs/src/getting_started/hidden_markov_model.md @@ -0,0 +1,23 @@ +# [Hidden Markov Models] (@id hmm) + +```@example hmm_tutorial +using Gen + +@gen function kernel(n, latent) + x ~ normal(latent, 1.0) + y ~ normal(x, 1.0) + return y +end + +hmm_model = Unfold(kernel) +``` + +```@example hmm_tutorial +tr = simulate(hmm_model, (2,0)) +get_choices(tr) +``` + +# Plot this trajectory multiple times... +```@example hmm_tutorial +traces = [simulate(hmm_mode, (0,0)) for i in range(1000)] +``` From 3f58b32c47e9901d16ac5d529480685ea6865db6 Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Wed, 17 Apr 2024 10:53:04 -0700 Subject: [PATCH 04/19] HMM example --- .../getting_started/hidden_markov_model.md | 107 +++++++++++++++++- 1 file changed, 103 insertions(+), 4 deletions(-) diff --git a/docs/src/getting_started/hidden_markov_model.md b/docs/src/getting_started/hidden_markov_model.md index 13a31504d..9088d6fb2 100644 --- a/docs/src/getting_started/hidden_markov_model.md +++ b/docs/src/getting_started/hidden_markov_model.md @@ -2,7 +2,13 @@ ```@example hmm_tutorial using Gen +using GLMakie +import GLMakie: scatter!, plot +using LinearAlgebra +GLMakie.activate!(inline=true) +``` +```@example hmm_tutorial @gen function kernel(n, latent) x ~ normal(latent, 1.0) y ~ normal(x, 1.0) @@ -12,12 +18,105 @@ end hmm_model = Unfold(kernel) ``` +# Plot this trajectory multiple times... ```@example hmm_tutorial -tr = simulate(hmm_model, (2,0)) -get_choices(tr) +traces = [simulate(hmm_mode, (0,0)) for i in range(1000)] ``` -# Plot this trajectory multiple times... ```@example hmm_tutorial -traces = [simulate(hmm_mode, (0,0)) for i in range(1000)] +struct EField + positions::Matrix{Float64} + charges::Vector{Float64} + k::Float64 + function EField(positions, charges, k=1.0) + size(positions, 1) != length(charges) && throw("Number of positions does not match number of charges.") + new(positions, charges, k) + end +end + +function force(field::EField, charge::Float64, particle::AbstractVector) + r_vec = field.positions .- particle' + r = sqrt.(sum(r_vec .^ 2; dims=2)).^3 + forces = field.charges * charge + forces = r_vec .* (forces ./ r ) + total = sum(forces; dims=1) + -vec(total) +end + +function (field::EField)(charge, position) + force(field, charge, position) +end +``` + +```@example hmm_tutorial +positions = [ + 1.0 0; + 0.0 1.0; + -1.0 0; + 0 -1.0 +] +charges = [ + 1.0, + 1.0, + 1.0, + 1.0 +] +field = EField(positions, charges) +forces = force(field, 1.0, [0.0, 1.000]) +field(1.0, [0,0]) ``` + +```@example hmm_tutorial +scatter!(ax, field::EField) = scatter!(ax, field.positions[:,1], field.positions[:,2], marker=:circle, markersize=20) + +function plot(ax,field::EField) + scatter!(ax, field) + if length(field.charges) == 1 # Set lengths manually + x, y = field.positions[1,:] + x_min, x_max = x-0.5, x+0.5 + y_min, y_max = y-0.5, y+0.5 + else + x_min, y_min = minimum(field.positions; dims=1) + x_max, y_max = maximum(field.positions; dims=1) + end + xs = range(x_min, x_max, 10) + ys = range(y_min, y_max, 10) + charges = [field(1.0, [x, y]) for x in xs, y in ys] + normalize!.(charges) + charges = stack(charges) ./ 10 + u = charges[1,:,:] + v = charges[2,:,:] + + arrows!(ax, xs, ys, u, v) +end +``` + +```@example hmm_tutorial +function trajectory(field, charge, position; N=1, dt=0.01, velocity = [0.00, 0.00]) + init = position + path = Matrix{Float64}(undef, N, 2) + path[1,:] = init + + for i in 2:N + a = field(charge, init) + # println("a ", a) + velocity = velocity + dt*a + # println("v ", v) + init = 0.5 * dt^2 * a + velocity * dt + init + path[i,:] = init + end + path +end +``` + +```@example hmm_tutorial +f = Figure(size=(500,500)) +ga = f[1,1] = GridLayout() +ax = Axis(ga[1,1]) +plot(ax, field) +# display(trajectory(field, 0.1, [0.0, 0.5],20)) +traj = trajectory(field, -1.0, [0.0, 0.0], N=300, velocity=[0.02, 0.01]) +# println(traj) +lines!(ax, traj, color=:red) +f +``` \ No newline at end of file From ad9760dde9a72125fb9c92b14b1635d5e5089954 Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Wed, 17 Apr 2024 11:37:03 -0700 Subject: [PATCH 05/19] Serialization as an extension --- Project.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Project.toml b/Project.toml index 72ff91c38..5dc5b533e 100644 --- a/Project.toml +++ b/Project.toml @@ -35,3 +35,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = ["Test"] + +[extensions] +GenSerialization = "Serialization" From 3728b917f5efadd0bfee05c9e0d5ea4518a5a72a Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Mon, 22 Apr 2024 09:04:04 -0500 Subject: [PATCH 06/19] Edit docs. Add attribution --- docs/{README => README.md} | 5 ++++- docs/make.jl | 2 -- docs/pages.jl | 10 +++++----- docs/src/assets/header.js | 2 ++ docs/src/misc/changelog.md | 0 docs/src/misc/contributing.md | 0 6 files changed, 11 insertions(+), 8 deletions(-) rename docs/{README => README.md} (84%) delete mode 100644 docs/src/misc/changelog.md delete mode 100644 docs/src/misc/contributing.md diff --git a/docs/README b/docs/README.md similarity index 84% rename from docs/README rename to docs/README.md index 16e69157d..7b83f87c4 100644 --- a/docs/README +++ b/docs/README.md @@ -2,7 +2,10 @@ - `pages.jl` to find skeleton of website. - `make.jl` to build the website index. -# How to add tutorials +# Build Docs Locally +Open the Julia REPL and include `make.jl`. + +# Add Tutorial Code Currently you must write the tutorial directly in the docs rather than a source file (e.g. Quarto). See `getting_started` or `tutorials` for examples. Code snippets must use the triple backtick with a label to run. The environment carries over so long as the labels match. Example: diff --git a/docs/make.jl b/docs/make.jl index 6807a0c4c..d2c6d6442 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -4,11 +4,9 @@ using Documenter, Gen include("pages.jl") makedocs( modules = [Gen], - # linkcheck = true, doctest = false, clean = true, warnonly = true, - # warnonly = [:missing_docs], format = Documenter.HTML(; assets = String["assets/header.js", "assets/header.css", "assets/theme.css"] ), diff --git a/docs/pages.jl b/docs/pages.jl index 1c2460fd3..99afcddf6 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -45,9 +45,9 @@ pages = [ "Optimizing Trainable Parameters" => "api/internals/parameter_optimization.md", "Modeling Language Implementation" => "api/internals/language_implementation.md" ], - "Miscellanea" => [ - "Changelog" => "misc/changelog.md", - "Contributing" => "misc/contributing.md", - "References" => "misc/references.md" - ] + # "Miscellanea" => [ + # "Changelog" => "misc/changelog.md", + # "Contributing" => "misc/contributing.md", + # "References" => "misc/references.md" + # ] ] diff --git a/docs/src/assets/header.js b/docs/src/assets/header.js index b25445a32..2f211b385 100644 --- a/docs/src/assets/header.js +++ b/docs/src/assets/header.js @@ -1,3 +1,5 @@ +// Source: +// https://github.com/ReactiveBayes/RxInfer.jl/blob/246b196a3ea29d0b5744ce241a923c7a3b30eaf4/docs/src/assets/header.js#L4 // We add a simple `onload` hook to inject the custom header for our `HTML`-generated pages window.onload = function() { diff --git a/docs/src/misc/changelog.md b/docs/src/misc/changelog.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/misc/contributing.md b/docs/src/misc/contributing.md deleted file mode 100644 index e69de29bb..000000000 From dad5ec4a4726d25d2eaa4a9e13386789cbec8940 Mon Sep 17 00:00:00 2001 From: Ian Limarta Date: Mon, 29 Apr 2024 11:04:00 -0500 Subject: [PATCH 07/19] API listings --- Project.toml | 5 +- docs/pages.jl | 36 ++++++------- docs/src/api/choice_maps.md | 5 ++ docs/src/api/distributions.md | 2 +- docs/src/api/gfi.md | 8 ++- .../api/internals/language_implementation.md | 4 +- docs/src/api/learning.md | 2 +- docs/src/api/modeling.md | 34 ++++++------- docs/src/api/parameter_optimization.md | 2 +- docs/src/api/trace_translators.md | 1 + .../src/getting_started/gaussian_processes.md | 0 .../getting_started/hidden_markov_model.md | 2 +- .../infinite_mixture_models.md | 0 docs/src/getting_started/linear_regression.md | 2 +- docs/src/getting_started/mcmc_cipher.md | 0 .../getting_started/variational_inference.md | 0 docs/src/misc/references.md | 0 docs/src/tutorials/combinators.md | 50 +++++-------------- docs/src/tutorials/gfi.md | 48 ++++++++---------- docs/src/tutorials/languages/extending.md | 2 +- docs/src/tutorials/mcmc.md | 4 +- docs/src/tutorials/modeling_in_gen.md | 2 +- docs/src/tutorials/trace_translators.md | 2 +- src/address.jl | 5 ++ src/choice_map.jl | 29 ++++++++--- src/diff.jl | 11 ++-- src/inference/trace_translators.jl | 9 +++- src/modeling_library/recurse/recurse.jl | 2 +- src/modeling_library/vector.jl | 1 + 29 files changed, 137 insertions(+), 131 deletions(-) delete mode 100644 docs/src/getting_started/gaussian_processes.md delete mode 100644 docs/src/getting_started/infinite_mixture_models.md delete mode 100644 docs/src/getting_started/mcmc_cipher.md delete mode 100644 docs/src/getting_started/variational_inference.md delete mode 100644 docs/src/misc/references.md diff --git a/Project.toml b/Project.toml index 5dc5b533e..8e428b1cc 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,4 @@ julia = "1.3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] - -[extensions] -GenSerialization = "Serialization" +test = ["Test"] \ No newline at end of file diff --git a/docs/pages.jl b/docs/pages.jl index 99afcddf6..e321b42b4 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -22,32 +22,32 @@ pages = [ "tutorials/languages/extending.md" ], ], - "Modeling Languages and APIs" => [ + "How-to Guides" => [ + + ], + "API Reference" => [ + "Modeling Language" => [ "Generative Functions" => "api/gfi.md", "Probability Distributions" => "api/distributions.md", - "Built-in Modeling Language" => "api/modeling.md", + "Built-in Modeling Languages" => "api/modeling.md", "Generative Function Combinators" => "api/combinators.md", "Choice Maps" => "api/choice_maps.md", "Selections" => "api/selections.md", "Optimizing Trainable Parameters" => "api/parameter_optimization.md", "Trace Translators" => "api/trace_translators.md", - ], - "Standard Inference Library" => [ - "Importance Sampling" => "api/importance.md", - "MAP Optimization" => "api/map.md", - "Markov chain Monte Carlo" => "api/mcmc.md", - "MAP Optimization" => "api/map.md", - "Particle Filtering" => "api/pf.md", - "Variational Inference" => "api/vi.md", - "Learning Generative Functions" => "api/learning.md" + ], + "Standard Inference Library" => [ + "Importance Sampling" => "api/importance.md", + "MAP Optimization" => "api/map.md", + "Markov chain Monte Carlo" => "api/mcmc.md", + "MAP Optimization" => "api/map.md", + "Particle Filtering" => "api/pf.md", + "Variational Inference" => "api/vi.md", + "Learning Generative Functions" => "api/learning.md" + ], ], - "Internals" => [ + "Explanation and Internals" => [ "Optimizing Trainable Parameters" => "api/internals/parameter_optimization.md", "Modeling Language Implementation" => "api/internals/language_implementation.md" - ], - # "Miscellanea" => [ - # "Changelog" => "misc/changelog.md", - # "Contributing" => "misc/contributing.md", - # "References" => "misc/references.md" - # ] + ] ] diff --git a/docs/src/api/choice_maps.md b/docs/src/api/choice_maps.md index c065b1b32..89b59e831 100644 --- a/docs/src/api/choice_maps.md +++ b/docs/src/api/choice_maps.md @@ -50,3 +50,8 @@ choicemap set_value! set_submap! ``` + +```@docs +DynamicChoiceMap +EmptyChoiceMap +``` diff --git a/docs/src/api/distributions.md b/docs/src/api/distributions.md index 478cdb1e4..4646f60f5 100644 --- a/docs/src/api/distributions.md +++ b/docs/src/api/distributions.md @@ -1,4 +1,4 @@ -# Probability Distributions +# [Probability Distributions](@id distributions) Gen provides a library of built-in probability distributions, and three ways of defining custom distributions, each of which are explained below: diff --git a/docs/src/api/gfi.md b/docs/src/api/gfi.md index 805efb97b..aa8db7cad 100644 --- a/docs/src/api/gfi.md +++ b/docs/src/api/gfi.md @@ -1,7 +1,8 @@ -## Generative Functions +## [Generative Functions](@id gfi_api) ```@docs GenerativeFunction +Trace ``` add example @@ -31,3 +32,8 @@ accumulate_param_gradients! choice_gradients get_params ``` + +```@docs +NoChange +UnknownChange +``` \ No newline at end of file diff --git a/docs/src/api/internals/language_implementation.md b/docs/src/api/internals/language_implementation.md index 03f45cc9e..8381f36fa 100644 --- a/docs/src/api/internals/language_implementation.md +++ b/docs/src/api/internals/language_implementation.md @@ -2,11 +2,11 @@ ## Parsing `@gen` functions -Gen's built-in modeling languages are designed to preserve Julia's syntax as far as possible, apart from the [Tilde syntax](@ref) for calling generative functions, and the restrictions imposed on the [Static Modeling Language](@ref). In order to preserve that syntax, including the use of non-Gen macros within `@gen` functions, we relegate as much of the parsing of `@gen` functions as possible to Julia's macro-expander and parser. +Gen's built-in modeling languages are designed to preserve Julia's syntax as far as possible, apart from the [Tilde syntax](@ref) for calling generative functions, and the restrictions imposed on the [Static Modeling Language](@ref sml). In order to preserve that syntax, including the use of non-Gen macros within `@gen` functions, we relegate as much of the parsing of `@gen` functions as possible to Julia's macro-expander and parser. In particular, we adopt an implementation strategy that enforces a separation between the surface syntax associated with Gen-specific macros (i.e., `@trace` and `@param`) and their corresponding implementations, which differ across the Dynamic Modeling Language (DML) and the Static Modeling Language (SML). We do this by introducing the custom expressions `Expr(:gentrace, call, addr)` and `Expr(:genparam, name, type)`, which serve as intermediate representations in the macro-expanded [abstract syntax tree](https://docs.julialang.org/en/v1/manual/metaprogramming/#Program-representation-1). -Each modeling language can then handle these custom expressions in their own manner, either by parsing them to nodes in the [Static Computation Graph](@ref) (for the SML), or by substituting them with their implementations (for the DML). This effectively allows the SML and DML to have separate implementations of `@trace` and `@param`. +Each modeling language can then handle these custom expressions in their own manner, either by parsing them to nodes in the [Static Computation Graphs](@ref) (for the SML), or by substituting them with their implementations (for the DML). This effectively allows the SML and DML to have separate implementations of `@trace` and `@param`. For clarity, below is a procedural description of how the `@gen` macro processes Julia function syntax: diff --git a/docs/src/api/learning.md b/docs/src/api/learning.md index 2192793fc..daa62eb85 100644 --- a/docs/src/api/learning.md +++ b/docs/src/api/learning.md @@ -209,7 +209,7 @@ Then, the traces of the model can be obtained by simulating from the variational Instead of fitting the variational approximation from scratch for each observation, it is possible to fit an *inference model* instead, that takes as input the observation, and generates a distribution on latent variables as output (as in the wake sleep algorithm). When we train the variational approximation by minimizing the evidence lower bound (ELBO) this is called amortized variational inference. Variational autencoders are an example. -It is possible to perform amortized variational inference using [`black_box_vi`](@ref) or [`black_box_vimco!`](@ref). +It is possible to perform amortized variational inference using [`black_box_vi!`](@ref) or [`black_box_vimco!`](@ref). ## References diff --git a/docs/src/api/modeling.md b/docs/src/api/modeling.md index a89367fd7..ae985e745 100644 --- a/docs/src/api/modeling.md +++ b/docs/src/api/modeling.md @@ -1,4 +1,4 @@ -# Built-in Modeling Language +# [The Dynamic Modeling Language](@id dynamic_modeling_language) Gen provides a built-in embedded modeling language for defining generative functions. The language uses a syntax that extends Julia's syntax for defining regular Julia functions, and is also referred to as the **Dynamic Modeling Language**. @@ -29,14 +29,14 @@ We can also trace its execution: ``` Optional arguments can be left out of the above operations, and default values will be filled in automatically: ```julia -julia> (trace, _) = generate(foo, (,)); +julia> (trace, _) = generate(foo, ()) julia> get_args(trace) (0.1,) ``` -See [Generative Functions](@ref) for the full set of operations supported by a generative function. +See [Generative Functions](@ref gfi_api) for the full set of operations supported by a generative function. Note that the built-in modeling language described in this section is only one of many ways of defining a generative function -- generative functions can also be constructed using other embedded languages, or by directly implementing the methods of the generative function interface. However, the built-in modeling language is intended to being flexible enough cover a wide range of use cases. -In the remainder of this section, we refer to generative functions defined using the built-in modeling language as `@gen` functions. Details about the implementation of `@gen` functions can be found in the [Modeling Language Implementation](@ref) section. +In the remainder of this section, we refer to generative functions defined using the built-in modeling language as `@gen` functions. Details about the implementation of `@gen` functions can be found in the [Modeling Language Implementation](@ref language-implementation) section. ## Annotations @@ -57,7 +57,7 @@ Each argument can have the following different syntactic forms: Currently, the possible argument annotations are: -- `grad` (see [Differentiable programming](@ref)). +- `grad` (see [Differentiable programming](@ref differentiable_modeling)). **Function annotations.** The `@gen` function itself can also be optionally associated with zero or more annotations, which are separate from the per-argument annotations. Function-level annotations use the following different syntactic forms: @@ -70,11 +70,11 @@ Function-level annotations use the following different syntactic forms: Currently the possible function annotations are: -- `grad` (see [Differentiable programming](@ref)). +- `grad` (see [Differentiable programming](@ref differentiable_modeling)). -- `static` (see [Static Modeling Language](@ref)). +- `static` (see [Static Modeling Language](@ref sml)). -- `nojuliacache` (see [Static Modeling Language](@ref)). +- `nojuliacache` (see [Static Modeling Language](@ref sml)). ## Making random choices @@ -82,7 +82,7 @@ Random choices are made by calling a probability distribution on some arguments: ```julia val::Bool = bernoulli(0.5) ``` -See [Probability Distributions](@ref) for the set of built-in probability distributions, and for information on implementing new probability distributions. +See [Probability Distributions](@ref distributions) for the set of built-in probability distributions, and for information on implementing new probability distributions. In the body of a `@gen` function, wrapping a call to a random choice with an `@trace` expression associates the random choice with an *address*, and evaluates to the value of the random choice. The syntax is: @@ -145,7 +145,7 @@ It is recommended to write disciplined generative functions when possible. **Untraced call**: If `foo` is a generative function, we can invoke `foo` from within the body of a `@gen` function using regular call syntax. -The random choices made within the call are not given addresses in our trace, and are therefore *untraced* random choices (see [Generative Function Interface](@ref) for details on untraced random choices). +The random choices made within the call are not given addresses in our trace, and are therefore *untraced* random choices (see [Generative Function Interface](@ref gfi) for details on untraced random choices). ```julia val = foo(0.5) ``` @@ -247,10 +247,10 @@ Note that `~` is also defined in `Base` as a unary operator that performs the bi Like regular Julia functions, `@gen` functions return either the expression used in a `return` keyword, or by evaluating the last expression in the function body. Note that the return value of a `@gen` function is different from a trace of `@gen` function, which contains the return value associated with an execution as well as the assignment to each random choice made during the execution. -See [Generative Function Interface](@ref) for more information about traces. +See [Generative Function Interface](@ref gfi) for more information about traces. -## Trainable parameters +## [Trainable Parameters](@id trainable_parameters_modeling) A `@gen` function may begin with an optional block of *trainable parameter declarations*. The block consists of a sequence of statements, beginning with `@param`, that declare the name and Julia type for each trainable parameter. @@ -281,7 +281,7 @@ zero_param_grad! Trainable parameters are designed to be trained using gradient-based methods. This is discussed in the next section. -## Differentiable programming +## [Differentiable Programming](@id differentiable_modeling) Given a trace of a `@gen` function, Gen supports automatic differentiation of the log probability (density) of all of the random choices made in the trace with respect to the following types of inputs: @@ -371,7 +371,7 @@ See [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl) for more details. When making a random choice, each argument is either a tracked value or not. If the argument is a tracked value, then the probability distribution must support differentiation of the log probability (density) with respect to that argument. Otherwise, an error is thrown. -The [`has_argument_grads`](@ref) function indicates which arguments support differentiation for a given distribution (see [Probability Distributions](@ref)). +The [`has_argument_grads`](@ref) function indicates which arguments support differentiation for a given distribution (see [Probability Distributions](@ref distributions)). If the gradient is required for the *value* of a random choice, the distribution must support differentiation of the log probability (density) with respect to the value. This is indicated by the [`has_output_grad`](@ref) function. @@ -381,7 +381,7 @@ It is an error if a tracked value is passed as an argument of a generative funct If a generative function `gen_fn` has `accepts_output_grad(gen_fn) = true`, then the return value of the generative function call will be tracked and will propagate further through the caller `@gen` function's computation. -## Static Modeling Language +## [Static Modeling Language](@id sml) The *static modeling language* is a restricted variant of the built-in modeling language. Models written in the static modeling language can result in better inference performance (more inference operations per second and less memory consumption), than the full built-in modeling language, especially for models used with iterative inference algorithms like Markov chain Monte Carlo. @@ -399,7 +399,7 @@ end ``` After running this code, `foo` is a Julia value whose type is a subtype of `StaticIRGenerativeFunction`, which is a subtype of [`GenerativeFunction`](@ref). -### Static computation graph +### Static Computation Graphs Using the `static` annotation instructs Gen to statically construct a directed acyclic graph for the computation represented by the body of the function. For the function `foo` above, the static graph looks like: ```@raw html @@ -431,7 +431,7 @@ First, the definition of a `(static)` generative function is always expected to Next, in order to be able to construct the static graph, Gen restricts the permitted syntax that can be used in functions annotated with `static`. In particular, each statement in the body must be one of the following: -- A `@param` statement specifying any [Trainable parameters](@ref), e.g.: +- A `@param` statement specifying any [trainable parameters](@ref trainable_parameters_modeling), e.g.: ```julia @param theta::Float64 diff --git a/docs/src/api/parameter_optimization.md b/docs/src/api/parameter_optimization.md index 60e05f41d..91c6410ea 100644 --- a/docs/src/api/parameter_optimization.md +++ b/docs/src/api/parameter_optimization.md @@ -1,4 +1,4 @@ -# Optimizing Trainable Parameters +# Trainable Parameters(@trainable_parameter_optimization) Trainable parameters of generative functions are initialized differently depending on the type of generative function. Trainable parameters of the built-in modeling language are initialized with [`init_param!`](@ref). diff --git a/docs/src/api/trace_translators.md b/docs/src/api/trace_translators.md index 30209823e..59e1b7563 100644 --- a/docs/src/api/trace_translators.md +++ b/docs/src/api/trace_translators.md @@ -5,6 +5,7 @@ @read @write @copy +@tcall pair_bijections! is_involution! inverse diff --git a/docs/src/getting_started/gaussian_processes.md b/docs/src/getting_started/gaussian_processes.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/getting_started/hidden_markov_model.md b/docs/src/getting_started/hidden_markov_model.md index 9088d6fb2..c28ad4c71 100644 --- a/docs/src/getting_started/hidden_markov_model.md +++ b/docs/src/getting_started/hidden_markov_model.md @@ -20,7 +20,7 @@ hmm_model = Unfold(kernel) # Plot this trajectory multiple times... ```@example hmm_tutorial -traces = [simulate(hmm_mode, (0,0)) for i in range(1000)] +traces = [simulate(hmm_model, (0,0)) for i in 1:1000] ``` ```@example hmm_tutorial diff --git a/docs/src/getting_started/infinite_mixture_models.md b/docs/src/getting_started/infinite_mixture_models.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/getting_started/linear_regression.md b/docs/src/getting_started/linear_regression.md index 5a936127c..73db75845 100644 --- a/docs/src/getting_started/linear_regression.md +++ b/docs/src/getting_started/linear_regression.md @@ -1,4 +1,4 @@ -# [Getting Started with Gen.jl: Linear Regression](@id gen_linear_regression) +# Getting Started with Gen.jl: Linear Regression Let's write a short Gen program that does Bayesian linear regression - that is, given a set of observation points in the xy-plane, we want to find a line that fits them "well". What does "well" mean in Bayesian linear regression? Well one way of interpreting this is by proposing a line that makes the data highly likely - as in a high probability of occuring. diff --git a/docs/src/getting_started/mcmc_cipher.md b/docs/src/getting_started/mcmc_cipher.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/getting_started/variational_inference.md b/docs/src/getting_started/variational_inference.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/misc/references.md b/docs/src/misc/references.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/tutorials/combinators.md b/docs/src/tutorials/combinators.md index 3f31498fe..77186f793 100644 --- a/docs/src/tutorials/combinators.md +++ b/docs/src/tutorials/combinators.md @@ -87,54 +87,28 @@ get_retval(trace) Consider the following constructions: -```@example switch_combinator +```@setup switch_combinator using Gen -@gen function bang((grad)(x::Float64), (grad)(y::Float64)) - std::Float64 = 3.0 - z = @trace(normal(x + y, std), :z) +``` + +```@example switch_combinator +@gen function line(x) + z ~ normal(3*x+1,1.0) return z end -@gen function fuzz((grad)(x::Float64), (grad)(y::Float64)) - std::Float64 = 3.0 - z = @trace(normal(x + 2 * y, std), :z) +@gen function outlier(x) + z ~ normal(3*x+1, 10.0) return z end -sc = Switch(bang, fuzz) -``` - -This creates a new generative function `sc`. We can then obtain the trace of `sc`: - -```@example switch_combinator -trace, _ = simulate(sc, (2, 5.0, 3.0)) -trace +switch_model = Switch(line, outlier) ``` -The resulting trace contains the subtrace from the branch with index `2` - in this case, a call to `fuzz`: +This creates a new generative function `switch_model` whose arguments take the form `(branch, args...)`. By default, +branch is an integer indicating which generative function to execute. For example, branch `2` corresponds to `outlier`: ```@example switch_combinator +trace = simulate(switch_model, (2, 5.0)) get_choices(trace) ``` -``` -│ -└── :z : 13.552870875213735 -``` - -## Recurse combinator - -TODO: document me - -```@raw html -
- schematic of recurse combinatokr -
-``` - -TODO: document me - -```@raw html -
- schematic of recurse combinatokr -
-``` \ No newline at end of file diff --git a/docs/src/tutorials/gfi.md b/docs/src/tutorials/gfi.md index 28a2d30f5..c5f1c27e9 100644 --- a/docs/src/tutorials/gfi.md +++ b/docs/src/tutorials/gfi.md @@ -8,8 +8,8 @@ Generative functions are used to represent a variety of different types of proba Generative functions are represented by the following abstact type: There are various kinds of generative functions, which are represented by concrete subtypes of [`GenerativeFunction`](@ref). -For example, the [Built-in Modeling Language](@ref) allows generative functions to be constructed using Julia function definition syntax: -```@example gfi_example +For example, the [Built-in Modeling Language](@ref dynamic_modeling_language) allows generative functions to be constructed using Julia function definition syntax: +```@example gfi_tutorial using Gen # hide @gen function foo(a, b=0) if @trace(bernoulli(0.5), :z) @@ -19,11 +19,11 @@ using Gen # hide end end ``` -Users can also extend Gen by implementing their own [Custom generative function types](@ref), which can be new modeling languages, or just specialized optimized implementations of a fragment of a specific model. +Users can also extend Gen by implementing their own [custom generative functions](@ref custom_generative_functions), which can be new modeling languages, or just specialized optimized implementations of a fragment of a specific model. Generative functions behave like Julia functions in some respects. For example, we can call a generative function `foo` on arguments and get an output value using regular Julia call syntax: -```@example gfi_example +```@example gfi_tutorial foo(2, 4) ``` @@ -50,7 +50,7 @@ The set of valid argument tuples to the function, denoted ``X``. A family of probability distributions ``p(t, r; x)`` on maps ``t`` from random choice addresses to their values, and non-addressable randomness ``r``, indexed by arguments ``x``, for all ``x \in X``. Note that the distribution must be normalized: ```math -\sum_{t, r} p(t, r; x) = 1 \;\; \mbox{for all} \;\; x \in X +\sum_{t, r} p(t, r; x) = 1 \;\; \text{for all} \;\; x \in X ``` This corresponds to a requirement that the function terminate with probabability 1 for all valid arguments. We use ``p(t; x)`` to denote the marginal distribution on the map ``t``: @@ -78,17 +78,17 @@ It is typically used by inference programs, for logging and for caching the resu A family of probability distributions ``q(t; x, u)`` on maps ``t`` from random choice addresses to their values, indexed by tuples ``(x, u)`` where ``u`` is a map from random choice addresses to values, and where ``x`` are the arguments to the function. It must satisfy the following conditions: ```math -\sum_{t} q(t; x, u) = 1 \;\; \mbox{for all} \;\; x \in X, u +\sum_{t} q(t; x, u) = 1 \;\; \text{for all} \;\; x \in X, u ``` ```math -p(t; x) > 0 \mbox{ if and only if } q(t; x, u) > 0 \mbox{ for all } u \mbox{ where } u \mbox{ and } t \mbox{ agree } +p(t; x) > 0 \text{ if and only if } q(t; x, u) > 0 \text{ for all } u \text{ where } u \text{ and } t \text{ agree } ``` ```math -q(t; x, u) > 0 \mbox{ implies that } u \mbox{ and } t \mbox{ agree }. +q(t; x, u) > 0 \text{ implies that } u \text{ and } t \text{ agree }. ``` There is also a family of probability distributions ``q(r; x, t)`` on non-addressable randomness, that satisfies: ```math -q(r; x, t) > 0 \mbox{ if and only if } p(r | t, x) > 0 +q(r; x, t) > 0 \text{ if and only if } p(r | t, x) > 0 ``` @@ -113,9 +113,7 @@ Traces contain: Different concrete types of generative functions use different data structures and different Julia types for their traces, but traces are subtypes of [`Trace`](@ref). -```@docs -Trace -``` + The concrete trace type that a generative function uses is the second type parameter of the [`GenerativeFunction`](@ref) abstract type. For example, the trace type of [`DynamicDSLFunction`](@ref) is `DynamicDSLTrace`. @@ -147,10 +145,9 @@ z = trace[:z] ``` When a generative function has default values specified for trailing arguments, those arguments can be left out when calling [`simulate`](@ref), [`generate`](@ref), and other functions provided by the generative function interface. The default values will automatically be filled in: -```julia -julia> trace = simulate(foo, (2,)); -julia> get_args(trace) -(2, 0) +```@example gfi_tutorial +trace = simulate(foo, (2,)); +get_args(trace) ``` ## Updating traces @@ -159,7 +156,7 @@ It is often important to incrementally modify the trace of a generative function In Gen, traces are **functional data structures**, meaning they can be treated as immutable values. There are several methods that take a trace of a generative function as input and return a new trace of the generative function based on adjustments to the execution history of the function. We will illustrate these methods using the following generative function: -```@example gfi_example +```@example gfi_tutorial @gen function bar() val = @trace(bernoulli(0.3), :a) if @trace(bernoulli(0.4), :b) @@ -172,7 +169,7 @@ We will illustrate these methods using the following generative function: end ``` Suppose we have a trace (`trace`) of `bar` with initial choices: -```@example gfi_example +```@example gfi_tutorial a = 1 ``` ``` @@ -234,7 +231,7 @@ Suppose we run [`update`](@ref) on the example `trace`, with the following const │ └── :b : false ``` -```@example gfi_example +```@example gfi_tutorial constraints = choicemap((:b, false)) (new_trace, w, _, discard) = update(trace, (), (), constraints) ``` @@ -282,7 +279,7 @@ The [`regenerate`](@ref) method takes a trace and generates an adjusted trace th **Example.** Suppose we run [`regenerate`](@ref) on the example `trace`, with selection `:a` and `:b`: -```@example gfi_example +```@example gfi_tutorial (new_trace, w, _) = regenerate(trace, (), (), select(:a, :b)) ``` Then, a new value for `:a` will be sampled from `bernoulli(0.3)`, and a new value for `:b` will be sampled from `bernoulli(0.4)`. @@ -291,7 +288,7 @@ If the new value for `:b` is `false`, then a new value for `:d` will be sampled The previous value for `:c` will always be retained. Suppose the new value for `:a` is `true`, and the new value for `:b` is `true`. Then `get_choices(new_trace)` will be: -```@example gfi_example +```@example gfi_tutorial get_choices(new_trace) ``` + + +## Speeding Up Inference Using the `Unfold` Combinator +For the particle filtering algorithms above, within an update step it is only +necessary to revisit the most recent state (or the most recent 5 states if +the rejuvenation moves are used) because the initial states are never +updated, and the contribution of these states to the weight computation +cancel. + +However, each update step of the particle filter inference programs above +scales *linearly* in the size of the trace because it visits every state when +computing the weight update. This is because the built-in modeling DSL by +default always performs an end-to-end execution of the generative function +body whenever performing a trace update. This allows the built-in modeling +DSL to be very flexible and to have a simple implementation, at the cost of +performance. There are several ways of improving performance after one has a +prototype written in the built-in modeling DSL. One of these is [Generative +Function Combinators](https://www.gen.dev/docs/dev/ref/combinators/), which make +the flow of information through the generative process more explicit to Gen, +and enable asymptotically more efficient inference programs. + +To exploit the opportunity for incremental computation, and improve the +scaling behavior of our particle filter inference programs, we will write a +new model using a generative function combinator to replaces the following +Julia `for` loop in our model. + +```julia + # generate successive states and measurements + for t=1:T + + # update the state of the point + vx = {(:vx, t)} ~ normal(vx, sqrt(velocity_var)) + vy = {(:vy, t)} ~ normal(vy, sqrt(velocity_var)) + x += vx + y += vy + + # bearing measurement + {(:z, t)} ~ normal(bearing(x, y), measurement_noise) + + # record position + xs[t+1] = x + ys[t+1] = y + end +``` + +This `for` loop has a very specific pattern of information flow—there is a +sequence of states (represented by `x`, `y`, `vx`, and `vy`), and each state is +generated from the previous state. This is exactly the pattern that the +[Unfold](https://www.gen.dev/docs/dev/ref/combinators/#Unfold-combinator-1) +generative function combinator is designed to handle. + +Below, we re-express the Julia `for` loop over the state sequence using the +Unfold combinator. Specifically, we define a generative function (`kernel`) +that takes the prevous state as its second argument, and returns the new +state. The Unfold combinator takes the kernel and returns a new generative +function (`chain`) that applies kernel repeatedly. Read the Unfold combinator +documentation for details on the behavior of the resulting generative +function (`chain`). + + +```@example particle_filter_tutorial +struct State + x::Float64 + y::Float64 + vx::Float64 + vy::Float64 +end + +@gen (static) function kernel(t::Int, prev_state::State, + velocity_var::Float64, measurement_noise::Float64) + vx ~ normal(prev_state.vx, sqrt(velocity_var)) + vy ~ normal(prev_state.vy, sqrt(velocity_var)) + x = prev_state.x + vx + y = prev_state.y + vy + z ~ normal(bearing(x, y), measurement_noise) + next_state = State(x, y, vx, vy) + return next_state +end + +chain = Gen.Unfold(kernel) + +Gen.@load_generated_functions # To allow use of the generative function written in the static modeling language above. +``` + +We can understand the behavior of `chain` by getting a trace of it and printing the random choices: + +trace = Gen.simulate(chain, (4, State(0., 0., 0., 0.), 0.01, 0.01)) +Gen.get_choices(trace) + +We now write a new version of the generative model that invokes `chain` instead of using the Julia `for` loop: + +```@example particle_fitler_tutorial +@gen (static) function unfold_model(T::Int) + + # parameters + measurement_noise = 0.005 + velocity_var = 1e-6 + + # initial conditions + x0 ~ normal(0.01, 0.01) + y0 ~ normal(0.95, 0.01) + vx0 ~ normal(0.002, 0.01) + vy0 ~ normal(-0.013, 0.01) + + # initial measurement + z0 ~ normal(bearing(x0, y0), measurement_noise) + + # record initial state + init_state = State(x0, y0, vx0, vy0) + + # run `chain` function under address namespace `:chain`, producing a vector of states + chain ~ chain(T, init_state, velocity_var, measurement_noise) + + result = (init_state, chain) + return result +end + +Gen.@load_generated_functions +``` + +Let's generate a trace of this new model program to understand its structure: + +```@example particle_filter_tutorial +(trace, _) = Gen.generate(unfold_model, (4,)) +Gen.get_choices(trace) +``` + +We can now run a particle filter on the Unfold model and see a speedup: + +```@example particle_filter_tutorial +function unfold_particle_filter(num_particles::Int, zs::Vector{Float64}, num_samples::Int) + init_obs = Gen.choicemap((:z0, zs[1])) + state = Gen.initialize_particle_filter(unfold_model, (0,), init_obs, num_particles) + + for t=1:length(zs)-1 + maybe_resample!(state, ess_threshold=num_particles/2) + obs = Gen.choicemap((:chain => t => :z, zs[t+1])) + Gen.particle_filter_step!(state, (t,), (UnknownChange(),), obs) + end + + # return a sample of traces from the weighted collection: + return Gen.sample_unweighted_traces(state, num_samples) +end + +@time unfold_pf_traces = unfold_particle_filter(5000, zs, 200); + +function unfold_render(trace; show_data=true, max_T=get_args(trace)[1], overlay=false) + (T,) = Gen.get_args(trace) + choices = Gen.get_choices(trace) + (init_state, states) = Gen.get_retval(trace) + xs = Vector{Float64}(undef, T+1) + ys = Vector{Float64}(undef, T+1) + zs = Vector{Float64}(undef, T+1) + xs[1] = init_state.x + ys[1] = init_state.y + zs[1] = choices[:z0] + for t=1:T + xs[t+1] = states[t].x + ys[t+1] = states[t].y + zs[t+1] = choices[:chain => t => :z] + end + f = overlay ? scatter! : scatter + fig = f(xs[1:max_T+1], ys[1:max_T+1], msize=3, msw=1, label=nothing) + if show_data + for z in zs[1:max_T+1] + dx = cos(z) * 0.5 + dy = sin(z) * 0.5 + plot!([0., dx], [0., dy], color="red", alpha=0.3, label=nothing) + end + end +end +``` + +Let's check that the results are reasonable: + +```@example particle_filter_tutorial +overlay(unfold_render, unfold_pf_traces, same_data=true) +``` + +We now empirically investigate the scaling behavior of +1. the inference program that uses the Julia `for` loop (`particle_filter`), +and +2. the equivalent inference program that uses `Unfold` +(`unfold_particle_filter`). + +We will use a synthetic long vector of z data, and we will investigate how the running time depends on the number of observations. + +```@example particle_filter_tutorial +fake_zs = rand(1000); + +function timing_experiment(num_observations_list::Vector{Int}, num_particles::Int, num_samples::Int) + times = Vector{Float64}() + times_unfold = Vector{Float64}() + for num_observations in num_observations_list + println("evaluating inference programs for num_observations: $num_observations") + tstart = time_ns() + traces = particle_filter(num_particles, fake_zs[1:num_observations], num_samples) + push!(times, (time_ns() - tstart) / 1e9) + + tstart = time_ns() + traces = unfold_particle_filter(num_particles, fake_zs[1:num_observations], num_samples) + push!(times_unfold, (time_ns() - tstart) / 1e9) + + end + (times, times_unfold) +end; + +num_observations_list = [1, 3, 10, 30, 50, 100, 150, 200, 500] +(times, times_unfold) = timing_experiment(num_observations_list, 100, 20); +``` + +Notice that the running time of the inference program without unfold appears to be quadratic in the number of observations, whereas the inference program that uses unfold appears to scale linearly: + +```@example particle_filter_tutorial +plot(num_observations_list, times, color="blue", + xlabel="# observations", ylabel="running time (sec.)", label="for loop") +plot!(num_observations_list, times_unfold, color="red", label="unfold") +``` \ No newline at end of file diff --git a/docs/src/tutorials/data_driven_proposals.md b/docs/src/tutorials/data_driven_proposals.md new file mode 100644 index 000000000..c346a83ea --- /dev/null +++ b/docs/src/tutorials/data_driven_proposals.md @@ -0,0 +1,1201 @@ +# [Improving Inference using Data-Driven Custom Proposals _(with applications to Inverse Planning)_](@id custom_proposal_tutorial) + +In our [introduction to modeling tutorial](@ref introduction_to_modeling_in_gen), we used **importance sampling** for inference: the inference engine _proposed_ many possible explanations for a dataset, and then chose one. Importance sampling can be difficult to scale to more complex problems, because it is essentially "guessing and checking." If we ran importance sampling with 1000 particles, for example, the method would fail unless those 1000 proposed solutions (blind guesses, essentially) contained _something_ close to the true answer. In complex problems, it is difficult +to "guess" (or "propose") an entire solution all at once. + +This tutorial addresses the scalability problem while staying inside +the importance sampling framework: the inference in this notebook is all importance +sampling, but with **customized "data-driven" proposals**. Such proposals +can be used to accelerate Monte Carlo inference, making importance sampling +feasible for a broader class of models. Data-driven proposals work by incorporating +information from the observed data set to make better proposals for the +latent variables in a generative model. Data-driven proposals can be based on heuristics, +or general-purpose _discriminative_ models, such as neural networks or random forests. +Many data-driven proposals have **trainable parameters**, which can be learned via gradient +descent using synthetic data simulated from the generative model itself. This training +process is sometimes called 'amortized inference' or 'inference compilation'. + +Although we focus on using data-driven proposals with importance sampling in this notebook, +data-driven proposals can also be used with Markov Chain Monte Carlo (MCMC) and +[sequential Monte Carlo](@ref object_tracking_with_particle_filter_tutorial) (SMC), covered in [other](Iterative%20Inference%20in%20Gen.jl). + +### Application to Inverse Planning +This notebook begins by introducing a probabilistic model for the motion of +an autonomous agent. The model itself demonstrates an important feature of +Gen: because it is embedded in Julia, we can use complex, black-box programs +as sub-routines in our models. The model we develop here uses an _RRT path planner_ +to model the goal-directed motion of the agent. + +After developing the model, we set out to improve the efficiency of inference. +We show that we can improve the efficiency of inference in this model using +a custom proposal for the destination of the agent. + +*Outline* + +**Section 1.** [A generative model of an autonomous agent](#model) + +**Section 2.** [Writing a data-driven proposal as a generative function](#custom-proposal) + +**Section 3.** [Using a data-driven proposal within importance sampling](#using) + +**Section 4.** [Training the parameters of a data-driven proposal](#training) + + +``` +using Gen, Distributions +``` + +## 1: A generative model of an autonomous agent + +We begin by writing a generative probabilistic model of the motion of an +intelligent agent that is navigating a two-dimensional scene. The model will +be *algorithmic* --- it will invoke a path planning algorithm implemented in +regular Julia code to generate the agent's motion plan from its destination +and its map of the scene. + +First, we load some basic geometric primitives for a two-dimensional scene. We +implemented these already in an auxiliary Julia file: + +``` +include("../inverse-planning/geometric_primitives.jl"); +``` + +The file we loaded gives us the `Point` data type (imported from [Luxor.jl](https://juliagraphics.github.io/Luxor.jl/)), which has fields `x` and `y`. + +``` +point = Point(1.0, 2.0) +println(point.x) +println(point.y) +``` + +The file also defines an `Obstacle` data type, which represents a polygonal +obstacle in a two-dimensional scene, that is constructed from a list of +vertices. Here, we construct a square: + +``` +obstacle = Obstacle([Point(0.0, 0.0), Point(0.5, 0.0), Point(0.5, 1.0), Point(0.0, 1.0)]) +``` + +Next, we load the definition of a `Scene` data type that represents a +two-dimensional scene. + +include("../inverse-planning/scene.jl"); + +The `Scene` data type represents a rectangle on the two-dimensional x-y plane +with a list of obstacles, which is initially empty. + +``` +scene = Scene(xmin=0, xmax=1, ymin=0, ymax=1) +``` + +Obstacles can be added to a scene with the `add_obstacle!` function: + +``` +add_obstacle!(scene, obstacle); +``` + +The file we imported also defines functions that simplify the construction of obstacles: + +`make_square(center::Point, size::Float64)` constructs a square-shaped +obstacle centered at the given point with the given side length: + +``` +square_obstacle = make_square(Point(.25, .75), 0.1) +``` + +`make_line(vertical::Bool, start::Point, length::Float64, +thickness::Float64)` constructs an axis-aligned line (either vertical or +horizontal) with given thickness that extends from a given strating point for +a certain length: + +``` +vertical_wall_obstacle = make_line(false, Point(0.20, 0.40), 0.40, 0.02) +``` + +We now construct a scene value that we will use in the rest of the notebook: + +scene = Scene(xmin=0, xmax=1, ymin=0, ymax=1) +add_obstacle!(scene, make_square(Point(0.30, 0.20), 0.1)) +add_obstacle!(scene, make_square(Point(0.83, 0.80), 0.1)) +add_obstacle!(scene, make_square(Point(0.80, 0.40), 0.1)) +horizontal = false +vertical = true +wall_thickness = 0.02 +add_obstacle!(scene, make_line(horizontal, Point(0.20, 0.40), 0.40, wall_thickness)) +add_obstacle!(scene, make_line(vertical, Point(0.60, 0.40), 0.40, wall_thickness)) +add_obstacle!(scene, make_line(horizontal, Point(0.60 - 0.15, 0.80), 0.15 + wall_thickness, wall_thickness)) +add_obstacle!(scene, make_line(horizontal, Point(0.20, 0.80), 0.15, wall_thickness)) +add_obstacle!(scene, make_line(vertical, Point(0.20, 0.40), 0.40, wall_thickness)); + +# We visualize the scene below. +include("../inverse-planning/viz.jl") +visualize() do + draw_scene(scene) +end + +Next, we load a file that defines a `Path` data type (a sequence of +`Point`s), and a `plan_path` method, which uses a path planning algorithm +based on rapidly exploring random tree (RRT, [1]) to find a sequence of +`Point`s beginning with `start` and ending in `dest` such that the line +segment between each consecutive pair of points does not intersect any +obstacles in the scene. The planning algorithm may fail to find a valid path, +in which case it will return a value of type `Nothing`. + +`path::Union{Path,Nothing} = plan_path(start::Point, dest::Point, +scene::Scene, planner_params::PlannerParams)` + +[1] [_Rapidly-exploring random trees: A new tool for path planning._](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.35.1853&rep=rep1&type=pdf) +# S. M. LaValle. TR 98-11, Computer Science Dept., Iowa State University, October 1998. + +include("../inverse-planning/planning.jl"); + +Let's use `plan_path` to plan a path from the lower-left corner of the scene +into the interior of the box. + +start = Point(0.1, 0.1) +dest = Point(0.5, 0.5) +planner_params = PlannerParams(rrt_iters=600, rrt_dt=3.0, + refine_iters=3500, refine_std=1.) +example_path = plan_path(start, dest, scene, planner_params) + +We visualize the path below with the function `visualize` (defined in the +external file we loaded), which will visualize the path in the scene. +The start location is shown as a blue circle, the +destination as a red rhombus, and the path in orange. Run the cell above followed by +the cell below a few times to see the variability in the paths generated by +`plan_path` for these inputs. + +visualize() do + draw_trace(Dict(:start => start, + :dest => dest, + :scene => scene, + :path => example_path.points); + should_draw_measurements=false) +end + +We also need a model for how the agent moves along its path. +We will assume that the agent moves along its path a constant speed. The file +loaded above also defines a method `walk_path(path, speed, dt, num_ticks)` +that computes the locations of the agent at a set of timepoints +(a vector of `Point`s sampled at time intervals of `dt` starting at time `0.`), +given the path and the speed of the agent: + +speed = 1. +dt = 0.1 +num_ticks = 10; +example_locations = walk_path(example_path, speed, dt, num_ticks) +println(example_locations) + +# Now, we are prepated to write our generative model for the motion of the agent. + +@gen function agent_model( + scene::Scene, dt::Float64, num_ticks::Int, + planner_params::PlannerParams) + + # sample the start point of the agent from the prior + start_x ~ uniform(0, 1) + start_y ~ uniform(0, 1) + start = Point(start_x, start_y) + + # sample the destination point of the agent from the prior + dest_x ~ uniform(0, 1) + dest_y ~ uniform(0, 1) + dest = Point(dest_x, dest_y) + + # plan a path that avoids obstacles in the scene + maybe_path = plan_path(start, dest, scene, planner_params) + planning_failed = maybe_path === nothing + + # sample the speed from the prior + speed ~ uniform(0.3, 1) + + if planning_failed + # path planning failed; assume agent stays at start location indefinitely + locations = fill(start, num_ticks) + else + # path planning succeeded; move along the path at constant speed + locations = walk_path(maybe_path, speed, dt, num_ticks) + end + + # generate noisy measurements of the agent's location at each time point + noise = 0.01 + for (i, point) in enumerate(locations) + x = {:meas => (i, :x)} ~ normal(point.x, noise) + y = {:meas => (i, :y)} ~ normal(point.y, noise) + end + + return (planning_failed, maybe_path) +end; + +We can now perform a traced execution of `agent_model` and examine the +random choices it made: + +### Exercise + +Using `simulate` (or `generate`, without any constraints) to sample a trace, +get the random choices made by this model. Parameterize the planner using +`PlannerParams` with the same parameters as above. + + +
+ +Next we explore the assumptions of the model by sampling many traces from the +generative function and visualizing them. We have created a visualization +specialized for this generative function in the file we included above. +It also defined a `visualize_grid` method to plot traces in a grid. + ++ +Let's visualize several traces of the function, with the start location fixed to +a point in the lower-left corner: +constraints = Gen.choicemap() +constraints[:start_x] = 0.1 +constraints[:start_y] = 0.1 + +traces = [Gen.generate( + agent_model, (scene, dt, num_ticks, planner_params), constraints + )[1] for i in 1:12]; +visualize_grid(traces, 4, 600; separators="gray") do trace, frame + draw_trace(trace, frame; draw_measurements=true, markersize=6) +end +- + +In this visualization, the measured coordinates at each +time point are represented by black $\times$ marks. The path, if path planning was +succesful, is shown as an orange line from the start point to the destination +point. Notice that the speed of the agent is different in each case. Also note that +the we observe the agent for a fixed amount of time, in which they may or may not +finish walking their planned path. + +
+ +### Exercise + +Edit the constraints passed to the inference algorithm: + +1. Constrain the start of the agent to be at $x = 0.9$, $y = 0.1$. +2. Constrain the destination of the agent to be at $x = 0.9$, $y = 0.8$. + + +Visualize the resulting prior. We have provided some starter code. + ++ +< put your code here> +traces_constrained = [] +for i in 1:12 + # Modify the following line: + (trace_constrained, _) = Gen.generate(agent_model, (scene, dt, num_ticks, planner_params)) + push!(traces_constrained, trace_constrained) +end + +# Visualize: +visualize_grid(traces_constrained, 4, 600; separators="gray") do trace, frame + draw_trace(trace, frame; draw_measurements=true) +end +- + + +
+ +### Exercise +The `rrt_iters` field of `PlannerParams` is the number of iterations of the RRT +algorithm to use. The `refine_iters` field of `PlannerParams` is the number of +iterations of path refinement. These parameters affect the distribution on +paths of the agent. Visualize traces of the `agent_model` with a couple of +different settings of these two parameters to the path planning algorithm for +fixed starting point and destination point. Try setting them to smaller +values. Discuss. + +We have provided starter code. + +constraints = Gen.choicemap() +constraints[:start_x] = 0.1 +constraints[:start_y] = 0.1; + +Modify the `PlannerParams` in the cell below. + +# + +planner_params = PlannerParams( + rrt_iters=300, rrt_dt=3.0, refine_iters=2000, refine_std=1.) # < change this line> + +traces = [Gen.generate(agent_model, (scene, dt, num_ticks, planner_params), constraints)[1] for i in 1:12]; +visualize_grid(traces, 4, 600; separators="gray") do trace, frame + draw_trace(trace, frame; draw_measurements=true) +end +- + +
+ +For the next few sections of the notebook, let's reset any variables that may have changed during your exploration with the model. + +start = Point(0.1, 0.1) +dt = 0.1 +num_ticks = 10 +planner_params = PlannerParams(rrt_iters=600, rrt_dt=3.0, + refine_iters=3500, refine_std=1.); + +We will infer the destination of the agent for the given sequence of observed locations: + +measurements = [ + Point(0.0980245, 0.104775), + Point(0.113734, 0.150773), + Point(0.100412, 0.195499), + Point(0.114794, 0.237386), + Point(0.0957668, 0.277711), + Point(0.140181, 0.31304), + Point(0.124384, 0.356242), + Point(0.122272, 0.414463), + Point(0.124597, 0.462056), + Point(0.126227, 0.498338)]; + +
+ +### Exercise + +Run inference using Gen's built-in importance resampling implementation. Use +50 importance samples (`amt_computation`). + +To see how to use the built-in importance resampling function, run +```?Gen.importance_resampling``` or check out the +[documentation](https://www.gen.dev/docs/dev/ref/importance/#Gen.importance_resampling). + +We have provided some starter code. + +function do_inference( + scene::Scene, dt::Float64, num_ticks::Int, + planner_params::PlannerParams, + start::Point, measurements::Vector{Point}, amount_of_computation::Int) + + # Constrain the observed measurements. + observations = Gen.choicemap() + observations[:start_x] = start.x + observations[:start_y] = start.y + for (i, m) in enumerate(measurements) + observations[:meas => (i, :x)] = m.x + observations[:meas => (i, :y)] = m.y + end + + # < put your code here> + + return trace +end; + + +#### Visualize your answer +Below, we run this algorithm 500 times, to generate 500 approximate samples +from the posterior distribution on the destination. The inferred destinations +should appear as red rhombuses on the map. The following function +visualizes the paths overlaid on the scene. + +function visualize_inference(measurements, scene, start; computation_amt=50, samples=1000) + visualize() do + for i in 1:samples + trace = do_inference(scene, dt, num_ticks, planner_params, start, measurements, computation_amt) + draw_trace(trace; draw_measurements=true, draw_path=false) + end + end +end; + + +And now we run it! Note that this might take a while. + +visualize_inference( + measurements, scene, start, computation_amt=100, samples=500) + +The algorithm has made reasonable inferences about where the agent was likely +trying to go. + +Note that the above illustration takes a while to produce. This is +because computing each endpoint requires sampling 100 times from the default proposal (which +runs the RRT planner). When our models contain more expensive components, like +the path-planner, the computational demands of inference increase accordingly. +This motivates us to find more efficient inference algorithms, that will +require fewer model evaluations to produce good results. + +
+ +### Exercise + +In this problem, you'll explore the effect of changing the _scene_ on the +inferences we make about the agent. Below, we've reproduced the code for +constructing the scene in which we performed inference above. Modify the scene +so that there is an opening into the "room" along the _bottom_ wall, in +addition to the already-existing door along top wall. Otherwise, the scene +should be identical to the one above. + +Rerun inference. The results should be qualitatively different from the +results generated above, even though the observed movements of the agent are +identical. **Write a one- or two-sentence description of how the results are +different, and why.** Please address: + +1. Why would a _human_ make different inferences about the agent's likely + destination in the two different scenes? +2. To what extent does the _model_ succeed in producing qualitiatively + different results in the two scenes? Why? (Concretely, why are certain + proposals more often rejected by importance sampling in the two-door scene + than in the one-door scene?) + +# + +scene_2doors = Scene(xmin=0, xmax=1, ymin=0, ymax=1) + +Add the three blocks. +``` +add_obstacle!(scene_2doors, make_square(Point(0.30, 0.20), 0.1)) +add_obstacle!(scene_2doors, make_square(Point(0.83, 0.80), 0.1)) +add_obstacle!(scene_2doors, make_square(Point(0.80, 0.40), 0.1)) +``` + +Add the walls. You will need to change this code. In particular, you will need to edit +one of these lines (the one that constructs the bottom wall of the room) and add one new line +(because in order to create the "door", you will actually need to represent the bottom wall +as two separate rectangular obstacles -- as the sample code already does for the top wall). +``` +horizontal = false +vertical = true +wall_thickness = 0.02 +add_obstacle!(scene_2doors, make_line(horizontal, Point(0.20, 0.40), 0.40, wall_thickness)) +add_obstacle!(scene_2doors, make_line(vertical, Point(0.60, 0.40), 0.40, wall_thickness)) +add_obstacle!(scene_2doors, make_line(horizontal, Point(0.60 - 0.15, 0.80), 0.15 + wall_thickness, wall_thickness)) +add_obstacle!(scene_2doors, make_line(horizontal, Point(0.20, 0.80), 0.15, wall_thickness)) +add_obstacle!(scene_2doors, make_line(vertical, Point(0.20, 0.40), 0.40, wall_thickness)); +# - +``` + +Perform and visualize inference: +``` +visualize_inference(measurements, scene_2doors, start, computation_amt=100, samples=500) +``` + +**Free response:** What changed about the inferences when you changed the scene, +and why? You might address: + +1. Why would a _human_ make different inferences about the agent's likely + destination in the two different scenes? +2. To what extent does the _model_ succeed in producing qualitiatively + different results in the two scenes? (Concretely, why are certain proposals + more often rejected by importance sampling in the two-door scene than in + the one-door scene?) + +
+ + +## 2. Writing a data-driven proposal as a generative function + +The inference algorithm above used a variant of +[`Gen.importance_resampling`](https://www.gen.dev/docs/stable/ref/importance/#Gen.importance_resampling) +that does not take a custom proposal distribution. It uses the default +proposal distribution associated with the generative model. For generative +functions defined using the built-in modeling DSL, the default proposal +distribution is based on *ancestral sampling*, which involves sampling +unconstrained random choices from the distributions specified in the +generative model. Put more simply, each "guess" the inference algorithm +makes about the possible destination of the agent is totally uninformed +by the observed measurements; it is sampled using the prior generative + model's `dest_x` and `dest_y` sampling statements. + +We can visualize this default proposal distribution by sampling from it, + using `Gen.generate` (note, we also could use `Gen.simulate` for the same purpose, since we are not passing any constraints). The cell below shows samples of the agent's destination + drawn from this distribution. + ++ +include("../inverse-planning/viz.jl"); + +traces = [Gen.generate(agent_model, (scene, dt, num_ticks, planner_params))[1] for i in 1:1000] +visualize() do + + for i in 1:1000 + trace, = Gen.generate(agent_model, (scene, dt, num_ticks, planner_params)) + draw_dest(scene, Point(trace[:dest_x], trace[:dest_y])) + end + + draw_scene(scene) + draw_start(scene, start) + draw_measurements(scene, measurements) +end +- + +Intuitively, if we see the data set above (where blue is the starting +location, and the measurements are in black), we might guess that the +agent is more likely to be heading into the upper part of the scene. This +is because we don't expect the agent to unecessarily backtrack on its route +to its destnation. A simple heuristic for biasing the proposal distribution +of the destination using just the first measurement and the last measurement might be: + +- If the x-coordinate of the last measurement is greater than the + x-coordinate of the first measurement, we think the agent is probably + headed generally to the right. Make values `:dest_x` that are greater than + the x-coordinate of the last measurement more probable. + +- If the x-coordinate of the last measurment is less than the x-coordinate of + the first measurement, we think the agent is probably headed generally to + the left. Make values `:dest_x` that are smaller than the x-coordinate of + the last measurement more probable. + +We can apply the same heuristic separately for the y-coordinate. + +To implement this idea, we discretize the x-axis and y-axis of the scene into +bins: + +num_x_bins = 5 +num_y_bins = 5; + +We will propose the x-coordinate of the destination from a +[piecewise_uniform](https://www.gen.dev/docs/dev/ref/distributions/#Gen.piecewise_uniform) +distribution, where we set higher probability for certain bins based on the +heuristic described above and use a uniform continuous distribution for the +coordinate within a bin. The `compute_bin_probs` function below computes the +probability for each bin. The bounds of the scene are given by `min` and +`max`. The coordinates of the first and last measured points respectively are +given by `first` and `last`. We compute the probability by assigning a +"score" to each bin based on the heuristic above --- if the bin should +receive lower probability, it gets a score of 1., and if it should receive +higher probability, it gets a bin of `score_high`, where `score_high` is some +value greater than 1. + ++ +function compute_bin_prob(first::Float64, last::Float64, bin::Int, last_bin::Int, score_high) + last >= first && bin >= last_bin && return score_high + last < first && bin <= last_bin && return score_high + return 1. +end + +function compute_bin_probs(num_bins::Int, min::Float64, max::Float64, first::Float64, last::Float64, score_high) + bin_len = (max - min) / num_bins + last_bin = Int(floor(last / bin_len)) + 1 + probs = [compute_bin_prob(first, last, bin, last_bin, score_high) for bin=1:num_bins] + total = sum(probs) + return [p / total for p in probs] +end; +- + +We will see how to automatically tune the value of `score_high` shortly. For +now, we will use a value of 5. Below, we see that for the data set of +measurements, shown above the probabilities of higher bins are indeed 5x +larger than those of lower bins, becuase the agent seems to be headed up. + +compute_bin_probs(num_y_bins, scene.ymin, scene.ymax, measurements[1].y, measurements[end].y, 5.) + +Next, we write a generative function that applies this heuristic for both the +x-coordinate and y-coordinate, and samples the destination coordinates at +addresses `:dest_x` and `:dest_y`. + +@gen function custom_dest_proposal(measurements::Vector{Point}, scene::Scene) + + score_high = 5. + + x_first = measurements[1].x + x_last = measurements[end].x + y_first = measurements[1].y + y_last = measurements[end].y + + # sample dest_x + x_probs = compute_bin_probs(num_x_bins, scene.xmin, scene.xmax, x_first, x_last, score_high) + x_bounds = collect(range(scene.xmin, stop=scene.xmax, length=num_x_bins+1)) + dest_x ~ piecewise_uniform(x_bounds, x_probs) + + # sample dest_y + y_probs = compute_bin_probs(num_y_bins, scene.ymin, scene.ymax, y_first, y_last, score_high) + y_bounds = collect(range(scene.ymin, stop=scene.ymax, length=num_y_bins+1)) + dest_y ~ piecewise_uniform(y_bounds, y_probs) + + return nothing +end; + +We can propose values of random choices from the proposal function using +[`Gen.propose`](https://www.gen.dev/docs/stable/ref/gfi/#Gen.propose). +This method returns the choices, as well as some other information, which we +won't need for our purposes. For now, you can think of `Gen.propose` as +similar to `Gen.generate` except that it does not produce a full execution +trace (only the choices), and it does not accept constraints. We can see the +random choices for one run of the proposal on our data set: + +(proposed_choices, _, _) = Gen.propose(custom_dest_proposal, (measurements, scene)) +proposed_choices + +The function below runs the proposal 1000 times. For each run, it generates a +trace of the model where the `:dest_x` and `:dest_y` choices are constrained +to the proposed values, and then visualizes the resulting traces. We make the +proposal a parameter of the function because we will be visualizing the +output distribution of various proposals later in the notebook. + +function visualize_custom_destination_proposal(measurements, start, dest_proposal; num_samples=100) + visualize() do + for i=1:num_samples + (proposed_choices, _) = Gen.propose(dest_proposal, (measurements, scene)) + constraints = choicemap(proposed_choices) + constraints[:start_x] = start.x + constraints[:start_y] = start.y + (trace, _) = Gen.generate(agent_model, (scene, dt, num_ticks, planner_params), constraints) + draw_dest(scene, Point(trace[:dest_x], trace[:dest_y])) + end + draw_scene(scene) + draw_start(scene, start) + draw_measurements(scene, measurements) + end +end; + +Let's visualize the output distribution of `custom_dest_proposal` for our +data set: + +visualize_custom_destination_proposal(measurements, start, custom_dest_proposal, num_samples=1000) + +We see that the proposal distribution indeed samples destinations in the top +half of the scene with higher probability than destinations in the bottom +half. + +Alone, this is just a heuristic. But we can use it as a proposal for importance sampling, turning it into an asymptotically valid Bayesian inference algorithm. Alternatively, we can view it as a tool for speeding up our naive importance sampler, by focusing computation on regions of the space that are more likely. + +## 3. Using a data-driven proposal within importance sampling + +We now use our data-driven proposal within an inference algorithm. There is a +second variant of +[`Gen.importance_resampling`](https://www.gen.dev/docs/stable/ref/importance/#Gen.importance_resampling) +that accepts a generative function representing a custom proposal. This +proposal generative function makes traced random choices at the addresses of +a subset of the unobserved random choices made by the generative model. In +our case, these addresses are `:dest_x` and `:dest_y`. + +
+ +### Exercise + +Implement an inference program that uses this second variant of importance resampling. + + +Because we will experiment with different data-driven proposals, we make the +proposal into an agument of our inference program. We assume that the +proposal accepts arguments `(measurements, scene)`. + +This time, use only 5 importance samples (`amt_computation`). You can run +`?Gen.importance_resampling` or check out the +[documentation](https://www.gen.dev/docs/stable/ref/inference/#Importance-Sampling-1) +to understand how to supply the arguments to invoke this second version of of +importance resampling. + +We have provided some starter code. + +# + +function do_inference_data_driven( + dest_proposal::GenerativeFunction, + scene::Scene, dt::Float64, + num_ticks::Int, planner_params::PlannerParams, + start::Point, measurements::Vector{Point}, + amount_of_computation::Int) + + observations = Gen.choicemap((:start_x, start.x), (:start_y, start.y)) + for (i, m) in enumerate(measurements) + observations[:meas => (i, :x)] = m.x + observations[:meas => (i, :y)] = m.y + end + + # trace = ... < put your code here> + + return trace +end; + +function visualize_data_driven_inference(measurements, scene, start, proposal; amt_computation=50, samples=1000) + visualize() do + for i=1:samples + trace = do_inference_data_driven(proposal, + scene, dt, num_ticks, planner_params, start, measurements, amt_computation) + draw_dest(scene, Point(trace[:dest_x], trace[:dest_y])) + end + draw_scene(scene) + draw_start(scene, start) + draw_measurements(scene, measurements) + end +end; +# - + +visualize_data_driven_inference(measurements, scene, start, custom_dest_proposal; amt_computation=5, samples=1000) + +The code executes much more quickly than before, because we are only taking five proposal samples to generate each. + +We compare this to the original algorithm that used the default proposal, for +the same "amount of computation" of 5. + +visualize_inference(measurements, scene, start, computation_amt=5, samples=1000) + +We should see that the results are somewhat more accurate using the +data-driven proposal. In particular, there is less probability mass in the +lower left corner when using the data-driven proposal. + +
+ + + +## 4. Training the parameters of a data-driven proposal + +Our choice of the `score_high` value of 5. was somewhat arbitrary. To use +more informed value, we can make `score_high` into a [*trainable +parameter*](https://www.gen.dev/docs/dev/ref/gfi/#Trainable-parameters-1) +of the generative function. Below, we write a new version of the proposal +function that makes `score_high` trainable. However, the optimization +algorithms we will use for training work best with *unconstrained* parameters +(parameters that can take any value on the real line), but `score_high` must +be positive. Therefore, we introduce an unconstrained trainable parameter +mamed `log_score_high`, and use `exp()` to ensure that `score_high` is +positive: + +@gen function custom_dest_proposal_trainable(measurements::Vector{Point}, scene::Scene) + + @param log_score_high::Float64 + + x_first = measurements[1].x + x_last = measurements[end].x + y_first = measurements[1].y + y_last = measurements[end].y + + # sample dest_x + x_probs = compute_bin_probs(num_x_bins, scene.xmin, scene.xmax, x_first, x_last, exp(log_score_high)) + x_bounds = collect(range(scene.xmin, stop=scene.xmax, length=num_x_bins+1)) + dest_x ~ piecewise_uniform(x_bounds, x_probs) + + # sample dest_y + y_probs = compute_bin_probs(num_y_bins, scene.ymin, scene.ymax, y_first, y_last, exp(log_score_high)) + y_bounds = collect(range(scene.ymin, stop=scene.ymax, length=num_y_bins+1)) + dest_y ~ piecewise_uniform(y_bounds, y_probs) + + return nothing +end; + +We initialize the value of `score_high` to 1. For this value, our custom +proposal gives a uniform distribution, and is the same as the default +proposal. + +Gen.init_param!(custom_dest_proposal_trainable, :log_score_high, 0.); + +Let's visualize the proposed distribution prior to training to confirm that +it is a uniform distribution. + +visualize_custom_destination_proposal(measurements, start, custom_dest_proposal_trainable, num_samples=1000) + +Now, we train the generative function. First, we will require a +data-generator that generates the training data. The data-generator is a +function of no arguments that returns a tuple of the form `(inputs, constraints)`. +The `inputs` are the arguments to the generative function +being trained, and the `constraints` contains the desired values of random +choices made by the function for those arguments. For the training +distribution, we will use the distribution induced by the generative model +(`agent_model`), restricted to cases where planning actually succeeded. When +planning failed, the agent just stays at the same location for all time, and +we won't worry about tuning our proposal for that case. The training +procedure will attempt to maximize the expected conditional log probablity +(density) that the proposal function generates the constrained values, +when run on the arguments. +Note that this is an *average case* objective function --- the resulting proposal +distribution may perform better on some data sets than others. + +function data_generator() + + # since these names are used in the global scope, explicitly declare it + # local to avoid overwriting the global variable + local measurements + local choices + + # obtain an execution of the model where planning succeeded + done = false + while !done + (choices, _, retval) = Gen.propose(agent_model, (scene, dt, num_ticks, planner_params)) + (planning_failed, maybe_path) = retval + done = !planning_failed + end + + # construct arguments to the proposal function being trained + measurements = [Point(choices[:meas => (i, :x)], choices[:meas => (i, :y)]) for i=1:num_ticks] + inputs = (measurements, scene) + + # construct constraints for the proposal function being trained + constraints = Gen.choicemap() + constraints[:dest_x] = choices[:dest_x] + constraints[:dest_y] = choices[:dest_y] + + return (inputs, constraints) +end; + +Next, we choose type of optimization algorithm we will use for training. Gen +supports a set of gradient-based optimization algorithms (see [Optimizing +Trainable +Parameters](https://www.gen.dev/docs/dev/ref/parameter_optimization/#Optimizing-Trainable-Parameters-1)). +Here we will use gradient descent with a fixed step size of 0.001. + +update = Gen.ParamUpdate(Gen.FixedStepGradientDescent(0.001), custom_dest_proposal_trainable); + +Finally, we use the +[`Gen.train!`](https://www.gen.dev/docs/stable/ref/inference/#Gen.train!) +method to actually do the training. + +For each epoch, `Gen.train!` makes `epoch_size` calls to the data-generator +to construct a batch of training data for that epoch. Then, it iteratively +selects `num_minibatch` subsets of the epoch training data, each of size +`100`, and applies the update once per minibatch. At the end of the epoch, it +generates another batch of evaluation data (of size `evaluation_size`) which +it uses to estimate the objective function (the expected conditional log +likelihood under the data-generating distribution). + +Here, we are running 200 gradient-descent updates, where each update is using +a gradient estimate obtained from 100 training examples. The method prints +the estimate of the objective function after each epoch. + +@time scores = Gen.train!(custom_dest_proposal_trainable, data_generator, update, + num_epoch=200, epoch_size=100, num_minibatch=1, minibatch_size=100, evaluation_size=100, verbose=true); + +using Plots +plot(scores, + xlabel="Iterations of stochastic gradient descent", + ylabel="Estimate of expected\nconditional log probability density", + label=nothing) + +We can read out the new value for `score_high`: + +println(exp(Gen.get_param(custom_dest_proposal_trainable, :log_score_high))) + +We see that the optimal value of the parameter is indeed larger than our +initial guess. This validates that the heuristic is indeed a useful one. We +visualize the proposal distribution below: + +visualize_custom_destination_proposal( + measurements, start, custom_dest_proposal_trainable, + num_samples=1000) + +We can visualize the results of inference, using this newly trained proposal: + +visualize_data_driven_inference( + measurements, scene, start, custom_dest_proposal_trainable, + amt_computation=5, samples=1000) + +visualize_data_driven_inference( + measurements, scene, start, custom_dest_proposal_trainable, + amt_computation=10, samples=1000) + +------------ + +### Exercise + +Can you devise a data-driven proposal for the speed of the agent? If so, would you +expect it to work equally well on all data sets? + +
+ +## 5. Writing and training a deep learning based data-driven proposal + +The heuristic data-driven proposal above gave some improvement in efficiency, +but it was very simple. One way of constructing complex data-driven proposals +is to parametrize the proposal with a deep neural network or use another +class of high-capacity machine learning model (e.g. random forest). Here, we +will will write a data-driven proposal for the destination of the agent that +uses deep neural networks. In this section, we do everything manually, without +the aid of neural network libraries. We also provide an [extension to the tutorial](#) that +shows how to use PyTorch to make this process a lot easier. + +First, we define a sigmoid function for the nonlinearity in our networks. + +nonlinearity(x) = 1.7159 * tanh.(x * 0.66666); + +We will use a deep neural network with two hidden layers that takes as input +x- and y- coordinates of the first and last measurement (4 values) and +produces as output a vector of un-normalized probabilities, one for each bin +of the x-dimension. We will later sample `:dest_x` from this distribution. + +function dest_x_neural_net(nn_params, x_first::Real, y_first::Real, x_last::Real, y_last::Real) + (W1, b1, W2, b2, W3, b3) = nn_params + input_layer = [x_first, y_first, x_last, y_last] + hidden_layer_1 = nonlinearity(W1 * input_layer .+ b1) + hidden_layer_2 = nonlinearity(W2 * hidden_layer_1 .+ b2) + output_layer = exp.(W3 * hidden_layer_2 .+ b3) + return output_layer +end; + +After sampling the value of `:dest_x`, we will use a second deep neural +network with the same structure to sample `:dest_y`. + +function dest_y_neural_net(nn_params, x_first::Real, y_first::Real, x_last::Real, y_last::Real)#, dest_x::Real) + (W1, b1, W2, b2, W3, b3) = nn_params + input_layer = [x_first, y_first, x_last, y_last] + hidden_layer_1 = nonlinearity(W1 * input_layer .+ b1) + hidden_layer_2 = nonlinearity(W2 * hidden_layer_1 .+ b2) + output_layer = exp.(W3 * hidden_layer_2 .+ b3) + return output_layer +end; + +Now that we have defined our neural networks, we define our new proposal. +This generative function has a number of parameters. + ++ +scale_coord(coord, min, max) = (coord / (max - min)) - 0.5 + +@gen function custom_dest_proposal_neural(measurements::Vector{Point}, scene::Scene) + @param x_W1::Matrix{Float64} + @param x_b1::Vector{Float64} + @param x_W2::Matrix{Float64} + @param x_b2::Vector{Float64} + @param x_W3::Matrix{Float64} + @param x_b3::Vector{Float64} + + @param y_W1::Matrix{Float64} + @param y_b1::Vector{Float64} + @param y_W2::Matrix{Float64} + @param y_b2::Vector{Float64} + @param y_W3::Matrix{Float64} + @param y_b3::Vector{Float64} + + num_x_bins = length(x_b3) + num_y_bins = length(y_b3) + + # scale inputs to be in the range [-0.5, 0.5] + x_first = scale_coord(measurements[1].x, scene.xmin, scene.xmax) + x_last = scale_coord(measurements[end].x, scene.xmin, scene.xmax) + y_first = scale_coord(measurements[1].y, scene.ymin, scene.ymax) + y_last = scale_coord(measurements[end].y, scene.ymin, scene.ymax) + + # sample dest_x + x_bounds = collect(range(scene.xmin, stop=scene.xmax, length=num_x_bins+1)) + x_probs = dest_x_neural_net((x_W1, x_b1, x_W2, x_b2, x_W3, x_b3), x_first, y_first, x_last, y_last) + dest_x ~ piecewise_uniform(x_bounds, x_probs / sum(x_probs)) + + # sample dest_y + y_bounds = collect(range(scene.xmin, stop=scene.xmax, length=num_y_bins+1)) + y_probs = dest_y_neural_net((y_W1, y_b1, y_W2, y_b2, y_W3, y_b3), x_first, y_first, x_last, y_last) + dest_y ~ Gen.piecewise_uniform(y_bounds, y_probs / sum(y_probs)) + + return nothing +end; +- + +We will use 50 hidden units in each of the layers of the two networks: + +num_hidden_1 = 50 +num_hidden_2 = 50; + +Next, we initialize the parameters: + ++ +import Random +Random.seed!(3) + +init_weight(shape...) = (1. / sqrt(shape[2])) * randn(shape...) + +init_x_W1 = init_weight(num_hidden_1, 4) +init_x_W2 = init_weight(num_hidden_2, num_hidden_1) +init_x_W3 = init_weight(num_x_bins, num_hidden_2) + +set parameters for dest_x_neural_net predictor network +init_param!(custom_dest_proposal_neural, :x_W1, init_x_W1) +init_param!(custom_dest_proposal_neural, :x_b1, zeros(num_hidden_1)) +init_param!(custom_dest_proposal_neural, :x_W2, init_x_W2) +init_param!(custom_dest_proposal_neural, :x_b2, zeros(num_hidden_2)) +init_param!(custom_dest_proposal_neural, :x_W3, init_x_W3) +init_param!(custom_dest_proposal_neural, :x_b3, zeros(num_x_bins)) + +init_y_W1 = init_weight(num_hidden_1, 4) +init_y_W2 = init_weight(num_hidden_2, num_hidden_1) +init_y_W3 = init_weight(num_x_bins, num_hidden_2) + +set parameters for dest_y_neural_net predictor network +init_param!(custom_dest_proposal_neural, :y_W1, init_y_W1) +init_param!(custom_dest_proposal_neural, :y_b1, zeros(num_hidden_1)) +init_param!(custom_dest_proposal_neural, :y_W2, init_y_W2) +init_param!(custom_dest_proposal_neural, :y_b2, zeros(num_hidden_2)) +init_param!(custom_dest_proposal_neural, :y_W3, init_y_W3) +init_param!(custom_dest_proposal_neural, :y_b3, zeros(num_y_bins)); +- + +Now, we visualize the proposal distribution prior to training: + +visualize_custom_destination_proposal(measurements, start, custom_dest_proposal_neural, num_samples=1000) + +It looks like the initial distribution is roughly uniform, like the default +proposal. + +Now we train the network stochastic gradient descent with a fixed step size +of 0.001 that is shared among all of the trainable parameters. + +update = Gen.ParamUpdate(Gen.FixedStepGradientDescent(0.001), custom_dest_proposal_neural); + +We use 50 epochs of training. In each epoch, we generate 100 training +examples, and we apply 100 gradient updates, where each update is based on +the gradient estimate obtained from a random set of 100 of the trainable +examples. At the end of each epoch, we estimate the objective function value +using 10000 freshly sampled examples. This process takes about 10 minutes to +run on a typical laptop CPU, so we have precomputed the results for you. + +```julia +using JLD2 +@time scores = Gen.train!(custom_dest_proposal_neural, data_generator, update, + num_epoch=50, epoch_size=100, num_minibatch=100, minibatch_size=100, + evaluation_size=1000, verbose=true); + +let data = Dict() + for name in [:x_W1, :x_b1, :x_W2, :x_b2, :x_W3, :x_b3, :y_W1, :y_b1, :y_W2, :y_b2, :y_W3, :y_b3] + data[(:param, name)] = Gen.get_param(custom_dest_proposal_neural, name) + end + data[:scores] = scores + save("params/custom_dest_proposal_neural_trained.jld2", "data", data) +end +``` + +We load the results here: + +``` +using JLD2 +scores = let data = JLD2.load("params/custom_dest_proposal_neural_trained.jld2", "data") + for name in [:x_W1, :x_b1, :x_W2, :x_b2, :x_W3, :x_b3, :y_W1, :y_b1, :y_W2, :y_b2, :y_W3, :y_b3] + Gen.init_param!(custom_dest_proposal_neural, name, data[(:param, name)]) + end + data[:scores] +end; +``` + +We plot the estimate of the objective function over epochs: + +plot(scores, + xlabel="Epochs", + ylabel="Estimate of expected\nconditional log probability density", + label=nothing) + +Below, we visualize the trained proposal distribution for our data set: + +visualize_custom_destination_proposal(measurements, start, custom_dest_proposal_neural, num_samples=1000) + +If we run inference with `amt_computation` set to 5, we see that the inferred distribution reflects the bias of the proposal: + +visualize_data_driven_inference(measurements, scene, start, custom_dest_proposal_neural, + amt_computation=5, samples=1000) + +As we increase the amount of computation, the effect of the proposal's bias +is reduced: + +visualize_data_driven_inference(measurements, scene, start, custom_dest_proposal_neural, + amt_computation=50, samples=1000) + +This bias-correction is more noticeable the more computation we use (though here we only draw 100 approximate posterior samples): + +visualize_data_driven_inference(measurements, scene, start, custom_dest_proposal_neural, + amt_computation=1000, samples=100) + +This example highlights an important aspect of importance sampling: it not +only _upweights_ guesses that explain the data well; it also _downweights_ guesses +that are too high-probability under the proposal distribution. That is, if a proposal +is heavily biased toward one region of the state space, all guesses in that region will +be downweighted accordingly. That's why, even though (a) guesses in the left and right halves +of the room are equally likely, and (b) the +proposal stongly prefers the left half of the room, the importance sampling algorithm +samples roughly the same number of points in each half of the room. + +In the limit of infinite computation, the distribution induced by importance sampling +converges to the true posterior, independent of the proposal. Indeed, using the +generic proposal with a high amount of computation produces very similar results: + +visualize_inference(measurements, scene, start; computation_amt=1000, samples=100) \ No newline at end of file diff --git a/docs/src/tutorials/model_optimizations/scaling_with_sml.md b/docs/src/tutorials/model_optimizations/scaling_with_sml.md index 668fd9fb8..29ea8f438 100644 --- a/docs/src/tutorials/model_optimizations/scaling_with_sml.md +++ b/docs/src/tutorials/model_optimizations/scaling_with_sml.md @@ -1,10 +1,10 @@ -# Scaling with the Static Modeling Language +# [Scaling with the Static Modeling Language](@id scaling_with_sml_tutorial) ## Introduction For prototyping models and working with dynamic structures, Gen's [Dynamic Modeling Language](@ref dynamic_modeling_language) is a great (and the default) way of writing probabilistic programs in nearly pure Julia. However, better performance and scaling characteristics can be obtained using specialized modeling languages or modeling constructs. This notebook introduces a more specialized modeling language known as the [Static Modeling Language](@ref sml) (SML) which is also built into Gen. The SML provides model speedups by carefully analyzing what work is necessary during inference. **Prerequisites for this tutorial** -- [Using Generative Combinators](@ref combinators_tutorial) +- [Generative Combinators](@ref combinators_tutorial) This tutorial will take the robust regression model used to introduce iterative inference in [an earlier tutorial] and optimize the speed of inference using the SML. diff --git a/docs/src/tutorials/reversible_jump.md b/docs/src/tutorials/reversible_jump.md new file mode 100644 index 000000000..e69de29bb diff --git a/src/address.jl b/src/address.jl index 0d0fa45c3..8236a7e9c 100644 --- a/src/address.jl +++ b/src/address.jl @@ -29,7 +29,7 @@ export AllAddressSchema ###################### """ - abstract type Selection end + Selection Abstract type for selections of addresses. diff --git a/src/modeling_library/switch/update.jl b/src/modeling_library/switch/update.jl index d14af1278..49cf48cbb 100644 --- a/src/modeling_library/switch/update.jl +++ b/src/modeling_library/switch/update.jl @@ -10,6 +10,11 @@ mutable struct SwitchUpdateState{T} SwitchUpdateState{T}(weight::Float64, score::Float64, noise::Float64, prev_trace::Trace) where T = new{T}(weight, score, noise, prev_trace) end +""" + update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap) + +Returns choices that are in constraints, merged with all choices in the previous trace that do not have the same address as some choice in the constraints." +""" function update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap) prev_choice_submap_iterator = get_submaps_shallow(prev_choices) prev_choice_value_iterator = get_values_shallow(prev_choices) @@ -48,13 +53,14 @@ function update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap) return new_choices end -@doc( -""" -update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap) -Returns choices that are in constraints, merged with all choices in the previous trace that do not have the same address as some choice in the constraints." -""", update_recurse_merge) +""" + update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices::ChoiceMap) +Returns choices from previous trace that: + 1. have an address which does not appear in the new trace. + 2. have an address which does appear in the constraints. +""" function update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices::ChoiceMap) discard = choicemap() for (k, v) in get_submaps_shallow(prev_choices) @@ -71,14 +77,6 @@ function update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices discard end -@doc( -""" -update_discard(prev_choices::ChoiceMap, choices::ChoiceMap, new_choices::ChoiceMap) - -Returns choices from previous trace that: - 1. have an address which does not appear in the new trace. - 2. have an address which does appear in the constraints. -""", update_discard) @inline update_discard(prev_trace::Trace, choices::ChoiceMap, new_trace::Trace) = update_discard(get_choices(prev_trace), choices, get_choices(new_trace))