diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py index 18a178a41..9ddc244fb 100644 --- a/bilby/core/sampler/dynesty.py +++ b/bilby/core/sampler/dynesty.py @@ -17,6 +17,7 @@ logger, safe_file_dump, ) +from . import dynesty_utils from .base_sampler import NestedSampler, Sampler, _SamplingContainer, signal_wrapper @@ -196,15 +197,6 @@ def default_kwargs(self): kwargs["seed"] = None return kwargs - @property - def new_dynesty_api(self): - try: - import dynesty.internal_samplers # noqa - - return True - except ImportError: - return False - def __init__( self, likelihood, @@ -281,61 +273,54 @@ def sampler_init_kwargs(self): # if we're using a Bilby implemented sampling method we need to register the # method. If we aren't we need to make sure the default "live" isn't set as # the bounding method - if self.new_dynesty_api: - internal_kwargs = dict( - ndim=self.ndim, - nonbounded=self.kwargs.get("nonbounded", None), - periodic=self.kwargs.get("periodic", None), - reflective=self.kwargs.get("reflective", None), - maxmcmc=self.maxmcmc, - ) - - from . import dynesty3_utils as dynesty_utils + internal_kwargs = dict( + ndim=self.ndim, + nonbounded=self.kwargs.get("nonbounded", None), + periodic=self.kwargs.get("periodic", None), + reflective=self.kwargs.get("reflective", None), + maxmcmc=self.maxmcmc, + ) - if kwargs["sample"] == "act-walk": - internal_kwargs["nact"] = self.nact - internal_sampler = dynesty_utils.ACTTrackingEnsembleWalk( - **internal_kwargs - ) - bound = "none" - logger.info( - f"Using the bilby-implemented ensemble rwalk sampling tracking the " - f"autocorrelation function and thinning by {internal_sampler.thin} with " - f"maximum length {internal_sampler.thin * internal_sampler.maxmcmc}." - ) - elif kwargs["sample"] == "acceptance-walk": - internal_kwargs["naccept"] = self.naccept - internal_kwargs["walks"] = self.kwargs["walks"] - internal_sampler = dynesty_utils.EnsembleWalkSampler(**internal_kwargs) - bound = "none" - logger.info( - f"Using the bilby-implemented ensemble rwalk sampling method with an " - f"average of {internal_sampler.naccept} accepted steps up to chain " - f"length {internal_sampler.maxmcmc}." - ) - elif kwargs["sample"] == "rwalk": - internal_kwargs["nact"] = self.nact - internal_sampler = dynesty_utils.AcceptanceTrackingRWalk( - **internal_kwargs - ) - bound = "none" - logger.info( - f"Using the bilby-implemented ensemble rwalk sampling method with ACT " - f"estimated chain length. An average of {2 * internal_sampler.nact} " - f"steps will be accepted up to chain length {internal_sampler.maxmcmc}." - ) - elif kwargs["bound"] == "live": - logger.info( - "Live-point based bound method requested with dynesty sample " - f"'{kwargs['sample']}', overwriting to 'multi'" - ) - internal_sampler = kwargs["sample"] - bound = "multi" - else: - internal_sampler = kwargs["sample"] - bound = kwargs["bound"] - kwargs["sample"] = internal_sampler - kwargs["bound"] = bound + if kwargs["sample"] == "act-walk": + internal_kwargs["nact"] = self.nact + internal_sampler = dynesty_utils.ACTTrackingEnsembleWalk(**internal_kwargs) + bound = "none" + logger.info( + f"Using the bilby-implemented ensemble rwalk sampling tracking the " + f"autocorrelation function and thinning by {internal_sampler.thin} with " + f"maximum length {internal_sampler.thin * internal_sampler.maxmcmc}." + ) + elif kwargs["sample"] == "acceptance-walk": + internal_kwargs["naccept"] = self.naccept + internal_kwargs["walks"] = self.kwargs["walks"] + internal_sampler = dynesty_utils.EnsembleWalkSampler(**internal_kwargs) + bound = "none" + logger.info( + f"Using the bilby-implemented ensemble rwalk sampling method with an " + f"average of {internal_sampler.naccept} accepted steps up to chain " + f"length {internal_sampler.maxmcmc}." + ) + elif kwargs["sample"] == "rwalk": + internal_kwargs["nact"] = self.nact + internal_sampler = dynesty_utils.AcceptanceTrackingRWalk(**internal_kwargs) + bound = "none" + logger.info( + f"Using the bilby-implemented ensemble rwalk sampling method with ACT " + f"estimated chain length. An average of {2 * internal_sampler.nact} " + f"steps will be accepted up to chain length {internal_sampler.maxmcmc}." + ) + elif kwargs["bound"] == "live": + logger.info( + "Live-point based bound method requested with dynesty sample " + f"'{kwargs['sample']}', overwriting to 'multi'" + ) + internal_sampler = kwargs["sample"] + bound = "multi" + else: + internal_sampler = kwargs["sample"] + bound = kwargs["bound"] + kwargs["sample"] = internal_sampler + kwargs["bound"] = bound return kwargs def _translate_kwargs(self, kwargs): @@ -514,107 +499,12 @@ def sampler_class(self): return Sampler - def _set_sampling_method(self): - """ - Resolve the sampling method and sampler to use from the provided - :code:`bound` and :code:`sample` arguments. - - This requires registering the :code:`bilby` specific methods in the - appropriate locations within :code:`dynesty`. - - Additionally, some combinations of bound/sample/proposals are not - compatible and so we either warn the user or raise an error. - """ - if self.new_dynesty_api: - return - - import dynesty - - _set_sampling_kwargs((self.nact, self.maxmcmc, self.proposals, self.naccept)) - - sample = self.kwargs["sample"] - bound = self.kwargs["bound"] - - if sample not in ["rwalk", "act-walk", "acceptance-walk"] and bound in [ - "live", - "live-multi", - ]: - logger.info( - "Live-point based bound method requested with dynesty sample " - f"'{sample}', overwriting to 'multi'" - ) - self.kwargs["bound"] = "multi" - elif bound == "live": - from .dynesty_utils import LivePointSampler - - dynesty.dynamicsampler._SAMPLERS["live"] = LivePointSampler - elif bound == "live-multi": - from .dynesty_utils import MultiEllipsoidLivePointSampler - - dynesty.dynamicsampler._SAMPLERS[ - "live-multi" - ] = MultiEllipsoidLivePointSampler - elif sample == "acceptance-walk": - raise DynestySetupError( - "bound must be set to live or live-multi for sample=acceptance-walk" - ) - elif self.proposals is None: - logger.warning( - "No proposals specified using dynesty sampling, defaulting " - "to 'volumetric'." - ) - self.proposals = ["volumetric"] - _SamplingContainer.proposals = self.proposals - elif "diff" in self.proposals: - raise DynestySetupError( - "bound must be set to live or live-multi to use differential " - "evolution proposals" - ) - - if sample == "rwalk": - logger.info( - f"Using the bilby-implemented {sample} sample method with ACT estimated walks. " - f"An average of {2 * self.nact} steps will be accepted up to chain length " - f"{self.maxmcmc}." - ) - from .dynesty_utils import AcceptanceTrackingRWalk - - if self.kwargs["walks"] > self.maxmcmc: - raise DynestySetupError("You have maxmcmc < walks (minimum mcmc)") - if self.nact < 1: - raise DynestySetupError("Unable to run with nact < 1") - AcceptanceTrackingRWalk.old_act = None - dynesty.nestedsamplers._SAMPLING["rwalk"] = AcceptanceTrackingRWalk() - elif sample == "acceptance-walk": - logger.info( - f"Using the bilby-implemented {sample} sampling with an average of " - f"{self.naccept} accepted steps per MCMC and maximum length {self.maxmcmc}" - ) - from .dynesty_utils import FixedRWalk - - dynesty.nestedsamplers._SAMPLING["acceptance-walk"] = FixedRWalk() - elif sample == "act-walk": - logger.info( - f"Using the bilby-implemented {sample} sampling tracking the " - f"autocorrelation function and thinning by " - f"{self.nact} with maximum length {self.nact * self.maxmcmc}" - ) - from .dynesty_utils import ACTTrackingRWalk - - ACTTrackingRWalk._cache = list() - dynesty.nestedsamplers._SAMPLING["act-walk"] = ACTTrackingRWalk() - elif sample == "rwalk_dynesty": - sample = sample.strip("_dynesty") - self.kwargs["sample"] = sample - logger.info(f"Using the dynesty-implemented {sample} sample method") - @signal_wrapper def run_sampler(self): import dynesty logger.info(f"Using dynesty version {dynesty.__version__}") - self._set_sampling_method() self._setup_pool() if self.resume: @@ -666,25 +556,6 @@ def run_sampler(self): return self.result - def _setup_pool(self): - """ - In addition to the usual steps, we need to set the sampling kwargs on - every process. To make sure we get every process, run the kwarg setting - more times than we have processes. - """ - super(Dynesty, self)._setup_pool() - - if self.new_dynesty_api: - return - - if self.pool is not None: - args = ( - [(self.nact, self.maxmcmc, self.proposals, self.naccept)] - * self.npool - * 10 - ) - self.pool.map(_set_sampling_kwargs, args) - def _generate_result(self, out): """ Extract the information we need from the dynesty output. This includes @@ -895,10 +766,7 @@ def read_saved_state(self, continuing=False): mapper = self.pool.map else: mapper = map - if self.new_dynesty_api: - self.sampler.mapper = mapper - else: - self.sampler.M = mapper + self.sampler.mapper = mapper return True else: logger.info(f"Resume file {self.resume_file} does not exist.") @@ -941,10 +809,7 @@ def write_current_state(self): metadata = dict() versions = dict(bilby=bilby_version, dynesty=dynesty_version) self.sampler.pool = None - if self.new_dynesty_api: - self.sampler.mapper = map - else: - self.sampler.M = map + self.sampler.mapper = map if dill.pickles(self.sampler): safe_file_dump((self.sampler, versions, metadata), self.resume_file, dill) logger.info(f"Written checkpoint file {self.resume_file}") @@ -955,10 +820,7 @@ def write_current_state(self): ) self.sampler.pool = self.pool if self.sampler.pool is not None: - if self.new_dynesty_api: - self.sampler.mapper = self.sampler.pool.map - else: - self.sampler.M = self.sampler.pool.map + self.sampler.mapper = self.sampler.pool.map def dump_samples_to_dat(self): """ @@ -1090,7 +952,6 @@ def _run_test(self): """Run the sampler very briefly as a sanity test that it works.""" import pandas as pd - self._set_sampling_method() self._setup_pool() self.sampler = self.sampler_init( loglikelihood=_log_likelihood_wrapper, diff --git a/bilby/core/sampler/dynesty3_utils.py b/bilby/core/sampler/dynesty3_utils.py deleted file mode 100644 index 7c8f56b9f..000000000 --- a/bilby/core/sampler/dynesty3_utils.py +++ /dev/null @@ -1,942 +0,0 @@ -import warnings -from collections import namedtuple - -import numpy as np -from dynesty.internal_samplers import InternalSampler, SamplerReturn -from dynesty.utils import SamplerHistoryItem, apply_reflect, get_random_generator - -from ...bilby_mcmc.chain import calculate_tau -from ..utils.log import logger - -EnsembleSamplerArgument = namedtuple( - "EnsembleSamplerArgument", - [ - "u", - "loglstar", - "live_points", - "prior_transform", - "loglikelihood", - "rseed", - "kwargs", - ], -) -EnsembleAxisSamplerArgument = namedtuple( - "EnsembleAxisSamplerArgument", - [ - "u", - "loglstar", - "axes", - "live_points", - "prior_transform", - "loglikelihood", - "rseed", - "kwargs", - ], -) - - -class BaseEnsembleSampler(InternalSampler): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.ncdim = kwargs.get("ncdim") - self.sampler_kwargs["ncdim"] = self.ncdim - self.sampler_kwargs["proposals"] = kwargs.get("proposals", ["diff"]) - - def prepare_sampler( - self, - loglstar=None, - points=None, - axes=None, - seeds=None, - prior_transform=None, - loglikelihood=None, - nested_sampler=None, - ): - """ - Prepare the list of arguments for sampling. - - Parameters - ---------- - loglstar : float - Ln(likelihood) bound. - points : `~numpy.ndarray` with shape (n, ndim) - Initial sample points. - axes : `~numpy.ndarray` with shape (ndim, ndim) - Axes used to propose new points. - seeds : `~numpy.ndarray` with shape (n,) - Random number generator seeds. - prior_transform : function - Function transforming a sample from the a unit cube to the - parameter space of interest according to the prior. - loglikelihood : function - Function returning ln(likelihood) given parameters as a 1-d - `~numpy` array of length `ndim`. - nested_sampler : `~dynesty.samplers.Sampler` - The nested sampler object used to sample. - - Returns - ------- - arglist: - List of `SamplerArgument` objects containing the parameters - needed for sampling. - """ - arg_list = [] - kwargs = self.sampler_kwargs.copy() - self.nlive = nested_sampler.nlive - for curp, curaxes, curseed in zip(points, axes, seeds): - vals = dict( - u=curp, - loglstar=loglstar, - live_points=nested_sampler.live_u, - prior_transform=prior_transform, - loglikelihood=loglikelihood, - rseed=curseed, - kwargs=kwargs, - ) - if "volumetric" in kwargs["proposals"]: - vals["axes"] = curaxes - curarg = EnsembleAxisSamplerArgument(**vals) - else: - curarg = EnsembleSamplerArgument(**vals) - arg_list.append(curarg) - return arg_list - - -class EnsembleWalkSampler(BaseEnsembleSampler): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.walks = max(2, kwargs.get("walks", 25)) - self.sampler_kwargs["walks"] = self.walks - self.naccept = kwargs.get("naccept", 10) - self.maxmcmc = kwargs.get("maxmcmc", 5000) - - def tune(self, tuning_info, update=True): - """ - Update the proposal parameters based on the number of accepted steps - and MCMC chain length. - - The :code:`walks` parameter to asymptotically approaches the desired number - of accepted steps. Update :code:`self.scale` for inclusion in the state plot. - """ - # update walks to match target naccept - accept_prob = max(0.5, tuning_info["accept"]) / self.sampler_kwargs["walks"] - delay = max(self.nlive // 10 - 1, 0) - self.scale = tuning_info["accept"] - self.walks = (self.walks * delay + self.naccept / accept_prob) / (delay + 1) - self.sampler_kwargs["walks"] = min(int(np.ceil(self.walks)), self.maxmcmc) - - @staticmethod - def sample(args): - """ - Return a new live point proposed by random walking away from an - existing live point. - - Parameters - ---------- - u : `~numpy.ndarray` with shape (ndim,) - Position of the initial sample. **This is a copy of an existing - live point.** - - loglstar : float - Ln(likelihood) bound. - - axes : `~numpy.ndarray` with shape (ndim, ndim) - Axes used to propose new points. For random walks new positions are - proposed using the :class:`~dynesty.bounding.Ellipsoid` whose - shape is defined by axes. - - scale : float - Value used to scale the provided axes. - - prior_transform : function - Function transforming a sample from the a unit cube to the - parameter space of interest according to the prior. - - loglikelihood : function - Function returning ln(likelihood) given parameters as a 1-d - `~numpy` array of length `ndim`. - - kwargs : dict - A dictionary of additional method-specific parameters. - - Returns - ------- - u : `~numpy.ndarray` with shape (ndim,) - Position of the final proposed point within the unit cube. - - v : `~numpy.ndarray` with shape (ndim,) - Position of the final proposed point in the target parameter space. - - logl : float - Ln(likelihood) of the final proposed point. - - nc : int - Number of function calls used to generate the sample. - - sampling_info : dict - Collection of ancillary quantities used to tune :data:`scale`. - - """ - current_u = args.u - naccept = 0 - ncall = 0 - - periodic = args.kwargs["periodic"] - reflective = args.kwargs["reflective"] - boundary_kwargs = dict(periodic=periodic, reflective=reflective) - - proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args) - walks = len(proposals) - evaluation_history = list() - - for prop in proposals: - u_prop = proposal_funcs[prop]( - u=current_u, **common_kwargs, **proposal_kwargs[prop] - ) - u_prop = apply_boundaries_(u_prop=u_prop, **boundary_kwargs) - if u_prop is None: - continue - - v_prop = args.prior_transform(u_prop) - logl_prop = args.loglikelihood(v_prop) - evaluation_history.append( - SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) - ) - ncall += 1 - - if logl_prop > args.loglstar: - current_u = u_prop - current_v = v_prop - logl = logl_prop - naccept += 1 - - if naccept == 0: - logger.debug( - "Unable to find a new point using walk: returning a random point. " - "If this warning occurs often, increase naccept." - ) - # Technically we can find out the likelihood value - # stored somewhere - # But I'm currently recomputing it - current_u = common_kwargs["rstate"].uniform(0, 1, len(current_u)) - current_v = args.prior_transform(current_u) - logl = args.loglikelihood(current_v) - - sampling_info = { - "accept": naccept, - "reject": walks - naccept, - } - - return SamplerReturn( - u=current_u, - v=current_v, - logl=logl, - tuning_info=sampling_info, - ncalls=ncall, - proposal_stats=sampling_info, - evaluation_history=evaluation_history, - ) - - # return current_u, current_v, logl, ncall, sampling_info - - -class ACTTrackingEnsembleWalk(BaseEnsembleSampler): - """ - Run the MCMC sampler for many iterations in order to reliably estimate - the autocorrelation time. - - This builds a cache of :code:`nact / thin` points consistent with the - likelihood constraint. - While this approach is the most robust, it is not well optimized for - parallelised sampling as the length of the MCMC will be different for each - parallel process. - """ - - # the _cache is a class level variable to avoid being forgotten at every - # iteration when using multiprocessing - _cache = list() - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.act = 1 - self.thin = kwargs.get("nact", 2) - self.maxmcmc = kwargs.get("maxmcmc", 5000) * 50 - self.sampler_kwargs["rebuild"] = True - self.sampler_kwargs["thin"] = self.thin - self.sampler_kwargs["act"] = self.act - self.sampler_kwargs["maxmcmc"] = self.maxmcmc - # reset the cache at instantiation to avoid contamination from - # previous analyses - self.__class__._cache = list() - - def prepare_sampler( - self, - loglstar=None, - points=None, - axes=None, - seeds=None, - prior_transform=None, - loglikelihood=None, - nested_sampler=None, - ): - """ - Prepare the list of arguments for sampling. - - Parameters - ---------- - loglstar : float - Ln(likelihood) bound. - points : `~numpy.ndarray` with shape (n, ndim) - Initial sample points. - axes : `~numpy.ndarray` with shape (ndim, ndim) - Axes used to propose new points. - seeds : `~numpy.ndarray` with shape (n,) - Random number generator seeds. - prior_transform : function - Function transforming a sample from the a unit cube to the - parameter space of interest according to the prior. - loglikelihood : function - Function returning ln(likelihood) given parameters as a 1-d - `~numpy` array of length `ndim`. - nested_sampler : `~dynesty.samplers.Sampler` - The nested sampler object used to sample. - - Returns - ------- - arglist: - List of `SamplerArgument` objects containing the parameters - needed for sampling. - """ - arg_list = super().prepare_sampler( - loglstar=loglstar, - points=points, - axes=axes, - seeds=seeds, - prior_transform=prior_transform, - loglikelihood=loglikelihood, - nested_sampler=nested_sampler, - ) - self.sampler_kwargs["rebuild"] = False - return arg_list - - def tune(self, tuning_info, update=True): - """ - Update the proposal parameters based on the number of accepted steps - and MCMC chain length. - - The :code:`walks` parameter to asymptotically approach the - desired number of accepted steps. - """ - if tuning_info.get("remaining", 0) == 0: - self.sampler_kwargs["rebuild"] = True - self.scale = tuning_info["accept"] - self.sampler_kwargs["act"] = tuning_info["act"] - - @staticmethod - def sample(args): - cache = ACTTrackingEnsembleWalk._cache - if args.kwargs.get("rebuild", False): - logger.debug(f"Force rebuilding cache with {len(cache)}.") - cache.clear() - if len(cache) == 0: - ACTTrackingEnsembleWalk.build_cache(args) - - while len(cache) > 0 and cache[0][2] < args.loglstar: - state = cache.pop(0) - - if len(cache) == 0: - current_u, current_v, logl, ncall, blob, evaluation_history = state - else: - current_u, current_v, logl, ncall, blob, evaluation_history = cache.pop(0) - - blob["remaining"] = len(cache) - logger.debug(f"Returning point from cache with {len(cache)} remaining") - return SamplerReturn( - u=current_u, - v=current_v, - logl=logl, - tuning_info=blob, - ncalls=ncall, - proposal_stats=blob, - evaluation_history=evaluation_history, - ) - - @staticmethod - def build_cache(args): - # Bounds - periodic = args.kwargs.get("periodic", None) - reflective = args.kwargs.get("reflective", None) - boundary_kwargs = dict(periodic=periodic, reflective=reflective) - - proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args) - - # Setup - current_u = args.u - old_act = args.kwargs.get("act", 100) - check_interval = ACTTrackingEnsembleWalk.integer_act(old_act) - target_nact = 50 - next_check = check_interval - n_checks = 0 - maxmcmc = args.kwargs.get("maxmcmc", 5000) - evaluation_history = list() - - # Initialize internal variables - current_v = args.prior_transform(np.array(current_u)) - logl = args.loglikelihood(np.array(current_v)) - evaluation_history.append( - SamplerHistoryItem(u=current_u, v=current_v, logl=logl) - ) - accept = 0 - reject = 0 - nfail = 0 - ncall = 0 - act = np.inf - u_list = list() - v_list = list() - logl_list = list() - most_failures = 0 - current_failures = 0 - - iteration = 0 - while iteration < min(target_nact * act, maxmcmc): - iteration += 1 - - prop = proposals[iteration % len(proposals)] - u_prop = proposal_funcs[prop]( - u=current_u, **common_kwargs, **proposal_kwargs[prop] - ) - u_prop = apply_boundaries_(u_prop=u_prop, **boundary_kwargs) - success = False - if u_prop is not None: - v_prop = args.prior_transform(np.array(u_prop)) - logl_prop = args.loglikelihood(np.array(v_prop)) - evaluation_history.append( - SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) - ) - ncall += 1 - if logl_prop > args.loglstar: - success = True - current_u = u_prop - current_v = v_prop - logl = logl_prop - - u_list.append(current_u) - v_list.append(current_v) - logl_list.append(logl) - - if success: - accept += 1 - most_failures = max(most_failures, current_failures) - current_failures = 0 - else: - nfail += 1 - current_failures += 1 - - # If we've taken the minimum number of steps, calculate the ACT - if iteration > next_check and accept > target_nact: - n_checks += 1 - most_failures = max(most_failures, current_failures) - act = ACTTrackingEnsembleWalk._calculate_act( - accept=accept, - iteration=iteration, - samples=np.array(u_list), - most_failures=most_failures, - ) - to_next_update = (act * target_nact - iteration) // 2 - to_next_update = max(to_next_update, iteration // 100) - next_check += to_next_update - logger.debug( - f"it: {iteration}, accept: {accept}, reject: {reject}, " - f"fail: {nfail}, act: {act:.2f}, nact: {iteration / act:.2f} " - ) - elif iteration > next_check: - next_check += check_interval - - most_failures = max(most_failures, current_failures) - act = ACTTrackingEnsembleWalk._calculate_act( - accept=accept, - iteration=iteration, - samples=np.array(u_list), - most_failures=most_failures, - ) - reject += nfail - blob = {"accept": accept, "reject": reject, "act": act} - iact = ACTTrackingEnsembleWalk.integer_act(act) - thin = args.kwargs.get("thin", 2) * iact - - cache = ACTTrackingEnsembleWalk._cache - - if accept == 0: - logger.warning( - "Unable to find a new point using walk: returning a random point" - ) - u = common_kwargs["rstate"].uniform(size=len(current_u)) - v = args.prior_transform(u) - logl = args.loglikelihood(v) - evaluation_history = [ - SamplerHistoryItem(u=current_u, v=current_v, logl=logl) - ] - cache.append((u, v, logl, ncall, blob, evaluation_history)) - elif not np.isfinite(act): - logger.warning( - "Unable to find a new point using walk: try increasing maxmcmc" - ) - cache.append((current_u, current_v, logl, ncall, blob, evaluation_history)) - elif (thin == -1) or (len(u_list) <= thin): - cache.append((current_u, current_v, logl, ncall, blob, evaluation_history)) - else: - u_list = u_list[thin::thin] - v_list = v_list[thin::thin] - logl_list = logl_list[thin::thin] - evaluation_history_list = ( - evaluation_history[thin * ii : thin * (ii + 1)] - for ii in range(len(u_list)) - ) - n_found = len(u_list) - accept = max(accept // n_found, 1) - reject //= n_found - nfail //= n_found - ncall_list = [ncall // n_found] * n_found - blob_list = [ - dict(accept=accept, reject=reject, fail=nfail, act=act) - ] * n_found - cache.extend( - zip( - u_list, - v_list, - logl_list, - ncall_list, - blob_list, - evaluation_history_list, - ) - ) - logger.debug( - f"act: {act:.2f}, max failures: {most_failures}, thin: {thin}, " - f"iteration: {iteration}, n_found: {n_found}" - ) - logger.debug( - f"Finished building cache with length {len(cache)} after " - f"{iteration} iterations with {ncall} likelihood calls and ACT={act:.2f}" - ) - - @staticmethod - def _calculate_act(accept, iteration, samples, most_failures): - """ - Take the maximum of three ACT estimates, leading to a conservative estimate: - - - a full ACT estimate as done in :code:`bilby_mcmc`. This is almost always the - longest estimator, and is the most computationally expensive. The other - methods mostly catch cases where the estimated ACT is very small. - - the naive ACT used for the acceptance tracking walk. - - the most failed proposals between any pair of accepted steps. This is a strong - lower bound, because we know that if we thin by less than this, there will be - duplicate points. - """ - if accept > 0: - naive_act = 2 / accept * iteration - 1 - else: - return np.inf - return max(calculate_tau(samples), naive_act, most_failures) - - @staticmethod - def integer_act(act): - if np.isinf(act): - return act - else: - return int(np.ceil(act)) - - -class AcceptanceTrackingRWalk(EnsembleWalkSampler): - """ - This is a modified version of dynesty.sampling.sample_rwalk that runs the - MCMC random walk for a user-specified number of a crude approximation to - the autocorrelation time. - - This is the proposal method used by default for :code:`Bilby<2` and - corresponds to specifying :code:`sample="rwalk"` - """ - - # to retain state between calls to pool.Map, this needs to be a class - # level attribute - old_act = None - - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.nact = kwargs.get("nact", 40) - self.sampler_kwargs["nact"] = self.nact - self.sampler_kwargs["maxmcmc"] = self.maxmcmc - - def tune(self, tuning_info, update=True): - """ - The tuning all happens inside the MCMC for this proposal - so all we do is extract the number of accepted steps for plotting. - """ - self.scale = tuning_info["accept"] - - @staticmethod - def sample(args): - rstate = get_random_generator(args.rseed) - - periodic = args.kwargs.get("periodic", None) - reflective = args.kwargs.get("reflective", None) - boundary_kwargs = dict(periodic=periodic, reflective=reflective) - - current_u = args.u - nlive = args.kwargs.get("nlive", args.kwargs.get("walks", 100)) - - proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args) - - accept = 0 - reject = 0 - nfail = 0 - act = np.inf - nact = args.kwargs.get("nact", 40) - maxmcmc = args.kwargs.get("maxmcmc", 5000) - evaluation_history = list() - - iteration = 0 - while iteration < nact * act: - iteration += 1 - - prop = proposals[iteration % len(proposals)] - u_prop = proposal_funcs[prop]( - current_u, **common_kwargs, **proposal_kwargs[prop] - ) - u_prop = apply_boundaries_(u_prop, **boundary_kwargs) - - if u_prop is None: - nfail += 1 - continue - - # Check proposed point. - v_prop = args.prior_transform(np.array(u_prop)) - logl_prop = args.loglikelihood(np.array(v_prop)) - evaluation_history.append( - SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) - ) - if logl_prop > args.loglstar: - current_u = u_prop - current_v = v_prop - logl = logl_prop - accept += 1 - else: - reject += 1 - - # If we've taken the minimum number of steps, calculate the ACT - if iteration > nact: - act = AcceptanceTrackingRWalk.estimate_nmcmc( - accept_ratio=accept / (accept + reject + nfail), - safety=1, - tau=nlive, - maxmcmc=maxmcmc, - old_act=AcceptanceTrackingRWalk.old_act, - ) - - # If we've taken too many likelihood evaluations then break - if accept + reject > maxmcmc: - warnings.warn( - f"Hit maximum number of walks {maxmcmc} with accept={accept}," - f" reject={reject}, and nfail={nfail} try increasing maxmcmc" - ) - break - - if not (np.isfinite(act) and accept > 0): - logger.debug( - "Unable to find a new point using walk: returning a random point" - ) - current_u = rstate.uniform(size=len(current_u)) - current_v = args.prior_transform(current_u) - logl = args.loglikelihood(current_v) - - sampling_info = {"accept": accept, "reject": reject + nfail} - AcceptanceTrackingRWalk.old_act = act - - ncall = accept + reject - return SamplerReturn( - u=current_u, - v=current_v, - logl=logl, - tuning_info=sampling_info, - ncalls=ncall, - proposal_stats=sampling_info, - evaluation_history=evaluation_history, - ) - - @staticmethod - def estimate_nmcmc(accept_ratio, safety=5, tau=None, maxmcmc=5000, old_act=None): - """Estimate autocorrelation length of chain using acceptance fraction - - Using ACL = (2/acc) - 1 multiplied by a safety margin. Code adapted from: - - - https://github.com/farr/Ensemble.jl - - https://github.com/johnveitch/cpnest/blob/master/cpnest/sampler.py - - Parameters - ========== - accept_ratio: float [0, 1] - Ratio of the number of accepted points to the total number of points - old_act: int - The ACT of the last iteration - maxmcmc: int - The maximum length of the MCMC chain to use - safety: int - A safety factor applied in the calculation - tau: int (optional) - The ACT, if given, otherwise estimated. - - Notes - ===== - This method does not compute a reliable estimate of the autocorrelation - length for our proposal distributions. - """ - if tau is None: - tau = maxmcmc / safety - - if accept_ratio == 0.0: - if old_act is None: - Nmcmc_exact = np.inf - else: - Nmcmc_exact = (1 + 1 / tau) * old_act - else: - estimated_act = 2 / accept_ratio - 1 - Nmcmc_exact = safety * estimated_act - if old_act is not None: - Nmcmc_exact = (1 - 1 / tau) * old_act + Nmcmc_exact / tau - Nmcmc_exact = float(min(Nmcmc_exact, maxmcmc)) - return max(safety, Nmcmc_exact) - - -def _get_proposal_kwargs(args): - """ - Resolve the proposal cycle from the provided keyword arguments. - - The steps involved are: - - - extract the requested proposal types from the :code:`_SamplingContainer`. - If none are specified, only differential evolution will be used. - - differential evolution requires the live points to be passed. If they are - not present, raise an error. - - if a dictionary, e.g., :code:`dict(diff=5, volumetric=1)` is specified, - the keys will be used weighted by the values, e.g., 5:1 differential - evolution to volumetric. - - each proposal needs different keyword arguments, see the specific functions - for what requires which. - - - Parameters - ========== - args: dynesty.sampler.SamplerArgument - Object that carries around various pieces of information about the - analysis. - """ - rstate = get_random_generator(args.rseed) - walks = args.kwargs.get("walks", 100) - current_u = args.u - - proposals = args.kwargs.get("proposals", None) - if proposals is None: - proposals = ["diff"] - - n_cluster = args.kwargs.get("ncdim", None) - if n_cluster is None: - if hasattr(args, "live_points"): - n_cluster = args.live_points.shape[1] - elif hasattr(args, "axes"): - n_cluster = args.axes.shape[0] - - if "diff" in proposals: - live = args.live_points - if live is None: - raise ValueError( - "Live points not passed for differential evolution, specify " - "bound='live' to use differential evolution proposals." - ) - live = np.unique(live, axis=0) - matches = np.where(np.equal(current_u, live).all(axis=1))[0] - np.delete(live, matches, 0) - - if isinstance(proposals, (list, set, tuple)): - proposals = rstate.choice(proposals, int(walks)) - elif isinstance(proposals, dict): - props, weights = zip(*proposals.items()) - weights = np.array(weights) / sum(weights) - proposals = rstate.choice(list(props), int(walks), p=weights) - - common_kwargs = dict( - n=len(current_u), - n_cluster=n_cluster, - rstate=rstate, - ) - proposal_kwargs = dict() - if "diff" in proposals: - proposal_kwargs["diff"] = dict( - live=live[:, :n_cluster], - mix=0.5, - scale=2.38 / (2 * n_cluster) ** 0.5, - ) - if "volumetric" in proposals: - proposal_kwargs["volumetric"] = dict( - axes=args.axes, - scale=args.scale, - ) - return proposals, common_kwargs, proposal_kwargs - - -def propose_differential_evolution( - u, - live, - n, - n_cluster, - rstate, - mix=0.5, - scale=1, -): - r""" - Propose a new point using ensemble differential evolution - (`ter Braak + (2006) `_). - - .. math:: - - u_{\rm prop} = u + \gamma (v_{a} - v_{b}) - - We consider two choices for :math:`\gamma`: weighted by :code:`mix`. - - - :math:`\gamma = 1`: this is a mode-hopping mode for efficiently - exploring multi-modal spaces - - :math:`\gamma \sim \Gamma\left(\gamma; k=4, \theta=\frac{\kappa}{4}\right)` - - Here :math:`\kappa = 2.38 / \sqrt{2 n}` unless specified by the user and - we scale by a random draw from a Gamma distribution. The specific - distribution was chosen somewhat arbitrarily to have mean and mode close to - :math:`\kappa` and give good acceptance and autocorrelation times on a subset - of problems. - - Parameters - ---------- - u: np.ndarray - The current point. - live: np.ndarray - The ensemble of live points to select :math:`v` from. - n: int - The number of dimensions being explored - n_cluster: int - The number of dimensions to run the differential evolution over, the - first :code:`n_cluster` dimensions are used. The rest are randomly - sampled from the prior. - rstate: numpy.random.Generator - The numpy generator instance used for random number generation. - Consider using built in `bilby.core.utils.random.rng`. - mix: float - The fraction of proposed points that should follow the specified scale - rather than mode hopping. :code:`default=0.5` - scale: float - The amount to scale the difference vector by. - :code:`default = 2.38 / (2 * n_cluster)**0.5)` - - Returns - ------- - u_prop: np.ndarray - The proposed point. - """ - delta = np.diff(rstate.choice(live, 2, replace=False), axis=0)[0] - if rstate.uniform(0, 1) < mix: - if scale is None: - scale = 2.38 / (2 * n_cluster) ** 0.5 - scale *= rstate.gamma(4, 0.25) - else: - scale = 1 - u_prop = u.copy() - u_prop[:n_cluster] += delta * scale - u_prop[n_cluster:] = rstate.uniform(0, 1, n - n_cluster) - return u_prop - - -def propose_volumetric( - u, - axes, - scale, - n, - n_cluster, - rstate, -): - """ - Propose a new point using the default :code:`dynesty` proposal. - - The new difference vector is scaled by a vector isotropically drawn - from an ellipsoid centered on zero with covariance given by the - provided axis. Note that the magnitude of this proposal is heavily - skewed to the size of the ellipsoid. - - Parameters - ---------- - u: np.ndarray - The current point. - n: int - The number of dimensions being explored. - scale: float - The amount to scale the proposed point by. - n_cluster: int - The number of dimensions to run the differential evolution over, the - first :code:`n_cluster` dimensions are used. The rest are randomly - sampled from the prior. - rstate: numpy.random.Generator - The numpy generator instance used for random number generation. - Consider using built in `bilby.core.utils.random.rng`. - - Returns - ------- - u_prop: np.ndarray - The proposed point. - """ - # Propose a direction on the unit n-sphere. - drhat = rstate.normal(0, 1, n_cluster) - drhat /= np.linalg.norm(drhat) - - # Scale based on dimensionality. - dr = drhat * rstate.uniform(0, 1) ** (1.0 / n_cluster) - - # Transform to proposal distribution. - delta = scale * np.dot(axes, dr) - u_prop = u.copy() - u_prop[:n_cluster] += delta - u_prop[n_cluster:] = rstate.uniform(0, 1, n - n_cluster) - return u_prop - - -def apply_boundaries_(u_prop, periodic, reflective): - """ - Apply the periodic and reflective boundaries and test if we are inside the - unit cube. - - Parameters - ---------- - u_prop: np.ndarray - The proposed point in the unit hypercube space. - periodic: np.ndarray - Indices of the parameters with periodic boundaries. - reflective: np.ndarray - Indices of the parameters with reflective boundaries. - - Returns - ======= - [np.ndarray, None]: - Either the remapped proposed point, or None if the proposed point - lies outside the unit cube. - """ - # Wrap periodic parameters - if periodic is not None: - u_prop[periodic] = np.mod(u_prop[periodic], 1) - - # Reflect - if reflective is not None: - u_prop[reflective] = apply_reflect(u_prop[reflective]) - - if u_prop.min() < 0 or u_prop.max() > 1: - return None - else: - return u_prop - - -proposal_funcs = dict( - diff=propose_differential_evolution, volumetric=propose_volumetric -) diff --git a/bilby/core/sampler/dynesty_utils.py b/bilby/core/sampler/dynesty_utils.py index 22e725aaf..7c8f56b9f 100644 --- a/bilby/core/sampler/dynesty_utils.py +++ b/bilby/core/sampler/dynesty_utils.py @@ -1,118 +1,182 @@ import warnings +from collections import namedtuple import numpy as np -from dynesty.nestedsamplers import MultiEllipsoidSampler, UnitCubeSampler -from dynesty.utils import apply_reflect, get_random_generator +from dynesty.internal_samplers import InternalSampler, SamplerReturn +from dynesty.utils import SamplerHistoryItem, apply_reflect, get_random_generator from ...bilby_mcmc.chain import calculate_tau from ..utils.log import logger -from .base_sampler import _SamplingContainer +EnsembleSamplerArgument = namedtuple( + "EnsembleSamplerArgument", + [ + "u", + "loglstar", + "live_points", + "prior_transform", + "loglikelihood", + "rseed", + "kwargs", + ], +) +EnsembleAxisSamplerArgument = namedtuple( + "EnsembleAxisSamplerArgument", + [ + "u", + "loglstar", + "axes", + "live_points", + "prior_transform", + "loglikelihood", + "rseed", + "kwargs", + ], +) + + +class BaseEnsembleSampler(InternalSampler): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.ncdim = kwargs.get("ncdim") + self.sampler_kwargs["ncdim"] = self.ncdim + self.sampler_kwargs["proposals"] = kwargs.get("proposals", ["diff"]) + + def prepare_sampler( + self, + loglstar=None, + points=None, + axes=None, + seeds=None, + prior_transform=None, + loglikelihood=None, + nested_sampler=None, + ): + """ + Prepare the list of arguments for sampling. -class LivePointSampler(UnitCubeSampler): - """ - Modified version of dynesty UnitCubeSampler that adapts the MCMC - length in addition to the proposal scale, this corresponds to - :code:`bound=live`. - - In order to support live-point based proposals, e.g., differential - evolution (:code:`diff`), the live points are added to the - :code:`kwargs` passed to the evolve method. - - Note that this does not perform ellipsoid clustering as with the - :code:`bound=multi` option, if ellipsoid-based proposals are used, e.g., - :code:`volumetric`, consider using the - :code:`MultiEllipsoidLivePointSampler` (:code:`sample=live-multi`). - """ + Parameters + ---------- + loglstar : float + Ln(likelihood) bound. + points : `~numpy.ndarray` with shape (n, ndim) + Initial sample points. + axes : `~numpy.ndarray` with shape (ndim, ndim) + Axes used to propose new points. + seeds : `~numpy.ndarray` with shape (n,) + Random number generator seeds. + prior_transform : function + Function transforming a sample from the a unit cube to the + parameter space of interest according to the prior. + loglikelihood : function + Function returning ln(likelihood) given parameters as a 1-d + `~numpy` array of length `ndim`. + nested_sampler : `~dynesty.samplers.Sampler` + The nested sampler object used to sample. + + Returns + ------- + arglist: + List of `SamplerArgument` objects containing the parameters + needed for sampling. + """ + arg_list = [] + kwargs = self.sampler_kwargs.copy() + self.nlive = nested_sampler.nlive + for curp, curaxes, curseed in zip(points, axes, seeds): + vals = dict( + u=curp, + loglstar=loglstar, + live_points=nested_sampler.live_u, + prior_transform=prior_transform, + loglikelihood=loglikelihood, + rseed=curseed, + kwargs=kwargs, + ) + if "volumetric" in kwargs["proposals"]: + vals["axes"] = curaxes + curarg = EnsembleAxisSamplerArgument(**vals) + else: + curarg = EnsembleSamplerArgument(**vals) + arg_list.append(curarg) + return arg_list - rebuild = False - def update_user(self, blob, update=True): +class EnsembleWalkSampler(BaseEnsembleSampler): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.walks = max(2, kwargs.get("walks", 25)) + self.sampler_kwargs["walks"] = self.walks + self.naccept = kwargs.get("naccept", 10) + self.maxmcmc = kwargs.get("maxmcmc", 5000) + + def tune(self, tuning_info, update=True): """ Update the proposal parameters based on the number of accepted steps and MCMC chain length. - There are a number of logical checks performed: - - if the ACT tracking rwalk method is being used and any parallel - process has an empty cache, set the :code:`rebuild` flag to force - the cache to rebuild at the next call. This improves the efficiency - when using parallelisation. - - update the :code:`walks` parameter to asymptotically approach the - desired number of accepted steps for the :code:`FixedRWalk` proposal. - - update the ellipsoid scale if the ellipsoid proposals are being used. + The :code:`walks` parameter to asymptotically approaches the desired number + of accepted steps. Update :code:`self.scale` for inclusion in the state plot. """ - # do we need to trigger rebuilding the cache - if blob.get("remaining", 0) == 1: - self.rebuild = True - if update: - self.kwargs["rebuild"] = self.rebuild - self.rebuild = False - # update walks to match target naccept - accept_prob = max(0.5, blob["accept"]) / self.kwargs["walks"] + accept_prob = max(0.5, tuning_info["accept"]) / self.sampler_kwargs["walks"] delay = max(self.nlive // 10 - 1, 0) - n_target = getattr(_SamplingContainer, "naccept", 60) - self.walks = (self.walks * delay + n_target / accept_prob) / (delay + 1) - self.kwargs["walks"] = min(int(np.ceil(self.walks)), _SamplingContainer.maxmcmc) + self.scale = tuning_info["accept"] + self.walks = (self.walks * delay + self.naccept / accept_prob) / (delay + 1) + self.sampler_kwargs["walks"] = min(int(np.ceil(self.walks)), self.maxmcmc) - self.scale = blob["accept"] + @staticmethod + def sample(args): + """ + Return a new live point proposed by random walking away from an + existing live point. - update_rwalk = update_user + Parameters + ---------- + u : `~numpy.ndarray` with shape (ndim,) + Position of the initial sample. **This is a copy of an existing + live point.** - def propose_live(self, *args): - """ - We need to make sure the live points are passed to the proposal - function if we are using live point-based proposals. - """ - self.kwargs["nlive"] = self.nlive - self.kwargs["live"] = self.live_u - i = self.rstate.integers(self.nlive) - u = self.live_u[i, :] - return u, np.identity(self.ncdim) + loglstar : float + Ln(likelihood) bound. + axes : `~numpy.ndarray` with shape (ndim, ndim) + Axes used to propose new points. For random walks new positions are + proposed using the :class:`~dynesty.bounding.Ellipsoid` whose + shape is defined by axes. -class MultiEllipsoidLivePointSampler(MultiEllipsoidSampler): - """ - Modified version of dynesty MultiEllipsoidSampler that adapts the MCMC - length in addition to the proposal scale, this corresponds to - :code:`bound=live-multi`. + scale : float + Value used to scale the provided axes. - Additionally, in order to support live point-based proposals, e.g., - differential evolution (:code:`diff`), the live points are added to the - :code:`kwargs` passed to the evolve method. + prior_transform : function + Function transforming a sample from the a unit cube to the + parameter space of interest according to the prior. - When just using the :code:`diff` proposal method, consider using the - :code:`LivePointSampler` (:code:`sample=live`). - """ + loglikelihood : function + Function returning ln(likelihood) given parameters as a 1-d + `~numpy` array of length `ndim`. - rebuild = False + kwargs : dict + A dictionary of additional method-specific parameters. - def update_user(self, blob, update=True): - LivePointSampler.update_user(self, blob=blob, update=update) - super(MultiEllipsoidLivePointSampler, self).update_rwalk( - blob=blob, update=update - ) + Returns + ------- + u : `~numpy.ndarray` with shape (ndim,) + Position of the final proposed point within the unit cube. - update_rwalk = update_user + v : `~numpy.ndarray` with shape (ndim,) + Position of the final proposed point in the target parameter space. - def propose_live(self, *args): - """ - We need to make sure the live points are passed to the proposal - function if we are using ensemble proposals. - """ - self.kwargs["nlive"] = self.nlive - self.kwargs["live"] = self.live_u - return super(MultiEllipsoidLivePointSampler, self).propose_live(*args) + logl : float + Ln(likelihood) of the final proposed point. + nc : int + Number of function calls used to generate the sample. -class FixedRWalk: - """ - Run the MCMC walk for a fixed length. This is nearly equivalent to - :code:`bilby.sampling.sample_rwalk` except that different proposal - distributions can be used. - """ + sampling_info : dict + Collection of ancillary quantities used to tune :data:`scale`. - def __call__(self, args): + """ current_u = args.u naccept = 0 ncall = 0 @@ -123,8 +187,7 @@ def __call__(self, args): proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args) walks = len(proposals) - - accepted = list() + evaluation_history = list() for prop in proposals: u_prop = proposal_funcs[prop]( @@ -132,11 +195,13 @@ def __call__(self, args): ) u_prop = apply_boundaries_(u_prop=u_prop, **boundary_kwargs) if u_prop is None: - accepted.append(0) continue v_prop = args.prior_transform(u_prop) logl_prop = args.loglikelihood(v_prop) + evaluation_history.append( + SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) + ) ncall += 1 if logl_prop > args.loglstar: @@ -144,9 +209,6 @@ def __call__(self, args): current_v = v_prop logl = logl_prop naccept += 1 - accepted.append(1) - else: - accepted.append(0) if naccept == 0: logger.debug( @@ -160,16 +222,25 @@ def __call__(self, args): current_v = args.prior_transform(current_u) logl = args.loglikelihood(current_v) - blob = { + sampling_info = { "accept": naccept, "reject": walks - naccept, - "scale": args.scale, } - return current_u, current_v, logl, ncall, blob + return SamplerReturn( + u=current_u, + v=current_v, + logl=logl, + tuning_info=sampling_info, + ncalls=ncall, + proposal_stats=sampling_info, + evaluation_history=evaluation_history, + ) + # return current_u, current_v, logl, ncall, sampling_info -class ACTTrackingRWalk: + +class ACTTrackingEnsembleWalk(BaseEnsembleSampler): """ Run the MCMC sampler for many iterations in order to reliably estimate the autocorrelation time. @@ -185,32 +256,113 @@ class ACTTrackingRWalk: # iteration when using multiprocessing _cache = list() - def __init__(self): + def __init__(self, **kwargs): + super().__init__(**kwargs) self.act = 1 - self.thin = getattr(_SamplingContainer, "nact", 2) - self.maxmcmc = getattr(_SamplingContainer, "maxmcmc", 5000) * 50 + self.thin = kwargs.get("nact", 2) + self.maxmcmc = kwargs.get("maxmcmc", 5000) * 50 + self.sampler_kwargs["rebuild"] = True + self.sampler_kwargs["thin"] = self.thin + self.sampler_kwargs["act"] = self.act + self.sampler_kwargs["maxmcmc"] = self.maxmcmc + # reset the cache at instantiation to avoid contamination from + # previous analyses + self.__class__._cache = list() + + def prepare_sampler( + self, + loglstar=None, + points=None, + axes=None, + seeds=None, + prior_transform=None, + loglikelihood=None, + nested_sampler=None, + ): + """ + Prepare the list of arguments for sampling. + + Parameters + ---------- + loglstar : float + Ln(likelihood) bound. + points : `~numpy.ndarray` with shape (n, ndim) + Initial sample points. + axes : `~numpy.ndarray` with shape (ndim, ndim) + Axes used to propose new points. + seeds : `~numpy.ndarray` with shape (n,) + Random number generator seeds. + prior_transform : function + Function transforming a sample from the a unit cube to the + parameter space of interest according to the prior. + loglikelihood : function + Function returning ln(likelihood) given parameters as a 1-d + `~numpy` array of length `ndim`. + nested_sampler : `~dynesty.samplers.Sampler` + The nested sampler object used to sample. + + Returns + ------- + arglist: + List of `SamplerArgument` objects containing the parameters + needed for sampling. + """ + arg_list = super().prepare_sampler( + loglstar=loglstar, + points=points, + axes=axes, + seeds=seeds, + prior_transform=prior_transform, + loglikelihood=loglikelihood, + nested_sampler=nested_sampler, + ) + self.sampler_kwargs["rebuild"] = False + return arg_list + + def tune(self, tuning_info, update=True): + """ + Update the proposal parameters based on the number of accepted steps + and MCMC chain length. - def __call__(self, args): - self.args = args + The :code:`walks` parameter to asymptotically approach the + desired number of accepted steps. + """ + if tuning_info.get("remaining", 0) == 0: + self.sampler_kwargs["rebuild"] = True + self.scale = tuning_info["accept"] + self.sampler_kwargs["act"] = tuning_info["act"] + + @staticmethod + def sample(args): + cache = ACTTrackingEnsembleWalk._cache if args.kwargs.get("rebuild", False): - logger.debug("Force rebuilding cache") - self.build_cache() - while self.cache[0][2] < args.loglstar: - self.cache.pop(0) - current_u, current_v, logl, ncall, blob = self.cache.pop(0) - blob["remaining"] = len(self.cache) - return current_u, current_v, logl, ncall, blob - - @property - def cache(self): - if len(self._cache) == 0: - self.build_cache() + logger.debug(f"Force rebuilding cache with {len(cache)}.") + cache.clear() + if len(cache) == 0: + ACTTrackingEnsembleWalk.build_cache(args) + + while len(cache) > 0 and cache[0][2] < args.loglstar: + state = cache.pop(0) + + if len(cache) == 0: + current_u, current_v, logl, ncall, blob, evaluation_history = state else: - logger.debug(f"Not rebuilding cache, remaining size {len(self._cache)}") - return self._cache + current_u, current_v, logl, ncall, blob, evaluation_history = cache.pop(0) + + blob["remaining"] = len(cache) + logger.debug(f"Returning point from cache with {len(cache)} remaining") + return SamplerReturn( + u=current_u, + v=current_v, + logl=logl, + tuning_info=blob, + ncalls=ncall, + proposal_stats=blob, + evaluation_history=evaluation_history, + ) - def build_cache(self): - args = self.args + @staticmethod + def build_cache(args): # Bounds periodic = args.kwargs.get("periodic", None) reflective = args.kwargs.get("reflective", None) @@ -220,14 +372,20 @@ def build_cache(self): # Setup current_u = args.u - check_interval = self.integer_act + old_act = args.kwargs.get("act", 100) + check_interval = ACTTrackingEnsembleWalk.integer_act(old_act) target_nact = 50 next_check = check_interval n_checks = 0 + maxmcmc = args.kwargs.get("maxmcmc", 5000) + evaluation_history = list() # Initialize internal variables current_v = args.prior_transform(np.array(current_u)) logl = args.loglikelihood(np.array(current_v)) + evaluation_history.append( + SamplerHistoryItem(u=current_u, v=current_v, logl=logl) + ) accept = 0 reject = 0 nfail = 0 @@ -240,7 +398,7 @@ def build_cache(self): current_failures = 0 iteration = 0 - while iteration < min(target_nact * act, self.maxmcmc): + while iteration < min(target_nact * act, maxmcmc): iteration += 1 prop = proposals[iteration % len(proposals)] @@ -252,6 +410,9 @@ def build_cache(self): if u_prop is not None: v_prop = args.prior_transform(np.array(u_prop)) logl_prop = args.loglikelihood(np.array(v_prop)) + evaluation_history.append( + SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) + ) ncall += 1 if logl_prop > args.loglstar: success = True @@ -275,7 +436,7 @@ def build_cache(self): if iteration > next_check and accept > target_nact: n_checks += 1 most_failures = max(most_failures, current_failures) - act = self._calculate_act( + act = ACTTrackingEnsembleWalk._calculate_act( accept=accept, iteration=iteration, samples=np.array(u_list), @@ -292,16 +453,18 @@ def build_cache(self): next_check += check_interval most_failures = max(most_failures, current_failures) - self.act = self._calculate_act( + act = ACTTrackingEnsembleWalk._calculate_act( accept=accept, iteration=iteration, samples=np.array(u_list), most_failures=most_failures, ) reject += nfail - blob = {"accept": accept, "reject": reject, "scale": args.scale} - iact = self.integer_act - thin = self.thin * iact + blob = {"accept": accept, "reject": reject, "act": act} + iact = ACTTrackingEnsembleWalk.integer_act(act) + thin = args.kwargs.get("thin", 2) * iact + + cache = ACTTrackingEnsembleWalk._cache if accept == 0: logger.warning( @@ -310,34 +473,50 @@ def build_cache(self): u = common_kwargs["rstate"].uniform(size=len(current_u)) v = args.prior_transform(u) logl = args.loglikelihood(v) - self._cache.append((u, v, logl, ncall, blob)) + evaluation_history = [ + SamplerHistoryItem(u=current_u, v=current_v, logl=logl) + ] + cache.append((u, v, logl, ncall, blob, evaluation_history)) elif not np.isfinite(act): logger.warning( "Unable to find a new point using walk: try increasing maxmcmc" ) - self._cache.append((current_u, current_v, logl, ncall, blob)) - elif (self.thin == -1) or (len(u_list) <= thin): - self._cache.append((current_u, current_v, logl, ncall, blob)) + cache.append((current_u, current_v, logl, ncall, blob, evaluation_history)) + elif (thin == -1) or (len(u_list) <= thin): + cache.append((current_u, current_v, logl, ncall, blob, evaluation_history)) else: u_list = u_list[thin::thin] v_list = v_list[thin::thin] logl_list = logl_list[thin::thin] + evaluation_history_list = ( + evaluation_history[thin * ii : thin * (ii + 1)] + for ii in range(len(u_list)) + ) n_found = len(u_list) accept = max(accept // n_found, 1) reject //= n_found nfail //= n_found ncall_list = [ncall // n_found] * n_found blob_list = [ - dict(accept=accept, reject=reject, fail=nfail, scale=args.scale) + dict(accept=accept, reject=reject, fail=nfail, act=act) ] * n_found - self._cache.extend(zip(u_list, v_list, logl_list, ncall_list, blob_list)) + cache.extend( + zip( + u_list, + v_list, + logl_list, + ncall_list, + blob_list, + evaluation_history_list, + ) + ) logger.debug( - f"act: {self.act:.2f}, max failures: {most_failures}, thin: {thin}, " + f"act: {act:.2f}, max failures: {most_failures}, thin: {thin}, " f"iteration: {iteration}, n_found: {n_found}" ) logger.debug( - f"Finished building cache with length {len(self._cache)} after " - f"{iteration} iterations with {ncall} likelihood calls and ACT={self.act:.2f}" + f"Finished building cache with length {len(cache)} after " + f"{iteration} iterations with {ncall} likelihood calls and ACT={act:.2f}" ) @staticmethod @@ -359,15 +538,15 @@ def _calculate_act(accept, iteration, samples, most_failures): return np.inf return max(calculate_tau(samples), naive_act, most_failures) - @property - def integer_act(self): - if np.isinf(self.act): - return self.act + @staticmethod + def integer_act(act): + if np.isinf(act): + return act else: - return int(np.ceil(self.act)) + return int(np.ceil(act)) -class AcceptanceTrackingRWalk: +class AcceptanceTrackingRWalk(EnsembleWalkSampler): """ This is a modified version of dynesty.sampling.sample_rwalk that runs the MCMC random walk for a user-specified number of a crude approximation to @@ -381,18 +560,28 @@ class AcceptanceTrackingRWalk: # level attribute old_act = None - def __init__(self, old_act=None): - self.maxmcmc = getattr(_SamplingContainer, "maxmcmc", 5000) - self.nact = getattr(_SamplingContainer, "nact", 40) + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.nact = kwargs.get("nact", 40) + self.sampler_kwargs["nact"] = self.nact + self.sampler_kwargs["maxmcmc"] = self.maxmcmc + + def tune(self, tuning_info, update=True): + """ + The tuning all happens inside the MCMC for this proposal + so all we do is extract the number of accepted steps for plotting. + """ + self.scale = tuning_info["accept"] - def __call__(self, args): + @staticmethod + def sample(args): rstate = get_random_generator(args.rseed) periodic = args.kwargs.get("periodic", None) reflective = args.kwargs.get("reflective", None) boundary_kwargs = dict(periodic=periodic, reflective=reflective) - u = args.u + current_u = args.u nlive = args.kwargs.get("nlive", args.kwargs.get("walks", 100)) proposals, common_kwargs, proposal_kwargs = _get_proposal_kwargs(args) @@ -401,13 +590,18 @@ def __call__(self, args): reject = 0 nfail = 0 act = np.inf + nact = args.kwargs.get("nact", 40) + maxmcmc = args.kwargs.get("maxmcmc", 5000) + evaluation_history = list() iteration = 0 - while iteration < self.nact * act: + while iteration < nact * act: iteration += 1 prop = proposals[iteration % len(proposals)] - u_prop = proposal_funcs[prop](u, **common_kwargs, **proposal_kwargs[prop]) + u_prop = proposal_funcs[prop]( + current_u, **common_kwargs, **proposal_kwargs[prop] + ) u_prop = apply_boundaries_(u_prop, **boundary_kwargs) if u_prop is None: @@ -417,26 +611,31 @@ def __call__(self, args): # Check proposed point. v_prop = args.prior_transform(np.array(u_prop)) logl_prop = args.loglikelihood(np.array(v_prop)) + evaluation_history.append( + SamplerHistoryItem(u=v_prop, v=u_prop, logl=logl_prop) + ) if logl_prop > args.loglstar: - u = u_prop - v = v_prop + current_u = u_prop + current_v = v_prop logl = logl_prop accept += 1 else: reject += 1 # If we've taken the minimum number of steps, calculate the ACT - if iteration > self.nact: - act = self.estimate_nmcmc( + if iteration > nact: + act = AcceptanceTrackingRWalk.estimate_nmcmc( accept_ratio=accept / (accept + reject + nfail), safety=1, tau=nlive, + maxmcmc=maxmcmc, + old_act=AcceptanceTrackingRWalk.old_act, ) # If we've taken too many likelihood evaluations then break - if accept + reject > self.maxmcmc: + if accept + reject > maxmcmc: warnings.warn( - f"Hit maximum number of walks {self.maxmcmc} with accept={accept}," + f"Hit maximum number of walks {maxmcmc} with accept={accept}," f" reject={reject}, and nfail={nfail} try increasing maxmcmc" ) break @@ -445,17 +644,26 @@ def __call__(self, args): logger.debug( "Unable to find a new point using walk: returning a random point" ) - u = rstate.uniform(size=len(u)) - v = args.prior_transform(u) - logl = args.loglikelihood(v) + current_u = rstate.uniform(size=len(current_u)) + current_v = args.prior_transform(current_u) + logl = args.loglikelihood(current_v) - blob = {"accept": accept, "reject": reject + nfail, "scale": args.scale} + sampling_info = {"accept": accept, "reject": reject + nfail} AcceptanceTrackingRWalk.old_act = act ncall = accept + reject - return u, v, logl, ncall, blob + return SamplerReturn( + u=current_u, + v=current_v, + logl=logl, + tuning_info=sampling_info, + ncalls=ncall, + proposal_stats=sampling_info, + evaluation_history=evaluation_history, + ) - def estimate_nmcmc(self, accept_ratio, safety=5, tau=None): + @staticmethod + def estimate_nmcmc(accept_ratio, safety=5, tau=None, maxmcmc=5000, old_act=None): """Estimate autocorrelation length of chain using acceptance fraction Using ACL = (2/acc) - 1 multiplied by a safety margin. Code adapted from: @@ -482,19 +690,19 @@ def estimate_nmcmc(self, accept_ratio, safety=5, tau=None): length for our proposal distributions. """ if tau is None: - tau = self.maxmcmc / safety + tau = maxmcmc / safety if accept_ratio == 0.0: - if self.old_act is None: + if old_act is None: Nmcmc_exact = np.inf else: - Nmcmc_exact = (1 + 1 / tau) * self.old_act + Nmcmc_exact = (1 + 1 / tau) * old_act else: estimated_act = 2 / accept_ratio - 1 Nmcmc_exact = safety * estimated_act - if self.old_act is not None: - Nmcmc_exact = (1 - 1 / tau) * self.old_act + Nmcmc_exact / tau - Nmcmc_exact = float(min(Nmcmc_exact, self.maxmcmc)) + if old_act is not None: + Nmcmc_exact = (1 - 1 / tau) * old_act + Nmcmc_exact / tau + Nmcmc_exact = float(min(Nmcmc_exact, maxmcmc)) return max(safety, Nmcmc_exact) @@ -524,13 +732,20 @@ def _get_proposal_kwargs(args): rstate = get_random_generator(args.rseed) walks = args.kwargs.get("walks", 100) current_u = args.u - n_cluster = args.axes.shape[0] - proposals = getattr(_SamplingContainer, "proposals", None) + proposals = args.kwargs.get("proposals", None) if proposals is None: proposals = ["diff"] + + n_cluster = args.kwargs.get("ncdim", None) + if n_cluster is None: + if hasattr(args, "live_points"): + n_cluster = args.live_points.shape[1] + elif hasattr(args, "axes"): + n_cluster = args.axes.shape[0] + if "diff" in proposals: - live = args.kwargs.get("live", None) + live = args.live_points if live is None: raise ValueError( "Live points not passed for differential evolution, specify " @@ -567,7 +782,7 @@ def _get_proposal_kwargs(args): return proposals, common_kwargs, proposal_kwargs -def propose_differetial_evolution( +def propose_differential_evolution( u, live, n, @@ -722,4 +937,6 @@ def apply_boundaries_(u_prop, periodic, reflective): return u_prop -proposal_funcs = dict(diff=propose_differetial_evolution, volumetric=propose_volumetric) +proposal_funcs = dict( + diff=propose_differential_evolution, volumetric=propose_volumetric +) diff --git a/containers/env-template.yml b/containers/env-template.yml index 42b3f3f8a..925ea9f56 100644 --- a/containers/env-template.yml +++ b/containers/env-template.yml @@ -44,7 +44,7 @@ dependencies: - nbsphinx - sphinx_rtd_theme - sphinx-design - - dynesty + - dynesty>=3.0.0 - emcee - nestle - ptemcee diff --git a/requirements.txt b/requirements.txt index b045db212..84a0bee61 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ bilby.cython>=0.3.0 -dynesty>=2.0.1 +dynesty>=3.0.0 emcee corner numpy diff --git a/test/core/sampler/dynesty_test.py b/test/core/sampler/dynesty_test.py index 0be239f19..342e6ca3b 100644 --- a/test/core/sampler/dynesty_test.py +++ b/test/core/sampler/dynesty_test.py @@ -7,20 +7,10 @@ import bilby.core.sampler.dynesty import numpy as np import parameterized -import pytest from attr import define from scipy.stats import gamma, ks_1samp, uniform, powerlaw - -try: - import dynesty.internal_samplers # noqa - from bilby.core.sampler import dynesty3_utils as dynesty_utils # noqa - - NEW_DYNESTY_API = True -except ImportError: - from bilby.core.sampler import dynesty_utils # noqa - - NEW_DYNESTY_API = False +from bilby.core.sampler import dynesty_utils @define @@ -70,27 +60,25 @@ def init_sampler(self, **kwargs): skip_import_verification=True, **kwargs, ) - if NEW_DYNESTY_API: - bilby.core.sampler.base_sampler._initialize_global_variables( - self.likelihood, - self.priors, - self.priors.keys(), - False, - self.priors.sample(), - ) - self.dysampler = self.sampler.sampler_init( - loglikelihood=bilby.core.sampler.dynesty._log_likelihood_wrapper, - prior_transform=bilby.core.sampler.dynesty._prior_transform_wrapper, - ndim=self.sampler.ndim, - **self.sampler.sampler_init_kwargs, - ) + bilby.core.sampler.base_sampler._initialize_global_variables( + self.likelihood, + self.priors, + self.priors.keys(), + False, + self.priors.sample(), + ) + self.dysampler = self.sampler.sampler_init( + loglikelihood=bilby.core.sampler.dynesty._log_likelihood_wrapper, + prior_transform=bilby.core.sampler.dynesty._prior_transform_wrapper, + ndim=self.sampler.ndim, + **self.sampler.sampler_init_kwargs, + ) def tearDown(self): del self.likelihood del self.priors del self.sampler - if NEW_DYNESTY_API: - del self.dysampler + del self.dysampler def test_default_kwargs(self): """Only test the kwargs where we specify different defaults to dynesty""" @@ -121,58 +109,39 @@ def test_prior_boundary(self): self.priors["d"] = bilby.core.prior.Prior(boundary="reflective") self.priors["e"] = bilby.core.prior.Prior(boundary="periodic") self.init_sampler() - if NEW_DYNESTY_API: - self.assertEqual( - [0, 4], self.dysampler.internal_sampler_next.sampler_kwargs["periodic"] - ) - self.assertEqual( - [1, 3], - self.dysampler.internal_sampler_next.sampler_kwargs["reflective"], - ) - else: - self.assertEqual([0, 4], self.sampler.kwargs["periodic"]) - self.assertEqual([1, 3], self.sampler.kwargs["reflective"]) + self.assertEqual( + [0, 4], self.dysampler.internal_sampler_next.sampler_kwargs["periodic"] + ) + self.assertEqual( + [1, 3], + self.dysampler.internal_sampler_next.sampler_kwargs["reflective"], + ) def test_sampler_kwargs_act(self): self.init_sampler(sample="act-walk", nact=5, maxmcmc=200) - if NEW_DYNESTY_API: - self.assertIsInstance( - self.dysampler.internal_sampler_next, - dynesty_utils.ACTTrackingEnsembleWalk, - ) - self.assertEqual(self.dysampler.internal_sampler_next.thin, 5) - self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200 * 50) - else: - self.assertEqual(self.sampler.kwargs["sample"], "act-walk") - self.assertEqual(self.sampler.nact, 5) - self.assertEqual(self.sampler.maxmcmc, 200) + self.assertIsInstance( + self.dysampler.internal_sampler_next, + dynesty_utils.ACTTrackingEnsembleWalk, + ) + self.assertEqual(self.dysampler.internal_sampler_next.thin, 5) + self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200 * 50) def test_sampler_kwargs_rwalk(self): self.init_sampler(sample="rwalk", nact=5, maxmcmc=200) - if NEW_DYNESTY_API: - self.assertIsInstance( - self.dysampler.internal_sampler_next, - dynesty_utils.AcceptanceTrackingRWalk, - ) - self.assertEqual(self.dysampler.internal_sampler_next.nact, 5) - self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200) - else: - self.assertEqual(self.sampler.kwargs["sample"], "rwalk") - self.assertEqual(self.sampler.nact, 5) - self.assertEqual(self.sampler.maxmcmc, 200) + self.assertIsInstance( + self.dysampler.internal_sampler_next, + dynesty_utils.AcceptanceTrackingRWalk, + ) + self.assertEqual(self.dysampler.internal_sampler_next.nact, 5) + self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200) def test_sampler_kwargs_acceptance_walk(self): self.init_sampler(sample="acceptance-walk", naccept=5, maxmcmc=200) - if NEW_DYNESTY_API: - self.assertIsInstance( - self.dysampler.internal_sampler_next, dynesty_utils.EnsembleWalkSampler - ) - self.assertEqual(self.dysampler.internal_sampler_next.naccept, 5) - self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200) - else: - self.assertEqual(self.sampler.kwargs["sample"], "acceptance-walk") - self.assertEqual(self.sampler.naccept, 5) - self.assertEqual(self.sampler.maxmcmc, 200) + self.assertIsInstance( + self.dysampler.internal_sampler_next, dynesty_utils.EnsembleWalkSampler + ) + self.assertEqual(self.dysampler.internal_sampler_next.naccept, 5) + self.assertEqual(self.dysampler.internal_sampler_next.maxmcmc, 200) def test_run_test_runs(self): self.sampler._run_test() @@ -283,11 +252,8 @@ def test_propose_differential_evolution(self, scale): def test_get_proposal_kwargs_diff(self): args = Dummy(u=-np.ones(4), axes=np.zeros((2, 2)), scale=4) - if NEW_DYNESTY_API: - args.kwargs["proposals"] = ["diff"] - args.kwargs["ncdim"] = 2 - else: - dynesty_utils._SamplingContainer.proposals = ["diff"] + args.kwargs["proposals"] = ["diff"] + args.kwargs["ncdim"] = 2 proposals, common, specific = dynesty_utils._get_proposal_kwargs(args) del common["rstate"] self.assertTrue( @@ -300,11 +266,8 @@ def test_get_proposal_kwargs_diff(self): def test_get_proposal_kwargs_volumetric(self): args = Dummy(u=-np.ones(4), axes=np.zeros((2, 2)), scale=4) - if NEW_DYNESTY_API: - args.kwargs["proposals"] = ["volumetric"] - args.live_points = np.zeros((2, 2)) - else: - dynesty_utils._SamplingContainer.proposals = ["volumetric"] + args.kwargs["proposals"] = ["volumetric"] + args.live_points = np.zeros((2, 2)) proposals, common, specific = dynesty_utils._get_proposal_kwargs(args) del common["rstate"] self.assertTrue( @@ -315,22 +278,6 @@ def test_get_proposal_kwargs_volumetric(self): specific, dict(volumetric=dict(axes=args.axes, scale=args.scale)) ) - @pytest.mark.skipif(NEW_DYNESTY_API, reason="Invalid for new dynesty API") - def test_proposal_functions_run_old(self): - args = Dummy(u=np.ones(4) / 2, axes=np.ones((2, 2))) - args.kwargs["live"][0] += 1 - for proposals in [ - ["diff"], - ["volumetric"], - ["diff", "volumetric"], - {"diff": 5, "volumetric": 1}, - ]: - dynesty_utils._SamplingContainer.proposals = proposals - dynesty_utils.FixedRWalk()(args) - dynesty_utils.AcceptanceTrackingRWalk()(args) - dynesty_utils.ACTTrackingRWalk()(args) - - @pytest.mark.skipif(not NEW_DYNESTY_API, reason="Invalid for old dynesty API") def test_proposal_functions_run_new(self): args = Dummy(u=np.ones(4) / 2, axes=np.ones((2, 2))) args.live_points[0] += 1 @@ -348,69 +295,6 @@ def test_proposal_functions_run_new(self): dynesty_utils.ACTTrackingEnsembleWalk().sample(args) -@pytest.mark.skipif(NEW_DYNESTY_API, reason="Invalid for new dynesty API") -@parameterized.parameterized_class(("kind",), [("live",), ("live-multi",)]) -class TestCustomSampler(unittest.TestCase): - def setUp(self): - if self.kind == "live": - cls = dynesty_utils.LivePointSampler - elif self.kind == "live-multi": - cls = dynesty_utils.MultiEllipsoidLivePointSampler - else: - raise ValueError(f"Unknown sampler class {self.kind}") - - self.sampler = cls( - loglikelihood=lambda x: 1, - prior_transform=lambda x: x, - ndim=4, - live_points=(np.zeros((1000, 4)), np.zeros((1000, 4)), np.zeros(1000)), - update_interval=None, - first_update=dict(), - queue_size=1, - pool=None, - use_pool=dict(), - ncdim=2, - method="rwalk", - rstate=np.random.default_rng(1234), - kwargs=dict(walks=100), - ) - self.blob = dict(accept=5, reject=35, scale=1) - - def tearDown(self): - dynesty_utils._SamplingContainer.proposals = None - - def test_update_with_update(self): - """ - If there is only one element left in the cache for ACT tracking we - need to add the flag to rebuild the cache. After rebuilding, this is - then reset to False. - """ - self.sampler.rebuild = False - self.blob["remaining"] = 1 - self.sampler.update_user(blob=self.blob, update=True) - self.assertEqual(self.sampler.kwargs["rebuild"], True) - self.blob["remaining"] = 2 - self.sampler.update_user(blob=self.blob, update=True) - self.assertEqual(self.sampler.kwargs["rebuild"], False) - - def test_diff_update(self): - """ - Sampler updates do different things depending on whether ellipsoid - bounding is used. For the `live-multi` case we reproduce the dynesty - behaviour. For the `live` case, we overwrite the scale attribute to - store acceptance information. - """ - dynesty_utils._SamplingContainer.proposals = ["diff"] - dynesty_utils._SamplingContainer.maxmcmc = 10000 - dynesty_utils._SamplingContainer.naccept = 10 - self.sampler.update_user(blob=self.blob, update=False) - if self.kind == "live": - self.assertEqual(self.sampler.scale, self.blob["accept"]) - else: - self.assertEqual(self.sampler.scale, self.blob["scale"]) - self.assertEqual(self.sampler.walks, 101) - - class TestEstimateNMCMC(unittest.TestCase): def test_converges_to_correct_value(self): """