diff --git a/py/dynesty/dynesty.py b/py/dynesty/dynesty.py index 60fab7603..88e5a5adb 100644 --- a/py/dynesty/dynesty.py +++ b/py/dynesty/dynesty.py @@ -252,7 +252,9 @@ def NestedSampler(loglikelihood, update_func=None, ncdim=None, save_history=False, - history_filename=None): + history_filename=None, + proposals=None, + ): """ Initializes and returns a sampler object for Static Nested Sampling. @@ -472,6 +474,10 @@ def prior_transform(u): just sample uniformly from the prior distribution. If this is `None` (default), this will default to npdim. + proposals : iterable, optional + A list of proposal functions to use for rwalk evolution. This will be + used as a fixed proposal cycle. + Returns ------- sampler : sampler from :mod:`~dynesty.nestedsamplers` @@ -519,6 +525,7 @@ def prior_transform(u): kwargs['nonbounded'] = nonbounded kwargs['periodic'] = periodic kwargs['reflective'] = reflective + kwargs['proposals'] = proposals # Keyword arguments controlling the first update. if first_update is None: @@ -663,7 +670,9 @@ def DynamicNestedSampler(loglikelihood, update_func=None, ncdim=None, save_history=False, - history_filename=None): + history_filename=None, + proposals=None, + ): """ Initializes and returns a sampler object for Dynamic Nested Sampling. @@ -868,6 +877,10 @@ def prior_transform(u): just sample uniformly from the prior distribution. If this is `None` (default), this will default to npdim. + proposals : iterable, optional + A list of proposal functions to use for rwalk evolution. This will be + used as a fixed proposal cycle. + Returns ------- sampler : a :class:`dynesty.DynamicSampler` instance @@ -918,6 +931,7 @@ def prior_transform(u): kwargs['nonbounded'] = nonbounded kwargs['periodic'] = periodic kwargs['reflective'] = reflective + kwargs['proposals'] = proposals # Keyword arguments controlling the first update. if first_update is None: diff --git a/py/dynesty/nestedsamplers.py b/py/dynesty/nestedsamplers.py index 2987f707e..d946d9d27 100644 --- a/py/dynesty/nestedsamplers.py +++ b/py/dynesty/nestedsamplers.py @@ -30,6 +30,7 @@ from .sampler import Sampler from .bounding import (UnitCube, Ellipsoid, MultiEllipsoid, RadFriends, SupFriends, rand_choice) +from .proposals import DEFAULT_PROPOSALS, ENSEMBLE_PROPOSALS from .sampling import (sample_unif, sample_rwalk, sample_slice, sample_rslice, sample_hslice) from .utils import unitcheck, get_enlarge_bootstrap @@ -137,6 +138,13 @@ def __init__(self, self.fmove = self.kwargs.get('fmove', 0.9) self.max_move = self.kwargs.get('max_move', 100) + @property + def proposals(self): + proposals = self.kwargs.get("proposals", None) + if proposals is None: + proposals = DEFAULT_PROPOSALS + return proposals + def propose_unif(self, *args): pass @@ -327,6 +335,9 @@ def propose_live(self, *args): u = self.live_u[i, :] ax = np.identity(self.npdim) + if not ENSEMBLE_PROPOSALS.isdisjoint(self.proposals): + self.kwargs["live"] = copy.deepcopy(self.live_u) + return u, ax @@ -473,6 +484,9 @@ def propose_live(self, *args): else: ax = np.identity(self.ncdim) + if not ENSEMBLE_PROPOSALS.isdisjoint(self.proposals): + self.kwargs["live"] = copy.deepcopy(self.live_u) + return u, ax @@ -649,6 +663,9 @@ def propose_live(self, *args): else: ax = np.identity(self.npdim) + if not ENSEMBLE_PROPOSALS.isdisjoint(self.proposals): + self.kwargs["live"] = copy.deepcopy(self.live_u) + return u, ax @@ -799,6 +816,9 @@ def propose_live(self, *args): u = self.live_u[i, :] ax = self.radfriends.axes + if not ENSEMBLE_PROPOSALS.isdisjoint(self.proposals): + self.kwargs["live"] = copy.deepcopy(self.live_u) + return u, ax @@ -949,4 +969,7 @@ def propose_live(self, *args): u = self.live_u[i, :] ax = self.supfriends.axes + if not ENSEMBLE_PROPOSALS.isdisjoint(self.proposals): + self.kwargs["live"] = copy.deepcopy(self.live_u) + return u, ax diff --git a/py/dynesty/proposals.py b/py/dynesty/proposals.py new file mode 100644 index 000000000..cdfd2bd9e --- /dev/null +++ b/py/dynesty/proposals.py @@ -0,0 +1,322 @@ +import numpy as np + +ELLIPSOID_PROPOSALS = {"normal", "volumetric", "axis", "chi"} +ENSEMBLE_PROPOSALS = {"diff", "stretch", "walk", "snooker"} +DEFAULT_PROPOSALS = {"diff", "stretch", "normal"} + + +def register_proposal(name, proposal, kind=None): + f""" + Add a new proposal function to the known set. + + Parameters + ---------- + name: str + Name of the proposal, this is what is passed to :code:`dynesty`, + e.g., :code:`diff` + proposal: callable + The callable that takes in the standard proposal inputs + :code:`u, rstate, axis, live, scale` and returns either + proposed point and possibly a natural log Jacobian associated with the + step. + kind: str, optional + Either :code:`ellipsoid`, :code:`ensemble` or :code:`None`. This tells + the sampler whether bounding ellipsoids need to be computed for the + proposal or the current live points need to be passed respectively. + + Raises + ------ + ValueError + A value error is raised if the new proposal name already exists in the + known registered proposals. + The pre-defined proposal names are {", ".join(_proposal_map.keys())}. + """ + if name in _proposal_map: + raise ValueError(f"Proposal {name} already exists.") + if kind == "ellipsoid": + ELLIPSOID_PROPOSALS.add(name) + elif kind == "ensemble": + ENSEMBLE_PROPOSALS.add(name) + _proposal_map[name] = proposal + + +_acceptances = dict() +_failures = dict() + + +def propose_diff_evo(u, live, rstate, **kwargs): + r""" + Propose a new point using ensemble differential evolution. + + .. math:: + + u_{\rm prop} = u + \gamma (v_{a} - v_{b}) + + Parameters + ---------- + u: np.ndarray + The current point. + live: np.ndarray + The ensemble of live points to select :math:`v` from. + rstate: RandomState + The random state to use to generate random numbers. + kwargs: dict, unused + Required for flexibility of the interface + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + nlive, n = live.shape + first, second = rstate.choice(nlive, 2, replace=False) + diff = live[second] - live[first] + if rstate.uniform(0, 1) < 0.5: + diff *= 2.38 / n**0.5 + diff *= (100**rstate.uniform(0, 1)) / 10 + u_prop = u + diff + return u_prop + + +def propose_ensemble_snooker(u, live, rstate, **kwargs): + r""" + Propose a new point using ensemble differential evolution. + + .. math:: + + u_{\rm prop} = u + \gamma (v_{a} - v_{b}) + + Parameters + ---------- + u: np.ndarray + The current point. + live: np.ndarray + The ensemble of live points to select :math:`v` from. + rstate: RandomState + The random state to use to generate random numbers. + kwargs: dict, unused + Required for flexibility of the interface + + Returns + ------- + u_prop: np.ndarray + The proposed point. + ln_jacobian: float + The natural log transition asymmetry for the proposed step. + """ + nlive, n = live.shape + choices = rstate.choice(nlive, 3, replace=False) + z = live[choices[0]] + z1 = live[choices[1]] + z2 = live[choices[2]] + delta = u - z + norm = np.linalg.norm(delta) + delta /= norm + u_prop = u + 1.7 * delta * (np.dot(u, z1) - np.dot(u, z2)) + ln_jacobian = (n - 1.0) * np.log(np.linalg.norm(u_prop - z) / norm) + return u_prop, ln_jacobian + + +def propose_ensemble_stretch(u, live, rstate, **kwargs): + r""" + Propose a new point using ensemble differential evolution. + + .. math:: + + u_{\rm prop} = v + \gamma (u - v) + + Parameters + ---------- + u: np.ndarray + The current point. + live: np.ndarray + The ensemble of live points to select :math:`v` from. + rstate: RandomState + The random state to use to generate random numbers. + kwargs: dict, unused + Required for flexibility of the interface + + Returns + ------- + u_prop: np.ndarray + The proposed point. + ln_jacobian: float + The natural log transition asymmetry for the proposed step. + """ + nlive, n = live.shape + max_scale = 3 + scale = ((max_scale - 1.0) * rstate.uniform() + 1) ** 2.0 / max_scale + + other = rstate.choice(nlive) + other = live[other] + + u_prop = other + scale * (u - other) + ln_jacobian = np.log(scale) * (n - 1) + return u_prop, ln_jacobian + + +def propose_ensemble_walk(u, live, rstate, nsamples=3, **kwargs): + r""" + Propose a new point using ensemble differential evolution. + + .. math:: + + u_{\rm prop} = u + \frac{1}{N_{\rm samples}} + \sum^{N_{\rm samples}}_{i=1} \gamma_{i} (v_{i} - \hat{v}) + + Parameters + ---------- + u: np.ndarray + The current point. + live: np.ndarray + The ensemble of live points to select :math:`v` from. + rstate: RandomState + The random state to use to generate random numbers. + kwargs: dict, unused + Required for flexibility of the interface + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + nlive, n = live.shape + choices = rstate.choice(nlive, nsamples, replace=False) + center_of_mass = np.mean(live[choices], axis=0) + scales = rstate.normal(0, 1, nsamples)[:, np.newaxis] + diff = np.sum(scales * (live[choices] - center_of_mass), axis=0) + return u + diff + + +def propose_along_axis(u, axes, rstate, scale=1, **kwargs): + """ + Propose a new point along one of the principal axes of the bounding + ellipsoid. + + Parameters + ---------- + u: np.ndarray + The current point. + axes: np.ndarray + The principal component decomposition of the bounding ellipsoid. + rstate: RandomState + The random state to use to generate random numbers. + scale: float + The scale of the step to take. + kwargs: dict, unused + Required for flexibility of the interface. + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + n = len(u) + idx = rstate.choice(n) + axis = axes[idx] + scale *= rstate.normal(0, 1) + return u + axis * scale + + +def propose_chi(u, axes, rstate, scale=1, **kwargs): + """ + Propose a new point drawn using the principal axes of the bounding + ellipsoid. The scale of the proposal is drawn from a chi distribution + with n_dim degrees of freedom. + + Parameters + ---------- + u: np.ndarray + The current point. + axes: np.ndarray + The principal component decomposition of the bounding ellipsoid. + rstate: RandomState + The random state to use to generate random numbers. + scale: float + The scale of the step to take. + kwargs: dict, unused + Required for flexibility of the interface. + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + scale *= rstate.chisquare(len(u)) / len(u) + return _axes_proposal(u=u, axes=axes, rstate=rstate, scale=scale) + + +def propose_normal(u, axes, rstate, scale=1, **kwargs): + """ + Propose a new point drawn using the principal axes of the bounding + ellipsoid. The scale of the proposal is drawn from a unit normal + distribution. + + Parameters + ---------- + u: np.ndarray + The current point. + axes: np.ndarray + The principal component decomposition of the bounding ellipsoid. + rstate: RandomState + The random state to use to generate random numbers. + scale: float + The scale of the step to take. + kwargs: dict, unused + Required for flexibility of the interface. + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + scale *= rstate.normal(0, 1) + return _axes_proposal(u=u, axes=axes, rstate=rstate, scale=scale) + + +def propose_volumetric(u, axes, rstate, scale=1, **kwargs): + """ + Propose a new point drawn using the principal axes of the bounding + ellipsoid. The scale of the proposal is drawn uniformly from the ellipsoid. + + Parameters + ---------- + u: np.ndarray + The current point. + axes: np.ndarray + The principal component decomposition of the bounding ellipsoid. + rstate: RandomState + The random state to use to generate random numbers. + scale: float + The scale of the step to take. + kwargs: dict, unused + Required for flexibility of the interface. + + Returns + ------- + u_prop: np.ndarray + The proposed point. + """ + n = len(u) + scale *= rstate.uniform(0, 1) ** (1.0 / n) + return _axes_proposal(u=u, axes=axes, rstate=rstate, scale=scale) + + +def _axes_proposal(u, axes, rstate, scale): + n = len(u) + drhat = rstate.normal(0, 1, n) + drhat /= np.linalg.norm(drhat) + du = np.dot(axes, drhat) + return u + scale * du + + +_proposal_map = dict( + diff=propose_diff_evo, + stretch=propose_ensemble_stretch, + walk=propose_ensemble_walk, + snooker=propose_ensemble_snooker, + axis=propose_along_axis, + volumetric=propose_volumetric, + chi=propose_chi, + normal=propose_normal, +) diff --git a/py/dynesty/sampling.py b/py/dynesty/sampling.py index f9cfa8101..37c021e7b 100644 --- a/py/dynesty/sampling.py +++ b/py/dynesty/sampling.py @@ -10,11 +10,20 @@ import warnings import math +from copy import deepcopy + import numpy as np from numpy import linalg -from .utils import unitcheck, apply_reflect, get_random_generator -from .bounding import randsphere +from .utils import apply_boundary, get_random_generator, unitcheck +from .proposals import ( + _acceptances, + DEFAULT_PROPOSALS, + ELLIPSOID_PROPOSALS, + ENSEMBLE_PROPOSALS, + _failures, + _proposal_map, +) __all__ = [ "sample_unif", "sample_rwalk", "sample_slice", "sample_rslice", @@ -209,10 +218,42 @@ def generic_random_walk(u, loglstar, axes, scale, prior_transform, reflective = kwargs.get('reflective') # Setup. - n = len(u) n_cluster = axes.shape[0] walks = kwargs.get('walks', 25) # number of steps + proposal_cycle = kwargs.get("proposals", None) + proposal_kwargs = dict(rstate=rstate) + if proposal_cycle is None: + proposal_cycle = DEFAULT_PROPOSALS + proposal_cycle = list(proposal_cycle) + + for prop in proposal_cycle: + # make sure rolling count of per proposal efficiency is set up + if prop not in _acceptances: + _acceptances[prop] = 0 + _failures[prop] = 0 + + if "live" in kwargs: + # before starting the sampling we remove from the ensemble: + # - duplicate live points + # - the starting point + live = np.unique(deepcopy(kwargs["live"]), axis=0)[:, :n_cluster] + live = np.unique(live, axis=0) + matches = np.where(np.equal(u, live).all(axis=1))[0] + np.delete(live, matches, 0) + else: + live = None + + if not ELLIPSOID_PROPOSALS.isdisjoint(proposal_cycle): + proposal_kwargs["scale"] = scale + proposal_kwargs["axes"] = axes + if not ENSEMBLE_PROPOSALS.isdisjoint(proposal_cycle) and live is not None: + proposal_kwargs["live"] = live + + for prop in proposal_cycle: + if prop not in _proposal_map: + raise ValueError(f"Unknown proposal method {prop}.") + naccept = 0 # Total number of accepted points with L>L* @@ -224,38 +265,56 @@ def generic_random_walk(u, loglstar, axes, scale, prior_transform, # Total number of Likelihood calls (proposals evaluated) # Here we loop for exactly walks iterations. - while ncall < walks: - - # This proposes a new point within the ellipsoid - # This also potentially modifies the scale - u_prop, fail = propose_ball_point(u, - scale, - axes, - n, - n_cluster, - rstate=rstate, - periodic=periodic, - reflective=reflective, - nonbounded=nonbounded) - # If generation of points within an ellipsoid was - # highly inefficient we adjust the scale + # We loop through the specified proposals in order. + for ii in range(walks): + prop = proposal_cycle[ii % len(proposal_cycle)] + u_prop = rstate.uniform(0, 1, len(u)) + proposed = _proposal_map[prop](u=u[:n_cluster], **proposal_kwargs) + if isinstance(proposed, tuple): + u_prop[:n_cluster], ln_jacobian = proposed + else: + u_prop[:n_cluster] = proposed + ln_jacobian = 0 + + if ln_jacobian < np.log(rstate.uniform(0, 1)): + _failures[prop] += 1 + nreject += 1 + continue + + # Only apply boundary conditions if there is no jacobian for the + # proposal. This may be overly conservative, but some proposals + # that require a Jacobian, e.g., the stretch, break detailed + # balance if using a periodic or reflective boundary. + if ln_jacobian == 0: + u_prop, fail = apply_boundary( + u_prop=u_prop, + periodic=periodic, + reflective=reflective, + nonbounded=nonbounded + ) + else: + fail = (np.min(u_prop) < 0) or (np.max(u_prop) > 1) + if fail: + _failures[prop] += 1 nreject += 1 - ncall += 1 continue # Check proposed point. - v_prop = prior_transform(u_prop) - logl_prop = loglikelihood(v_prop) + v_prop = prior_transform(np.array(u_prop)) + logl_prop = loglikelihood(np.array(v_prop)) ncall += 1 if logl_prop > loglstar: u = u_prop v = v_prop logl = logl_prop + _acceptances[prop] += 1 naccept += 1 else: + _failures[prop] += 1 nreject += 1 + if naccept == 0: # Technically we can find out the likelihood value # stored somewhere @@ -268,55 +327,6 @@ def generic_random_walk(u, loglstar, axes, scale, prior_transform, return u, v, logl, ncall, blob -def propose_ball_point(u, - scale, - axes, - n, - n_cluster, - rstate=None, - periodic=None, - reflective=None, - nonbounded=None): - """ - Here we are proposing points uniformly within an n-d ellipsoid. - We are only trying once. - We return the tuple with - 1) proposed point or None - 2) failure flag (if True, the generated point was outside bounds) - """ - - # starting point for clustered dimensions - u_cluster = u[:n_cluster] - - # draw random point for non clustering parameters - # we only need to generate them once - u_non_cluster = rstate.uniform(0, 1, n - n_cluster) - u_prop = np.zeros(n) - u_prop[n_cluster:] = u_non_cluster - - # Propose a direction on the unit n-sphere. - dr = randsphere(n_cluster, rstate=rstate) - # This generates uniform distribution within n-d ball - - # Transform to proposal distribution. - du = np.dot(axes, dr) - u_prop[:n_cluster] = u_cluster + scale * du - - # 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]) - - # Check unit cube constraints. - if unitcheck(u_prop, nonbounded): - return u_prop, False - else: - return None, True - - def generic_slice_step(u, direction, nonperiodic, loglstar, loglikelihood, prior_transform, rstate): """ diff --git a/py/dynesty/utils.py b/py/dynesty/utils.py index f27be7a0e..14ee9159d 100644 --- a/py/dynesty/utils.py +++ b/py/dynesty/utils.py @@ -340,6 +340,22 @@ def unitcheck(u, nonbounded=None): and ub.max() < 1.5) +def apply_boundary(u_prop, periodic, reflective, nonbounded): + """Apply periodic and reflective boundary conditions to the proposed point""" + 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]) + + # Check unit cube constraints. + if unitcheck(u_prop, nonbounded): + return u_prop, False + else: + return u_prop, True + + def apply_reflect(u): """ Iteratively reflect a number until it is contained in [0, 1]. diff --git a/tests/test_proposals.py b/tests/test_proposals.py new file mode 100644 index 000000000..ac58992e0 --- /dev/null +++ b/tests/test_proposals.py @@ -0,0 +1,79 @@ +import itertools + +import numpy as np +import dynesty +from dynesty.proposals import register_proposal, _proposal_map, ELLIPSOID_PROPOSALS, ENSEMBLE_PROPOSALS +import pytest +from scipy.special import erf +from utils import get_rstate, get_printing + +nlive = 100 +printing = get_printing() +win = 100 +ndim = 2 + + +def loglike(x): + return -0.5 * x[1]**2 + + +def prior_transform(x): + return (2 * x - 1) * win + + +def dummy_proposal(x): + return x + 1 + + +@pytest.mark.parametrize( + "proposal,dynamic", + itertools.product( + ["diff", "stretch", "walk", "snooker", "axis", "volumetric", "chi", "normal"], + [True, False] + ) +) +def test_proposals(proposal, dynamic): + # hard test of dynamic sampler with high dlogz_init and small number + # of live points + logz_true = np.log(np.sqrt(2 * np.pi) * erf(win / np.sqrt(2)) / (2 * win)) + sampler = "rwalk" + rstate = get_rstate() + if dynamic: + dns = dynesty.DynamicNestedSampler(loglike, + prior_transform, + ndim, + nlive=nlive, + periodic=[0], + rstate=rstate, + proposals=[proposal], + sample=sampler) + dns.run_nested(dlogz_init=1, print_progress=printing) + else: + dns = dynesty.NestedSampler(loglike, + prior_transform, + ndim, + nlive=nlive, + periodic=[0], + rstate=rstate, + proposals=[proposal], + sample=sampler) + dns.run_nested(dlogz=1, print_progress=printing) + assert (np.abs(logz_true - dns.results.logz[-1]) < + 5. * dns.results.logzerr[-1]) + + +@pytest.mark.parametrize( + "name,kind", [("ellipse", "ellipsoid"), ("ensemble", "ensemble"), ("neither", None)] +) +def test_register_proposal(name, kind): + register_proposal(name, proposal=dummy_proposal, kind=kind) + assert _proposal_map[name](1) == 2 + ellipse = False + ensemble = False + if kind == "ellipsoid": + ellipse = True + assert name in ELLIPSOID_PROPOSALS + elif kind == "ensemble": + ensemble = True + assert ellipse == (name in ELLIPSOID_PROPOSALS) + assert ensemble == (name in ENSEMBLE_PROPOSALS)