<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A kernel-density-based, embarrassingly parallel ensemble sampler.
"""

from __future__ import (division, print_function, absolute_import, unicode_literals)
from .pool import choose_pool
import numpy as np
import numpy.ma as ma
from scipy.stats import chisquare
from .clustered_kde import optimized_kde, TransdimensionalKDE


def dynamic_pbar_update(iter, test_size, acc, pbar=None):
    if pbar is not None:
        pbar.set_postfix_str('| Single-Step Acceptence Rate: {0} | Test Stepsize: {1} >= 16'.format(acc, test_size))
        pbar.update(iter - pbar.n)


def in_notebook():
    import __main__ as main
    return not hasattr(main, '__file__')


class _GetLnProbWrapper(object):
    """Convenience class for evaluating multiple probability densities at a single point."""
    def __init__(self, lnpost, kde, *args):
        self.lnpost = lnpost
        self.kde = kde
        self.args = args

    def update_kde(self, kde):
        self.kde = kde

    def lnprobs(self, p):
        """
        Evaluate the log probability density of the stored target distribution fuction
        and KDE at `p`.

        :param p: Location to evaluate probability densties at.

        :returns: ``lnpost(p)``, ``kde(p)``
        """
        result = self.lnpost(p, *self.args)
        kde = self.kde(p)

        # allow posterior function to optionally
        # return additional metadata
        try:
            lnpost = result[0]
            blob = result[1]
            return lnpost, kde, blob
        except (IndexError, TypeError):
            lnpost = result
            return lnpost, kde

    __call__ = lnprobs


_lnprob_wrapper = None


def _set_global_lnprob_wrapper(wrapper_instance):
    """Sets `_lnprob_wrapper` to a global variable equal to the provided
    instance of `_GetLnProbWrapper`. Making the wrapper a global variable
    allows chilren processes to access their own copy of the wrapper without
    needing to push out the same data on every call. For this to work it must
    be done *before* a multiprocessing pool is initialized. See
    http://stackoverflow.com/a/10118250 for details.
    """
    global _lnprob_wrapper
    _lnprob_wrapper = wrapper_instance


def _get_lnprob_from_wrapper(p):
    return _lnprob_wrapper.lnprobs(p)

def _update_wrapper_kde(kde):
    _lnprob_wrapper.update_kde(kde)


class Sampler(object):
    """
    An Ensemble sampler.

    The :attr:`chain` member of this object has the shape: `(nsteps, nwalkers, ndim)` where
    `nsteps` is the stored number of steps taken thus far.

    :param nwalkers:
        The number of individual MCMC chains to include in the ensemble.

    :param ndim:
        Number of dimensions in the parameter space.  If `transd` is ``True`` this is the number of
        unique dimensions across the parameter spaces.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and returns the natural
        logarithm of the posterior probability for that position.

    :param transd:
        If ``True``, the sampler will operate across parameter spaces using a
        :class:`.clustered_kde.TransdimensionalKDE` proposal distribution. In this mode a masked
        array with samples in each of the possible sets of dimensions must be given for the initial
        ensemble distribution.

    :param processes: (optional)
        The number of processes to use with :mod:`multiprocessing`.  If ``None``, all available
        cores are used.

    :param pool: (optional)
        A pre-constructed pool with a map method. If ``None`` a pool will be created using
        :mod:`multiprocessing`.

    """
    def __init__(self, nwalkers, ndim, lnpostfn, transd=False,
                 processes=None, pool=None, mpi=False, args=[]):
        self.nwalkers = nwalkers
        self.dim = ndim
        self._kde = None
        self._kde_size = self.nwalkers
        self._updates = []
        self._burnin_length = None

        self._get_lnpost = lnpostfn
        self._lnpost_args = args

        self.iterations = 0
        self.stored_iterations = 0
        self.mpi = mpi
        self.processes = processes
        if processes is None and pool is not None:
            raise ValueError("Please provide the number of processes if "
                             " if also providing a pool instance.")
        self._managing_pool = False
        if pool is not None:
            self._pool = pool
        else:
            self._managing_pool = True
            self._pool = choose_pool(processes, mpi=mpi)

        if not hasattr(self._pool, 'map'):
            raise AttributeError("Pool object must have a map() method.")

        self._transd = transd
        if self._transd:
            self._chain = ma.masked_all((0, self.nwalkers, self.dim))
        else:
            self._chain = np.zeros((0, self.nwalkers, self.dim))

        self._lnpost = np.empty((0, self.nwalkers))
        self._lnprop = np.empty((0, self.nwalkers))
        self._acceptance = np.zeros((0, self.nwalkers))
        self._blobs = []
        self.tqdm = None

        self._last_run_mcmc_result = None
        self._burnin_spaces = None
        self._failed_p = None
        self._set_wrapper()

    def _set_wrapper(self):
        """ Determine the lnprob wrapper """
        # If the pool can guarantee a call to all processes we can update
        # the data explicitly, which is much faster
        if hasattr(self._pool, 'broadcast'):
            wrapper = _GetLnProbWrapper(self._get_lnpost, self._kde,

                                        *self._lnpost_args)
            self._pool.broadcast(_set_global_lnprob_wrapper, wrapper)

    def _get_wrapper(self):
        """ Retunr the lnprob wrapper function call """
        if hasattr(self._pool, 'broadcast'):
            return _get_lnprob_from_wrapper
        else:
            return _GetLnProbWrapper(self._get_lnpost, self._kde,

                                     *self._lnpost_args)

    def _get_dynamci_pbar(self, progress):
        pbar = None
        if progress:
            import tqdm
            self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(desc="Burning in until Constant Acceptence Rate", position=0, leave=True)
            else:
                pbar = self.tqdm.tqdm(desc="Burning in until Constant Acceptence Rate", position=0, leave=True)
        return pbar, dynamic_pbar_update

    def _get_finite_pbar(self, progress, N):
        pbar = None
        if progress:
            if self.tqdm is None:
                import tqdm
                self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(total=N, desc="Running MCMC")
            else:
                pbar = self.tqdm.tqdm(total=N, desc="Running MCMC")
        return pbar

    def burnin(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,
               test_steps=16, critical_pval=0.05, max_steps=None,
               verbose=False, callback=None, progress=False, **kwargs):
        """
        Evolve an ensemble until the acceptance rate becomes roughly constant.  This is done by
        splitting acceptances in half and checking for statistical consistency.  This isn't
        guaranteed to return a fully burned-in ensemble, but usually does.

        :param p0: (optional)
            A list of the initial walker positions.  It should have the shape `(nwalkers, ndim)`.
            If ``None`` and the sampler has been run previously, it'll pick up where it left off.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at positions `p0`. If ``lnpost0
            is None``, the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param lnprop0: (optional)
            List of log proposal densities for walkers at positions `p0`. If ``lnprop0 is None``,
            the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param blob0: (optional)
            The list of blob data for walkers at positions `p0`.

        :param test_steps: (optional)
            The (rough) number of accepted steps over which to check for acceptance rate
            consistency. If you find burnin repeatedly ending prematurely try increasing this.

        :param critical_pval: (optional)
            The critial p-value for considering the acceptance distribution
            consistent.  If the calculated p-value is over this, then the acceptance rate
            is considered to be stationary.  Lower this if you want burnin criteria to be
            less strict.

        :param max_steps: (optional)
            An absolute maximum number of steps to take, in case burnin is too painful.

        :param verbose: (optional)
            Print status messages each time a milestone is reached in the burnin.

        :param kwargs: (optional)
            The rest is passed to :meth:`run_mcmc`.

        After burnin...

        :returns:
            * ``p`` - A list of the current walker positions with shape `(nwalkers, ndim)`.

            * ``lnpost`` - Array of log posterior probabilities for walkers at positions `p`; has
              shape `(nwalkers, )`.

            * ``lnprop`` - Array of log proposal densities for walkers at positions `p`; has shape
              `(nwalkers, )`.

            * ``blob`` - (if `lnprobfn` returns blobs) The list of blob data for the walkers at
              positions `p`.
        """
        if p0 is not None:
            p0 = np.atleast_2d(p0)

        # Determine the maximum iteration to look for
        start = self.iterations
        # Confine walkers to their space during burnin
        freeze_transd = False
        if self._transd:
            freeze_transd = True
            self._burnin_spaces = ~p0.mask
        pbar, pbar_status = self._get_dynamci_pbar(progress)
        max_iter = np.inf
        if max_steps is not None:
            max_iter = start + max_steps
        if progress:
            verbose = False

        sampling_ct = 0
        step_size = 2
        while step_size <= test_steps:
            # Update the proposal
            if p0 is not None:
                self.update_proposal(p0, max_samples=self.nwalkers)
                lnprop0 = self._kde(p0)
            if verbose:
                print('Updated proposal')

            # Take one step to estimate acceptance rate
            test_interval = 1
            results = self.run_mcmc(test_interval, p0, lnpost0, lnprop0, blob0,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            # Use the fraction of acceptances in the last step to estimate acceptance rate
            #   Bottom out at 1% if acceptances are really bad
            last_acc_rate = max(np.mean(self.acceptance[-1]), 0.01)

            # Estimate ACT based on acceptance
            act = int(np.ceil(2.0/last_acc_rate - 1.0))

            if verbose:
                print('Single-step acceptance rate is ', last_acc_rate)
                print('Producing ACT of ', act)

            # Use the ACT to set the new test interval, but avoid
            # overstepping a specified max.  We throw away the first
            # 2*act worth of steps as an initial burnin when comparing
            # acceptance rates
            test_interval = min((step_size+2)*act, max_iter - self.iterations)

            # Make sure we're taking at least one step
            test_interval = max(test_interval, 1)
            sampling_ct += 1
            pbar_status(sampling_ct, step_size, last_acc_rate, pbar)

            if verbose:
                if ~np.isinf(max_iter):
                    print("Time {} running mcmc during burnin with {}/{} Iterations and test_interval = {}:".format(
                        sampling_ct, self.iterations, max_iter, test_interval))
                else:
                    print("Time {} running mcmc during burnin with test_interval = {}:".format(
                        sampling_ct, test_interval))

            results = self.run_mcmc(test_interval, p, lnpost, lnprop, blob,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces, progress=progress,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            if callback is not None:
                callback(self)

            # Quit if we hit the max
            if self.iterations >= max_iter:
                print("Burnin hit {} iterations before completing.".format(max_iter))
                break

            # Only check for consistency past the burn-in stage of 2*act
            if self.consistent_acceptance_rate(window_size=step_size*act, critical_pval=critical_pval):
                if verbose:
                    print('Acceptance rate constant over ', step_size, ' ACTs')
                step_size *= 2
            else:
                if verbose:
                    print('Acceptance rate varies, trying again')

            if verbose:
                print('') # Newline

            p0, lnpost0, lnprop0, blob0 = p, lnpost, lnprop, blob

        self._burnin_length = self.updates[-1]
        if pbar is not None:
            pbar.close()

        if blob is None:
            return p, lnpost, lnprop
        else:
            return p, lnpost, lnprop, blob

    def sample(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,<commit_msg>change outer burnin pbar to show progress from step_size -> test_steps<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
A kernel-density-based, embarrassingly parallel ensemble sampler.
"""

from __future__ import (division, print_function, absolute_import, unicode_literals)
from .pool import choose_pool
import numpy as np
import numpy.ma as ma
from scipy.stats import chisquare
from .clustered_kde import optimized_kde, TransdimensionalKDE


def dynamic_pbar_update(step_size, last_step_size, acc, pbar=None):
    if pbar is not None:
        pbar.set_postfix_str('| Single-Step Acceptence Rate: {0}'.format(acc))
        pbar.update(step_size-last_step_size)


def in_notebook():
    import __main__ as main
    return not hasattr(main, '__file__')


class _GetLnProbWrapper(object):
    """Convenience class for evaluating multiple probability densities at a single point."""
    def __init__(self, lnpost, kde, *args):
        self.lnpost = lnpost
        self.kde = kde
        self.args = args

    def update_kde(self, kde):
        self.kde = kde

    def lnprobs(self, p):
        """
        Evaluate the log probability density of the stored target distribution fuction
        and KDE at `p`.

        :param p: Location to evaluate probability densties at.

        :returns: ``lnpost(p)``, ``kde(p)``
        """
        result = self.lnpost(p, *self.args)
        kde = self.kde(p)

        # allow posterior function to optionally
        # return additional metadata
        try:
            lnpost = result[0]
            blob = result[1]
            return lnpost, kde, blob
        except (IndexError, TypeError):
            lnpost = result
            return lnpost, kde

    __call__ = lnprobs


_lnprob_wrapper = None


def _set_global_lnprob_wrapper(wrapper_instance):
    """Sets `_lnprob_wrapper` to a global variable equal to the provided
    instance of `_GetLnProbWrapper`. Making the wrapper a global variable
    allows chilren processes to access their own copy of the wrapper without
    needing to push out the same data on every call. For this to work it must
    be done *before* a multiprocessing pool is initialized. See
    http://stackoverflow.com/a/10118250 for details.
    """
    global _lnprob_wrapper
    _lnprob_wrapper = wrapper_instance


def _get_lnprob_from_wrapper(p):
    return _lnprob_wrapper.lnprobs(p)

def _update_wrapper_kde(kde):
    _lnprob_wrapper.update_kde(kde)


class Sampler(object):
    """
    An Ensemble sampler.

    The :attr:`chain` member of this object has the shape: `(nsteps, nwalkers, ndim)` where
    `nsteps` is the stored number of steps taken thus far.

    :param nwalkers:
        The number of individual MCMC chains to include in the ensemble.

    :param ndim:
        Number of dimensions in the parameter space.  If `transd` is ``True`` this is the number of
        unique dimensions across the parameter spaces.

    :param lnpostfn:
        A function that takes a vector in the parameter space as input and returns the natural
        logarithm of the posterior probability for that position.

    :param transd:
        If ``True``, the sampler will operate across parameter spaces using a
        :class:`.clustered_kde.TransdimensionalKDE` proposal distribution. In this mode a masked
        array with samples in each of the possible sets of dimensions must be given for the initial
        ensemble distribution.

    :param processes: (optional)
        The number of processes to use with :mod:`multiprocessing`.  If ``None``, all available
        cores are used.

    :param pool: (optional)
        A pre-constructed pool with a map method. If ``None`` a pool will be created using
        :mod:`multiprocessing`.

    """
    def __init__(self, nwalkers, ndim, lnpostfn, transd=False,
                 processes=None, pool=None, mpi=False, args=[]):
        self.nwalkers = nwalkers
        self.dim = ndim
        self._kde = None
        self._kde_size = self.nwalkers
        self._updates = []
        self._burnin_length = None

        self._get_lnpost = lnpostfn
        self._lnpost_args = args

        self.iterations = 0
        self.stored_iterations = 0
        self.mpi = mpi
        self.processes = processes
        if processes is None and pool is not None:
            raise ValueError("Please provide the number of processes if "
                             " if also providing a pool instance.")
        self._managing_pool = False
        if pool is not None:
            self._pool = pool
        else:
            self._managing_pool = True
            self._pool = choose_pool(processes, mpi=mpi)

        if not hasattr(self._pool, 'map'):
            raise AttributeError("Pool object must have a map() method.")

        self._transd = transd
        if self._transd:
            self._chain = ma.masked_all((0, self.nwalkers, self.dim))
        else:
            self._chain = np.zeros((0, self.nwalkers, self.dim))

        self._lnpost = np.empty((0, self.nwalkers))
        self._lnprop = np.empty((0, self.nwalkers))
        self._acceptance = np.zeros((0, self.nwalkers))
        self._blobs = []
        self.tqdm = None

        self._last_run_mcmc_result = None
        self._burnin_spaces = None
        self._failed_p = None
        self._set_wrapper()

    def _set_wrapper(self):
        """ Determine the lnprob wrapper """
        # If the pool can guarantee a call to all processes we can update
        # the data explicitly, which is much faster
        if hasattr(self._pool, 'broadcast'):
            wrapper = _GetLnProbWrapper(self._get_lnpost, self._kde,

                                        *self._lnpost_args)
            self._pool.broadcast(_set_global_lnprob_wrapper, wrapper)

    def _get_wrapper(self):
        """ Retunr the lnprob wrapper function call """
        if hasattr(self._pool, 'broadcast'):
            return _get_lnprob_from_wrapper
        else:
            return _GetLnProbWrapper(self._get_lnpost, self._kde,

                                     *self._lnpost_args)

    def _get_dynamci_pbar(self, progress, t):
        pbar = None
        if progress:
            import tqdm
            self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(desc="Burning in until Constant Acceptence Rate", position=0, leave=True,
                                               total=t)
            else:
                pbar = self.tqdm.tqdm(desc="Burning in until Constant Acceptence Rate", position=0, leave=True,
                                      total=t)
        return pbar, dynamic_pbar_update

    def _get_finite_pbar(self, progress, N):
        pbar = None
        if progress:
            if self.tqdm is None:
                import tqdm
                self.tqdm = tqdm
            if in_notebook():
                pbar = self.tqdm.tqdm_notebook(total=N, desc="Running MCMC")
            else:
                pbar = self.tqdm.tqdm(total=N, desc="Running MCMC")
        return pbar

    def burnin(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,
               test_steps=16, critical_pval=0.05, max_steps=None,
               verbose=False, callback=None, progress=False, **kwargs):
        """
        Evolve an ensemble until the acceptance rate becomes roughly constant.  This is done by
        splitting acceptances in half and checking for statistical consistency.  This isn't
        guaranteed to return a fully burned-in ensemble, but usually does.

        :param p0: (optional)
            A list of the initial walker positions.  It should have the shape `(nwalkers, ndim)`.
            If ``None`` and the sampler has been run previously, it'll pick up where it left off.

        :param lnpost0: (optional)
            The list of log posterior probabilities for the walkers at positions `p0`. If ``lnpost0
            is None``, the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param lnprop0: (optional)
            List of log proposal densities for walkers at positions `p0`. If ``lnprop0 is None``,
            the initial values are calculated. It should have the shape `(nwalkers, )`.

        :param blob0: (optional)
            The list of blob data for walkers at positions `p0`.

        :param test_steps: (optional)
            The (rough) number of accepted steps over which to check for acceptance rate
            consistency. If you find burnin repeatedly ending prematurely try increasing this.

        :param critical_pval: (optional)
            The critial p-value for considering the acceptance distribution
            consistent.  If the calculated p-value is over this, then the acceptance rate
            is considered to be stationary.  Lower this if you want burnin criteria to be
            less strict.

        :param max_steps: (optional)
            An absolute maximum number of steps to take, in case burnin is too painful.

        :param verbose: (optional)
            Print status messages each time a milestone is reached in the burnin.

        :param kwargs: (optional)
            The rest is passed to :meth:`run_mcmc`.

        After burnin...

        :returns:
            * ``p`` - A list of the current walker positions with shape `(nwalkers, ndim)`.

            * ``lnpost`` - Array of log posterior probabilities for walkers at positions `p`; has
              shape `(nwalkers, )`.

            * ``lnprop`` - Array of log proposal densities for walkers at positions `p`; has shape
              `(nwalkers, )`.

            * ``blob`` - (if `lnprobfn` returns blobs) The list of blob data for the walkers at
              positions `p`.
        """
        if p0 is not None:
            p0 = np.atleast_2d(p0)

        # Determine the maximum iteration to look for
        start = self.iterations
        # Confine walkers to their space during burnin
        freeze_transd = False
        if self._transd:
            freeze_transd = True
            self._burnin_spaces = ~p0.mask
        pbar, pbar_status = self._get_dynamci_pbar(progress, test_steps)
        max_iter = np.inf
        if max_steps is not None:
            max_iter = start + max_steps
        if progress:
            verbose = False

        sampling_ct = 0
        step_size = 2
        last_step_size = 0
        while step_size <= test_steps:
            # Update the proposal
            if p0 is not None:
                self.update_proposal(p0, max_samples=self.nwalkers)
                lnprop0 = self._kde(p0)
            if verbose:
                print('Updated proposal')

            # Take one step to estimate acceptance rate
            test_interval = 1
            results = self.run_mcmc(test_interval, p0, lnpost0, lnprop0, blob0,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            # Use the fraction of acceptances in the last step to estimate acceptance rate
            #   Bottom out at 1% if acceptances are really bad
            last_acc_rate = max(np.mean(self.acceptance[-1]), 0.01)

            # Estimate ACT based on acceptance
            act = int(np.ceil(2.0/last_acc_rate - 1.0))

            if verbose:
                print('Single-step acceptance rate is ', last_acc_rate)
                print('Producing ACT of ', act)

            # Use the ACT to set the new test interval, but avoid
            # overstepping a specified max.  We throw away the first
            # 2*act worth of steps as an initial burnin when comparing
            # acceptance rates
            test_interval = min((step_size+2)*act, max_iter - self.iterations)

            # Make sure we're taking at least one step
            test_interval = max(test_interval, 1)
            sampling_ct += 1
            pbar_status(step_size, last_step_size, last_acc_rate, pbar)

            if verbose:
                if ~np.isinf(max_iter):
                    print("Time {} running mcmc during burnin with {}/{} Iterations and test_interval = {}:".format(
                        sampling_ct, self.iterations, max_iter, test_interval))
                else:
                    print("Time {} running mcmc during burnin with test_interval = {}:".format(
                        sampling_ct, test_interval))

            results = self.run_mcmc(test_interval, p, lnpost, lnprop, blob,
                                    freeze_transd=freeze_transd, spaces=self._burnin_spaces, progress=progress,
                                    **kwargs)
            try:
                p, lnpost, lnprop, blob = results
            except ValueError:
                p, lnpost, lnprop = results
                blob = None

            if callback is not None:
                callback(self)

            # Quit if we hit the max
            if self.iterations >= max_iter:
                print("Burnin hit {} iterations before completing.".format(max_iter))
                break

            # Only check for consistency past the burn-in stage of 2*act
            if self.consistent_acceptance_rate(window_size=step_size*act, critical_pval=critical_pval):
                if verbose:
                    print('Acceptance rate constant over ', step_size, ' ACTs')
                last_step_size = step_size
                step_size *= 2
            else:
                if verbose:
                    print('Acceptance rate varies, trying again')

            if verbose:
                print('') # Newline

            p0, lnpost0, lnprop0, blob0 = p, lnpost, lnprop, blob

        self._burnin_length = self.updates[-1]
        if pbar is not None:
            pbar.close()

        if blob is None:
            return p, lnpost, lnprop
        else:
            return p, lnpost, lnprop, blob

    def sample(self, p0=None, lnpost0=None, lnprop0=None, blob0=None,
{'fraction_non_alphanumeric': 0.0770799785292539, 'fraction_numerical': 0.005797101449275362, 'mean_word_length': 1.9739278493136108, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 15, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3260952', 'n_tokens_mistral': 7991, 'n_tokens_neox': 7573, 'n_words': 2604}