From 614a10a31cdad027014bc531d51efdbaf3250880 Mon Sep 17 00:00:00 2001 From: Xuan Date: Sat, 5 Oct 2024 17:23:06 -0400 Subject: [PATCH] Reorganize documentation. --- docs/README.md | 21 +- docs/make.jl | 1 - docs/pages.jl | 76 +-- docs/src/api/inference/map.md | 4 - docs/src/api/inference/mcmc.md | 19 - docs/src/api/inference/pf.md | 10 - docs/src/api/inference/vi.md | 7 - docs/src/api/model/gfi.md | 55 -- docs/src/api/model/trace_translators.md | 20 - docs/src/assets/header.css | 138 ---- docs/src/assets/header.js | 98 --- docs/src/assets/theme.css | 27 - docs/src/explanations/combinator_design.md | 9 - docs/src/how_to/custom_derivatives.md | 25 - docs/src/how_to/custom_distributions.md | 163 +++-- docs/src/how_to/custom_dsl.md | 16 - docs/src/how_to/custom_gen_fns.md | 83 +++ docs/src/how_to/custom_gradients.md | 56 ++ .../how_to/custom_incremental_computation.md | 7 +- docs/src/how_to/extending_gen.md | 12 + docs/src/how_to/how_to.md | 3 - docs/src/index.md | 83 +-- docs/src/ref/core/change_hints.md | 30 + .../{api/model => ref/core}/choice_maps.md | 8 +- .../src/{tutorials/basics => ref/core}/gfi.md | 61 +- .../src/{api/model => ref/core}/selections.md | 6 +- docs/src/{api => ref}/inference/importance.md | 1 + docs/src/ref/inference/map.md | 9 + .../mcmc_kernels.md => ref/inference/mcmc.md} | 33 +- .../inference}/parameter_optimization.md | 27 +- docs/src/ref/inference/pf.md | 37 ++ .../inference}/trace_translators.md | 31 +- docs/src/ref/inference/vi.md | 8 + docs/src/ref/inference/wake_sleep.md | 10 + .../internals}/language_implementation.md | 8 +- .../model => ref/modeling}/combinators.md | 20 +- docs/src/ref/modeling/custom_gen_fns.md | 29 + .../model => ref/modeling}/distributions.md | 71 +- .../model/modeling.md => ref/modeling/dml.md} | 135 +--- docs/src/ref/modeling/sml.md | 115 ++++ docs/src/tutorials/basics/combinators.md | 114 ---- docs/src/tutorials/basics/vi.md | 10 - .../getting_started.md} | 5 +- .../learning_gen_fns.md} | 10 +- .../tutorials/{basics => }/modeling_in_gen.md | 605 ++++++------------ docs/src/tutorials/reversible_jump.md | 0 .../scaling_with_sml.md | 68 +- .../{basics/particle_filter.md => smc.md} | 146 ++--- src/inference/map_optimize.jl | 1 - src/inference/particle_filter.jl | 20 +- src/modeling_library/modeling_library.jl | 7 +- 51 files changed, 1070 insertions(+), 1488 deletions(-) delete mode 100644 docs/src/api/inference/map.md delete mode 100644 docs/src/api/inference/mcmc.md delete mode 100644 docs/src/api/inference/pf.md delete mode 100644 docs/src/api/inference/vi.md delete mode 100644 docs/src/api/model/gfi.md delete mode 100644 docs/src/api/model/trace_translators.md delete mode 100644 docs/src/assets/header.css delete mode 100644 docs/src/assets/header.js delete mode 100644 docs/src/assets/theme.css delete mode 100644 docs/src/explanations/combinator_design.md delete mode 100644 docs/src/how_to/custom_derivatives.md delete mode 100644 docs/src/how_to/custom_dsl.md create mode 100644 docs/src/how_to/custom_gen_fns.md create mode 100644 docs/src/how_to/custom_gradients.md create mode 100644 docs/src/how_to/extending_gen.md delete mode 100644 docs/src/how_to/how_to.md create mode 100644 docs/src/ref/core/change_hints.md rename docs/src/{api/model => ref/core}/choice_maps.md (96%) rename docs/src/{tutorials/basics => ref/core}/gfi.md (93%) rename docs/src/{api/model => ref/core}/selections.md (94%) rename docs/src/{api => ref}/inference/importance.md (98%) create mode 100644 docs/src/ref/inference/map.md rename docs/src/{how_to/mcmc_kernels.md => ref/inference/mcmc.md} (97%) rename docs/src/{api/model => ref/inference}/parameter_optimization.md (64%) create mode 100644 docs/src/ref/inference/pf.md rename docs/src/{tutorials => ref/inference}/trace_translators.md (97%) create mode 100644 docs/src/ref/inference/vi.md create mode 100644 docs/src/ref/inference/wake_sleep.md rename docs/src/{explanations => ref/internals}/language_implementation.md (92%) rename docs/src/{api/model => ref/modeling}/combinators.md (87%) create mode 100644 docs/src/ref/modeling/custom_gen_fns.md rename docs/src/{api/model => ref/modeling}/distributions.md (88%) rename docs/src/{api/model/modeling.md => ref/modeling/dml.md} (75%) create mode 100644 docs/src/ref/modeling/sml.md delete mode 100644 docs/src/tutorials/basics/combinators.md delete mode 100644 docs/src/tutorials/basics/vi.md rename docs/src/{getting_started/linear_regression.md => tutorials/getting_started.md} (97%) rename docs/src/{api/inference/learning.md => tutorials/learning_gen_fns.md} (99%) rename docs/src/tutorials/{basics => }/modeling_in_gen.md (90%) delete mode 100644 docs/src/tutorials/reversible_jump.md rename docs/src/tutorials/{model_optimizations => }/scaling_with_sml.md (84%) rename docs/src/tutorials/{basics/particle_filter.md => smc.md} (91%) diff --git a/docs/README.md b/docs/README.md index 78ddd3f69..fcda114d5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,17 +1,18 @@ -# Website Docs -- `pages.jl` to find skeleton of website. -- `make.jl` to build the website index. +# Gen.jl Documentation -The docs are divided in roughly four sections: +- `pages.jl` to find skeleton of the Gen.jl documentation. +- `make.jl` to build the documentation website index. + +The documentation is divided into three sections: - Getting Started + Tutorials - How-to Guides -- API = Modeling API + Inference API -- Explanations + Internals - +- Reference Guides # Developing -To build the docs, run `julia --make.jl` or alternatively startup the Julia REPL and include `make.jl`. For debugging, consider setting `draft=true` in the `makedocs` function found in `make.jl`. -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. + +To build the docs, run `julia --make.jl`. Alternatively, start the Julia REPL, activate the `Project.toml` in this directory, then include `make.jl`. For debugging, consider setting `draft=true` in the `makedocs` function found in `make.jl`. This will avoid running the `@example` blocks when generating the tutorials. + +Currently you must write the tutorial directly in the docs rather than in a source file. See `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: @@ -21,4 +22,4 @@ x = rand() ```@example tutorial_1 print(x) -``` \ No newline at end of file +``` diff --git a/docs/make.jl b/docs/make.jl index 92a7cd065..e7b6b45c3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -8,7 +8,6 @@ makedocs( clean = true, warnonly = true, format = Documenter.HTML(; - assets = String["assets/header.js", "assets/header.css", "assets/theme.css"], collapselevel=1, ), sitename = "Gen.jl", diff --git a/docs/pages.jl b/docs/pages.jl index dc340eb7e..d6b1f6c4f 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -1,53 +1,45 @@ pages = [ - "Home" => "index.md", - "Getting Started" => [ - "Example 1: Linear Regression" => "getting_started/linear_regression.md", - ], + "Gen.jl" => "index.md", "Tutorials" => [ - "Basics" => [ - "tutorials/basics/modeling_in_gen.md", - "tutorials/basics/gfi.md", - "tutorials/basics/combinators.md", - "tutorials/basics/particle_filter.md", - "tutorials/basics/vi.md", - ], - "Advanced" => [ - "tutorials/trace_translators.md", - ], - "Model Optmizations" => [ - "Speeding Inference with the Static Modeling Language" => "tutorials/model_optimizations/scaling_with_sml.md", - ], + "Getting Started" => "tutorials/getting_started.md", + "Introduction to Modeling in Gen" => "tutorials/modeling_in_gen.md", + "Object Tracking with SMC" => "tutorials/smc.md", + "Learning Generative Functions" => "tutorials/learning_gen_fns.md", + "Speeding Up Inference with the SML" => "tutorials/scaling_with_sml.md", ], "How-to Guides" => [ - "MCMC Kernels" => "how_to/mcmc_kernels.md", - "Custom Distributions" => "how_to/custom_distributions.md", - "Custom Modeling Languages" => "how_to/custom_dsl.md", - "Custom Gradients" => "how_to/custom_derivatives.md", - "Incremental Computation" => "how_to/custom_incremental_computation.md", + "Extending Gen" => "how_to/extending_gen.md", + "Adding New Distributions" => "how_to/custom_distributions.md", + "Adding New Generative Functions" => "how_to/custom_gen_fns.md", + "Custom Gradients" => "how_to/custom_gradients.md", + "Custom Incremental Computation" => "how_to/custom_incremental_computation.md", ], - "API Reference" => [ + "Reference" => [ + "Core Interfaces" => [ + "Generative Function Interface" => "ref/core/gfi.md", + "Choice Maps" => "ref/core/choice_maps.md", + "Selections" => "ref/core/selections.md", + "Change Hints" => "ref/core/change_hints.md", + ], "Modeling Library" => [ - "Generative Functions" => "api/model/gfi.md", - "Probability Distributions" => "api/model/distributions.md", - "Choice Maps" => "api/model/choice_maps.md", - "Built-in Modeling Languages" => "api/model/modeling.md", - "Combinators" => "api/model/combinators.md", - "Selections" => "api/model/selections.md", - "Optimizing Trainable Parameters" => "api/model/parameter_optimization.md", - "Trace Translators" => "api/model/trace_translators.md", + "Built-In Modeling Language" => "ref/modeling/dml.md", + "Static Modeling Language" => "ref/modeling/sml.md", + "Probability Distributions" => "ref/modeling/distributions.md", + "Combinators" => "ref/modeling/combinators.md", + "Custom Generative Functions" => "ref/modeling/custom_gen_fns.md", ], "Inference Library" => [ - "Importance Sampling" => "api/inference/importance.md", - "MAP Optimization" => "api/inference/map.md", - "Markov chain Monte Carlo" => "api/inference/mcmc.md", - "MAP Optimization" => "api/inference/map.md", - "Particle Filtering" => "api/inference/pf.md", - "Variational Inference" => "api/inference/vi.md", - "Learning Generative Functions" => "api/inference/learning.md" + "Importance Sampling" => "ref/inference/importance.md", + "Markov Chain Monte Carlo" => "ref/inference/mcmc.md", + "Particle Filtering & SMC" => "ref/inference/pf.md", + "Trace Translators" => "ref/inference/trace_translators.md", + "Parameter Optimization" => "ref/inference/parameter_optimization.md", + "MAP Optimization" => "ref/inference/map.md", + "Variational Inference" => "ref/inference/vi.md", + "Wake-Sleep Learning" => "ref/inference/wake_sleep.md", ], + "Internals" => [ + "Modeling Language Implementation" => "ref/internals/language_implementation.md", + ] ], - "Explanation and Internals" => [ - "Modeling Language Implementation" => "explanations/language_implementation.md", - "explanations/combinator_design.md" - ] ] diff --git a/docs/src/api/inference/map.md b/docs/src/api/inference/map.md deleted file mode 100644 index fcd1ba880..000000000 --- a/docs/src/api/inference/map.md +++ /dev/null @@ -1,4 +0,0 @@ -# MAP Optimization -```@docs -map_optimize -``` diff --git a/docs/src/api/inference/mcmc.md b/docs/src/api/inference/mcmc.md deleted file mode 100644 index bd6df6d00..000000000 --- a/docs/src/api/inference/mcmc.md +++ /dev/null @@ -1,19 +0,0 @@ -# Markov chain Monte Carlo (MCMC) - -Gen supports standard Markov Chain Monte Carlo algorithms and allows users to write their own custom kernels. -```@index -Pages = ["mcmc.md"] -``` - -```@docs -metropolis_hastings -mh -mala -hmc -elliptical_slice -@pkern -@kern -@rkern -reversal -involutive_mcmc -``` diff --git a/docs/src/api/inference/pf.md b/docs/src/api/inference/pf.md deleted file mode 100644 index a356a76d9..000000000 --- a/docs/src/api/inference/pf.md +++ /dev/null @@ -1,10 +0,0 @@ -# Particle Filtering -```@docs -initialize_particle_filter -particle_filter_step! -maybe_resample! -log_ml_estimate -get_traces -get_log_weights -sample_unweighted_traces -``` diff --git a/docs/src/api/inference/vi.md b/docs/src/api/inference/vi.md deleted file mode 100644 index 4d55fb43c..000000000 --- a/docs/src/api/inference/vi.md +++ /dev/null @@ -1,7 +0,0 @@ -## 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! -``` diff --git a/docs/src/api/model/gfi.md b/docs/src/api/model/gfi.md deleted file mode 100644 index 0385988be..000000000 --- a/docs/src/api/model/gfi.md +++ /dev/null @@ -1,55 +0,0 @@ -## [Generative Functions](@id gfi_api) - -```@docs -GenerativeFunction -Trace -``` - -The complete set of methods in the generative function interface (GFI) is: - -```@docs -simulate -generate -update -regenerate -get_args -get_retval -get_choices -get_score -get_gen_fn -Base.getindex -project -propose -assess -has_argument_grads -has_submap -accepts_output_grad -accumulate_param_gradients! -choice_gradients -get_params -``` - -```@docs -Diff -NoChange -UnknownChange -SetDiff -Diffed -``` - -```@docs -CustomUpdateGF -apply_with_state -update_with_state -``` - -```@docs -CustomGradientGF -apply -gradient -``` - -```@docs -Gen.init_update_state -Gen.apply_update! -``` \ No newline at end of file diff --git a/docs/src/api/model/trace_translators.md b/docs/src/api/model/trace_translators.md deleted file mode 100644 index 14de30ec9..000000000 --- a/docs/src/api/model/trace_translators.md +++ /dev/null @@ -1,20 +0,0 @@ -## Trace Translators - -```@docs -@transform -@read -@write -@copy -@tcall -pair_bijections! -is_involution! -inverse -TraceTranslator -DeterministicTraceTranslator -GeneralTraceTranslator -SimpleExtendingTraceTranslator -SymmetricTraceTranslator -``` -```@docs -TraceTransformDSLProgram -``` diff --git a/docs/src/assets/header.css b/docs/src/assets/header.css deleted file mode 100644 index 155b10a3b..000000000 --- a/docs/src/assets/header.css +++ /dev/null @@ -1,138 +0,0 @@ - -@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 deleted file mode 100644 index 2f211b385..000000000 --- a/docs/src/assets/header.js +++ /dev/null @@ -1,98 +0,0 @@ -// 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() { - //
- 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 deleted file mode 100644 index 67017e105..000000000 --- a/docs/src/assets/theme.css +++ /dev/null @@ -1,27 +0,0 @@ -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/explanations/combinator_design.md b/docs/src/explanations/combinator_design.md deleted file mode 100644 index 2d999e330..000000000 --- a/docs/src/explanations/combinator_design.md +++ /dev/null @@ -1,9 +0,0 @@ -# Combinator Design and Implementation - -Internally, the Combinators use custom trace types. -```@docs -Gen.VectorTrace -Gen.process_all_new! -Gen.update_recurse_merge -Gen.update_discard -``` \ No newline at end of file diff --git a/docs/src/how_to/custom_derivatives.md b/docs/src/how_to/custom_derivatives.md deleted file mode 100644 index e52f92e9c..000000000 --- a/docs/src/how_to/custom_derivatives.md +++ /dev/null @@ -1,25 +0,0 @@ -# How to Write Custom Gradients - -To add a custom gradient for a differentiable deterministic computation, define a concrete subtype of [`CustomGradientGF`](@ref) with the following methods: - -- [`apply`](@ref) - -- [`gradient`](@ref) - -- [`has_argument_grads`](@ref) - -For example: - -```julia -struct MyPlus <: CustomGradientGF{Float64} end - -Gen.apply(::MyPlus, args) = args[1] + args[2] -Gen.gradient(::MyPlus, args, retval, retgrad) = (retgrad, retgrad) -Gen.has_argument_grads(::MyPlus) = (true, true) -``` - -# [Optimizing Trainable Parameters](@id optimizing-internal) - -To add support for a new type of gradient-based parameter update, create a new type with the following methods defined for the types of generative functions that are to be supported. -- [`Gen.init_update_state`](@ref) -- [`Gen.apply_update!`](@ref) diff --git a/docs/src/how_to/custom_distributions.md b/docs/src/how_to/custom_distributions.md index 4722bfa29..573472c64 100644 --- a/docs/src/how_to/custom_distributions.md +++ b/docs/src/how_to/custom_distributions.md @@ -1,126 +1,123 @@ -# How to Write Custom Gen Distributions +# [Adding New Distributions](@id custom_distributions_howto) -## [Custom distributions](@id custom_distributions) +In addition to the built-in distributions, mixture distributions, and product distributions, +Gen provides two primary ways of adding new distributions: -Users can extend Gen with new probability distributions, which can then be used -to make random choices within generative functions. Simple transformations of -existing distributions can be created using the [`@dist` DSL](@ref dist_dsl). -For arbitrary distributions, including distributions that cannot be expressed -in the `@dist` DSL, users can define a custom distribution by implementing -Gen's Distribution interface directly, as defined below. +## [Defining New Distributions Inline with the `@dist` DSL](@id dist_dsl_howto) -Probability distributions are singleton types whose supertype is `Distribution{T}`, where `T` indicates the data type of the random sample. +The `@dist` DSL allows users to concisely define a distribution, as long as +that distribution can be expressed as a certain type of deterministic +transformation of an existing distribution: ```julia -abstract type Distribution{T} end +@dist name(arg1, arg2, ..., argN) = body +``` +or +```julia +@dist function name(arg1, arg2, ..., argN) + body +end ``` -A new `Distribution` type must implement the following methods: - -- [`random`](@ref) - -- [`logpdf`](@ref) - -- [`has_output_grad`](@ref) - -- [`logpdf_grad`](@ref) +Here `body` is ordinary Julia code, with the constraint that `body` must +contain exactly one random choice. The resulting distribution is called `name`, +parameterized by `arg1, ..., argN`, and represents a distribution over +_return values_ of `body`. -- [`has_argument_grads`](@ref) +### Common Use Cases - -By convention, distributions have a global constant lower-case name for the singleton value. -For example: +The `@dist` DSL makes it easy to implement labeled uniform or categorical +distributions, instead of having to use integers to refer to categories: ```julia -struct Bernoulli <: Distribution{Bool} end -const bernoulli = Bernoulli() -``` -Distribution values should also be callable, which is a syntactic sugar with the same behavior of calling `random`: +"""Labeled uniform distribution""" +@dist labeled_uniform(labels) = labels[uniform_discrete(1, length(labels))] -```julia -bernoulli(0.5) # identical to random(bernoulli, 0.5) and random(Bernoulli(), 0.5) +"""Labeled categorical distribution""" +@dist labeled_categorical(labels, probs) = labels[categorical(probs)] ``` -## [Custom Generative Functions](@id custom_generative_functions) +Note however that there is a slight overhead incurred by having to detect the +possibility of duplicate labels. -We recommend the following steps for implementing a new type of generative function, and also looking at the implementation for the [`DynamicDSLFunction`](@ref) type as an example. +It is also possible to implement a distribution that takes a Boolean input +and flips it with some probability: -##### Define a trace data type ```julia -struct MyTraceType <: Trace - .. -end +"Bit-flip distribution." +@dist bit_flip(x::Bool, p::Float64) = bernoulli((1-x) * p + x * (1-p)) ``` -##### Decide the return type for the generative function -Suppose our return type is `Vector{Float64}`. +Finally, it is possible to distributions that are defined in terms of shifting, +scaling, exponentiation, or taking the logarithm of another distribution: -##### Define a data type for your generative function -This should be a subtype of [`GenerativeFunction`](@ref), with the appropriate type parameters. ```julia -struct MyGenerativeFunction <: GenerativeFunction{Vector{Float64},MyTraceType} -.. -end -``` -Note that your generative function may not need to have any fields. -You can create a constructor for it, e.g.: -``` -function MyGenerativeFunction(...) -.. -end -``` - -##### Decide what the arguments to a generative function should be -For example, our generative functions might take two arguments, `a` (of type `Int`) and `b` (of type `Float64`). -Then, the argument tuple passed to e.g. [`generate`](@ref) will have two elements. - -NOTE: Be careful to distinguish between arguments to the generative function itself, and arguments to the constructor of the generative function. -For example, if you have a generative function type that is parametrized by, for example, modeling DSL code, this DSL code would be a parameter of the generative function constructor. - -##### Decide what the traced random choices (if any) will be -Remember that each random choice is assigned a unique address in (possibly) hierarchical address space. -You are free to design this address space as you wish, although you should document it for users of your generative function type. +"Symmetric Binomial distribution." +@dist sym_binom(mean::Int, scale::Int) = binom(2*scale, 0.5) - scale + mean -##### Implement methods of the Generative Function Interface +"Shifted geometric distribution." +@dist shifted_geometric(p::Real, shift::Int) = geometric(p) + shift -At minimum, you need to implement the following methods: +"Log-normal distribution." +@dist log_normal(mu::Real, sigma::Real) = exp(normal(mu, sigma)) -- [`simulate`](@ref) - -- [`has_argument_grads`](@ref) +"Gumbel distribution." +@dist gumbel(mu::Real, beta::Real) = mu - beta * log(0.0 - log(uniform(0, 1))) +``` -- [`accepts_output_grad`](@ref) +### Restrictions -- [`get_args`](@ref) +There are a number of restrictions imposed by the `@dist` DSL, which are +explained further in the [reference docmumentation](@ref dist_dsl). +Most importantly, only a limited set of deterministic transformations are currently +supported (`+`, `-`, `*`, `/`, `exp`, `log`, `getindex`), and only *one* random +choice can be used in the body of the definition. -- [`get_retval`](@ref) +## Defining New Distributions From Scratch -- [`get_choices`](@ref) +For distributions that cannot be expressed in the `@dist` DSL, users can define +a custom distribution by defining an (ordinary Julia) subtype of +`Gen.Distribution` and implementing the methods of the [Distribution API](@ref +distributions). This method requires more custom code than using the +`@dist` DSL, but also affords more flexibility: arbitrary user-defined logic +for sampling, PDF evaluation, etc. -- [`get_score`](@ref) +Probability distributions are singleton types whose supertype is `Distribution{T}`, where `T` indicates the data type of the random sample. -- [`get_gen_fn`](@ref) +```julia +abstract type Distribution{T} end +``` -- [`project`](@ref) +A new `Distribution` type must implement the following methods: -If you want to use the generative function within models, you should implement: +- [`random`](@ref) -- [`generate`](@ref) +- [`logpdf`](@ref) -If you want to use MCMC on models that call your generative function, then implement: +- [`logpdf_grad`](@ref) -- [`update`](@ref) +- [`has_argument_grads`](@ref) -- [`regenerate`](@ref) +- [`has_output_grad`](@ref) -If you want to use gradient-based inference techniques on models that call your generative function, then implement: +- [`is_discrete`](@ref) -- [`choice_gradients`](@ref) +By convention, distributions have a global constant lower-case name for the singleton value. +For example: -- [`update`](@ref) +```julia +struct Bernoulli <: Distribution{Bool} end +const bernoulli = Bernoulli() +``` -If your generative function has trainable parameters, then implement: +Distribution values should also be callable, which is a syntactic sugar with the same behavior of calling `random`: -- [`accumulate_param_gradients!`](@ref) +```julia +bernoulli(0.5) # identical to random(bernoulli, 0.5) and random(Bernoulli(), 0.5) +``` +For example, this can be done by adding a method definition for `Bernoulli`: +```julia +(::Bernoulli)(prob) = random(Bernoulli(), prob) +``` diff --git a/docs/src/how_to/custom_dsl.md b/docs/src/how_to/custom_dsl.md deleted file mode 100644 index 7e503041b..000000000 --- a/docs/src/how_to/custom_dsl.md +++ /dev/null @@ -1,16 +0,0 @@ -# How to Write a Custom Modeling Languages - -Gen can be extended with new modeling languages by implementing new generative function types, and constructors for these types that take models as input. -This typically requires implementing the entire generative function interface, and is advanced usage of Gen. - - -Gen is designed for extensibility. -To implement behaviors that are not directly supported by the existing modeling languages, users can implement `black-box' generative functions directly, without using built-in modeling language. -These generative functions can then be invoked by generative functions defined using the built-in modeling language. -This includes several special cases: - -- Extending Gen with custom gradient computations - -- Extending Gen with custom incremental computation of return values - -- Extending Gen with new modeling languages. \ No newline at end of file diff --git a/docs/src/how_to/custom_gen_fns.md b/docs/src/how_to/custom_gen_fns.md new file mode 100644 index 000000000..f172bf009 --- /dev/null +++ b/docs/src/how_to/custom_gen_fns.md @@ -0,0 +1,83 @@ +# [Adding New Types of Generative Functions](@id custom_gen_fns_howto) + +We recommend the following steps for implementing a new type of generative function, and also looking at the implementation for the [`DynamicDSLFunction`](@ref) type as an example. + +#### Define a trace data type +```julia +struct MyTraceType <: Trace + .. +end +``` + +#### Decide the return type for the generative function +Suppose our return type is `Vector{Float64}`. + +#### Define a data type for your generative function +This should be a subtype of [`GenerativeFunction`](@ref), with the appropriate type parameters. +```julia +struct MyGenerativeFunction <: GenerativeFunction{Vector{Float64},MyTraceType} +.. +end +``` +Note that your generative function may not need to have any fields. +You can create a constructor for it, e.g.: +``` +function MyGenerativeFunction(...) +.. +end +``` + +#### Decide what the arguments to a generative function should be +For example, our generative functions might take two arguments, `a` (of type `Int`) and `b` (of type `Float64`). +Then, the argument tuple passed to e.g. [`generate`](@ref) will have two elements. + +NOTE: Be careful to distinguish between arguments to the generative function itself, and arguments to the constructor of the generative function. +For example, if you have a generative function type that is parametrized by, for example, modeling DSL code, this DSL code would be a parameter of the generative function constructor. + +#### Decide what the traced random choices (if any) will be +Remember that each random choice is assigned a unique address in (possibly) hierarchical address space. +You are free to design this address space as you wish, although you should document it for users of your generative function type. + +#### Implement methods of the Generative Function Interface + +At minimum, you need to implement the following methods: + +- [`simulate`](@ref) + +- [`has_argument_grads`](@ref) + +- [`accepts_output_grad`](@ref) + +- [`get_args`](@ref) + +- [`get_retval`](@ref) + +- [`get_choices`](@ref) + +- [`get_score`](@ref) + +- [`get_gen_fn`](@ref) + +- [`project`](@ref) + +If you want to use the generative function within models, you should implement: + +- [`generate`](@ref) + +If you want to use MCMC on models that call your generative function, then implement: + +- [`update`](@ref) + +- [`regenerate`](@ref) + +If you want to use gradient-based inference techniques on models that call your generative function, then implement: + +- [`choice_gradients`](@ref) + +- [`update`](@ref) + +If your generative function has trainable parameters, then implement: + +- [`accumulate_param_gradients!`](@ref) + + diff --git a/docs/src/how_to/custom_gradients.md b/docs/src/how_to/custom_gradients.md new file mode 100644 index 000000000..333b889a7 --- /dev/null +++ b/docs/src/how_to/custom_gradients.md @@ -0,0 +1,56 @@ +# [Customizing Gradients](@id custom_gradients_howto) + +## Determistic Functions with Custom Gradients + +To add a custom gradient for a differentiable deterministic computation, define a concrete subtype of [`CustomGradientGF`](@ref) with the following methods: + +- [`apply`](@ref) + +- [`gradient`](@ref) + +- [`has_argument_grads`](@ref) + +For example, we can implement binary addition with a manually-defined gradient: + +```julia +struct MyPlus <: CustomGradientGF{Float64} end + +Gen.apply(::MyPlus, args) = args[1] + args[2] +Gen.gradient(::MyPlus, args, retval, retgrad) = (retgrad, retgrad) +Gen.has_argument_grads(::MyPlus) = (true, true) +``` + +## Customizing Parameter Updates + +To add support for a new type of gradient-based parameter update, create a new [parameter update configuration](@ref update_configurations) with the following methods defined for the types of generative functions that are to be supported. + +- [`Gen.init_update_state`](@ref) +- [`Gen.apply_update!`](@ref) + +As an example, the built-in update configuration, [`FixedStepGradientDescent`](@ref), is implemented as follows: + +```julia +struct FixedStepGradientDescent + step_size::Float64 +end + +mutable struct FixedStepGradientDescentBuiltinDSLState + step_size::Float64 + gen_fn::Union{Gen.DynamicDSLFunction,Gen.StaticIRGenerativeFunction} + param_list::Vector +end + +function Gen.init_update_state(conf::FixedStepGradientDescent, + gen_fn::Union{Gen.DynamicDSLFunction,Gen.StaticIRGenerativeFunction}, param_list::Vector) + FixedStepGradientDescentBuiltinDSLState(conf.step_size, gen_fn, param_list) +end + +function Gen.apply_update!(state::FixedStepGradientDescentBuiltinDSLState) + for param_name in state.param_list + value = Gen.get_param(state.gen_fn, param_name) + grad = Gen.get_param_grad(state.gen_fn, param_name) + Gen.set_param!(state.gen_fn, param_name, value + grad * state.step_size) + Gen.zero_param_grad!(state.gen_fn, param_name) + end +end +``` diff --git a/docs/src/how_to/custom_incremental_computation.md b/docs/src/how_to/custom_incremental_computation.md index b44cb91bd..d1318ce61 100644 --- a/docs/src/how_to/custom_incremental_computation.md +++ b/docs/src/how_to/custom_incremental_computation.md @@ -1,6 +1,7 @@ -# How to Customize Incremental Computation +# [Customizing Incremental Computation](@id custom_incremental_computation_howto) + +Iterative inference techniques like Markov Chain Monte Carlo and Sequential Monte Carlo involve repeatedly updating the execution traces of generative models. -Iterative inference techniques like Markov chain Monte Carlo involve repeatedly updating the execution traces of generative models. In some cases, the output of a deterministic computation within the model can be incrementally computed during each of these updates, instead of being computed from scratch. To add a custom incremental computation for a deterministic computation, define a concrete subtype of [`CustomUpdateGF`](@ref) with the following methods: @@ -50,4 +51,4 @@ function Gen.update_with_state(::MySum, state, args, argdiffs::Tuple{VectorDiff} end Gen.num_args(::MySum) = 1 -``` \ No newline at end of file +``` diff --git a/docs/src/how_to/extending_gen.md b/docs/src/how_to/extending_gen.md new file mode 100644 index 000000000..12fd93661 --- /dev/null +++ b/docs/src/how_to/extending_gen.md @@ -0,0 +1,12 @@ +# [Extending Gen](@id extending_gen_howto) + +Gen is designed for extensibility. To implement behaviors that are not directly supported by the existing modeling languages, users can implement `black-box' generative functions that directly implement the [generative function interface](@ref gfi). These generative functions can then be invoked by generative functions defined using the built-in modeling language. + +The following how-tos describe various ways of extending Gen: + +- [Adding custom distributions](@ref custom_distributions_howto) +- [Custom incremental computation of return values](@ref custom_incremental_computation_howto) +- [Custom gradient computations](@ref custom_gradients_howto) +- [Implementing custom generative functions from scratch](@ref custom_gen_fns_howto) + +Gen can also be extended with entirely new modeling languages by implementing new generative function types, and constructors for these types that take models as input. This typically requires implementing the entire generative function interface, and is advanced usage of Gen. diff --git a/docs/src/how_to/how_to.md b/docs/src/how_to/how_to.md deleted file mode 100644 index f2f389b42..000000000 --- a/docs/src/how_to/how_to.md +++ /dev/null @@ -1,3 +0,0 @@ -# "How-to" guides - -"How do I do..." but for Gen. \ No newline at end of file diff --git a/docs/src/index.md b/docs/src/index.md index cc741f22f..e02652e38 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,64 +1,67 @@ # Gen.jl -*A general-purpose probabilistic programming system with programmable inference, embedded in Julia* +A general-purpose probabilistic programming system with programmable inference, embedded in Julia. -- 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. +## Features + +- Multi-paradigm Bayesian inference via [Sequential Monte Carlo](ref/inference/pf.md), [variational inference](ref/inference/vi.md), [MCMC](ref/inference/mcmc.md), and more. +- Gradient-based training of generative models via [parameter optimization](ref/inference/parameter_optimization.md), [wake-sleep learning](ref/inference/wake_sleep.md), etc. +- An expressive and intuitive [modeling language](ref/modeling/dml.md) for writing and composing probabilistic programs. +- Inference algorithms are *programmable*: Write [custom proposals](https://www.gen.dev/tutorials/data-driven-proposals/tutorial), [variational families](tutorials/vi.md), [MCMC kernels](ref/inference/mcmc.md) or [SMC updates](ref/inference/trace_translators.md) without worrying about the math. +- Support for Bayesian structure learning via [involutive MCMC](@ref involutive_mcmc) and [SMCP³](@ref advanced-particle-filtering). +- [Specialized modeling constructs](tutorials/scaling_with_sml.md) that speed-up inference by supporting incremental computation. +- Well-defined APIs for implementing [custom generative models](how_to/custom_gen_fns.md), [distributions](how_to/custom_distributions.md), [gradients](how_to/custom_gradients.md), etc. ## 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 +add Gen ``` -!!! note - Alternatively, - ```julia - julia> import Pkg; Pkg.add("Gen") - ``` +To install the latest development version, you may instead run: + +``` +add https://github.com/probcomp/Gen.jl.git +``` + +Gen can now be used in the Julia REPL, or at the top of a script: + +```julia +using Gen +``` To test the installation locally, you can run the tests with: + ```julia using Pkg; Pkg.test("Gen") ``` -## Getting Started -!!! warning "Tutorial Redirect" - We are in the process of moving tutorial docs around. For more stable tutorial docs, see [these tutorials](https://www.gen.dev/tutorials/) instead. See more examples [here](https://github.com/probcomp/GenExamples.jl). +## Tutorials -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](@ref introduction_to_modeling_in_gen). +Learn how to use Gen by following these tutorials: ```@contents Pages = [ - "getting_started/linear_regression.md", - ] -Depth = 2 + "tutorials/getting_started.md", + "tutorials/modeling_in_gen.md", + "tutorials/smc.md", + "tutorials/scaling_with_sml.md", + "tutorials/learning_gen_fns.md" +] +Depth = 1 ``` -## Contributing -See the [Developer's Guide](https://gen.dev) on how to contribute to the Gen ecosystem. +More tutorials [can also be found on our website](https://www.gen.dev/tutorials/). +Additional examples of Gen usage can be found in [GenExamples.jl](https://github.com/probcomp/GenExamples.jl). + +## Questions and Contributions + +If you have questions about using Gen.jl, feel free to open a [discussion on GitHub](https://github.com/probcomp/Gen.jl/discussions). If you encounter a bug, please [open an issue](https://github.com/probcomp/Gen.jl/issues). We also welcome bug fixes and feature additions as [pull requests](https://github.com/probcomp/Gen.jl/pulls). ## Supporting and Citing -This repo is part of ongoing research at [ProbComp](http://probcomp.csail.mit.edu) and may later include new experimental (for the better)! If you use Gen for your work, please consider citing us: - -```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 + +Gen.jl is part of ongoing research at the [MIT Probabilistic Computing Project](http://probcomp.csail.mit.edu). If you use Gen for your work, please consider citing us: + +> *Gen: A General-Purpose Probabilistic Programming System with Programmable Inference.* Cusumano-Towner, M. F.; Saad, F. A.; Lew, A.; and Mansinghka, V. K. In Proceedings of the 40th ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI ‘19). ([pdf](https://dl.acm.org/doi/10.1145/3314221.3314642)) ([bibtex](https://www.gen.dev/assets/gen-pldi.txt)) diff --git a/docs/src/ref/core/change_hints.md b/docs/src/ref/core/change_hints.md new file mode 100644 index 000000000..57b8186ba --- /dev/null +++ b/docs/src/ref/core/change_hints.md @@ -0,0 +1,30 @@ +# Change Hints + +Gen defines a system of *change hints*, or [`Diff`](@ref) types, in order to +pass information to and from generative functions about whether their arguments +and return values have changed (and how). + +```@docs +Diff +``` + +Change hints are used to support incremental computation in +[generative function combinators](@ref combinators) and the +[static modeling language](@ref sml). The most important change hints are +[`NoChange`](@ref) and [`UnknownChange`](@ref), which respectively indicate that +a value has not changed and that a value may have changed in an unknown way. + +```@docs +NoChange +UnknownChange +``` + +A number of other change hints are also implemented in Gen, such as `IntDiff`, +`VectorDiff`, and `SetDiff` and `DictDiff`, which support incremental +computation for certain operations when applied to [`Diffed`](@ref) values. +These change hints are not documented, and are currently considered an +implementation detail. + +```@docs +Diffed +``` diff --git a/docs/src/api/model/choice_maps.md b/docs/src/ref/core/choice_maps.md similarity index 96% rename from docs/src/api/model/choice_maps.md rename to docs/src/ref/core/choice_maps.md index b5220c2ba..3210e0ca7 100644 --- a/docs/src/api/model/choice_maps.md +++ b/docs/src/ref/core/choice_maps.md @@ -1,6 +1,7 @@ # Choice Maps Maps from the addresses of random choices to their values are stored in associative tree-structured data structures that have the following abstract type: + ```@docs ChoiceMap ``` @@ -9,9 +10,11 @@ Choice maps are constructed by users to express observations and/or constraints Choice maps are also returned by certain Gen inference methods, and are used internally by various Gen inference methods. Choice maps provide the following methods: + ```@docs has_value get_value +has_submap get_submap get_values_shallow get_submaps_shallow @@ -19,6 +22,7 @@ to_array from_array get_selected ``` + Note that none of these methods mutate the choice map. Choice maps also implement: @@ -29,7 +33,7 @@ Choice maps also implement: - `==`, which tests if two choice maps have the same addresses and values at those addresses. - +Gen.jl defines the following concrete types for choice maps: ```@docs DynamicChoiceMap @@ -39,6 +43,7 @@ choicemap ``` A mutable choice map can be constructed with [`choicemap`](@ref), and then populated: + ```julia choices = choicemap() choices[:x] = true @@ -51,7 +56,6 @@ There is also a constructor that takes initial (address, value) pairs: choices = choicemap((:x, true), ("foo", 1.25), (:y => 1 => :z, -6.3)) ``` - ```@docs set_value! set_submap! diff --git a/docs/src/tutorials/basics/gfi.md b/docs/src/ref/core/gfi.md similarity index 93% rename from docs/src/tutorials/basics/gfi.md rename to docs/src/ref/core/gfi.md index 723f07865..fc41e6b13 100644 --- a/docs/src/tutorials/basics/gfi.md +++ b/docs/src/ref/core/gfi.md @@ -1,14 +1,13 @@ # [Generative Function Interface](@id gfi) -One of the core abstractions in Gen is the **generative function**. +One of the core abstractions in Gen is the **generative function**. The interface for interacting with generative functions is called the **generative function interface (GFI)** . 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 dynamic_modeling_language) allows generative functions to be constructed using Julia function definition syntax: + +For example, the [Built-in Modeling Language](@ref dml) allows generative functions to be constructed using Julia function definition syntax: ```@example gfi_tutorial using Gen # hide @gen function foo(a, b=0) @@ -17,9 +16,9 @@ using Gen # hide else return a + b end -end +end; ``` -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. +Users can also extend Gen by implementing their own [custom generative functions](@ref custom_gen_fns), 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: @@ -253,7 +252,6 @@ 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. @@ -268,10 +266,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_tutorial -get_choices(new_trace) -``` - +``` The weight (`w`) is ``\log 1 = 0``. @@ -305,8 +300,7 @@ Generative functions may also be able to process more specialized diff data type 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](@id differentiable_programming_tutorial) - +## 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). @@ -321,11 +315,50 @@ To compute gradients with respect to the arguments, and to increment a stateful 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). + +## API + +The following GFI methods should be implemented for a [`Trace`](@ref): + +```@docs +Trace +get_args +get_retval +get_choices +get_score +get_gen_fn +Base.getindex +``` + +The following GFI methods should be implemented for a [`GenerativeFunction`](@ref) and its associated [`Trace`](@ref) datatype: + +```@docs +GenerativeFunction +simulate +generate +update +regenerate +project +propose +assess +``` + +Generative functions that support gradient computation with respect to arguments or trainable parameters should implement the following static properties: + +```@docs +has_argument_grads +accepts_output_grad +choice_gradients +accumulate_param_gradients! +get_params +``` diff --git a/docs/src/api/model/selections.md b/docs/src/ref/core/selections.md similarity index 94% rename from docs/src/api/model/selections.md rename to docs/src/ref/core/selections.md index 86b746bbc..3692d48b4 100644 --- a/docs/src/api/model/selections.md +++ b/docs/src/ref/core/selections.md @@ -45,9 +45,10 @@ complement The [`select`](@ref) method returns a selection with concrete type [`DynamicSelection`](@ref). The [`selectall`](@ref) method returns a selection with concrete type [`AllSelection`](@ref). -The full list of concrete types of selections is shown below. -Most users need not worry about these types. + +The full list of concrete types of selections is shown below. Most users need not worry about these types. Note that only selections of type [`DynamicSelection`](@ref) are mutable (using `push!` and `set_subselection!`). + ```@docs EmptySelection AllSelection @@ -56,6 +57,7 @@ DynamicSelection StaticSelection ComplementSelection ``` + ```@docs Gen.get_address_schema ``` diff --git a/docs/src/api/inference/importance.md b/docs/src/ref/inference/importance.md similarity index 98% rename from docs/src/api/inference/importance.md rename to docs/src/ref/inference/importance.md index e1d69fc29..8cbb97ab5 100644 --- a/docs/src/api/inference/importance.md +++ b/docs/src/ref/inference/importance.md @@ -1,4 +1,5 @@ # Importance Sampling + ```@docs importance_sampling importance_resampling diff --git a/docs/src/ref/inference/map.md b/docs/src/ref/inference/map.md new file mode 100644 index 000000000..ef8e36c3f --- /dev/null +++ b/docs/src/ref/inference/map.md @@ -0,0 +1,9 @@ +# MAP Optimization + +In contrast to [parameter optimization](parameter_optimization.md), which optimizes parameters of a generative function that are not associated with any prior distribution, *maximum a posteriori* (MAP) optimization can be used to maximize the posterior probability of a selection of traced random variables: + +```@docs +map_optimize +``` + +To use `map_optimize`, a trace of a generative function should be first created using the [`generate`](@ref) method with the appropriate observations. Users may also implement more complex optimization algorithms beyond backtracking gradient ascent by using the gradients returned by [`choice_gradients`](@ref). Note that if the selected random variables have bounded support over the real line, errors may occur during gradient-based optimization. diff --git a/docs/src/how_to/mcmc_kernels.md b/docs/src/ref/inference/mcmc.md similarity index 97% rename from docs/src/how_to/mcmc_kernels.md rename to docs/src/ref/inference/mcmc.md index 12195d28d..a36332b28 100644 --- a/docs/src/how_to/mcmc_kernels.md +++ b/docs/src/ref/inference/mcmc.md @@ -1,4 +1,4 @@ -# How to Write Markov Chain Monte Carlo Kernels +# [Markov Chain Monte Carlo (MCMC)](@id 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.[^1] @@ -198,7 +198,7 @@ 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 +## [Involutive MCMC](@id 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. @@ -208,7 +208,7 @@ These kernels are particularly useful for inferring the structure (e.g. control 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 language](@ref dynamic_modeling_language), 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*. +- 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 language](@ref dml), 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``. @@ -345,13 +345,12 @@ 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) +![rjmcmc plot](../../assets/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 @@ -363,4 +362,26 @@ For custom primitive kernels declared with [`@pkern`](@ref), users can declare t ``` 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 +The reversal of a kernel (primitive or composite) can be obtained with [`reversal`](@ref). + +## API + +Built-in stationary kernels: + +```@docs +metropolis_hastings +mh +mala +hmc +elliptical_slice +involutive_mcmc +``` + +Macros and functions for defining custom stationary kernels: + +```@docs +@pkern +@kern +@rkern +reversal +``` diff --git a/docs/src/api/model/parameter_optimization.md b/docs/src/ref/inference/parameter_optimization.md similarity index 64% rename from docs/src/api/model/parameter_optimization.md rename to docs/src/ref/inference/parameter_optimization.md index 50f4b323c..8865a5c50 100644 --- a/docs/src/api/model/parameter_optimization.md +++ b/docs/src/ref/inference/parameter_optimization.md @@ -1,4 +1,6 @@ -# Trainable Parameters(@trainable_parameter_optimization) +# Parameter Optimization + +Gen provides support for gradient-based optimization of the trainable parameters of generative functions, e.g. for maximixum likelihood learning, expectation-maximization, or empirical Bayes. 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). @@ -12,17 +14,19 @@ Gradient-based optimization of the trainable parameters of generative functions ## Parameter update A *parameter update* reads from the gradient accumulators for certain trainable parameters, updates the values of those parameters, and resets the gradient accumulators to zero. + A paramter update is constructed by combining an *update configuration* with the set of trainable parameters to which the update should be applied: + ```@docs ParamUpdate ``` -The set of possible update configurations is described in [Update configurations](@ref). -An update is applied with: +The set of possible update configurations is described in [Update configurations](@ref update_configurations).An update is applied with: + ```@docs apply! ``` -## Update configurations +## [Update configurations](@id update_configurations) Gen has built-in support for the following types of update configurations. ```@docs @@ -30,5 +34,18 @@ FixedStepGradientDescent GradientDescent ADAM ``` -For adding new types of update configurations, see [Optimizing Trainable Parameters (Internal)](@ref optimizing-internal). +Custom update configurations should implement the following methods: + +```@docs +Gen.init_update_state +Gen.apply_update! +``` + +## Training generative functions + +The [`train!`](@ref) method can be used to train the parameters of a generative function to maximize the likelihood of a dataset defined by a `data_generator`: + +```@docs +train! +``` diff --git a/docs/src/ref/inference/pf.md b/docs/src/ref/inference/pf.md new file mode 100644 index 000000000..6a7ad8494 --- /dev/null +++ b/docs/src/ref/inference/pf.md @@ -0,0 +1,37 @@ +# [Particle Filtering and Sequential Monte Carlo](@id particle_filtering) + +Gen.jl provides support for Sequential Monte Carlo (SMC) inference in the form of particle filtering. +The state of a particle filter is a represented as a `ParticleFilterState` object. + +```@docs +Gen.ParticleFilterState +``` + +## Particle Filtering Steps + +The basic steps of particle filtering are *initialization* (via [`initialize_particle_filter`](@ref)), *updating* (via [`particle_filter_step!`](@ref)), and *resampling* (via [`maybe_resample!`](@ref)). The latter two operations are applied to a [`ParticleFilterState`](@ref), modifying it in place. + +```@docs +initialize_particle_filter +particle_filter_step! +maybe_resample! +``` + +## Accessors + +The following accessor functions can be used to return information about a [`ParticleFilterState`](@ref), or to sample traces from the distribution that the particle filter approximates. + +```@docs +log_ml_estimate +get_traces +get_log_weights +sample_unweighted_traces +``` + +## [Advanced Particle Filtering](@id advanced-particle-filtering) + +For a richer set of particle filtering techniques, including support for stratified sampling, multiple resampling methods, MCMC rejuvenation moves, particle filter resizing, users are recommended to use the [GenParticleFilters.jl](https://github.com/probcomp/GenParticleFilters.jl) extension library. + +To use the generalization of standard SMC known as Sequential Monte Carlo with Probabilistic Program Proposals (SMCP³), use the API provided by [GenSMCP3.jl](https://github.com/probcomp/GenSMCP3.jl), or implement an [`UpdatingTraceTranslator`](https://probcomp.github.io/GenParticleFilters.jl/stable/translate/#GenParticleFilters.UpdatingTraceTranslator) in GenParticleFilters.jl. + +Even more advanced SMC techniques (such as divide-and-conquer SMC) are not currently supported by Gen. diff --git a/docs/src/tutorials/trace_translators.md b/docs/src/ref/inference/trace_translators.md similarity index 97% rename from docs/src/tutorials/trace_translators.md rename to docs/src/ref/inference/trace_translators.md index 47d00dd1c..dba99168e 100644 --- a/docs/src/tutorials/trace_translators.md +++ b/docs/src/ref/inference/trace_translators.md @@ -1,7 +1,8 @@ # Trace Translators -While [generative functions](@ref gfi_api) 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). +While [generative functions](@ref gfi) 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 involutive_mcmc) or [SMCP³](@ref advanced-particle-filtering). 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). @@ -294,7 +295,7 @@ 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). +Symmetric trace translators are important because they form the basis of [Involutive MCMC](@ref involutive_mcmc). 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. @@ -310,7 +311,7 @@ This has two benefits when the previous and new traces have random choices that 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, +is the basic operation used in [Particle Filtering](pf.md). For example, we might have a model that sequentially samples new latent variables `(:z, t)` and observations `(:x, t)` for each timestep `t`: @@ -433,3 +434,25 @@ 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 + + +```@docs +@transform +@read +@write +@copy +@tcall +pair_bijections! +is_involution! +inverse +TraceTranslator +DeterministicTraceTranslator +GeneralTraceTranslator +SimpleExtendingTraceTranslator +SymmetricTraceTranslator +``` +```@docs +TraceTransformDSLProgram +``` diff --git a/docs/src/ref/inference/vi.md b/docs/src/ref/inference/vi.md new file mode 100644 index 000000000..9eca22602 --- /dev/null +++ b/docs/src/ref/inference/vi.md @@ -0,0 +1,8 @@ +# Variational Inference + +There are two procedures in the inference library for performing black box variational inference, including amortized 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! +``` diff --git a/docs/src/ref/inference/wake_sleep.md b/docs/src/ref/inference/wake_sleep.md new file mode 100644 index 000000000..5bf4680b5 --- /dev/null +++ b/docs/src/ref/inference/wake_sleep.md @@ -0,0 +1,10 @@ +# Wake-Sleep Learning + +Wake-sleep learning is algorithm that jointly learns a generative model, represented as a generative function `p` and an inference model, represented as an inference function `q`. + +During the wake phase, `p` is trained on complete traces generated by running `q` on the observed data, which can be done using the [`train!](@ref) method. During the sleep phase, `q` is trained on data generated by simulating from `p`, which can be done using [`lecture!`](@ref) or [`lecture_batched!`](@ref): + +```@docs +lecture! +lecture_batched! +``` diff --git a/docs/src/explanations/language_implementation.md b/docs/src/ref/internals/language_implementation.md similarity index 92% rename from docs/src/explanations/language_implementation.md rename to docs/src/ref/internals/language_implementation.md index a5434196c..2b9b7bc4f 100644 --- a/docs/src/explanations/language_implementation.md +++ b/docs/src/ref/internals/language_implementation.md @@ -18,6 +18,10 @@ For clarity, below is a procedural description of how the `@gen` macro processes 4. For static `@gen` functions, match `:gentrace` expressions when adding address nodes to the static computation graph, and match `:genparam` expressions when adding parameter nodes to the static computation graph. A `StaticIRGenerativeFunction` is then compiled from the static computation graph. 5. For dynamic `@gen` functions, rewrite any `:gentrace` expression with its implementation `dynamic_trace_impl`, and rewrite any `:genparam` expression with its implementation `dynamic_param_impl`. The rewritten syntax tree is then evaluated as a standard Julia function, which serves as the implementation of the constructed `DynamicDSLFunction`. +## Dynamic Modeling Language + +The following methods are used to implement the semantics of the DML via non-standard interpretation: + ```@docs Gen.make_dynamic_gen_function Gen.dynamic_param_impl @@ -28,7 +32,9 @@ Gen.escape_default Gen.state Gen.choice_or_call_at ``` -## Static Modeling Language Tooling +## Static Modeling Language + +The following methods are used to parse functions written in the SML into static computation graphs: ```@docs Gen.StaticIRGenerativeFunction diff --git a/docs/src/api/model/combinators.md b/docs/src/ref/modeling/combinators.md similarity index 87% rename from docs/src/api/model/combinators.md rename to docs/src/ref/modeling/combinators.md index 229896523..5d7967d72 100644 --- a/docs/src/api/model/combinators.md +++ b/docs/src/ref/modeling/combinators.md @@ -1,4 +1,4 @@ -# Generative Function Combinators +# [Generative Function Combinators](@id 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. @@ -14,7 +14,7 @@ Map In the schematic below, the kernel is denoted ``\mathcal{G}_{\mathrm{k}}``. ```@raw html
- schematic of map combinator + schematic of map combinator
``` @@ -61,7 +61,7 @@ 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 + schematic of unfold combinator
``` @@ -118,7 +118,7 @@ Recurse ```@raw html
- schematic of recurse combinatokr + schematic of recurse combinatokr
``` ## Switch combinator @@ -129,7 +129,7 @@ Switch ```@raw html
- schematic of switch combinator + schematic of switch combinator
``` @@ -164,3 +164,13 @@ The resulting trace contains the subtrace from the branch with index `2` - in th └── :z : 13.552870875213735 ``` +## Design and Implementation + +Internally, the Combinators use custom trace types such as [`Gen.VectorTrace`](@ref), and are implemented using the following methods: + +```@docs +Gen.VectorTrace +Gen.process_all_new! +Gen.update_recurse_merge +Gen.update_discard +``` diff --git a/docs/src/ref/modeling/custom_gen_fns.md b/docs/src/ref/modeling/custom_gen_fns.md new file mode 100644 index 000000000..099a2b4dc --- /dev/null +++ b/docs/src/ref/modeling/custom_gen_fns.md @@ -0,0 +1,29 @@ +# [Custom Generative Functions](@id custom_gen_fns) + +Gen provides scaffolding for custom (deterministic) generative functions that make use of either incremental computation or custom gradient updates. + +```@docs +CustomDetermGF +CustomDetermGFTrace +``` + +## Custom Incremental Computation + +A [`CustomUpdateGF`](@ref) is a generative function that allows for easier implementation of custom incremental computation. + +```@docs +CustomUpdateGF +apply_with_state +update_with_state +num_args +``` + +## Custom Gradient Computations + +A [`CustomGradientGF`](@ref) is a generative function that allows for easier implementation of custom gradients computations and updates. + +```@docs +CustomGradientGF +apply +gradient +``` diff --git a/docs/src/api/model/distributions.md b/docs/src/ref/modeling/distributions.md similarity index 88% rename from docs/src/api/model/distributions.md rename to docs/src/ref/modeling/distributions.md index d3828b59f..9dfb710bb 100644 --- a/docs/src/api/model/distributions.md +++ b/docs/src/ref/modeling/distributions.md @@ -1,28 +1,21 @@ # [Probability Distributions](@id distributions) -```@docs -random -logpdf -has_output_grad -logpdf_grad -``` - Gen provides a library of built-in probability distributions, and four ways of -defining custom distributions, each of which are explained below: - -1. The [`@dist` constructor](@ref dist_dsl), for a distribution that can be expressed as a - simple deterministic transformation (technically, a - [pushforward](https://en.wikipedia.org/wiki/Pushforward_measure)) of an - existing distribution. +constructing custom distributions, each of which are explained below: -2. The [`HeterogeneousMixture`](@ref) and [`HomogeneousMixture`](@ref) constructors +1. The [`HeterogeneousMixture`](@ref) and [`HomogeneousMixture`](@ref) constructors for distributions that are mixtures of other distributions. -3. The [`ProductDistribution`](@ref) constructor for distributions that are products of +2. The [`ProductDistribution`](@ref) constructor for distributions that are products of other distributions. +3. The [`@dist` constructor](@ref dist_dsl), for a distribution that can be expressed as a + simple deterministic transformation (technically, a + [pushforward](https://en.wikipedia.org/wiki/Pushforward_measure)) of an + existing distribution. + 4. An API for defining arbitrary [custom distributions](@ref - custom_distributions) in plain Julia code. + custom_distributions_howto) in plain Julia code. ## Built-In Distributions @@ -49,7 +42,22 @@ uniform_discrete broadcasted_normal ``` -## [Defining New Distributions Inline with the `@dist` DSL](@id dist_dsl) +## Mixture Distribution Constructors + +There are two built-in constructors for defining mixture distributions: +```@docs +HomogeneousMixture +HeterogeneousMixture +``` + +## Product Distribution Constructors + +There is a built-in constructor for defining product distributions: +```@docs +ProductDistribution +``` + +## [The `@dist` DSL](@id dist_dsl) The `@dist` DSL allows the user to concisely define a distribution, as long as that distribution can be expressed as a certain type of deterministic @@ -70,7 +78,8 @@ end Here `body` is ordinary Julia code, with the constraint that `body` must contain exactly one random choice. The value of the `@dist` expression is then a `Gen.Distribution` object called `name`, parameterized by `arg1, ..., argN`, -representing the distribution over _return values_ of `body`. +representing the distribution over _return values_ of `body`. Arguments are +optionally typed. This DSL is designed to address the issue that sometimes, values stored in the trace do not correspond to the most natural physical elements of the model @@ -223,26 +232,14 @@ log(normal(exp(x), exp(x))) :: RND (by rule 6) log(normal(exp(x), exp(x))) + (x * (2 + 3)) :: RND (by rule 6) ``` -## Mixture Distribution Constructors +## API -There are two built-in constructors for defining mixture distributions: ```@docs -HomogeneousMixture -HeterogeneousMixture -``` - -## Product Distribution Constructors - -There is a built-in constructor for defining product distributions: -```@docs -ProductDistribution +random +logpdf +logpdf_grad +has_output_grad +is_discrete ``` -## Defining New Distributions From Scratch - -For distributions that cannot be expressed in the `@dist` DSL, users can define -a custom distribution by defining an (ordinary Julia) subtype of -`Gen.Distribution` and implementing the methods of the [Distribution API](@ref -custom_distributions). This method requires more custom code than using the -`@dist` DSL, but also affords more flexibility: arbitrary user-defined logic -for sampling, PDF evaluation, etc. +The [`has_argument_grads`](@ref) function is also part of the distribution API. diff --git a/docs/src/api/model/modeling.md b/docs/src/ref/modeling/dml.md similarity index 75% rename from docs/src/api/model/modeling.md rename to docs/src/ref/modeling/dml.md index 2e3c319a2..d867f7f92 100644 --- a/docs/src/api/model/modeling.md +++ b/docs/src/ref/modeling/dml.md @@ -1,10 +1,10 @@ -# [The Dynamic Modeling Language](@id dynamic_modeling_language) +# [Built-in Modeling Language](@id dml) -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**. +Gen provides a built-in modeling language for defining generative functions, using a syntax that extends Julia's syntax for defining regular Julia functions. This language is also referred to as the **Dynamic Modeling Language (DML)**. -Generative functions in the modeling language are identified using the `@gen` keyword in front of a Julia function definition. +Generative functions in this modeling language are identified using the `@gen` keyword in front of a Julia function definition. Here is an example `@gen` function that samples two random choices: + ```julia @gen function foo(prob::Float64=0.1) z1 = @trace(bernoulli(prob), :a) @@ -12,7 +12,9 @@ Here is an example `@gen` function that samples two random choices: return z1 || z2 end ``` + After running this code, `foo` is a Julia value of type [`DynamicDSLFunction`](@ref): + ```@docs DynamicDSLFunction ``` @@ -20,20 +22,26 @@ DynamicDSLFunction Note that it is possible to provide default values for trailing positional arguments. However, keyword arguments are currently *not* supported. We can call the resulting generative function like we would a regular Julia function: + ```julia retval::Bool = foo(0.5) ``` + We can also trace its execution: + ```julia (trace, _) = generate(foo, (0.5,)) ``` + Optional arguments can be left out of the above operations, and default values will be filled in automatically: + ```julia julia> (trace, _) = generate(foo, ()) julia> get_args(trace) (0.1,) ``` -See [Generative Functions](@ref gfi_api) for the full set of operations supported by a generative function. + +See [Generative Functions](@ref gfi) 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 language-implementation) section. @@ -382,120 +390,3 @@ This is indicated by the [`has_output_grad`](@ref) function. Like distributions, generative functions indicate which of their arguments support differentiation, using the `has_argument_grads` function. It is an error if a tracked value is passed as an argument of a generative function, when differentiation is not supported by the generative function for that argument. 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](@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. - -A function is identified as using the static modeling language by adding the `static` annotation to the function. -For example: -```julia -@gen (static) function foo(prob::Float64) - z1 = @trace(bernoulli(prob), :a) - z2 = @trace(bernoulli(prob), :b) - z3 = z1 || z2 - z4 = !z3 - return z4 -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 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 -
- example static computation graph -
-``` -In this graph, oval nodes represent random choices, square nodes represent Julia computations, and diamond nodes represent arguments. -The light blue shaded node is the return value of the function. -Having access to the static graph allows Gen to generate specialized code for [Updating traces](@ref) that skips unecessary parts of the computation. -Specifically, when applying an update operation, the graph is analyzed, and each value in the graph identified as having possibly changed, or not. -Nodes in the graph do not need to be re-executed if none of their input values could have possibly changed. -Also, even if some inputs to a generative function node may have changed, knowledge that some of the inputs have not changed often allows the generative function being called to more efficiently perform its update operation. -This is the case for functions produced by [Generative Function Combinators](@ref). - -You can plot the graph for a function with the `static` annotation if you have PyCall installed, and a Python environment that contains the [graphviz](https://pypi.org/project/graphviz/) Python package, using, e.g.: -```julia -using PyCall -@pyimport graphviz -using Gen: draw_graph -draw_graph(foo, graphviz, "test") -``` -This will produce a file `test.pdf` in the current working directory containing the rendered graph. - -### Restrictions - -First, the definition of a `(static)` generative function is always expected to occur as a [top-level definition](https://docs.julialang.org/en/v1/manual/modules/) (aka global variable); usage in non–top-level scopes is unsupported and may result in incorrect behavior. - -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 trainable_parameters_modeling), e.g.: - -```julia -@param theta::Float64 -``` - -- An assignment, with a symbol or tuple of symbols on the left-hand side, and a Julia expression on the right-hand side, which may include `@trace` expressions, e.g.: - -```julia -mu, sigma = @trace(bernoulli(p), :x) ? (mu1, sigma1) : (mu2, sigma2) -``` - -- A top-level `@trace` expression, e.g.: - -```julia -@trace(bernoulli(1-prob_tails), :flip) -``` -All `@trace` expressions must use a literal Julia symbol for the first component in the address. Unlike the full built-in modeling-language, the address is not optional. - -- A `return` statement, with a Julia expression on the right-hand side, e.g.: - -```julia -return @trace(geometric(prob), :n_flips) + 1 -``` - -The functions are also subject to the following restrictions: - -- Default argument values are not supported. - -- Constructing named or anonymous Julia functions (and closures) is not allowed. - -- List comprehensions with internal `@trace` calls are not allowed. - -- Splatting within `@trace` calls is not supported - -- Generative functions that are passed in as arguments cannot be traced. - -- For composite addresses (e.g. `:a => 2 => :c`) the first component of the address must be a literal symbol, and there may only be one statement in the function body that uses this symbol for the first component of its address. - -- Julia control flow constructs (e.g. `if`, `for`, `while`) cannot be used as top-level statements in the function body. Control flow should be implemented inside either Julia functions that are called, or other generative functions. - -Certain loop constructs can be implemented using [Generative Function Combinators](@ref) instead. For example, the following loop: -```julia -for (i, prob) in enumerate(probs) - @trace(foo(prob), :foo => i) -end -``` -can instead be implemented as: -```julia -@trace(Map(foo)(probs), :foo) -``` - -### Loading generated functions - -Once a `(static)` generative function is defined, it can be used in the same way as a non-static generative function. In previous versions of Gen, the `@load_generated_functions` macro had to be called before a function with a `(static)` annotation could be used. This macro is no longer necessary, and will be removed in a future release. - -### Performance tips -For better performance when the arguments are simple data types like `Float64`, annotate the arguments with the concrete type. -This permits a more optimized trace data structure to be generated for the generative function. - -### Caching Julia values - -By default, the values of Julia computations (all calls that are not random choices or calls to generative functions) are cached as part of the trace, so that [Updating traces](@ref) can avoid unecessary re-execution of Julia code. -However, this cache may grow the memory footprint of a trace. -To disable caching of Julia values, use the function annotation `nojuliacache` (this annotation is ignored unless the `static` function annotation is also used). diff --git a/docs/src/ref/modeling/sml.md b/docs/src/ref/modeling/sml.md new file mode 100644 index 000000000..f10c1c714 --- /dev/null +++ b/docs/src/ref/modeling/sml.md @@ -0,0 +1,115 @@ +# [Static Modeling Language](@id sml) + +The **Static Modeling Language (SML)** is a restricted variant of the built-in [dynamic modeling language](@ref dml). +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](@ref mcmc) or [particle filtering](@ref particle_filtering). + +A function is identified as using the static modeling language by adding the `static` annotation to the function. +For example: +```julia +@gen (static) function foo(prob::Float64) + z1 = @trace(bernoulli(prob), :a) + z2 = @trace(bernoulli(prob), :b) + z3 = z1 || z2 + z4 = !z3 + return z4 +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 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 +
+ example static computation graph +
+``` +In this graph, oval nodes represent random choices, square nodes represent Julia computations, and diamond nodes represent arguments. +The light blue shaded node is the return value of the function. +Having access to the static graph allows Gen to generate specialized code for [Updating traces](@ref) that skips unecessary parts of the computation. +Specifically, when applying an update operation, the graph is analyzed, and each value in the graph identified as having possibly changed, or not. +Nodes in the graph do not need to be re-executed if none of their input values could have possibly changed. +Also, even if some inputs to a generative function node may have changed, knowledge that some of the inputs have not changed often allows the generative function being called to more efficiently perform its update operation. +This is the case for functions produced by [Generative Function Combinators](@ref combinators). + +You can plot the graph for a function with the `static` annotation if you have PyCall installed, and a Python environment that contains the [graphviz](https://pypi.org/project/graphviz/) Python package, using, e.g.: +```julia +using PyCall +@pyimport graphviz +using Gen: draw_graph +draw_graph(foo, graphviz, "test") +``` +This will produce a file `test.pdf` in the current working directory containing the rendered graph. + +## Restrictions + +First, the definition of a `(static)` generative function is always expected to occur as a [top-level definition](https://docs.julialang.org/en/v1/manual/modules/) (aka global variable); usage in non–top-level scopes is unsupported and may result in incorrect behavior. + +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 trainable_parameters_modeling), e.g.: + +```julia +@param theta::Float64 +``` + +- An assignment, with a symbol or tuple of symbols on the left-hand side, and a Julia expression on the right-hand side, which may include `@trace` expressions, e.g.: + +```julia +mu, sigma = @trace(bernoulli(p), :x) ? (mu1, sigma1) : (mu2, sigma2) +``` + +- A top-level `@trace` expression, e.g.: + +```julia +@trace(bernoulli(1-prob_tails), :flip) +``` +All `@trace` expressions must use a literal Julia symbol for the first component in the address. Unlike the full built-in modeling-language, the address is not optional. + +- A `return` statement, with a Julia expression on the right-hand side, e.g.: + +```julia +return @trace(geometric(prob), :n_flips) + 1 +``` + +The functions are also subject to the following restrictions: + +- Default argument values are not supported. + +- Constructing named or anonymous Julia functions (and closures) is not allowed. + +- List comprehensions with internal `@trace` calls are not allowed. + +- Splatting within `@trace` calls is not supported + +- Generative functions that are passed in as arguments cannot be traced. + +- For composite addresses (e.g. `:a => 2 => :c`) the first component of the address must be a literal symbol, and there may only be one statement in the function body that uses this symbol for the first component of its address. + +- Julia control flow constructs (e.g. `if`, `for`, `while`) cannot be used as top-level statements in the function body. Control flow should be implemented inside either Julia functions that are called, or other generative functions. + +Certain loop constructs can be implemented using [Generative Function Combinators](@ref combinators) instead. For example, the following loop: +```julia +for (i, prob) in enumerate(probs) + @trace(foo(prob), :foo => i) +end +``` +can instead be implemented as: +```julia +@trace(Map(foo)(probs), :foo) +``` + +## Loading generated functions + +Once a `(static)` generative function is defined, it can be used in the same way as a non-static generative function. In previous versions of Gen, the `@load_generated_functions` macro had to be called before a function with a `(static)` annotation could be used. This macro is no longer necessary, and will be removed in a future release. + +## Performance tips +For better performance when the arguments are simple data types like `Float64`, annotate the arguments with the concrete type. +This permits a more optimized trace data structure to be generated for the generative function. + +## Caching Julia values + +By default, the values of Julia computations (all calls that are not random choices or calls to generative functions) are cached as part of the trace, so that [Updating traces](@ref) can avoid unecessary re-execution of Julia code. +However, this cache may grow the memory footprint of a trace. +To disable caching of Julia values, use the function annotation `nojuliacache` (this annotation is ignored unless the `static` function annotation is also used). diff --git a/docs/src/tutorials/basics/combinators.md b/docs/src/tutorials/basics/combinators.md deleted file mode 100644 index ed7542c82..000000000 --- a/docs/src/tutorials/basics/combinators.md +++ /dev/null @@ -1,114 +0,0 @@ -# [Generative Combinators](@id combinators_tutorial) - -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: - -```@setup switch_combinator -using Gen -``` - -```@example switch_combinator -@gen function line(x) - z ~ normal(3*x+1,1.0) - return z -end - -@gen function outlier(x) - z ~ normal(3*x+1, 10.0) - return z -end - -switch_model = Switch(line, outlier) -``` - -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) -``` diff --git a/docs/src/tutorials/basics/vi.md b/docs/src/tutorials/basics/vi.md deleted file mode 100644 index ccc45158d..000000000 --- a/docs/src/tutorials/basics/vi.md +++ /dev/null @@ -1,10 +0,0 @@ -# [Variational Inference] (@id vi_tutorial) - -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. - -## 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 diff --git a/docs/src/getting_started/linear_regression.md b/docs/src/tutorials/getting_started.md similarity index 97% rename from docs/src/getting_started/linear_regression.md rename to docs/src/tutorials/getting_started.md index 73db75845..8ace2977c 100644 --- a/docs/src/getting_started/linear_regression.md +++ b/docs/src/tutorials/getting_started.md @@ -1,5 +1,4 @@ -# Getting Started with Gen.jl: Linear Regression - +# [Getting Started: Linear Regression](@id getting_started) 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. @@ -26,6 +25,7 @@ using Gen {"y-$i"} ~ normal(slope * x + intercept, 1) end end +nothing # hide ``` Second, we write an _inference program_ that implements an algorithm for manipulating the execution traces of the model. @@ -58,6 +58,7 @@ function my_inference_program(xs::Vector{Float64}, ys::Vector{Float64}, num_iter choices = get_choices(trace) return (choices[:slope], choices[:intercept]) end +nothing # hide ``` Finally, we run the inference program on some data, and get the results: diff --git a/docs/src/api/inference/learning.md b/docs/src/tutorials/learning_gen_fns.md similarity index 99% rename from docs/src/api/inference/learning.md rename to docs/src/tutorials/learning_gen_fns.md index daa62eb85..5ddc2f845 100644 --- a/docs/src/api/inference/learning.md +++ b/docs/src/tutorials/learning_gen_fns.md @@ -1,4 +1,4 @@ -# Learning Generative Functions +# [Learning Generative Functions](@id learning_gen_fns_tutorial) Learning and inference are closely related concepts, and the distinction between the two is not always clear. Often, **learning** refers to inferring long-lived unobserved quantities that will be reused across many problem instances (like a dynamics model for an entity that we are trying to track), whereas **inference** refers to inferring shorter-lived quantities (like a specific trajectory of a specific entity). @@ -227,11 +227,3 @@ It is possible to perform amortized variational inference using [`black_box_vi!` [7] Diederik P. Kingma, Max Welling: Auto-Encoding Variational Bayes. ICLR 2014 [Link](https://arxiv.org/pdf/1312.6114.pdf) - -## API - -```@docs -lecture! -lecture_batched! -train! -``` diff --git a/docs/src/tutorials/basics/modeling_in_gen.md b/docs/src/tutorials/modeling_in_gen.md similarity index 90% rename from docs/src/tutorials/basics/modeling_in_gen.md rename to docs/src/tutorials/modeling_in_gen.md index 7ec032a98..1a0c7b724 100644 --- a/docs/src/tutorials/basics/modeling_in_gen.md +++ b/docs/src/tutorials/modeling_in_gen.md @@ -26,22 +26,19 @@ 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** + **This tutorial 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 + solving the problem. In this tutorial, 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**, + **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 +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 @@ -62,14 +59,13 @@ enable improved performance, but are not covered here. Gen is a package for the Julia language. The package can be loaded with: - -```@example tutorial_1 +```@example modeling_tutorial 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. +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 tutorial uses the [Plots.jl](https://github.com/JuliaPlots/Plots.jl) Julia package for plotting. -```@example tutorial_1 +```@example modeling_tutorial using Plots ``` @@ -78,17 +74,15 @@ using Plots 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 +custom proposal distributions, and variational approximations (see the +[Gen documentation](../index.md) 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. - +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 +modeling DSL](@ref dml). Generative functions written in the built-in modeling DSL are based on Julia function definition syntax, but are prefixed with the `@gen` macro: @@ -166,7 +160,7 @@ 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 +```@example modeling_tutorial @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 @@ -197,14 +191,16 @@ of a gas as a function of its measured temperature. # It can sometimems be useful to return something # meaningful, however; here, we return the function `y`. return y -end; +end +nothing # hide ``` 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.]; +```@example modeling_tutorial +xs = [-5., -4., -3., -2., -1., 0., 1., 2., 3., 4., 5.] +nothing # hide ``` Given this vector, the generative function samples a random choice @@ -219,18 +215,10 @@ 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 +```@example modeling_tutorial 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 @@ -242,74 +230,59 @@ 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 +generative function and obtain its trace using the [`simulate`](@ref) method from the Gen API: - -```@example tutorial_1 +```@example modeling_tutorial trace = Gen.simulate(line_model, (xs,)); +nothing # hide ``` - 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 +```@example modeling_tutorial 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`](@ref): -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 +```@example modeling_tutorial 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`](@ref): -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 +```@example modeling_tutorial Gen.get_choices(trace) ``` - - We can pull out individual values from this map using Julia's subscripting syntax `[...]`: - -```@example tutorial_1 +```@example modeling_tutorial 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: - +We can also read the value of a random choice directly from the trace, without having to use [`get_choices`](@ref) first: -```@example tutorial_1 +```@example modeling_tutorial trace[:slope] ``` -The return value is also recorded in the trace, and is accessible with the `get_retval` API method: +The return value is also recorded in the trace, and is accessible with the [`get_retval`](@ref) API method: - -```@example tutorial_1 +```@example modeling_tutorial Gen.get_retval(trace) ``` Or we can access the return value directly from the trace via the syntactic sugar `trace[]`: - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial function render_trace(trace; show_data=true) # Pull out xs from the trace @@ -334,44 +307,35 @@ function render_trace(trace; show_data=true) end return fig -end; +end +nothing # hide ``` -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial function grid(renderer::Function, traces) Plots.plot(map(renderer, traces)...) -end; +end +nothing # hide ``` - -```@example tutorial_1 +```@example modeling_tutorial 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 @@ -379,10 +343,10 @@ 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)). +[`Gen.gamma`](@ref)). 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)). +[`Gen.uniform`](@ref)). The sine wave should implement: @@ -398,8 +362,7 @@ When calling `trace = Gen.simulate(sine_model, (xs,))`, the following choices sh We have provided some starter code for the sine wave model: - -```@example tutorial_1 +```julia @gen function sine_model(xs::Vector{Float64}) # < your code here, for sampling a phase, period, and amplitude > @@ -413,29 +376,14 @@ We have provided some starter code for the sine wave model: end return y # We return the y function so it can be used for plotting, below. -end; +end ``` - -```@example tutorial_1 -traces = [Gen.simulate(sine_model, (xs,)) for _=1:12]; +```@raw html +
Solution ``` - -```@example tutorial_1 -grid(render_trace, traces) -``` - - - - - - - - - -
- +end +nothing # hide +``` + +```@example modeling_tutorial +traces = [Gen.simulate(sine_model, (xs,)) for _=1:12]; +grid(render_trace, traces) +``` + +```@raw html +
+``` ## Posterior Inference @@ -464,24 +423,16 @@ 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 +```@example modeling_tutorial ys = [6.75003, 6.1568, 4.26414, 1.84894, 3.09686, 1.94026, 1.36411, -0.83959, -0.976, -1.93363, -2.91303]; +nothing # hide ``` - -```@example tutorial_1 +```@example modeling_tutorial scatter(xs, ys, color="black", label=nothing, title="Observed data (linear)", xlabel="X", ylabel="Y") +nothing # hide ``` - - - - - - - - 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. @@ -499,8 +450,7 @@ 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 +```@example modeling_tutorial function do_inference(model, xs, ys, amount_of_computation) # Create a choice map that maps model addresses (:y, i) @@ -515,45 +465,27 @@ function do_inference(model, xs, ys, amount_of_computation) # with our observations. (trace, _) = Gen.importance_resampling(model, (xs,), observations, amount_of_computation); return trace -end; +end +nothing # hide ``` We can run the inference program to obtain a trace, and then visualize the result: - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial function overlay(renderer, traces; same_data=true, args...) fig = renderer(traces[1], show_data=true, args...) @@ -568,24 +500,16 @@ function overlay(renderer, traces; same_data=true, args...) xlim=(xmin, xmax), ylim=(xmin, xmax)) end return fig -end; +end +nothing # hide ``` - -```@example tutorial_1 +```@example modeling_tutorial traces = [do_inference(line_model, xs, ys, 100) for _=1:10]; overlay(render_trace, traces) ``` - - - - - - - - --------------- +------------------ ### Exercise @@ -596,42 +520,30 @@ The results above were obtained for `amount_of_computation = 100`. Run the algor ### Exercise Consider the following data set. - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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. +Using the API method [`generate`](@ref), 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 +```@example modeling_tutorial constraints = Gen.choicemap() constraints[:slope] = 0. constraints[:intercept] = 0. @@ -639,17 +551,8 @@ constraints[:intercept] = 0. 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 @@ -660,9 +563,7 @@ 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 +```@example modeling_tutorial function predict_new_data(model, trace, new_xs::Vector{Float64}, param_addrs) # Copy parameter values from the inferred trace (`trace`) @@ -679,34 +580,23 @@ function predict_new_data(model, trace, new_xs::Vector{Float64}, param_addrs) # Pull out the y-values and return them. ys = [new_trace[(:y, i)] for i=1:length(new_xs)] return ys -end; +end +nothing # hide ``` To illustrate, we call the function above given the previous trace (which constrained slope and intercept to be zero). - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial function infer_and_predict(model, xs, ys, new_xs, param_addrs, num_traces, amount_of_computation) pred_ys = [] for i=1:num_traces @@ -714,105 +604,69 @@ function infer_and_predict(model, xs, ys, new_xs, param_addrs, num_traces, amoun push!(pred_ys, predict_new_data(model, trace, new_xs, param_addrs)) end pred_ys -end; +end +nothing # hide ``` To illustrate, we generate predictions at `[1., 2., 3.]` given one (approximate) posterior trace. - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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; +end +nothing # hide ``` Recall the original dataset for the line model. The x-coordinates span the interval -5 to 5. - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial new_xs = collect(range(-5, stop=10, length=100)); +nothing # hide ``` We generate and plot the predicted data: - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial @gen function line_model_fancy(xs::Vector{Float64}) slope = ({:slope} ~ normal(0, 1)) intercept = ({:intercept} ~ normal(0, 2)) @@ -831,8 +685,7 @@ end; Then, we compare the predictions using inference of the unmodified and modified models on the `ys` data set: - -```@example tutorial_1 +```@example modeling_tutorial 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") @@ -842,38 +695,21 @@ inferred_noise_plot = plot_predictions(xs, ys, new_xs, pred_ys; title="inferred 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 +```@example modeling_tutorial 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. ------------------------- @@ -884,7 +720,7 @@ Write a modified version of the sine model that makes noise into a random choice We have provided you with starter code: -```@example tutorial_1 +```julia @gen function sine_model_fancy(xs::Vector{Float64}) # < your code here > @@ -893,11 +729,34 @@ We have provided you with starter code: {(:y, i)} ~ normal(0., 0.1) # < edit this line > end return nothing -end; +end ``` +```@raw html +
Solution +``` + +```@example modeling_tutorial +@gen function sine_model_fancy(xs::Vector{Float64}) + period = ({:period} ~ gamma(5, 1)) + amplitude = ({:amplitude} ~ gamma(1, 1)) + phase = ({:phase} ~ uniform(0, 2*pi)) + noise = ({:noise} ~ gamma(1, 1)) + + function y(x) + return amplitude * sin(x * (2 * pi / period) + phase) + end + + for (i, x) in enumerate(xs) + {(:y, i)} ~ normal(y(x), noise) + end -```@example tutorial_1 + return y +end +nothing # hide +``` + +```@example modeling_tutorial # 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") @@ -909,8 +768,7 @@ inferred_noise_plot = plot_predictions(xs, ys_sine, new_xs, pred_ys; title="Infe Plots.plot(fixed_noise_plot, inferred_noise_plot) ``` - -```@example tutorial_1 +```@example modeling_tutorial # 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") @@ -922,24 +780,11 @@ inferred_noise_plot = plot_predictions(xs, ys_noisy, new_xs, pred_ys; title="Inf Plots.plot(fixed_noise_plot, inferred_noise_plot) ``` -
- +------------------------- ## Calling other generative functions @@ -968,7 +813,7 @@ In this case, it is best to provide an address (`{addr} ~ f(x)`): `f`'s random c be placed under the _namespace_ `addr`. -```@example tutorial_1 +```@example modeling_tutorial @gen function foo() {:y} ~ normal(0, 1) end @@ -988,83 +833,41 @@ end # will appear in our trace at the # hierarchical address `:z => :y`. {:z} ~ foo() -end; +end +nothing # hide ``` We first show the addresses sampled by `bar`: - -```@example tutorial_1 +```@example modeling_tutorial trace = Gen.simulate(bar, ()) Gen.get_choices(trace) ``` - - - - │ - ├── :y : 0.6475752978557953 - │ - └── :x : true - - - - And the addresses sampled by `bar_using_namespace`: - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial @gen function baz() {:a} ~ bar_using_namespace() end @@ -1074,31 +877,15 @@ trace = simulate(baz, ()) trace[:a => :z => :y] ``` - - - - 0.5622991768999647 - - - Note that the `=>` operator associated right, so this is equivalent to: - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial @gen function combined_model(xs::Vector{Float64}) if ({:is_line} ~ bernoulli(0.5)) # Call line_model_fancy on xs, and import @@ -1109,22 +896,20 @@ Now, we write a generative function that combines the line and sine models. It m # its random choices directly into our trace return ({*} ~ sine_model_fancy(xs)) end -end; +end +nothing # hide ``` We visualize some traces, and see that sometimes it samples linear data and other times sinusoidal data. - -```@example tutorial_1 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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]; @@ -1132,7 +917,6 @@ 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. ------- @@ -1143,7 +927,7 @@ Construct a data set for which it is ambiguous whether the line or sine wave mod 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 @@ -1156,28 +940,19 @@ nonparametric* model. We will consider two data sets: - -```@example tutorial_1 +```@example modeling_tutorial 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; +nothing # hide ``` - -```@example tutorial_1 +```@example modeling_tutorial 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 @@ -1187,15 +962,16 @@ 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 +```@example modeling_tutorial struct Interval l::Float64 u::Float64 end +nothing # hide ``` -```@example tutorial_1 +```@example modeling_tutorial abstract type Node end struct InternalNode <: Node @@ -1208,12 +984,13 @@ struct LeafNode <: Node value::Float64 interval::Interval end +nothing # hide ``` 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 +```@example modeling_tutorial @gen function generate_segments(l::Float64, u::Float64) interval = Interval(l, u) if ({:isleaf} ~ bernoulli(0.7)) @@ -1230,13 +1007,14 @@ We now write a generative function that randomly creates such a tree. Note the u right = ({:right} ~ generate_segments(mid, u)) return InternalNode(left, right, interval) end -end; +end +nothing # hide ``` We also define some helper functions to visualize traces of the `generate_segments` function. -```@example tutorial_1 +```@example modeling_tutorial function render_node!(node::LeafNode) plot!([node.interval.l, node.interval.u], [node.value, node.value], label=nothing, linewidth=5) end @@ -1244,41 +1022,34 @@ end function render_node!(node::InternalNode) render_node!(node.left) render_node!(node.right) -end; +end +nothing # hide ``` - -```@example tutorial_1 +```@example modeling_tutorial 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; +end +nothing # hide ``` 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial # get_value_at searches a binary tree for # the leaf node containing some value. function get_value_at(x::Float64, node::LeafNode) @@ -1303,13 +1074,13 @@ end {(:y, i)} ~ normal(get_value_at(x, node), noise) end return node -end; +end +nothing # hide ``` We write a visualization for `changepoint_model` below: - -```@example tutorial_1 +```@example modeling_tutorial function render_changepoint_model_trace(trace; show_data=true) xs = Gen.get_args(trace)[1] node = Gen.get_retval(trace) @@ -1319,43 +1090,26 @@ function render_changepoint_model_trace(trace; show_data=true) ys = [trace[(:y, i)] for i=1:length(xs)] scatter!(xs, ys, c="gray", label=nothing, alpha=0.3, markersize=3) end -end; +end +nothing # hide ``` 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 +```@example modeling_tutorial 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 +```@example modeling_tutorial 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. @@ -1363,14 +1117,11 @@ For inference about the complex data set, we use more computation. You can exper !!! warning The following cell may run for 2-3 minutes. - -```@example tutorial_1 +```julia 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. ------ @@ -1385,4 +1136,4 @@ Hint: The return value of `changepoint_model` is the tree of `Node` values. Walk ### 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 +Hint: You will need to guarantee that all addresses are unique. How can you label each node in a binary tree using an integer? diff --git a/docs/src/tutorials/reversible_jump.md b/docs/src/tutorials/reversible_jump.md deleted file mode 100644 index e69de29bb..000000000 diff --git a/docs/src/tutorials/model_optimizations/scaling_with_sml.md b/docs/src/tutorials/scaling_with_sml.md similarity index 84% rename from docs/src/tutorials/model_optimizations/scaling_with_sml.md rename to docs/src/tutorials/scaling_with_sml.md index 29ea8f438..261173d43 100644 --- a/docs/src/tutorials/model_optimizations/scaling_with_sml.md +++ b/docs/src/tutorials/scaling_with_sml.md @@ -1,14 +1,12 @@ -# [Scaling with the Static Modeling Language](@id scaling_with_sml_tutorial) +# [Speeding Up Inference 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** -- [Generative Combinators](@ref combinators_tutorial) +For prototyping models and working with dynamic structures, Gen's [Dynamic Modeling Language](@ref dml) 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 tutorial 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. -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. +This tutorial will take a robust regression model with outliers and optimize the speed of inference using the SML. -## Slow Inference Program Case Study +## A Slow Inference Program ```@example sml_tutorial using Gen @@ -33,6 +31,7 @@ using Plots end ys end +nothing # hide ``` We wrote a Markov chain Monte Carlo inference update for this model that performs updates on each of the 'global' parameters (noise, slope, intercept, and prob_outlier), as well as the 'local' `is_outlier` variable associated with each data point. The update takes a trace as input, and returns the new trace as output. We reproduce this here: @@ -57,6 +56,7 @@ function block_resimulation_update(tr) # Return the updated trace tr end +nothing # hide ``` We write a helper function that takes a vector of y-coordinates and populates a constraints choice map: @@ -69,6 +69,7 @@ function make_constraints(ys::Vector{Float64}) end constraints end +nothing # hide ``` Finally, we package this into an inference program that takes the data set of all x- and y-coordinates ,and returns a trace. We will be experimenting with different variants of the model, so we make the model an argument to this function: @@ -82,6 +83,7 @@ function block_resimulation_inference(model, xs, ys) end tr end +nothing # hide ``` Let's see how the running time of this inference program changes as we increase the number of data points. We don't expect the running time to depend too much on the actual values of the data points, so we just construct a random data set for each run: @@ -98,13 +100,12 @@ for n in ns push!(times, (time_ns() - start) / 1e9) end end # hide -nothing +nothing # hide ``` We now plot the running time versus the number of data points: ```@example sml_tutorial - plot(ns, times, xlabel="number of data points", ylabel="running time (seconds)", label=nothing) ``` @@ -119,13 +120,14 @@ n = length(xs) for i=1:n (tr, _) = mh(tr, select(:data => i => :is_outlier)) end +nothing # hide ``` The reason for the quadratic scaling is that the running time of the call to `mh` inside this loop also grows in proportion to the number of data points. This is because the updates to a trace of a model written the generic built-in modeling language always involve re-running **the entire** model generative function. However, it should be possible for the algorithm to scale linearly in the number of data points. Briefly, deciding whether to update a given `is_outlier` variable can be done without referencing the other data points. This is because each `is_outiler` variable is conditionally independent of the outlier variables and y-coordinates of the other data points, conditioned on the parameters. -We can make this conditional independence structure explicit using the [Map generative function combinator](https://www.gen.dev/docs/stable/ref/combinators/#Map-combinator-1). Combinators like map encapsulate common modeling patterns (e.g., a loop in which each iteration is making independent choices), and when you use them, Gen can take advantage of the restrictions they enforce to implement performance optimizations automatically during inference. The `Map` combinator, like the `map` function in a functional programming language, helps to execute the same generative code repeatedly. +We can make this conditional independence structure explicit using the [Map generative function combinator](@ref Gen.Map). Combinators like map encapsulate common modeling patterns (e.g., a loop in which each iteration is making independent choices), and when you use them, Gen can take advantage of the restrictions they enforce to implement performance optimizations automatically during inference. The `Map` combinator, like the `map` function in a functional programming language, helps to execute the same generative code repeatedly. ## Rewriting the Program with Combinators @@ -139,13 +141,15 @@ To use the map combinator to express the conditional independences in our model, std = is_outlier ? 10. : noise y ~ normal(mu, std) return y -end; +end +nothing # hide ``` -We then apply the [`Map`](https://www.gen.dev/docs/stable/ref/combinators/#Map-combinator-1), which is a Julia function, to this generative function, to obtain a new generative function: +We then apply the [`Map`](@ref), which is a Julia function, to this generative function, to obtain a new generative function: ```@example sml_tutorial -generate_all_points = Map(generate_single_point); +generate_all_points = Map(generate_single_point) +nothing # hide ``` This new generative function has one argument for each argument of `generate_single_point`, except that these arguments are now vector-valued instead of scalar-valued. We can run the generative function on some fake data to test this: @@ -157,6 +161,7 @@ noises = fill(0.2, 5) slopes = fill(0.7, 5) intercepts = fill(-2.0, 5) trace = simulate(generate_all_points, (xs, prob_outliers, noises, slopes, intercepts)); +nothing # hide ``` We see that the `generate_all_points` function has traced 5 calls to `generate_single_point`, under namespaces `1` through `5`. The `Map` combinator automatically adds these indices to the trace address. @@ -176,7 +181,8 @@ Now, let's replace the Julia `for` loop in our model with a call to this new fun n = length(xs) data ~ generate_all_points(xs, fill(prob_outlier, n), fill(noise, n), fill(slope, n), fill(intercept, n)) return data -end; +end +nothing # hide ``` Note that this new model has the same address structure as our original model had, so our inference code will not need to change. For example, the 5th data point's $y$ coordinate will be stored at the address `:data => 5 => :y`, just as before. (The `:data` comes from our `data ~ ...` invocation in the `better_model` definition, and the `:y` comes from `generate_point`; only the `5` has been inserted automatically by `Map`.) @@ -197,6 +203,7 @@ for n in ns tr = block_resimulation_inference(model_with_map, xs, ys) push!(with_map_times, (time_ns() - start) / 1e9) end +nothing # hide ``` We plot the results and compare them to the original model, which used the Julia `for` loop: @@ -206,11 +213,9 @@ plot(ns, times, label="original", xlabel="number of data points", ylabel="runnin plot!(ns, with_map_times, label="with map") ``` -We see that the quadratic scaling did not improve. In fact, we actually got a that happed was a constant factor **slowdown**. +We see that the quadratic scaling did not improve. We can understand why we still have quadratic scaling, by examining the call to `generate_single_point`: -We can understand why we still have quadratic scaling, by examining the call to `generate_single_point`: - -```@example sml_tutorial +```julia data ~ generate_all_points(xs, fill(prob_outlier, n), fill(noise, n), fill(slope, n), fill(intercept, n)) ``` @@ -218,7 +223,7 @@ Even though the function `generate_all_points` knows that each of the calls to ` ## Rewriting in the Static Modeling Language -In order to provide `generate_all_points` with the knowledge that its arguments do not change during an update to the `is_outlier` variable, we need to write the top-level model generative function that calls `generate_all_points` in the [Static Modeling Language](https://www.gen.dev/docs/stable/ref/modeling/#Static-Modeling-Language-1), which is a restricted variant of the built-in modeling language that uses static analysis of the computation graph to generate specialized trace data structures and specialized implementations of trace operations. We indicate that a function is to be interpreted using the static language using the `static` annotation: +In order to provide `generate_all_points` with the knowledge that its arguments do not change during an update to the `is_outlier` variable, we need to write the top-level model generative function that calls `generate_all_points` in the [Static Modeling Language](@ref sml), which is a restricted variant of the built-in modeling language that uses static analysis of the computation graph to generate specialized trace data structures and specialized implementations of trace operations. We indicate that a function is to be interpreted using the static language using the `static` annotation: ```@example sml_tutorial @gen (static) function static_model_with_map(xs::Vector{Float64}) @@ -230,28 +235,27 @@ In order to provide `generate_all_points` with the knowledge that its arguments data ~ generate_all_points(xs, fill(prob_outlier, n), fill(noise, n), fill(slope, n), fill(intercept, n)) return data end +nothing # hide ``` The static language has a number of restrictions that make it more amenable to static analysis than the unrestricted modeling language. For example, we cannot use Julia `for` loops, and the return value needs to explicitly use the `return` keyword, followed by a symbol (e.g. `data`). Also, each symbol used on the left-hand side of an assignment must be unique. A more complete list of restrictions is given in the documentation. Below, we show the static dependency graph that Gen builds for this function. Arguments are shown as diamonds, Julia computations are shown as squares, random choices are shown as circles, and calls to other generative function are shown as stars. The call that produces the return value of the function is shaded in blue. - +```@raw html +
+ static dependency graph +
+``` Now, consider the update to the `is_outlier` variable: -```@example sml_tutorial +```julia (tr, _) = mh(tr, select(:data => i => :is_outlier)) ``` Because this update only causes values under address `:data` to change, the `static_model_with_map` function can use the graph above to infer that none of the arguments to `generate_all_point` could have possibly changed. This will allow us to obtain the linear scaling we expected. -However, before we can use a function written in the static modeling language, we need to run the following function (this is required for technical reasons, because functions written in the static modeling language use a staged programming feature of Julia called *generated functions*). - -``` -Gen.@load_generated_functions -``` - -Finally, we can re-run the experiment with our model that combines the map combinator with the static language: +Now we can re-run the experiment with our model that combines the map combinator with the static language: ```@example sml_tutorial static_with_map_times = [] @@ -262,7 +266,7 @@ for n in ns tr = block_resimulation_inference(static_model_with_map, xs, ys) push!(static_with_map_times, (time_ns() - start) / 1e9) end -nothing +nothing # hide ``` We compare the results to the results for the earlier models: @@ -302,8 +306,7 @@ static_generate_all_points = Map(static_generate_single_point); data ~ static_generate_all_points(xs, fill(prob_outlier, n), fill(noise, n), fill(slope, n), fill(intercept, n)) return data end; - -Gen.@load_generated_functions +nothing # hide ``` Now, we re-run the experiment with our new model: @@ -319,9 +322,10 @@ for n in ns push!(fully_static_with_map_times, (time_ns() - start) / 1e9) end end # hide +nothing # hide ``` -In earlier versions of Julia, we saw a modest improvement in running time, but here (running Julia 1.7.1) we see it makes little to no difference: +In earlier versions of Julia, we saw a modest improvement in running time, but here we see it makes little to no difference: ```@example sml_tutorial plot(ns, times, label="original", xlabel="number of data points", ylabel="running time (seconds)") @@ -370,6 +374,7 @@ function render_trace(trace, title) ys = [trace[:data => i => :y] for i=1:length(xs)] scatter!(xs, ys, label=nothing) end; +nothing # hide ``` Finally, we run the experiment. We will visualize just one trace produced by applying our inference program to each of the four variants of our model: @@ -391,5 +396,4 @@ fig4 = render_trace(tr, "fully static model with map") plot(fig1, fig2, fig3, fig4) ``` - It looks like inference in all the models seems to be working reasonably. \ No newline at end of file diff --git a/docs/src/tutorials/basics/particle_filter.md b/docs/src/tutorials/smc.md similarity index 91% rename from docs/src/tutorials/basics/particle_filter.md rename to docs/src/tutorials/smc.md index 4da7ac27d..7e68543c3 100644 --- a/docs/src/tutorials/basics/particle_filter.md +++ b/docs/src/tutorials/smc.md @@ -1,4 +1,4 @@ -# [Object Tracking with Particle Filters](@id object_tracking_with_particle_filter_tutorial) +# [Object Tracking with Sequential Monte Carlo](@id smc_tutorial) ## Introduction So far, we've seen two general classes of inference algorithm, importance sampling and MCMC. @@ -13,7 +13,7 @@ as follows: possible solutions. At every iteration, the current state is an entire proposed solution to the problem. -In this notebook, we will explore a third paradigm: **Sequential Monte Carlo**. SMC methods, +In this tutorial, we will explore a third paradigm: **Sequential Monte Carlo (SMC)**. SMC methods, such as particle filtering, iteratively solve a *sequence of inference problems* using techniques based on importance sampling and in some cases MCMC [^1][^2]. The solution to each problem in the sequence is represented as a collection of samples or *particles*. The particles for each problem are based on @@ -28,31 +28,29 @@ observations. We then consider a slightly more difficult inference problem: join inference of the first _two_ time steps' latent variables, given both time steps' observations. And so on, until the observations stop. -But SMC is a more general algorithm than the particle filter might suggest. +SMC is a more general algorithm than the particle filter might suggest. Sometimes, the sequence of problems does not arise from data arriving incrementally, but is rather constructed instrumentally to facilitate inference, as in annealed importance sampling [^3]. -However, this notebook focuses on particle filtering for a typical tracking problem. +However, this tutorial focuses on particle filtering for a typical tracking problem. We show how Gen's support for SMC integrates with its support for MCMC, enabling "rejuvenation" MCMC moves. Specifically, we will address the "bearings only tracking" problem described in [^4]. -This notebook will also introduce you to the -[`Unfold`](@ref Unfold) combinator, +This tutorial will also introduce you to the +[`Unfold`](@ref) combinator, which can be used to improve performance of SMC. `Unfold` is just one example of the levers that Gen provides for improving performance; once you understand it, you can check Gen's documentation to see how similar principles apply to the -[`Map`](https://www.gen.dev/docs/dev/ref/combinators/#Map-combinator-1) combinator +[`Map`](@ref) combinator and to the static DSL. (These features are also covered in the previous tutorial, [Scaling with the Static Modeling Language](@ref scaling_with_sml_tutorial)) **Prerequisites for this tutorial** -- [Generative Combinators](@ref combinators_tutorial) - [Static Modeling Language](@ref scaling_with_sml_tutorial) for the second half of the tutorial. - [^1]: Doucet, Arnaud, Nando De Freitas, and Neil Gordon. "An introduction to sequential Monte Carlo methods." Sequential Monte Carlo methods in practice. Springer, New York, NY, 2001. 3-14. [^2]: Del Moral, Pierre, Arnaud Doucet, and Ajay Jasra. "Sequential Monte Carlo samplers." Journal of the Royal Statistical Society: Series B (Statistical Methodology) 68.3 (2006): 411-436. @@ -83,7 +81,7 @@ then generates `T` successive states in a `for` loop. The argument to the model (`T`) is the number of time steps not including the initial state. -```@example particle_filter_tutorial +```@example smc_tutorial using Gen, Plots bearing(x, y) = atan(y, x) @@ -133,7 +131,8 @@ bearing(x, y) = atan(y, x) # return the sequence of positions return (xs, ys) -end; +end +nothing # hide ``` Note that the `model` function itself uses mutation to evolve the variables @@ -146,7 +145,7 @@ We generate a data set of positions, and observed bearings, by sampling from this model, with `T=50`: -```@example particle_filter_tutorial +```@example smc_tutorial import Random Random.seed!(3) @@ -154,11 +153,11 @@ Random.seed!(3) T = 50 constraints = Gen.choicemap((:x0, 0.01), (:y0, 0.95), (:vx0, 0.002), (:vy0, -0.013)) (trace, _) = Gen.generate(model, (T,), constraints) -nothing +nothing # hide ``` Let's extract the observed data `zs` from the trace -```@example particle_filter_tutorial +```@example smc_tutorial choices = Gen.get_choices(trace) zs = Vector{Float64}(undef, T+1) zs[1] = choices[:z0] @@ -172,7 +171,7 @@ We next write a visualization for full traces of this model. It shows the ship's positions (as dots) as well as the observed bearings (as fixed length line segments from the origin): -```@example particle_filter_tutorial +```@example smc_tutorial function render(trace; show_data=true, max_T=get_args(trace)[1], overlay=false) (T,) = Gen.get_args(trace) (xs, ys) = Gen.get_retval(trace) @@ -196,11 +195,12 @@ function render(trace; show_data=true, max_T=get_args(trace)[1], overlay=false) return fig end +nothing # hide ``` We visualize the synthetic trace below: -```@example particle_filter_tutorial +```@example smc_tutorial render(trace) title!("Observed bearings (lines) and positions (dots)") ``` @@ -225,27 +225,25 @@ sample of `num_samples` traces from the weighted collection that the particle filter produces. Gen provides methods for initializing and updating the state of a particle -filter, documented in [Particle Filtering](https://www.gen.dev/docs/dev/ref/pf/). +filter, documented in [Particle Filtering](@ref particle_filtering). -- `Gen.initialize_particle_filter` +- [`Gen.initialize_particle_filter`](@ref) -- `Gen.particle_filter_step!` +- [`Gen.particle_filter_step!`](@ref) Both of these methods can used either with the default proposal or a custom -proposal. In this -problem, -we will use the default proposal. There is also a method that resamples +proposal. In this problem, we will use the default proposal. There is also a method that resamples particles based on their weights, which serves to redistribute the particles to more promising parts of the latent space. -- `Gen.maybe_resample!` +- [`Gen.maybe_resample!`](@ref) Gen also provides a method for sampling a collection of unweighted traces from the current weighted collection in the particle filter state: -- `Gen.sample_unweighted_traces` +- [`Gen.sample_unweighted_traces`](@ref) -```@example particle_filter_tutorial +```@example smc_tutorial function particle_filter(num_particles::Int, zs::Vector{Float64}, num_samples::Int) # construct initial observations @@ -262,6 +260,7 @@ function particle_filter(num_particles::Int, zs::Vector{Float64}, num_samples::I # return a sample of unweighted traces from the weighted collection return Gen.sample_unweighted_traces(state, num_samples) end +nothing # hide ``` The initial state is obtained by providing the following to @@ -286,7 +285,7 @@ and then we introduce one additional bearing measurement by calling - The new arguments to the generative function for this step. In our case, this is the number of measurements beyond the first measurement. -- The [argdiff](https://www.gen.dev/docs/dev/ref/gfi/#Argdiffs-1) +- The [argdiff](../ref/core/change_hints.md) value, which provides detailed information about the change to the arguments between the previous step and this step. We will revisit this value later. For now, we indicate that we do not know how the `T::Int` @@ -300,14 +299,15 @@ We run this particle filter with 5000 particles, and return a sample of 200 particles. This will take 30-60 seconds. We will see one way of speeding up the particle filter in a later section. -```@example particle_filter_tutorial +```@example smc_tutorial @time pf_traces = particle_filter(5000, zs, 200); +nothing # hide ``` To render these traces, we first define a function that overlays many renderings: -```@example particle_filter_tutorial +```@example smc_tutorial function overlay(renderer, traces; same_data=true, args...) fig = plot(xlabel="X", ylabel="Y", title="Observed bearings (red) and \npositions of individual traces (one color per trace)") @@ -318,12 +318,12 @@ function overlay(renderer, traces; same_data=true, args...) end fig end +nothing # hide ``` - We then render the traces from the particle filter: -```@example particle_filter_tutorial +```@example smc_tutorial overlay(render, pf_traces) ``` @@ -340,16 +340,14 @@ explanation". !!! note Run the particle filter with fewer particles and visualize the results. - Run the particle filter without the `maybe_resample!` step, and visualize the results. What do you observe? Why do you think this is? Answer in the free response section below. -The code for particle_filter (from above) is copied in the body of the function below. Modify it so that it does NOT perform resampling after each -time step. +The code for particle_filter (from above) is copied in the body of the function below. Modify it so that it does NOT perform resampling after each time step. -```@example particle_filter_tutorial +```@example smc_tutorial function particle_filter_no_resampling(num_particles::Int, zs::Vector{Float64}, num_samples::Int) # construct initial observations @@ -374,7 +372,7 @@ overlay(render, pf_traces_no_resampling) !!! note Describe how the inference differs with and without resampling (based on the two plots above). Why do you think that is? Is it desirable? -```@example particle_filter_tutorial +```@example smc_tutorial function particle_filter_no_resampling(num_particles::Int, zs::Vector{Float64}, num_samples::Int) # construct initial observations @@ -390,16 +388,13 @@ function particle_filter_no_resampling(num_particles::Int, zs::Vector{Float64}, # return a sample of unweighted traces from the weighted collection return Gen.sample_unweighted_traces(state, num_samples) end +nothing # hide ``` - - Without resampling, we get particle collapse, because the model converges on one best guess. This is not desirable; we have lost diversity in our samples. -END ANSWER KEY - ## Adding Rejuvenation Moves The particle filter we developed above works as follows: @@ -455,7 +450,7 @@ First, the resimulation MH rejuvenation move (this function is the same as the previous, but with the addition of a rejuvenation move targeting the initial choices of each particle): -```@example particle_filter_tutorial +```@example smc_tutorial function particle_filter_rejuv_resim(num_particles::Int, zs::Vector{Float64}, num_samples::Int) init_obs = Gen.choicemap((:z0, zs[1])) state = Gen.initialize_particle_filter(model, (0,), init_obs, num_particles) @@ -492,27 +487,29 @@ You may notice slightly more variety in the initial state, compared to our first make sure you use the right _address_ --- **you want to use the same address in your proposal as was used in the model.** - We have provided starter code. -```@example particle_filter_tutorial +```@example smc_tutorial @gen function perturbation_proposal(prev_trace, a::Int, b::Int) + choices = get_choices(prev_trace) (T,) = get_args(prev_trace) speed = Array{Float64}(undef, 2, 1) for t=a:b - speed[1] = 0. # - speed[2] = 0. # + speed[1] = {(:vx, t)} ~ normal(choices[(:vx, t)], 1e-3) + speed[2] = {(:vy, t)} ~ normal(choices[(:vy, t)], 1e-3) end - return speed # Return an array of the most recent [vx, vy] for testing + return speed end function perturbation_move(trace, a::Int, b::Int) Gen.metropolis_hastings(trace, perturbation_proposal, (a, b)) -end; +end; +nothing # hide ``` + We add this into our particle filtering inference program below. We apply the rejuvenation move to adjust the velocities for the previous 5 time steps. -```@example particle_filter_tutorial +```@example smc_tutorial function particle_filter_rejuv(num_particles::Int, zs::Vector{Float64}, num_samples::Int) init_obs = Gen.choicemap((:z0, zs[1])) state = Gen.initialize_particle_filter(model, (0,), init_obs, num_particles) @@ -531,42 +528,27 @@ function particle_filter_rejuv(num_particles::Int, zs::Vector{Float64}, num_samp # return a sample of unweighted traces from the weighted collection return Gen.sample_unweighted_traces(state, num_samples) end +nothing # hide ``` We run the particle filter with rejuvenation below. This will take a minute or two. We will see one way of speeding up the particle filter in a later section. -```@example particle_filter_tutorial +```@example smc_tutorial @time pf_rejuv_traces = particle_filter_rejuv(5000, zs, 200); +nothing # hide ``` We render the traces: -```@example particle_filter_tutorial +```@example smc_tutorial overlay(render, pf_rejuv_traces) title!("Rejuvenation with resimulation MH on the starting points") ``` - - - ## 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 @@ -581,7 +563,7 @@ 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 +Function Combinators](@ref combinators), which make the flow of information through the generative process more explicit to Gen, and enable asymptotically more efficient inference programs. @@ -612,8 +594,7 @@ Julia `for` loop in our model. 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. +[`Unfold`](@ref) 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`) @@ -623,8 +604,7 @@ 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 +```@example smc_tutorial struct State x::Float64 y::Float64 @@ -644,8 +624,7 @@ end end chain = Gen.Unfold(kernel) - -Gen.@load_generated_functions # To allow use of the generative function written in the static modeling language above. +nothing # hide ``` We can understand the behavior of `chain` by getting a trace of it and printing the random choices: @@ -655,7 +634,7 @@ 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 +```@example smc_tutorial @gen (static) function unfold_model(T::Int) # parameters @@ -680,20 +659,19 @@ We now write a new version of the generative model that invokes `chain` instead result = (init_state, chain) return result end - -Gen.@load_generated_functions +nothing # hide ``` Let's generate a trace of this new model program to understand its structure: -```@example particle_filter_tutorial +```@example smc_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 +```@example smc_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) @@ -735,11 +713,12 @@ function unfold_render(trace; show_data=true, max_T=get_args(trace)[1], overlay= end end end +nothing # hide ``` Let's check that the results are reasonable: -```@example particle_filter_tutorial +```@example smc_tutorial overlay(unfold_render, unfold_pf_traces, same_data=true) ``` @@ -751,7 +730,7 @@ and 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 +```@example smc_tutorial fake_zs = rand(1000); function timing_experiment(num_observations_list::Vector{Int}, num_particles::Int, num_samples::Int) @@ -773,12 +752,13 @@ end; num_observations_list = [1, 3, 10, 30, 50, 100, 150, 200, 500] (times, times_unfold) = timing_experiment(num_observations_list, 100, 20); +nothing # hide ``` 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 +```@example smc_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/src/inference/map_optimize.jl b/src/inference/map_optimize.jl index 1ea837fd7..3250e9e38 100644 --- a/src/inference/map_optimize.jl +++ b/src/inference/map_optimize.jl @@ -3,7 +3,6 @@ max_step_size=0.1, tau=0.5, min_step_size=1e-16, verbose=false) Perform backtracking gradient ascent to optimize the log probability of the trace over selected continuous choices. - Selected random choices must have support on the entire real line. """ function map_optimize(trace, selection::Selection; diff --git a/src/inference/particle_filter.jl b/src/inference/particle_filter.jl index 4e2f1df79..ddedb08e2 100644 --- a/src/inference/particle_filter.jl +++ b/src/inference/particle_filter.jl @@ -12,9 +12,26 @@ function normalize_weights(log_weights::Vector{Float64}) end ####################################### -# building blocks for particle filter # +# building blocks for particle filters # ####################################### +""" + ParticleFilterState{U} + +Represents the state of a particle filter as a collection of weighted traces, +where the type of each trace is `U`. + +# Fields + +- `traces`: A vector of traces, one for each particle. +- `new_traces`: A preallocated vector for storing new traces. +- `log_weights`: A vector of log importance weights for each trace. +- `log_ml_est`: Estimate of the log marginal likelihood before the last resampling step. +- `parents`: The parent indices of each trace in `traces`. + +!!! note + The fields above are an implementation detail that are subject to future changes. +""" mutable struct ParticleFilterState{U} traces::Vector{U} new_traces::Vector{U} @@ -257,5 +274,6 @@ function maybe_resample!( return do_resample end +export ParticleFilterState export initialize_particle_filter, particle_filter_step!, maybe_resample! export get_traces, get_log_weights, log_ml_estimate, sample_unweighted_traces diff --git a/src/modeling_library/modeling_library.jl b/src/modeling_library/modeling_library.jl index 7e2a782fc..7bf691f2a 100644 --- a/src/modeling_library/modeling_library.jl +++ b/src/modeling_library/modeling_library.jl @@ -24,7 +24,7 @@ function logpdf end """ has::Bool = has_output_grad(dist::Distribution) -Return true of the gradient if the distribution computes the gradient of the logpdf with respect to the value of the random choice. +Return true if the distribution computes the gradient of the logpdf with respect to the value of the random choice. """ function has_output_grad end @@ -40,6 +40,11 @@ Otherwise, this element contains the gradient with respect to the `i`th argument """ function logpdf_grad end +""" + discrete::Bool = is_discrete(::Distribution) + +Return true if the distribution is discrete, false otherwise. +""" is_discrete(::Distribution) = false # default # NOTE: has_argument_grad is documented and exported in gen_fn_interface.jl