MCMCSpec

Contents

MCMCSpec#

class liesel.goose.MCMCSpec(kernel, kernel_kwargs=<factory>, kernel_group=None, jitter_dist=None, jitter_method='additive')[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 (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.

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

kernel

kernel_kwargs

Contents