"""
MCMC engine builder
The purpose of the engine builder is to provide a simple API to gradually assemble
the components needed by the MCMC engine. The builder is responsible of returning an
engine in a well-defined state. Furthermore, the builder can return different engine
implementations.
"""
from __future__ import annotations
import logging
import math
from collections.abc import Iterable, Mapping
from functools import partial
from typing import TYPE_CHECKING, cast
import jax
import jax.numpy as jnp
from liesel.option import Option
from .engine import Engine
from .epoch import EpochConfig, EpochManager, EpochType
from .kernel_sequence import KernelSequence
from .pytree import stack_leaves
from .types import (
JitterFunction,
JitterFunctions,
Kernel,
KeyArray,
ModelInterface,
ModelState,
Position,
QuantityGenerator,
)
from .warmup import stan_epochs
_EpochConfig = partial(EpochConfig, optional=None)
logger = logging.getLogger(__name__)
def _find_duplicate(xs: list[str]) -> Option[str]:
"""Checks if a list of strings contains any duplicates."""
set_ = set()
for x in xs:
if x in set_:
return Option(x)
else:
set_.add(x)
return Option(None)
[docs]
class EngineBuilder:
"""
The :class:`~liesel.goose.EngineBuilder` is used to construct an MCMC Engine.
.. rubric:: Liesel Workflow
The Liesel workflow usually looks something like this:
#. You intialize your :class:`~liesel.model.nodes.Var` objects with
:class:`~liesel.goose.MCMCSpec` that specify how each variable should be
sampled.
#. You initilize an EngineBuilder using the method
:meth:`~liesel.goose.LieselMCMC.get_engine_builder` of
:class:`~liesel.goose.LieselMCMC`.
#. Set the desired number of warmup and posterior samples with
:meth:`~liesel.goose.EngineBuilder.add_burnin`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`, and
:meth:`~liesel.goose.EngineBuilder.add_posterior`.
#. Build an :class:`~liesel.goose.Engine` with
:meth:`~liesel.goose.EngineBuilder.build`.
.. rubric:: General Workflow
The general workflow usually looks something like this:
#. Create a builder with :class:`~liesel.goose.EngineBuilder`.
#. Set the desired number of warmup and posterior samples with
:meth:`~liesel.goose.EngineBuilder.add_burnin`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`, and
:meth:`~liesel.goose.EngineBuilder.add_posterior`.
#. Set the model interface with :meth:`~liesel.goose.EngineBuilder.set_model`.
#. Set the initial values with
:meth:`~liesel.goose.EngineBuilder.set_initial_values`.
#. Add MCMC kernels with :meth:`~liesel.goose.EngineBuilder.add_kernel`.
#. Build an :class:`~liesel.goose.Engine` with
:meth:`~liesel.goose.EngineBuilder.build`.
Optionally, you can also:
- Add position keys to :attr:`~liesel.goose.EngineBuilder.positions_included` for
tracking. If you are using a
:class:`~liesel.model.Model` object, position keys are the names of variables or
nodes in the
model. Refer to :attr:`~liesel.goose.EngineBuilder.positions_included` for more
information.
- Add custom jittering for start values with
:meth:`~liesel.goose.EngineBuilder.set_jitter_fns`.
Parameters
----------
seed
Either an int or a key generated from jax.random.PRNGKey. Used for jittering
initial values and MCMC sampling.
num_chains
The number of chains to be used.
See Also
--------
liesel.goose.Engine : The MCMC engine returned by :meth:`build
<liesel.goose.EngineBuilder.build>`.
liesel.goose.LieselInterface : Interface for a :class:`~liesel.model.Model`.
liesel.goose.NUTSKernel : The NUTS kernel.
liesel.goose.HMCKernel : The HMC kernel.
liesel.goose.IWLSKernel : The IWLS kernel.
liesel.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
:attr:`~liesel.goose.EngineBuilder.positions_included` and
:attr:`~liesel.goose.EngineBuilder.positions_excluded`.
Examples
--------
.. rubric:: 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 :meth:`~liesel.goose.Engine.sample_all_epochs` to
draw
samples from your posterior distribution.
.. rubric:: 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 :meth:`~liesel.goose.Engine.sample_all_epochs` to
draw
samples from your posterior distribution.
"""
if TYPE_CHECKING:
minimize_transition_infos: bool
"""Whether to retain only minimal transition information."""
positions_excluded: list[str]
"""Position keys excluded from the sampling results."""
positions_included: list[str]
"""Additional position keys to retain in the sampling results."""
show_progress: bool
"""Whether to show progress bars during sampling."""
store_kernel_states: bool
"""Whether to retain kernel states in the sampling results."""
def __init__(self, seed: int | KeyArray, num_chains: int):
if isinstance(seed, int):
keys = jax.random.split(jax.random.PRNGKey(seed), 3)
elif isinstance(seed, jax.Array):
keys = jax.random.split(seed, 3)
else:
raise TypeError(
"Provide either an int or a key from jax.random.PRNGKey as seed."
)
initial_values = [
_EpochConfig(EpochType.INITIAL_VALUES, duration=1, thinning=1)
]
self._epochs = EpochManager(configs=initial_values)
self._prng_key: KeyArray = keys[0]
self._engine_key: KeyArray = keys[1]
self._jitter_key: KeyArray = keys[2]
self._num_chains: int = num_chains
self._kernels: list[Kernel] = []
self._quantity_generators: list[QuantityGenerator] = []
self._model_state: Option[ModelState] = Option(None)
self._model: Option[ModelInterface] = Option(None)
self._jitter_fns: Option[JitterFunctions] = Option(None)
# public fields, only simple states
self.store_kernel_states: bool = False
self.minimize_transition_infos: bool = False
self.show_progress: bool = True
"""Whether to show progress bars during sampling."""
self.positions_included: list[str] = []
"""
List of additional position keys that should be tracked.
If a position key is tracked that means the correspond element of the model
state will be saved and included in the :class:`~liesel.goose.SamplingResults`
and the
posterior samples returned by
:meth:`~liesel.goose.SamplingResults.get_posterior_samples`.
By default, only position keys associated with an MCMC kernel are tracked. You
can easily add additional position keys by appending to this list.
Examples
--------
For this example, we import ``tensorflow_probability`` as follows:
>>> import tensorflow_probability.substrates.jax.distributions as tfd
Consider the following simple model, in which we use the logarithm of the scale
parameter in a normal distribution and take the exponential value for including
the actual scale:
>>> log_scale = lsl.Var.new_param(0.0, name="log_scale")
>>> scale = lsl.Var.new_calc(jnp.exp, log_scale, name="scale")
>>> dist = lsl.Dist(tfd.Normal, loc=0.0, scale=scale)
>>> y = lsl.Var.new_obs(jnp.array([1.0, 2.0, 3.0]), dist, name="y")
>>> model = lsl.Model([y])
Now we might want to set up an engine builder with a NUTS kernel for the
parameter ``"log_scale"``:
>>> builder = gs.EngineBuilder(seed=1, num_chains=4)
>>> builder.set_model(gs.LieselInterface(model))
>>> builder.add_kernel(gs.NUTSKernel(["log_scale"]))
By default, only the position key ``"log_scale"`` is tracked and will be
included in the results. Now, if you also want the value of ``"scale"`` to be
included, you can add it to the list of included position keys:
>>> builder.positions_included.append("scale")
>>> builder.positions_included
['scale']
Beware however that including many intermediate position keys can lead to large
results. In some cases it may be preferable to keep the tracked positions to a
minimum and recompute the intermediate values from the posterior samples.
"""
self.positions_excluded: list[str] = []
"""List of position keys that should not be tracked. Excluded keys override
additional keys."""
[docs]
def set_engine_seed(self, seed: int | KeyArray):
"""Sets a seed used to initialize the MCMC engine."""
if jnp.isscalar(seed):
seed_int: int = cast(int, seed)
self._engine_key = jax.random.PRNGKey(seed_int)
else:
seed_keyarray = cast(KeyArray, seed)
self._engine_key = seed_keyarray
@property
def engine_seed(self) -> KeyArray:
"""The seed for the engine's pseudo-random number generation."""
return self._engine_key
[docs]
def add_kernel(self, kernel: Kernel):
"""Adds a :class:`~liesel.goose.Kernel`."""
self._kernels.append(kernel)
@property
def kernels(self) -> tuple[Kernel, ...]:
"""Tuple of all Kernels that are present in the builder."""
return tuple(self._kernels)
[docs]
def add_quantity_generator(self, generator: QuantityGenerator):
"""Adds a ``QuantityGenerator``."""
self._quantity_generators.append(generator)
@property
def quantity_generators(self) -> tuple[QuantityGenerator, ...]:
"""Tuple of all quantity generators present in the builder."""
return tuple(self._quantity_generators)
[docs]
def set_initial_values(self, model_state: ModelState, multiple_chains=False):
"""
Sets the initial model state.
If ``multiple_chains`` is true the
:attr:`~liesel.goose.EngineBuilder.model_state` will be used as is;
otherwise :attr:`~liesel.goose.EngineBuilder.model_state` will be used as the
initial values for each
chain. Note that if ``multiple_chains`` is true, the first axis of each
leaf of :attr:`~liesel.goose.EngineBuilder.model_state` refers to the chain.
"""
if not multiple_chains:
model_states = stack_leaves(model_state for _ in range(self._num_chains))
else:
model_states = model_state
self._model_state = Option(model_states)
[docs]
def set_jitter_fns(self, jitter_fns: Mapping[str, JitterFunction] | None):
"""
Set the jittering functions.
A jittering function is a function that takes as input a key and a value,
and applies a random jittering (noise) to the input value
based on the given key.
If no jitter function is provided a `Warning` will be raised and the
values won't be jittered.
Parameters
----------
jitter_fns
A mapping where a jittering function is assigned to each position key.
A shallow copy is stored. Later changes to the supplied mapping do not
change the builder's configuration; call this method again to update it.
Examples
--------
>>> import jax
>>> import jax.numpy as jnp
>>> import tensorflow_probability.substrates.jax.distributions as tfd
>>> key = jax.random.PRNGKey(42)
In this example, we show how to use the method
:meth:`EngineBuilder.set_jitter_fns <liesel.goose.EngineBuilder.set_jitter_fns>`
to apply jittering
to the initial values of each chain.
First, we sample 500 data points from a Normal Distrbution
with mean 2.0 and standard deviation 1.0.
>>> n = 500
>>> true_mu = 2.0
>>> true_sigma = 1.0
>>> x_vec = tfd.Normal(loc=true_mu, scale=true_sigma).sample((n,), key)
Then, we define the distribution we want to sample from, which is
parametrized by a single parameter `mu`.
>>> mu = lsl.Var.new_param(1.0, name="mu")
>>> x_dist = lsl.Dist(tfd.Normal, loc=true_mu, scale=true_sigma)
>>> x = lsl.Var(x_vec, dist=x_dist, name="x")
Now, we can create the model.
>>> model = lsl.Model([x])
Finally, we build the model with :class:`~liesel.goose.EngineBuilder`. We will
use 4
parallel chains and sample our varaible using a
:class:`~liesel.goose.NUTSKernel`.
>>> builder = gs.EngineBuilder(seed=1337, num_chains=4)
>>> builder.set_model(gs.LieselInterface(model))
>>> builder.set_initial_values(model.state)
>>> builder.add_kernel(gs.NUTSKernel(["mu"]))
A jitter function takes as input a key and value and applies random jittering
to the given value using the key. In this case, we apply a uniform noise with
a minimum value of -1.0 and maximum value of 1.0. Notice that the shape of `val`
is `(4, 1)`, where the first dimension corresponds to the number of chains.
>>> def jitter_fn(key, val):
... jitter = jax.random.uniform(key, val.shape, val.dtype, -1.0, 1.0)
... return val + jitter
The method takes as input a dictionary where a
jittering function is assigned to each position key.
>>> builder.set_jitter_fns({"mu": jitter_fn})
"""
self._jitter_fns = Option(dict(jitter_fns) if jitter_fns is not None else None)
@property
def jitter_fns(self) -> Option[JitterFunctions]:
"""Jittering functions."""
return self._jitter_fns
@property
def model_state(self) -> Option[ModelState]:
"""Model state."""
return self._model_state
[docs]
def set_epochs(self, epochs: Iterable[EpochConfig]):
"""Sets epochs."""
initial_values = [
_EpochConfig(EpochType.INITIAL_VALUES, duration=1, thinning=1)
]
self._epochs = EpochManager(initial_values + list(epochs))
[docs]
def append_epochs(self, epochs: Iterable[EpochConfig]):
"""Appends epochs."""
for epoch in epochs:
self._epochs.append(epoch)
[docs]
def add_adaptation(
self,
duration: int,
init: float = 0.1,
term: float = 0.2,
base: int = 25,
thinning: int = 1,
):
"""
Adds adaptation epochs, for tuning kernel hyperparameters.
This follows the Stan Development Team (Stan Reference Manual 2021, Chapter 15.2
[#stanmanual]_).
Note that the order in which you add epochs matters.
Adaptation epochs should usually come before all other epochs.
Parameters
----------
duration
The number of adaptation samples.
init
The share of samples in the *initial fast* adaptation epoch, relative to \
the total number of adaptation samples (if given as a float), or the total \
number of init samples (if given as an integer).
term
The share of samples in the *final fast* adaptation epoch, relative to the \
total number of adaptation samples (if given as a float), or the total \
number of term samples (if given as an integer).
base
The number of samples in the *first slow* adaptation epoch.
thinning
Thinning applied in each adaptation.
Warnings
--------
Kernels which rely on history tuning might not be able to deal with thinning \
during adaptation.
See Also
--------
.add_posterior : For adding posterior epochs.
.add_burnin : For adding burnin epochs.
.add_adaptation : For adding adaptation epochs.
.set_duration : An alternative method to specify adaptation and posterior
epochs, see notes for details.
Notes
-----
There are some differences in the details between the methods
:meth:`~liesel.goose.EngineBuilder.add_posterior`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`,
:meth:`~liesel.goose.EngineBuilder.add_burnin` and
:meth:`~liesel.goose.EngineBuilder.set_duration`.
- :meth:`~liesel.goose.EngineBuilder.set_duration` overwrites all existing
epochs with the epochs defined
in this method. In contrast, the ``add_`` methods only add additional epochs
that are appended to the potentially existing ones.
- :meth:`~liesel.goose.EngineBuilder.set_duration` sets both adaptation and
posterior epochs, while each
``add_`` method only adds epochs of one type.
- :meth:`~liesel.goose.EngineBuilder.set_duration` has an argument
``term_duration`` that defines the
number of samples drawn in the *final fast* adaptation epoch, and it defaults
to a fixed value of 50 samples. In contrast,
:meth:`~liesel.goose.EngineBuilder.add_adaptation` defines
the number of samples drawn in the *final fast* adaptation epoch as a fraction
of the total number of adaptation samples, and allows users to change this
fraction via the argument ``term``. As a result, if
:meth:`~liesel.goose.EngineBuilder.add_adaptation`
is used, the length of the final fast
adaptation phase increases automatically, if the total number of warmup
samples increases, while it remains constant if
:meth:`~liesel.goose.EngineBuilder.set_duration` is
used.
References
----------
.. [#stanmanual] https://mc-stan.org/docs/2_28/reference-manual/hmc-algorithm-parameters.html
"""
if duration < 20:
raise ValueError("duration too short (< 20)")
if isinstance(init, float):
init_samples = int(init * duration)
else:
init_samples = init
if isinstance(term, float):
term_samples = int(term * duration)
else:
term_samples = term
if duration < init_samples + term_samples + base:
raise ValueError(
"duration too short (< init_samples + term_samples + base)"
)
epochs = []
epochs.append(_EpochConfig(EpochType.FAST_ADAPTATION, init_samples, thinning))
time_left = duration - init_samples - term_samples
this_time = base
while 3 * this_time <= time_left:
epochs.append(_EpochConfig(EpochType.SLOW_ADAPTATION, this_time, thinning))
time_left -= this_time
this_time *= 2
epochs.append(_EpochConfig(EpochType.SLOW_ADAPTATION, time_left, thinning))
epochs.append(_EpochConfig(EpochType.FAST_ADAPTATION, term_samples, thinning))
self.append_epochs(epochs)
[docs]
def add_burnin(self, duration: int, thinning: int = 1):
"""
Adds burnin epochs.
Note that the order in which you add epochs matters.
Adaptation epochs should usually come before all other epochs.
Warnings
--------
If you use any MCMC kernel that expects hyperparameter tuning an adaptation
phase, like :class:`~liesel.goose.NUTSKernel`, :class:`~liesel.goose.HMCKernel`
or
our default implementation of :class:`~liesel.goose.IWLSKernel`, you *must* add
an adaptation phase via :meth:`~liesel.goose.EngineBuilder.add_adaptation`. Pure
burnin
epochs like the ones added with :meth:`~liesel.goose.EngineBuilder.add_burnin`
do not perform any
adaptation during warmup. They are only applicable if you *exclusively* use
MCMC kernels that do not require adaptation, such as
:class:`~liesel.goose.GibbsKernel`
or classic, untuned IWLS kernels via :meth:`IWLSKernel.untuned
<liesel.goose.IWLSKernel.untuned>`.
See Also
--------
.add_posterior : For adding posterior epochs.
.add_burnin : For adding burnin epochs.
.add_adaptation : For adding adaptation epochs.
.set_duration : An alternative method to specify adaptation and posterior
epochs, see notes for details.
Notes
-----
There are some differences in the details between the methods
:meth:`~liesel.goose.EngineBuilder.add_posterior`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`,
:meth:`~liesel.goose.EngineBuilder.add_burnin` and
:meth:`~liesel.goose.EngineBuilder.set_duration`.
- :meth:`~liesel.goose.EngineBuilder.set_duration` overwrites all existing
epochs with the epochs defined
in this method. In contrast, the ``add_`` methods only add additional epochs
that are appended to the potentially existing ones.
- :meth:`~liesel.goose.EngineBuilder.set_duration` sets both adaptation and
posterior epochs, while each
``add_`` method only adds epochs of one type.
- :meth:`~liesel.goose.EngineBuilder.set_duration` has an argument
``term_duration`` that defines the
number of samples drawn in the *final fast* adaptation epoch, and it defaults
to a fixed value of 50 samples. In contrast,
:meth:`~liesel.goose.EngineBuilder.add_adaptation` defines
the number of samples drawn in the *final fast* adaptation epoch as a fraction
of the total number of adaptation samples, and allows users to change this
fraction via the argument ``term``. As a result, if
:meth:`~liesel.goose.EngineBuilder.add_adaptation`
is used, the length of the final fast
adaptation phase increases automatically, if the total number of warmup
samples increases, while it remains constant if
:meth:`~liesel.goose.EngineBuilder.set_duration` is
used.
"""
epoch = _EpochConfig(EpochType.BURNIN, duration, thinning)
self.append_epochs([epoch])
[docs]
def add_posterior(self, duration: int, thinning: int = 1):
"""
Adds posterior epochs.
Note that the order in which you add epochs matters.
Adaptation epochs should usually come before all other epochs.
See Also
--------
.add_posterior : For adding posterior epochs.
.add_burnin : For adding burnin epochs.
.add_adaptation : For adding adaptation epochs.
.set_duration : An alternative method to specify adaptation and posterior
epochs, see notes for details.
Notes
-----
There are some differences in the details between the methods
:meth:`~liesel.goose.EngineBuilder.add_posterior`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`,
:meth:`~liesel.goose.EngineBuilder.add_burnin` and
:meth:`~liesel.goose.EngineBuilder.set_duration`.
- :meth:`~liesel.goose.EngineBuilder.set_duration` overwrites all existing
epochs with the epochs defined
in this method. In contrast, the ``add_`` methods only add additional epochs
that are appended to the potentially existing ones.
- :meth:`~liesel.goose.EngineBuilder.set_duration` sets both adaptation and
posterior epochs, while each
``add_`` method only adds epochs of one type.
- :meth:`~liesel.goose.EngineBuilder.set_duration` has an argument
``term_duration`` that defines the
number of samples drawn in the *final fast* adaptation epoch, and it defaults
to a fixed value of 50 samples. In contrast,
:meth:`~liesel.goose.EngineBuilder.add_adaptation` defines
the number of samples drawn in the *final fast* adaptation epoch as a fraction
of the total number of adaptation samples, and allows users to change this
fraction via the argument ``term``. As a result, if
:meth:`~liesel.goose.EngineBuilder.add_adaptation`
is used, the length of the final fast
adaptation phase increases automatically, if the total number of warmup
samples increases, while it remains constant if
:meth:`~liesel.goose.EngineBuilder.set_duration` is
used.
"""
epoch = _EpochConfig(EpochType.POSTERIOR, duration, thinning)
self.append_epochs([epoch])
[docs]
def set_duration_simple(
self,
burnin_duration: int = 10,
posterior_duration: int = 10,
thinning_posterior: int = 1,
):
"""
Defines a simple epoch configuration that does not include any adaptative
warmup. Includes only burnin and posterior samples.
Warnings
--------
The default values of samples to be drawn are extremely low and not intended
for practical use -- instead, they are intended for quickly checking whether
your model runs successfully.
If you use any MCMC kernel that expects hyperparameter tuning an adaptation
phase, like :class:`~liesel.goose.NUTSKernel`, :class:`~liesel.goose.HMCKernel`
or
our default implementation of :class:`~liesel.goose.IWLSKernel`, you *must* add
an adaptation phase via :meth:`~liesel.goose.EngineBuilder.add_adaptation`. Pure
burnin
epochs like the ones added with :meth:`~liesel.goose.EngineBuilder.add_burnin`
do not perform any
adaptation during warmup. They are only applicable if you *exclusively* use
MCMC kernels that do not require adaptation, such as
:class:`~liesel.goose.GibbsKernel`
or classic, untuned IWLS kernels via :meth:`IWLSKernel.untuned
<liesel.goose.IWLSKernel.untuned>`.
See Also
--------
.add_posterior : For adding posterior epochs.
.add_burnin : For adding burnin epochs.
.add_adaptation : For adding adaptation epochs.
.set_duration : An alternative method to specify adaptation and posterior
epochs, see notes for details.
Notes
-----
There are some differences in the details between the methods
:meth:`~liesel.goose.EngineBuilder.add_posterior`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`,
:meth:`~liesel.goose.EngineBuilder.add_burnin` and
:meth:`~liesel.goose.EngineBuilder.set_duration`.
- :meth:`~liesel.goose.EngineBuilder.set_duration` overwrites all existing
epochs with the epochs defined
in this method. In contrast, the ``add_`` methods only add additional epochs
that are appended to the potentially existing ones.
- :meth:`~liesel.goose.EngineBuilder.set_duration` sets both adaptation and
posterior epochs, while each
``add_`` method only adds epochs of one type.
- :meth:`~liesel.goose.EngineBuilder.set_duration` has an argument
``term_duration`` that defines the
number of samples drawn in the *final fast* adaptation epoch, and it defaults
to a fixed value of 50 samples. In contrast,
:meth:`~liesel.goose.EngineBuilder.add_adaptation` defines
the number of samples drawn in the *final fast* adaptation epoch as a fraction
of the total number of adaptation samples, and allows users to change this
fraction via the argument ``term``. As a result, if
:meth:`~liesel.goose.EngineBuilder.add_adaptation`
is used, the length of the final fast
adaptation phase increases automatically, if the total number of warmup
samples increases, while it remains constant if
:meth:`~liesel.goose.EngineBuilder.set_duration` is
used.
"""
epochs = []
epochs.append(_EpochConfig(EpochType.BURNIN, burnin_duration, 1))
epochs.append(
_EpochConfig(EpochType.POSTERIOR, posterior_duration, thinning_posterior)
)
self.set_epochs(epochs)
[docs]
def set_duration(
self,
warmup_duration: int,
posterior_duration: int,
term_duration: int = 50,
thinning_posterior: int = 1,
thinning_warmup: int = 1,
):
"""
Sets epochs using the ``liesel.goose.stan_epochs`` function.
Note that ``term_duration`` needs to be long enough that tuning algorithms
like dual averaging can converge.
See Also
--------
.add_posterior : For adding posterior epochs.
.add_burnin : For adding burnin epochs.
.add_adaptation : For adding adaptation epochs.
.set_duration : An alternative method to specify adaptation and posterior
epochs, see notes for details.
Notes
-----
There are some differences in the details between the methods
:meth:`~liesel.goose.EngineBuilder.add_posterior`,
:meth:`~liesel.goose.EngineBuilder.add_adaptation`,
:meth:`~liesel.goose.EngineBuilder.add_burnin` and
:meth:`~liesel.goose.EngineBuilder.set_duration`.
- :meth:`~liesel.goose.EngineBuilder.set_duration` overwrites all existing
epochs with the epochs defined
in this method. In contrast, the ``add_`` methods only add additional epochs
that are appended to the potentially existing ones.
- :meth:`~liesel.goose.EngineBuilder.set_duration` sets both adaptation and
posterior epochs, while each
``add_`` method only adds epochs of one type.
- :meth:`~liesel.goose.EngineBuilder.set_duration` has an argument
``term_duration`` that defines the
number of samples drawn in the *final fast* adaptation epoch, and it defaults
to a fixed value of 50 samples. In contrast,
:meth:`~liesel.goose.EngineBuilder.add_adaptation` defines
the number of samples drawn in the *final fast* adaptation epoch as a fraction
of the total number of adaptation samples, and allows users to change this
fraction via the argument ``term``. As a result, if
:meth:`~liesel.goose.EngineBuilder.add_adaptation`
is used, the length of the final fast
adaptation phase increases automatically, if the total number of warmup
samples increases, while it remains constant if
:meth:`~liesel.goose.EngineBuilder.set_duration` is
used.
"""
epochs = stan_epochs(
warmup_duration,
posterior_duration,
term_duration=term_duration,
thinning_posterior=thinning_posterior,
thinning_warmup=thinning_warmup,
)
self.set_epochs(epochs)
@property
def epochs(self) -> tuple[EpochConfig, ...]:
"""Tuple of epoch configurations."""
return tuple(self._epochs._configs)
[docs]
def set_model(self, model: ModelInterface):
"""Sets the model interface for all kernels and quantity generators."""
# avoid circular import
from liesel.model import Model as LieselModel
if isinstance(model, LieselModel):
raise TypeError(
f"{model=} is a `lsl.Model` instance. Please wrap it in a"
" `gs.LieselInterface`."
)
self._model = Option(model)
[docs]
def build(self) -> Engine:
"""Builds the MCMC engine with the provided setup."""
# build list of position keys
pos_keys: list[str] = []
for ker in self._kernels:
pos_keys.extend(ker.position_keys)
dupl = _find_duplicate(pos_keys)
if dupl.is_some():
raise RuntimeError(
f"The position key {dupl.unwrap()} is claimed by multiple kernels"
)
# find good jittable number
epochs = self._epochs._configs # FIXME: use of private field
durations = [e.duration for e in epochs[1:]]
jit_duration = math.gcd(*durations)
# seeds
seeds = self._engine_key
if seeds.shape == (2,):
# no multi-chain key
seeds = jax.random.split(seeds, self._num_chains)
if seeds.shape != (self._num_chains, 2):
raise RuntimeError(
f"MCMC seed has the wrong dimensions {seeds}. "
f"Expected is {(self._num_chains, 2)}"
)
# check for duplicated identifiers in self._quantity_generators
idents = []
for qg in self._quantity_generators:
idents.append(qg.identifier)
dupl = _find_duplicate(idents)
if dupl.is_some():
raise RuntimeError(
f"The identifier {dupl.unwrap()} is used by multiple "
"quantity generators"
)
# set model interface for all kernels and quantity generators
model = self._model.expect("Model interface must be set")
for ker in self.kernels:
if not ker.has_model():
ker.set_model(model)
for qg in self.quantity_generators:
if not qg.has_model():
qg.set_model(model)
# assign identifiers to kernels
for idx, ker in enumerate(self.kernels):
if not ker.identifier:
ker.identifier = f"kernel_{idx:02d}"
# set initial values
model_states = self._model_state.expect("Model state must be set")
if self._jitter_fns.is_some():
jitter_fns = self._jitter_fns.unwrap()
jitter_fns_pos_keys = list(jitter_fns.keys())
missing_keys = []
for pos_key in pos_keys:
if pos_key not in jitter_fns_pos_keys:
missing_keys.append(pos_key)
if missing_keys:
pretty_keys = str(missing_keys)[1:-1]
logger.warning(
f"No jitter functions provided for position keys {pretty_keys}. "
"The initial values for these keys won't be jittered"
)
jitter_keys = jax.random.split(self._jitter_key, len(jitter_fns_pos_keys))
current_position = model.extract_position(jitter_fns_pos_keys, model_states)
jittered_position = Position({})
for i, pos_key in enumerate(jitter_fns_pos_keys):
jittered_position[pos_key] = jax.vmap(jitter_fns[pos_key])(
jax.random.split(jitter_keys[i], self._num_chains),
current_position[pos_key],
)
model_states = jax.vmap(model.update_state)(jittered_position, model_states)
else:
logger.warning(
"No jitter functions provided. The initial values won't be jittered"
)
# extending position keys
pos_keys.extend(self.positions_included)
pos_keys = [key for key in pos_keys if key not in self.positions_excluded]
return Engine(
seeds=seeds,
model_states=model_states,
kernel_sequence=KernelSequence(self.kernels),
epoch_configs=epochs,
jitted_sample_duration=jit_duration,
model=model,
position_keys=pos_keys,
minimize_transition_infos=self.minimize_transition_infos,
store_kernel_states=self.store_kernel_states,
quantity_generators=self.quantity_generators,
show_progress=self.show_progress,
)
