EngineBuilder

class liesel.goose.EngineBuilder(seed, num_chains)

Bases: object

The EngineBuilder is used to construct an MCMC Engine.

Liesel Workflow

The Liesel workflow usually looks something like this:

1. You intialize your Var objects with

   MCMCSpec that specify how each variable should be sampled.

2. You initilize an EngineBuilder using the method

   get_engine_builder() of LieselMCMC.

3. Set the desired number of warmup and posterior samples with

   add_burnin(), add_adaptation(), and add_posterior().

4. Build an Engine with build().

General Workflow

The general workflow usually looks something like this:

1. Create a builder with EngineBuilder.

2. Set the desired number of warmup and posterior samples with

   add_burnin(), add_adaptation(), and add_posterior().

3. Set the model interface with set_model().

4. Set the initial values with set_initial_values().

5. Add MCMC kernels with add_kernel().

6. 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 (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

Liesel Workflow

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.Var.new_param(0.0, name="mu", inference=gs.MCMCSpec(gs.NUTSKernel))
>>> dist = lsl.Dist(tfd.Normal, loc=mu, scale=1.0)
>>> y = lsl.Var.new_obs(jnp.array([1.0, 2.0, 3.0]), dist, name="y")
>>> model = lsl.Model([y])

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

>>> builder = gs.LieselMCMC(model).get_engine_builder(seed=1, num_chains=4)
>>> builder.add_adaptation(1000)
>>> builder.add_posterior(1000)

Finally, we build the engine:

>>> engine = builder.build()

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

General Workflow

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.Var.new_param(0.0, name="mu")
>>> dist = lsl.Dist(tfd.Normal, loc=mu, scale=1.0)
>>> y = lsl.Var.new_obs(jnp.array([1.0, 2.0, 3.0]), dist, name="y")
>>> model = lsl.Model([y])

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_adaptation(duration[, init, term, base, ...])	Adds adaptation epochs, for tuning kernel hyperparameters.
add_burnin(duration[, thinning])	Adds burnin epochs.
add_kernel(kernel)	Adds a Kernel.
add_posterior(duration[, thinning])	Adds posterior epochs.
add_quantity_generator(generator)	Adds a QuantityGenerator.
append_epochs(epochs)	Appends epochs.
build()	Builds the MCMC engine with the provided setup.
set_duration(warmup_duration, posterior_duration)	Sets epochs using the stan_epochs() function.
set_duration_simple([burnin_duration, ...])	Defines a simple epoch configuration that does not include any adaptative warmup.
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 during sampling.
positions_included	List of additional position keys that should be tracked.
positions_excluded	List of position keys that should not be tracked.