EngineBuilder#

class liesel.goose.EngineBuilder(seed, num_chains)[source]#

Bases: object

The EngineBuilder is used to construct an MCMC Engine.

Workflow

The general workflow usually looks something like this:

Create a builder with EngineBuilder.
Set the desired number of warmup and posterior samples set_duration().
Set the model interface with set_model().
Set the initial values with set_initial_values().
Add MCMC kernels with add_kernel().
Build an Engine with build().

Optionally, you can also:

Add position keys to positions_included for tracking. If you are using a Model object, position keys are the names of variables or nodes in the model. Refer to positions_included for more information.
Add custom jittering for start values with set_jitter_fns().

Parameters:

seed (Union[int, Any]) – Either an int or a key generated from jax.random.PRNGKey. Used for jittering initial values and MCMC sampling.
num_chains (int) – The number of chains to be used.

See also

~.goose.Engine : The MCMC engine, output of build(). ~.goose.LieselInterface : Interface for a Model object. ~.goose.NUTSKernel : The NUTS kernel. ~.goose.HMCKernel : The HMC kernel. ~.goose.IWLSKernel : The IWLS kernel. ~.goose.RWKernel : The random walk kernel.

Notes

By default, only position keys associated with an MCMC kernel is tracked. This behavior can be adjusted with the fields positions_included and positions_excluded.

Examples

For this example, we import tensorflow_probability as follows:

>>> import tensorflow_probability.substrates.jax.distributions as tfd

First, we set up a minimal model:

>>> mu = lsl.param(0.0, name="mu")
>>> dist = lsl.Dist(tfd.Normal, loc=mu, scale=1.0)
>>> y = lsl.obs(jnp.array([1.0, 2.0, 3.0]), dist, name="y")
>>> model = lsl.GraphBuilder().add(y).build_model()

Now we initialize the EngineBuilder and set the desired number of warmup and posterior samples:

>>> builder = gs.EngineBuilder(seed=1, num_chains=4)
>>> builder.set_duration(warmup_duration=1000, posterior_duration=1000)

Next, we set the model interface and initial values:

>>> builder.set_model(gs.LieselInterface(model))
>>> builder.set_initial_values(model.state)

We add a NUTS kernel for the parameter "mu":

>>> builder.add_kernel(gs.NUTSKernel(["mu"]))

Finally, we build the engine:

>>> engine = builder.build()

From here, you can continue with sample_all_epochs() to draw samples from your posterior distribution.

Methods

`add_kernel`(kernel)	Adds a `Kernel`.
`add_quantity_generator`(generator)	Adds a `QuantityGenerator`.
`build`()	Builds the MCMC engine with the provided setup.
`set_duration`(warmup_duration, posterior_duration)	Sets epochs using the `stan_epochs()` function.
`set_engine_seed`(seed)	Sets a seed used to initialize the MCMC engine.
`set_epochs`(epochs)	Sets epochs.
`set_initial_values`(model_state[, ...])	Sets the initial model state.
`set_jitter_fns`(jitter_fns)	Set the jittering functions.
`set_model`(model)	Sets the model interface for all kernels and quantity generators.

Attributes

`engine_seed`	The seed for the engine's pseudo-random number generation.
`epochs`	Tuple of epoch configurations.
`jitter_fns`	Jittering functions.
`kernels`	Tuple of all Kernels that are present in the builder.
`model_state`	Model state.
`quantity_generators`	Tuple of all quantity generators present in the builder.
`show_progress`	Whether to show progress bars curing sampling.
`positions_included`	List of additional position keys that should be tracked.
`positions_excluded`	List of position keys that should not be tracked.

EngineBuilder

Contents

EngineBuilder#