MCMCSpec

Contents

MCMCSpec#

class liesel.goose.MCMCSpec(kernel, kernel_kwargs=<factory>, kernel_group=None, jitter_dist=None, jitter_method='additive', order=99)[source]#

Bases: object

Specification for the MCMC kernel and optional jitter distribution associated with a model variable.

Parameters:

  • kernel (KernelFactory) – A KernelFactory that returns a Kernel instance when provided with position keys and keyword arguments.

  • kernel_kwargs (dict[str, Any], default: <factory>) – Additional keyword arguments to be passed to the kernel callable.

  • kernel_group (str | None, default: None) – Name of the kernel group this variable belongs to. Variables in the same group must share the same kernel type and arguments.

  • jitter_dist (tfp.substrates.jax.distributions.Distribution | None, default: None) – A TensorFlow Probability distribution used to apply random jitter to the initial value of the variable.

  • jitter_method (Literal['additive', 'multiplicative', 'replacement'], default: 'additive') – The type of jitter to be applied. This can be one of the following: - none: No jitter is applied. - additive: Additive jitter is applied. - multiplicative: Multiplicative jitter is applied. - replacement: Value is replaced when jitter is applied.

  • order (int, default: 99) – If you want to change the order in which parameter blocks are sampled. Blocks will be ordered by default based on the topological order of the graph (from the bottom up; i.e. the kernels for sampling parameters closest to the graph’s leaf nodes/responses come first), which is often a sensible default. After that, blocks will be ordered based on the integer provided here. The kernel with the smallest order integer will be used first.

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()

Methods

apply_jitter(seed, value)

Apply random jitter to a given value using the specified jitter distribution.

Attributes

jitter_dist

jitter_method

kernel_group

order

kernel

kernel_kwargs

Contents