Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

Various docs tasks for consistency #271

Merged
merged 8 commits into from
Jan 9, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions docs/make.jl
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,6 @@ examples = [
],
]

design = ["AbstractMCMC sampling API" => "API/AbstractMCMC.md"]

api = [
"CalibrateEmulateSample" => [
"Emulators" => [
Expand All @@ -31,18 +29,21 @@ api = [
],
]

emulate = [
"Emulator" => "emulate.md",
"Gaussian Process" => "GaussianProcessEmulator.md",
"Random Features" => "random_feature_emulator.md",
]
pages = [
"Home" => "index.md",
"Installation instructions" => "installation_instructions.md",
"Contributing" => "contributing.md",
"Calibrate" => "calibrate.md",
"Emulate" => "emulate.md",
"Examples" => examples,
"Gaussian Process" => "GaussianProcessEmulator.md",
"Random Features" => "random_feature_emulator.md",
"Package Design" => design,
"API" => api,
"Calibrate" => "calibrate.md",
"Emulate" => emulate,
"Sample" => "sample.md",
"Glossary" => "glossary.md",
"API" => api,
]

#----------
Expand Down
45 changes: 33 additions & 12 deletions docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,25 +25,46 @@ Given an observation ``y``, the computer model ``\mathcal{G}``, the observationa

As the name suggests, `CalibrateEmulateSample.jl` breaks this problem into a sequence of three steps: calibration, emulation, and sampling. A comprehensive treatment of the calibrate-emulate-sample approach to Bayesian inverse problems can be found in [Cleary et al. (2020)](https://arxiv.org/pdf/2001.03689.pdf).

### The three steps of the algorithm:
### The three steps of the algorithm: see our walkthrough of the [Sinusoid Example](@ref)

The **calibrate** step of the algorithm consists of an application of [Ensemble Kalman Processes](https://github.com/CliMA/EnsembleKalmanProcesses.jl), which generates input-output pairs ``\{\theta, \mathcal{G}(\theta)\}`` in high density around an optimal parameter ``\theta^*``. This ``\theta^*`` will be near a mode of the posterior distribution (Note: This the only time we interface with the forward model ``\mathcal{G}``).
**Given some noisy observations...**
```@raw html
<img src="assets/sinusoid_true_vs_observed_signal.png" width="400">
```

1. The **calibrate** step of the algorithm consists of an application of [Ensemble Kalman Processes](https://github.com/CliMA/EnsembleKalmanProcesses.jl), which generates input-output pairs ``\{\theta, \mathcal{G}(\theta)\}`` in high density around an optimal parameter ``\theta^*``. This ``\theta^*`` will be near a mode of the posterior distribution (Note: This the only time we interface with the forward model ``\mathcal{G}``).
odunbar marked this conversation as resolved.
Show resolved Hide resolved

**calibrate with EKP to generate data pairs...**
```@raw html
<img src="assets/sinusoid_eki_pairs.png" width="400">
```

The **emulate** step takes these pairs ``\{\theta, \mathcal{G}(\theta)\}`` and trains a statistical surrogate model (e.g., a Gaussian process), emulating the forward map ``\mathcal{G}``.
2. The **emulate** step takes these pairs ``\{\theta, \mathcal{G}(\theta)\}`` and trains a statistical surrogate model (e.g., a Gaussian process), emulating the forward map ``\mathcal{G}``.
odunbar marked this conversation as resolved.
Show resolved Hide resolved

**emulate the map statistically from EKP pairs...**
```@raw html
<img src="assets/sinusoid_GP_emulator_contours.png" width="400">
```
3. The **sample** step uses this surrogate in place of ``\mathcal{G}`` in a sampling method (Markov chain Monte Carlo) to sample the posterior distribution of ``\theta``.

**sample the emulated map with MCMC...**
```@raw html
<img src="assets/sinusoid_MCMC_hist_GP.png" width="400">
```

The **sample** step uses this surrogate in place of ``\mathcal{G}`` in a sampling method (Markov chain Monte Carlo) to sample the posterior distribution of ``\theta``.
## Code Components

`CalibrateEmulateSample.jl` contains the following modules:

Module | Purpose
-----------------------------|--------------------------------------------------------
CalibrateEmulateSample.jl | Pulls in the [Ensemble Kalman Processes](https://github.com/CliMA/EnsembleKalmanProcesses.jl) package
Emulator.jl | Emulate: Modular template for emulators
GaussianProcess.jl | - A Gaussian process emulator
MarkovChainMonteCarlo.jl | Sample: Modular template for MCMC
Utilities.jl | Helper functions
Module | Purpose
---------------------------------------|--------------------------------------------------------
CalibrateEmulateSample.jl | A wrapper for the pipeline
Emulator.jl | Modular template for the emulators
GaussianProcess.jl | A Gaussian process emulator
Scalar/VectorRandomFeatureInterface.jl | A Scalar/Vector-output Random Feature emulator
MarkovChainMonteCarlo.jl | Modular template for Markov Chain Monte Carlo samplers
Utilities.jl | Helper functions

**The best way to get started is to have a look at the examples!**

## Authors

Expand Down
32 changes: 30 additions & 2 deletions docs/src/random_feature_emulator.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,34 @@ The `VectorRandomFeatureInterface`, when applied to multidimensional problems, d

Building a random feature interface is similar to building a Gaussian process: one defines a kernel to encode similarities between outputs ``(y_i,y_j)`` based on inputs ``(x_i,x_j)``. Additionally, one must specify the number of random feature samples to be taken to build the emulator.

# Recommended configuration

Below is listed a recommended configuration that is flexible and requires learning relatively few parameters. Users can increase `r` to balance flexibility against having more kernel hyperparameters to learn.

```julia
using CalibrateEmulateSample.Emulators
# given input_dim, output_dim, and a PairedDataContainer

# define number of features for prediction
n_features = 400 # number of features for prediction

# define kernel
nugget = 1e8*eps() # small nugget term
r = 1 # start with smallest rank
lr_perturbation = LowRankFactor(r, nugget)
nonsep_lrp_kernel = NonseparableKernel(lr_perturbation)

# configure optimizer
optimizer_options = Dict(
"verbose" => true, # print diagnostics for optimizer
"n_features_opt" => 100, # use less features during hyperparameter optimization/kernel learning
"cov_sample_multiplier" => 1.0, # use to reduce/increase number of samples in initial cov estimation stage
)

machine_learning_tool = VectorRandomFeatureInterface(n_features, input_dim, output_dim, optimizer_options = optimizer_options)
```
Users can change the kernel complexity with `r`, and the number of features for prediciton with `n_features` and optimization with `n_features_opt`.

# User Interface

`CalibrateEmulateSample.jl` allows the random feature emulator to be built using the external package [`RandomFeatures.jl`](https://github.com/CliMA/RandomFeatures.jl). In the notation of this package's documentation, our interface allows for families of `RandomFourierFeature` objects to be constructed with different Gaussian distributions of the "`xi`" a.k.a weight distribution, and with a learnable "`sigma`", a.k.a scaling parameter.
Expand All @@ -37,8 +65,8 @@ To adjust the expressivity of the random feature family one can define the keywo

We have two types,
```julia
separable_kernel = Separable(input_cov_structure, output_cov_structure)
nonseparable_kernel = Nonseparable(cov_structure)
separable_kernel = SeparableKernel(input_cov_structure, output_cov_structure)
nonseparable_kernel = NonseparableKernel(cov_structure)
```
where the `cov_structure` implies some imposed user structure on the covariance structure. The basic covariance structures are given by
```julia
Expand Down
64 changes: 58 additions & 6 deletions docs/src/API/AbstractMCMC.md → docs/src/sample.md
Original file line number Diff line number Diff line change
@@ -1,14 +1,66 @@
# AbstractMCMC sampling API
# The Sample stage

```@meta
CurrentModule = CalibrateEmulateSample.MarkovChainMonteCarlo
```

The "sample" part of CES refers to exact sampling from the emulated posterior via [Markov chain Monte
Carlo](https://en.wikipedia.org/wiki/Markov_chain_Monte_Carlo) (MCMC). Within this paradigm, we want to provide the
flexibility to use multiple sampling algorithms; the approach we take is to use the general-purpose
[AbstractMCMC.jl](https://turing.ml/dev/docs/for-developers/interface) API, provided by the
[Turing.jl](https://turing.ml/dev/) probabilistic programming framework.
The "sample" part of CES refers to exact sampling from the emulated posterior. In our current framework this is acheived with a [Markov chain Monte
Carlo algorithm](https://en.wikipedia.org/wiki/Markov_chain_Monte_Carlo) (MCMC). Within this paradigm, we want to provide the flexibility to use multiple sampling algorithms; the approach we take is to use the general-purpose [AbstractMCMC.jl](https://turing.ml/dev/docs/for-developers/interface) API, provided by the [Turing.jl](https://turing.ml/dev/) probabilistic programming framework.


## User interface

We briefly outline an instance of how one sets up and uses MCMC within the CES package. The user first provides a Protocol (i.e. how one wishes to generate proposals)

```julia
protocol = RWMHSampling() # Random-Walk algorithm
# protocol = pCNMHSampling() # preconditioned-Crank-Nicholson algorithm
```
Then one builds the MCMC by providing the standard Bayesian ingredients (prior and data) from the Calibrate stage, alongside the trained statistical emulator from the Emulate stage:
```julia
mcmc = MCMCWrapper(
protocol,
truth_sample,
prior,
emulator;
init_params=mean_u_final,
burnin=10_000,
)
```
The keyword arguments `init_params` give a starting step of the chain (often taken to be the mean of the final iteration of Calibrate stage), and a `burnin` gives a number of initial steps to be discarded when drawing statistics from the Sampling method.

For good efficiency, one often needs to run MCMC with a problem-dependent step size. We provide a simple utility to help choose this. Here the optimizer runs short chains (of length `N`), and adjusts the step-size until the MCMC acceptance rate falls within an acceptable range, returning this step size.
```julia
new_step = optimize_stepsize(
mcmc;
init_stepsize = 1,
N = 2000
)
```
To generate ``10^5`` samples with a given step size (and optional random number generator `rng`), one calls
```julia
chain = sample(rng, mcmc, 100_000; stepsize = new_step)
display(chain) # gives diagnostics
```
The return argument is stored in an `MCMCChains.Chains` object. To convert this back into a `ParameterDistribution` type (which contains e.g. the transformation maps) one can call
```julia
posterior = get_posterior(mcmc, chain)
constrained_posterior = transform_unconstrained_to_constrained(prior, get_distribution(posterior))
```

One can quickly plot the marginals of the prior and posterior distribution with
```julia
using Plots
plot(prior)
plot!(posterior)
```
or extract statistics of the (unconstrained) distribution with
```julia
mean_posterior = mean(posterior)
cov_postierior = cov(posterior)
```
odunbar marked this conversation as resolved.
Show resolved Hide resolved

# [Further details on the implementation](@id AbstractMCMC sampling API)

This page provides a summary of AbstractMCMC which augments the existing documentation
(\[[1](https://turing.ml/dev/docs/for-developers/interface)\],
Expand Down
Loading