diff --git a/.gitignore b/.gitignore index 8ef9160b..43d0f951 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,6 @@ -scratch_* -*ipynb_checkpoints -*.DS_Store -*__pycache__ -build/* -epsie.egg-info/* +build/ +epsie.egg-info/ +.DS_Store reinstall_epsie.sh +*.pyc +scratch_* diff --git a/epsie/diagnostic/__init__.py b/epsie/diagnostic/__init__.py new file mode 100644 index 00000000..e940070d --- /dev/null +++ b/epsie/diagnostic/__init__.py @@ -0,0 +1,17 @@ +# Copyright (C) 2019 Collin Capano +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + +from .acl import acl_1d, acl_chain, thinned_samples +from .convergence import gelman_rubin_test diff --git a/epsie/diagnostic/acl.py b/epsie/diagnostic/acl.py new file mode 100644 index 00000000..6aae3a24 --- /dev/null +++ b/epsie/diagnostic/acl.py @@ -0,0 +1,222 @@ +# Copyright (C) 2022 Richard Stiskalek, Collin Capano +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + +"""Autocorrelation time calculation.""" +import numpy + +from epsie.samplers import MetropolisHastingsSampler, ParallelTemperedSampler + + +def acl_1d(x, c=5.0): + """ + Calculate the autocorrelation length (ACL) of some series. + + Algorithm used is from: + N. Madras and A.D. Sokal, J. Stat. Phys. 50, 109 (1988). + + Implementation from: + https://emcee.readthedocs.io/en/stable/tutorials/autocorr/#autocorr + + Arguments + --------- + x : 1-dimensional array + Samples whose autocorrelation length is to be calculated. + c : float, optional + ACL hyperparameter. As per Sokal by default set to 5. + + Returns + ------- + ACL: int + Autocorrelation length. + """ + if x.ndim != 1: + raise TypeError("``x`` not a 1D array. Currently {}".format(x.ndim)) + # n is the next power of 2 for length of x + n = (2**numpy.ceil(numpy.log2(x.size))).astype(int) + # Compute the FFT and then (from that) the auto-correlation function + f = numpy.fft.fft(x - numpy.mean(x), n=2 * n) + acf = numpy.fft.ifft(f * numpy.conjugate(f))[: len(x)].real / (4 * n) + + acf /= acf[0] + + taus = 2.0 * numpy.cumsum(acf) - 1.0 + window = auto_window(taus, c) + acl = numpy.ceil(taus[window]).astype(int) + if acl < 1: + acl = 1 + return acl + + +def auto_window(taus, c): + """ + ACL automated windowing procedure following Sokal (1989). + + Implementation from: + https://emcee.readthedocs.io/en/stable/tutorials/autocorr/#autocorr + """ + m = numpy.arange(len(taus)) < c * taus + if numpy.any(m): + return numpy.argmin(m) + return len(taus) - 1 + + +def acl_chain(chain, burnin_iter=0, c=5.0, full=False): + """ + Calculate the ACL of :py:class:`epsie.chain.Chain` chain. + + Arguments + --------- + chain : :py:class:`epsie.chain.Chain` + Epsie chain. + burnin_iter : int, optional + Number of burnin iterations to be thrown away. + c : float, optional + ACL calculation hyperparameter. By default 5. + full : bool, optionall + Whether to return the ACL of every parameter. By default False, thus + returning only the maximum ACL. + + Returns + ------- + ACL : int or array of integers + Chain's ACL(s). + """ + acls = [acl_1d(chain.positions[p][burnin_iter:]) for p in chain.parameters] + if full: + return acls + return max(acls) + + +def thinned_samples(sampler, burnin_iter=0, c=5.0, temp_acls_method="coldest"): + """ + Parse a sampler and return its samples thinned by the ACL. + + Arguments + --------- + sampler : {:py:class:`epsie.sampler.MetropolisHastingsSampler`, + :py:class:`epsie.sampler.ParallelTemperedSampler`} + Epsie sampler whose samples to extract. + burnin_iter : int, optional + Number of burnin iterations to be thrown away. + c : float, optional + ACL calculation hyperparameter. By default 5. + temp_acls_method : string, optional + Tempered acls selection, either `coldest` or `max`. The former + corresponds to taking the ACL of the coldest chain, while the latter to + the maximum ACL across all temperatures. This single ACL is then used + to thin all temperatures. + + Returns + ------- + thinned samples : structured array + Array whose `dtype.names` are the sampler parameters. Each parameter's + data is 1-dimensional for the MH sampler and 2-dimensional for the PT + sampler. In the latter case the first and second axis represent the + temperature and sample index. + """ + if isinstance(sampler, MetropolisHastingsSampler): + return _thinned_mh_samples(sampler, burnin_iter, c) + elif isinstance(sampler, ParallelTemperedSampler): + return _thinned_pt_samples(sampler, burnin_iter, c, temp_acls_method) + else: + raise ValueError("Unknown sampler type ``{}``".format(type(sampler))) + + +def _get_raw_data(sampler, burnin_iter): + raw = {} + raw.update({p: sampler.positions[..., burnin_iter:][p] + for p in sampler.parameters}) + raw.update({p: sampler.stats[..., burnin_iter:][p] + for p in ("logl", "logp")}) + if sampler.blobs is not None: + raw.update({p: sampler.blobs[..., burnin_iter:][p] + for p in sampler.blobs.dtype.names}) + return raw + + +def _thinned_mh_samples(sampler, burnin_iter=0, c=5.0): + # Calculate the ACL for each chain + acls = [acl_chain(chain, burnin_iter, c) for chain in sampler.chains] + + keys = sampler.parameters + ("logl", "logp") + keys += sampler.blobs.dtype.names if sampler.blobs is not None else () + raw = _get_raw_data(sampler, burnin_iter) + + # Cycle over the chains and thin them + _thinned = {p: [] for p in keys} + for ii in range(sampler.nchains): + for key in keys: + _thinned[key].append(raw[key][ii, :][::-1][::acls[ii]][::-1]) + + # Put the thinned samples into a structured array + dtype = sampler.positions.dtype.descr + sampler.stats.dtype.descr + if sampler.blobs is not None: + dtype += sampler.blobs.dtype.descr + thinned = numpy.zeros(sum(x.size for x in _thinned[keys[0]]), + dtype=numpy.dtype(dtype)) + for p in keys: + thinned[p] = numpy.concatenate(_thinned[p]) + + return thinned + + +def _thinned_pt_samples(sampler, burnin_iter=0, c=5.0, + temp_acls_method="coldest"): + allowed_methods = ["coldest", "max"] + if temp_acls_method not in allowed_methods: + raise ValueError("`Invalid `temp_acls_method`. Allowed methods: `{}`" + .format(allowed_methods)) + # Calculate the ACL across temperatures for each chain + temp_acls = numpy.zeros((sampler.nchains, sampler.ntemps), dtype=int) + for ii in range(sampler.nchains): + for jj in range(sampler.ntemps): + temp_acls[ii, jj] = acl_chain( + sampler.chains[ii].chains[jj], burnin_iter, c=c) + # By default we only take the coldest chain. No need to calculate + # the hotter chains then. + if temp_acls_method == "coldest": + break + # Grab the ACL for each chain. By default ACL of the coldest chain + if temp_acls_method == "coldest": + acls = temp_acls[:, 0] + else: + acls = numpy.max(temp_acls, axis=1) + + keys = sampler.parameters + ("logl", "logp") + keys += sampler.blobs.dtype.names if sampler.blobs is not None else () + raw = _get_raw_data(sampler, burnin_iter) + + thinned = {p: [] for p in keys} + # Cycle over the chains, thinning them + for tk in range(sampler.ntemps): + _thinned = {p: [] for p in keys} + + for ii in range(sampler.nchains): + for key in keys: + _thinned[key].append( + raw[key][tk, ii, :][::-1][::acls[ii]][::-1]) + + for param in keys: + thinned[param].append(numpy.concatenate(_thinned[param])) + # Cast the thinned samples from a list of arrays into a 2D array for each + # parameter and return as a structured array + dtype = sampler.positions.dtype.descr + sampler.stats.dtype.descr + if sampler.blobs is not None: + dtype += sampler.blobs.dtype.descr + thinned_samples = numpy.zeros((sampler.ntemps, thinned[keys[0]][0].size), + dtype=numpy.dtype(dtype)) + for param in keys: + thinned_samples[param] = numpy.asarray(thinned[param]) + return thinned_samples diff --git a/epsie/diagnostic/convergence.py b/epsie/diagnostic/convergence.py new file mode 100644 index 00000000..8348e051 --- /dev/null +++ b/epsie/diagnostic/convergence.py @@ -0,0 +1,82 @@ +# Copyright (C) 2022 Richard Stiskalek, Collin Capano +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + +"""MCMC convergence tests.""" + +import numpy +from epsie.samplers import MetropolisHastingsSampler, ParallelTemperedSampler + + +def gelman_rubin_test(sampler, burnin_iter): + """ + Calculate the Gelman-Rubin (GR) convergence test outlined in [1] and + summarised in [2] for parallel, independent MCMC chains. + + For a PT sampler the GR statistic is calculated for each temperature level. + + Typically values of less than 1.1 or 1.2 are recommended. + + Arguments + --------- + sampler : {:py:class:`epsie.sampler.MetropolisHastingsSampler`, + :py:class:`epsie.sampler.ParallelTemperedSampler`} + Epsie sampler whose samples to extract. + burnin_iter : int, optional + Number of burnin iterations to be thrown away. + + Returns + ------- + Rs : array + GR statistic of each parameter. In case of a MH sampler 1-dimensional + array of length `len(sampler.paramaters)` and in case of a PT sampler + 2-dimensional array of shape `len(sampler.ntemps, len(sampler.params)`. + + References + ---------- + .. [1] Gelman, A. and Rubin, D.B. (1992). "Inference from Iterative + Simulation using Multiple Sequences". Statistical Science, 7, + p. 457–511. + .. [2] https://mc-stan.org/docs/2_18/reference-manual/ + notation-for-samples-chains-and-draws.html + """ + params = sampler.parameters + # Cut off burnin iterations and for PT take coldest chains + samples = sampler.positions[..., burnin_iter:] + if isinstance(sampler, MetropolisHastingsSampler): + return _gelman_rubin_at_temp(samples, params) + elif isinstance(sampler, ParallelTemperedSampler): + Rs = numpy.zeros((sampler.ntemps, len(params))) + for tk in range(sampler.ntemps): + Rs[tk, :] = _gelman_rubin_at_temp(samples[tk, ...], params) + return Rs + else: + raise ValueError("Unknown sampler type ``{}``".format(type(sampler))) + + +def _gelman_rubin_at_temp(samples, params): + """Calculate the Gelman Rubin statistic at a given temperature level.""" + Nchains, N = samples.shape + if Nchains == 1: + raise TypeError("GL statistic requires > 1 independent chains.") + # Calculate the GH statistic for each parameter independently + Rs = numpy.zeros(len(params)) + for i, param in enumerate(params): + # Between chains variance + B = N * numpy.var(numpy.mean(samples[param], axis=1), ddof=1) + # Within chains variance + W = numpy.mean(numpy.var(samples[param], axis=1, ddof=1)) + Rs[i] = ((N - 1) / N + 1 / N * B / W)**0.5 + + return Rs diff --git a/epsie/samplers/ptsampler.py b/epsie/samplers/ptsampler.py index 7df03fa1..5b1186f8 100644 --- a/epsie/samplers/ptsampler.py +++ b/epsie/samplers/ptsampler.py @@ -67,6 +67,7 @@ class ParallelTemperedSampler(BaseSampler): Specify a process pool to use for parallelization. Default is to use a single core. """ + def __init__(self, parameters, model, nchains, betas, swap_interval=1, proposals=None, adaptive_annealer=None, reset_after_swap=False, default_proposal=None, diff --git a/test/test_diagnostics.py b/test/test_diagnostics.py new file mode 100644 index 00000000..c49dde35 --- /dev/null +++ b/test/test_diagnostics.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python + +# Copyright (C) 2020 Collin Capano +# This program is free software; you can redistribute it and/or modify it +# under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 3 of the License, or (at your +# option) any later version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General +# Public License for more details. +# +# You should have received a copy of the GNU General Public License along +# with this program; if not, write to the Free Software Foundation, Inc., +# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. +"""Performs unit tests on the ParallelTempered sampler.""" + +import pytest +import numpy +from epsie import make_betas_ladder +from epsie import diagnostic +from _utils import Model + +from test_ptsampler import _create_sampler as _create_pt_sampler +from test_mhsampler import _create_sampler as _create_mh_sampler + + +NCHAINS = 5 +NTEMPS = 2 +BETAS = make_betas_ladder(NTEMPS, 1e5) +ITERINT = 256 +SEED = 2020 + + +@pytest.mark.parametrize("sampler_cls", ["mh", "pt"]) +def test_thinning(sampler_cls): + if sampler_cls == "mh": + _test_mh_thinning() + elif sampler_cls == "pt": + _test_pt_thinning("coldest") + _test_pt_thinning("max") + else: + raise ValueError("Unknown sampler class.") + + +def _test_mh_thinning(): + model = Model() + sampler = _create_mh_sampler(model, nprocs=1, nchains=NCHAINS, + seed=SEED) + sampler.start_position = model.prior_rvs(size=NCHAINS) + # run both for a few iterations + sampler.run(ITERINT) + + # Given a random seed test a comparison to a known value + acl = diagnostic.acl_chain(sampler.chains[0], full=True) + # Check that the thinnd arrays have the right shape + thinned = diagnostic.thinned_samples(sampler, burnin_iter=int(ITERINT/2)) + shape = None + for i, param in enumerate(sampler.parameters): + px = thinned[param] + assert px.ndim == 1 + if i == 0: + shape = px.shape + else: + assert px.shape == shape + + # Test the GR calculation + Rs = diagnostic.gelman_rubin_test( + sampler, burnin_iter=int(ITERINT/2)) + assert isinstance(Rs, numpy.ndarray) + assert Rs.ndim == 1 + assert ~numpy.any(Rs < 1) + + +def _test_pt_thinning(temp_acls_method): + model = Model() + sampler = _create_pt_sampler(model, 1, nchains=NCHAINS, betas=BETAS, + seed=SEED, swap_interval=1, proposals=None, + set_start=False) + sampler.start_position = model.prior_rvs(size=(NTEMPS, NCHAINS)) + + sampler.run(ITERINT) + + thinned = diagnostic.thinned_samples( + sampler, burnin_iter=int(ITERINT/2), temp_acls_method=temp_acls_method) + + shape = None + for i, param in enumerate(sampler.parameters): + px = thinned[param] + assert px.ndim == 2 + if i == 0: + shape = px.shape + else: + assert px.shape == shape + + # Test the GR calculation + Rs = diagnostic.gelman_rubin_test( + sampler, burnin_iter=int(ITERINT/2)) + assert isinstance(Rs, numpy.ndarray) + print(Rs) + assert Rs.ndim == 2 + assert ~numpy.any(Rs < 1)