Skip to content

Commit

Permalink
first try at tempering / parallel tempering docstrings
Browse files Browse the repository at this point in the history
  • Loading branch information
@AudOlivier
AudOlivier committed Feb 2, 2023
1 parent 8917c83 commit 5e9dbed
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 56 deletions.
97 changes: 52 additions & 45 deletions src/UQpy/sampling/mcmc/tempering_mcmc/ParallelTemperingMCMC.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,16 @@ class ParallelTemperingMCMC(TemperingMCMC):
Parallel-Tempering MCMC
This algorithm runs the chains sampling from various tempered distributions in parallel. Periodically during the
run, the different temperatures swap members of their ensemble in a way that
preserves detailed balance.The chains closer to the reference chain (hot chains) can sample from regions that have
low probability under the target and thus allow a better exploration of the parameter space, while the cold chains
can better explore the regions of high likelihood.
run, the different temperatures swap members of their ensemble in a way that preserves detailed balance.The chains
closer to the reference chain (hot chains) can sample from regions that have low probability under the target and
thus allow a better exploration of the parameter space, while the cold chains can better explore the regions of high
likelihood.
In parallel tempering, the normalizing constant :math:`Z_1` is evaluated via thermodynamic integration. Define
the potential function :math:`U_{\beta}(x)=\frac{\partial \log{q_{\beta}(x)}}{\partial \beta}`, then
:math:`\log{Z_{1}} = \log{Z_{0}} + \int_{0}^{1} E_{x~p_{beta}} \left[ U_{\beta}(x) \right] d\beta`
where the expectations are approximated via MC sampling using samples from the intermediate distributions (see
:method:`evaluate_normalization_constant`).
**References**
Expand All @@ -24,14 +30,18 @@ class ParallelTemperingMCMC(TemperingMCMC):
**Inputs:**
Many inputs are similar to MCMC algorithms. Additional inputs are:
* **niter_between_sweeps**
* **mcmc_class**
**Methods:**
Many inputs are similar to MCMC algorithms (`n_samples`, `n_samples_per_chain`, 'random_state')
:param save_log_pdf: boolean, see :class:MCMC documentation. Importantly, this needs to be set to True if one wants
to evaluate the normalization constant via thermodynamic integration.
:param n_iterations_between_sweeps: number of iterations (sampling steps) between sweeps between chains
:param tempering_parameters: list of :math:`\beta` values so that
:math:`0 < \beta_1 < \beta_2 < \cdots < \beta_N \leq 1`. Either `tempering_parameters` or `n_tempering_parameters`
should be provided.
:param n_tempering_parameters: number of tempering levels N, the tempering parameters are selected to follow a
geometric suite :math:`\frac{1}{\sqrt{2}^(N-n)}` for n in 1:N.
:param samplers: :class:`MCMC` object or list of such objects: MCMC samplers used to sample the parallel chains. If
only one object is provided, the same MCMC sampler is used for all chains. Default to running a simple MH algorithm,
where the proposal covariance for a given chain is :math:`\frac{1}{\beta}`.
"""

@beartype
Expand Down Expand Up @@ -63,9 +73,9 @@ def __init__(self, n_iterations_between_sweeps: PositiveInteger,
self.n_tempering_parameters = n_tempering_parameters
if self.tempering_parameters is None:
if self.n_tempering_parameters is None:
raise ValueError('UQpy: either input temper_param_list or n_temper_params should be provided.')
raise ValueError('UQpy: either input tempering_parameters or n_tempering_parameters should be provided.')
elif not (isinstance(self.n_tempering_parameters, int) and self.n_tempering_parameters >= 2):
raise ValueError('UQpy: input n_temper_params should be a integer >= 2.')
raise ValueError('UQpy: input n_tempering_parameters should be a integer >= 2.')
else:
self.tempering_parameters = [1. / np.sqrt(2) ** i for i in
range(self.n_tempering_parameters - 1, -1, -1)]
Expand All @@ -74,7 +84,7 @@ def __init__(self, n_iterations_between_sweeps: PositiveInteger,
# or float(self.temperatures[0]) != 1.
):
raise ValueError(
'UQpy: temper_param_list should be a list of floats in [0, 1], starting at 0. and increasing to 1.')
'UQpy: tempering_parameters should be a list of floats in [0, 1], starting at 0. and increasing to 1.')
else:
self.n_tempering_parameters = len(self.tempering_parameters)

Expand All @@ -89,8 +99,10 @@ def __init__(self, n_iterations_between_sweeps: PositiveInteger,

# Initialize algorithm specific inputs: target pdfs
self.thermodynamic_integration_results = None
"""Results of the thermodynamic integration (normalization constant). """

self.mcmc_samplers = []
"""List of MCMC samplers, one per tempering level. """
for i, temper_param in enumerate(self.tempering_parameters):
log_pdf_target = (lambda x, temper_param=temper_param: self.evaluate_log_reference(
x) + self.evaluate_log_intermediate(x, temper_param))
Expand All @@ -110,18 +122,13 @@ def run(self, nsamples: PositiveInteger = None, nsamples_per_chain: PositiveInte
Run the MCMC algorithm.
This function samples from the MCMC chains and appends samples to existing ones (if any). This method leverages
the ``run_iterations`` method that is specific to each algorithm.
**Inputs:**
* **nsamples** (`int`):
Number of samples to generate.
the :method:`run_iterations` method specific to each of the samplers.
* **nsamples_per_chain** (`int`)
Number of samples to generate per chain.
:param nsamples: Number of samples to generate from the target (the same number of samples will be generated
for all intermediate distributions).
Either `nsamples` or `nsamples_per_chain` must be provided (not both). Not that if `nsamples` is not a multiple
of `nchains`, `nsamples` is set to the next largest integer that is a multiple of `nchains`.
:param nsamples_per_chain: Number of samples to generate per chain of the target MCMC sampler. Either `nsamples`
or `nsamples_per_chain` must be provided (not both).
"""
current_state, current_log_pdf = [], []
Expand Down Expand Up @@ -187,33 +194,33 @@ def run(self, nsamples: PositiveInteger = None, nsamples_per_chain: PositiveInte

# Samples connect to posterior samples, i.e. the chain with beta=1.
self.intermediate_samples = [sampler.samples for sampler in self.mcmc_samplers]
"""List of samples for the intermediate tempering levels. """
self.samples = self.mcmc_samplers[-1].samples
"""Samples from the target distribution. """
if self.save_log_pdf:
self.log_pdf_values = self.mcmc_samplers[-1].log_pdf_values
"""Log pdf values for samples from the target. """

@beartype
def evaluate_normalization_constant(self, compute_potential, log_Z0: float = None, nsamples_from_p0: int = None):
"""
Evaluate new log free energy as
:math:`\log{Z_{1}} = \log{Z_{0}} + \int_{0}^{1} E_{x~p_{beta}} \left[ U_{\beta}(x) \right] d\beta`
References (for the Bayesian case):
* https://emcee.readthedocs.io/en/v2.2.1/user/pt/
**Inputs:**
* **compute_potential** (callable):
Function that takes three inputs (`x`, `log_factor_tempered_values`, `beta`) and computes the potential
:math:`U_{\beta}(x)`. `log_factor_tempered_values` are the values saved during sampling of
:math:`\log{p_{\beta}(x)}` at saved samples x.
* **log_Z0** (`float`):
Value of :math:`\log{Z_{0}}`
* **nsamples_from_p0** (`int`):
N samples from the reference distribution p0. Then :math:`\log{Z_{0}}` is evaluate via MC sampling
as :math:`\frac{1}{N} \sum{p_{\beta=0}(x)}`. Used only if input *log_Z0* is not provided.
Evaluate normalization constant :math:`Z_1` as:
:math:`\log{Z_{1}} \approx \log{Z_{\beta_{1}}} + \int_{\beta_{1}}^{\beta_{N}} \frac{1}{M} \sum_m U_{\beta}(x_m)
d\beta`
where :math:`x_m` are samples from the intermediate distributions. The integration is performed via a
trapezoidal rule. The function returns an approximation of :math:`Z_1`, and also saves intermediate results
(value of :math:`log_Z0`, list of tempering parameters used in integration, and values of the associated
expected potentials :math:`E_{x \sim p_{\beta} \left[ U_{\beta}(x) \right]}\frac{1}{M} \sum_m U_{\beta}(x_m)`)
:param compute_potential: Function that takes three inputs: :code:`x` (sample points where to evaluate the
potential), :code:`log_factor_tempered_values` (values of :math:`\log{q_{\beta}(x)}` at these points), `beta`
(tempering parameter) and evaluates the potential
:math:`U_{\beta}(x)=\frac{\partial \log{q_{\beta}(x)}}{\partial \beta}`.
:param log_Z0: Value of :math:`\log{Z_{0}}` (float), if unknwon, see `nsamples_from_p0`
:param nsamples_from_p0: number of samples :math:`M_0` from the reference distribution :math:`p_0` to be used
to evaluate :math:`\log{Z_{0}} \approx \frac{1}{M_0} \sum{p_{\beta=0}(x)}`. Used only if input `log_Z0` is not
provided.
"""
if not self.save_log_pdf:
Expand Down
44 changes: 33 additions & 11 deletions src/UQpy/sampling/mcmc/tempering_mcmc/baseclass/TemperingMCMC.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,37 @@ class TemperingMCMC(ABC):
"""
Parent class to parallel and sequential tempering MCMC algorithms.
To sample from the target distribution :math:`p(x)`, a sequence of intermediate densities
:math:`p(x, \beta) \propto q(x, \beta) p_{0}(x)` for values of the parameter :math:`\beta` between 0 and 1,
where :math:`p_{0}` is a reference distribution (often set as the prior in a Bayesian setting).
Setting :math:`\beta = 1` equates sampling from the target, while
:math:`\beta \rightarrow 0` samples from the reference distribution.
**Inputs:**
**Methods:**
These algorithms aim at sampling from a target distribution :math:`p_1(x)` of the form
:math:`p_1(x)=\frac{q_1(x)p_0(x)}{Z_1}` where the intermediate factor :math:`q_1(x)` and reference distribution
:math:`p_0(x)` can be evaluated. Additionally, these algorithms return an estimate of the normalization
constant :math:`Z_1=\int{q_{1}(x) p_{0}(x)dx}`.
The algorithms sample from a sequence of intermediate densities
:math:`p_{\beta}(x) \propto q_{\beta}(x) p_{0}(x)` for values of the parameter :math:`\beta` between 0 and 1
(:math:`\beta=\frac{1}{T}` where :math:`T` is sometimes called the temperature).
Setting :math:`\beta = 1` equates sampling from the target, while :math:`\beta \rightarrow 0` samples from the
reference distribution.
Parallel tempering samples from all distributions simultaneously, and the tempering parameters :math:`\beta` must be
chosen in advance by the user. Sequential tempering samples from the various distributions sequentially, starting
from the reference distribution, and the tempering parameters are selected adaptively by the algorithm.
Inputs that are common to both parallel and sequential tempering algorithms are:
:param pdf_intermediate: callable that computes the intermediate factor :math:`q_{\beta}(x)`. It should take at
least two inputs :code:`x` (ndarray, point(s) at which to evaluate the function), and :code:`temper_param` (float,
tempering parameter :math:`\beta`). Either `pdf_intermediate` or `log_pdf_intermediate` must be provided
(`log_pdf_intermediate` is preferred). Within the code, the `log_pdf_intermediate` is evaluated as:
:code:`log q_{\beta}(x) = log_pdf_intermediate(x, \beta, *args_pdf_intermediate)`
where `args_pdf_intermediate` are additional positional arguments that are provided to the class via its
`args_pdf_intermediate` input.
:param log_pdf_intermediate: see `pdf_intermediate`
:param args_pdf_intermediate: see `pdf_intermediate`
:param distribution_reference: reference pdf :math:`p_0` as a :class:`.Distribution` object
:param save_log_pdf: see same input in :class:`MCMC`
:param random_state
"""

def __init__(self, pdf_intermediate=None, log_pdf_intermediate=None, args_pdf_intermediate=(),
Expand Down Expand Up @@ -44,8 +66,8 @@ def run(self, nsamples):

@abstractmethod
def evaluate_normalization_constant(self, **kwargs):
""" Computes the normalization constant :math:`Z_{1}=\int{q_{1}(x) p_{0}(x)dx}` where p0 is the reference pdf
and q1 is the intermediate density with :math:`\beta=1`, thus q1 p0 is the target pdf."""
""" Computes the normalization constant :math:`Z_{1}=\int{q_{1}(x) p_{0}(x)dx}` where :math:`p_0` is the
reference pdf and :math:`q_1` is the target factor."""
pass

def _preprocess_reference(self, dist_, **kwargs):
Expand Down

0 comments on commit 5e9dbed

Please sign in to comment.