Adding an ACT estimate for the walks

• Improvements to the bilby implemented rwalk method

  • rwalk a new implementation which calculates an ACT to better estimate the number of steps needed

• The rwalk sample method now defaults to bilby.core.sample.dynesty.sample_rwalk_bilby_with_act which adds the maxmcmc, nact. These control the maximum number of walks, the number of ACT required before accepting a point.

• Uses the CPNest method to calculate the ACT

• Adds a guide to the documentation on using dynesty

• Updates the default vol_dec to 8 and enlarge to 1.5 which was found to improve boundary behaviour Note: during development, I experimented with adding a Jacobian and adaptive walks. I found in practise these didn't help sampling so I've removed them for simplicity.

• Updates the walks to default to 100. This is now the minimum number of steps.

bilby/core/sampler/dynesty.py

@@ -83,17 +83,18 @@ class Dynesty(NestedSampler):
     default_kwargs = dict(bound='multi', sample='rwalk',
                           verbose=True, periodic=None, reflective=None,
                           check_point_delta_t=600, nlive=1000,
-                          first_update=None, walks=None,
+                          first_update=None, walks=100,
                           npdim=None, rstate=None, queue_size=None, pool=None,
                           use_pool=None, live_points=None,
                           logl_args=None, logl_kwargs=None,
                           ptform_args=None, ptform_kwargs=None,
-                          enlarge=None, bootstrap=None, vol_dec=0.5, vol_check=2.0,
-                          facc=0.5, slices=5,
+                          enlarge=1.5, bootstrap=None, vol_dec=0.5, vol_check=8.0,
+                          facc=0.2, slices=5,
                           update_interval=None, print_func=None,
                           dlogz=0.1, maxiter=None, maxcall=None,
                           logl_max=np.inf, add_live=True, print_progress=True,
-                          save_bounds=False, n_effective=None)
+                          save_bounds=False, n_effective=None,
+                          maxmcmc=5000, nact=5)
 
     def __init__(self, likelihood, priors, outdir='outdir', label='label',
                  use_ratio=False, plot=False, skip_import_verification=False,
@@ -198,6 +199,7 @@ class Dynesty(NestedSampler):
         # Constructing output.
         string = []
         string.append("bound:{:d}".format(bounditer))
+        string.append("nc:{:d}".format(nc))
         string.append("ncall:{:d}".format(ncall))
         string.append("eff:{:0.1f}%".format(eff))
         string.append("{}={:0.2f}+/-{:0.2f}".format(key, logz, logzerr))
@@ -231,8 +233,23 @@ class Dynesty(NestedSampler):
                 self.get_initial_points_from_prior(
                     self.kwargs['nlive']))
 
-        dynesty.dynesty._SAMPLING["rwalk"] = sample_rwalk_bilby
-        dynesty.nestedsamplers._SAMPLING["rwalk"] = sample_rwalk_bilby
+        if self.kwargs.get("sample", "rwalk") == "rwalk":
+            logger.info(
+                "Using the bilby-implemented rwalk sample method with ACT estimated walks")
+            dynesty.dynesty._SAMPLING["rwalk"] = sample_rwalk_bilby
+            dynesty.nestedsamplers._SAMPLING["rwalk"] = sample_rwalk_bilby
+            if self.kwargs.get("walks", 25) > self.kwargs.get("maxmcmc"):
+                raise DynestySetupError("You have maxmcmc > walks (minimum mcmc)")
+            if self.kwargs.get("nact", 5) < 1:
+                raise DynestySetupError("Unable to run with nact < 1")
+        elif self.kwargs.get("sample") == "rwalk_dynesty":
+            self._kwargs["sample"] = "rwalk"
+            logger.info(
+                "Using the dynesty-implemented rwalk sample method")
+        elif self.kwargs.get("sample") == "rstagger_dynesty":
+            self._kwargs["sample"] = "rstagger"
+            logger.info(
+                "Using the dynesty-implemented rstagger sample method")
 
         self.sampler = dynesty.NestedSampler(
             loglikelihood=self.log_likelihood,
@@ -524,10 +541,7 @@ class Dynesty(NestedSampler):
 
 
 def sample_rwalk_bilby(args):
-    """
-    Modified version of dynesty.sampling.sample_rwalk
-
-    """
+    """ Modified bilby-implemented version of dynesty.sampling.sample_rwalk """
 
     # Unzipping.
     (u, loglstar, axes, scale,
@@ -541,66 +555,76 @@ def sample_rwalk_bilby(args):
 
     # Setup.
     n = len(u)
-    walks = kwargs.get('walks', 25)  # number of steps
+    walks = kwargs.get('walks', 25)  # minimum number of steps
+    maxmcmc = kwargs.get('maxmcmc', 2000)  # Maximum number of steps
+    nact = kwargs.get('nact', 5)  # Number of ACT
+
+    # Initialize internal variables
     accept = 0
     reject = 0
-    fail = 0
     nfail = 0
-    nc = 0
-    ncall = 0
+    act = np.inf
+    u_list = [u]
+    v_list = [prior_transform(u)]
+    logl_list = [loglikelihood(v_list[-1])]
+    max_walk_warning = True
 
     drhat, dr, du, u_prop, logl_prop = np.nan, np.nan, np.nan, np.nan, np.nan
-    while nc + nfail < walks or accept == 0:
-        while True:
-
-            # Check scale-factor.
-            if scale == 0.:
-                raise RuntimeError("The random walk sampling is stuck! "
-                                   "Some useful output quantities:\n"
-                                   "u: {0}\n"
-                                   "drhat: {1}\n"
-                                   "dr: {2}\n"
-                                   "du: {3}\n"
-                                   "u_prop: {4}\n"
-                                   "loglstar: {5}\n"
-                                   "logl_prop: {6}\n"
-                                   "axes: {7}\n"
-                                   "scale: {8}."
-                                   .format(u, drhat, dr, du, u_prop,
-                                           loglstar, logl_prop, axes, scale))
-
-            # Propose a direction on the unit n-sphere.
-            drhat = rstate.randn(n)
-            drhat /= linalg.norm(drhat)
-
-            # Scale based on dimensionality.
-            dr = drhat * rstate.rand()**(1. / n)
-
-            # Transform to proposal distribution.
-            du = np.dot(axes, dr)
-            u_prop = u + 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] = reflect(u_prop[reflective])
-
-            # Check unit cube constraints.
-            if unitcheck(u_prop, nonbounded):
-                break
-            else:
-                fail += 1
-                nfail += 1
-
-            # Check if we're stuck generating bad numbers.
-            if fail > 100 * walks:
-                warnings.warn("Random number generation appears to be "
-                              "extremely inefficient. Adjusting the "
-                              "scale-factor accordingly.")
-                fail = 0
-                scale *= math.exp(-1. / n)
+    while len(u_list) < nact * act:
+
+        if scale == 0.:
+            raise RuntimeError("The random walk sampling is stuck! "
+                               "Some useful output quantities:\n"
+                               "u: {0}\n"
+                               "drhat: {1}\n"
+                               "dr: {2}\n"
+                               "du: {3}\n"
+                               "u_prop: {4}\n"
+                               "loglstar: {5}\n"
+                               "logl_prop: {6}\n"
+                               "axes: {7}\n"
+                               "scale: {8}."
+                               .format(u, drhat, dr, du, u_prop,
+                                       loglstar, logl_prop, axes, scale))
+
+        # Propose a direction on the unit n-sphere.
+        drhat = rstate.randn(n)
+        drhat /= linalg.norm(drhat)
+
+        # Scale based on dimensionality.
+        # dr = drhat * rstate.rand()**(1. / n)  # CHANGED FROM DYNESTY 1.0
+        dr = drhat * rstate.rand(n)
+
+        # Transform to proposal distribution.
+        du = np.dot(axes, dr)
+        u_prop = u + 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] = reflect(u_prop[reflective])
+
+        # Check unit cube constraints.
+        if unitcheck(u_prop, nonbounded):
+            pass
+        else:
+            nfail += 1
+            # Only start appending to the chain once a single jump is made
+            if accept > 0:
+                u_list.append(u_list[-1])
+                v_list.append(v_list[-1])
+                logl_list.append(logl_list[-1])
+            continue
+
+        # Check if we're stuck generating bad numbers.
+        if nfail > 100 * walks:
+            warnings.warn("Random number generation appears to be "
+                          "extremely inefficient. Adjusting the "
+                          "scale-factor accordingly.")
+            nfail = 0
+            scale *= math.exp(-1. / n)
 
         # Check proposed point.
         v_prop = prior_transform(np.array(u_prop))
@@ -610,19 +634,103 @@ def sample_rwalk_bilby(args):
             v = v_prop
             logl = logl_prop
             accept += 1
+            u_list.append(u)
+            v_list.append(v)
+            logl_list.append(logl)
         else:
             reject += 1
-        nc += 1
-        ncall += 1
+            # Only start appending to the chain once a single jump is made
+            if accept > 0:
+                u_list.append(u_list[-1])
+                v_list.append(v_list[-1])
+                logl_list.append(logl_list[-1])
+
+        # If we've taken the minimum number of steps, calculate the ACT
+        if accept + reject > walks:
+            act = estimate_nmcmc(
+                accept / (accept + reject + nfail), walks, maxmcmc)
+
+        # If we've taken too many likelihood evaluations then break
+        if accept + reject > maxmcmc and accept > 0:
+            if max_walk_warning:
+                warnings.warn(
+                    "Hit maximum number of walks {} with accept={}, reject={}, "
+                    "and nfail={} try increasing maxmcmc"
+                    .format(maxmcmc, accept, reject, nfail))
+                max_walk_warning = False
+            if accept > 0:
+                # Break if we are above maxmcmc and have at least one accepted point
+                break
 
         # Check if we're stuck generating bad points.
-        if nc > 50 * walks:
+        if accept + reject > 50 * walks:
             scale *= math.exp(-1. / n)
             warnings.warn("Random walk proposals appear to be "
                           "extremely inefficient. Adjusting the "
                           "scale-factor accordingly.")
-            nc, accept, reject = 0, 0, 0  # reset values
+
+    # If the act is finite, pick randomly from within the chain
+    if np.isfinite(act) and act < len(u_list):
+        idx = np.random.randint(act, len(u_list))
+        u = u_list[idx]
+        v = v_list[idx]
+        logl = logl_list[idx]
+    elif len(u_list) == 1:
+        logger.warning("Returning the only point in the chain")
+        u = u_list[-1]
+        v = v_list[-1]
+        logl = logl_list[-1]
+    else:
+        idx = np.random.randint(int(len(u_list) / 2), len(u_list))
+        logger.warning("Returning random point in second half of the chain")
+        u = u_list[idx]
+        v = v_list[idx]
+        logl = logl_list[idx]
 
     blob = {'accept': accept, 'reject': reject, 'fail': nfail, 'scale': scale}
 
+    ncall = accept + reject
     return u, v, logl, ncall, blob
+
+
+def estimate_nmcmc(accept_ratio, minmcmc, maxmcmc, safety=5, tau=None):
+    """ Estimate autocorrelation length of chain using acceptance fraction
+
+    Using ACL = (2/acc) - 1 multiplied by a safety margin. Code adapated from
+    CPNest:
+        - https://github.com/johnveitch/cpnest/blob/master/cpnest/sampler.py
+        - http://github.com/farr/Ensemble.jl
+
+    Parameters
+    ----------
+    accept_ratio: float [0, 1]
+        Ratio of the number of accepted points to the total number of points
+    minmcmc: int
+        The minimum length of the MCMC chain to use
+    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.
+
+    """
+    if tau is None:
+        tau = maxmcmc / safety
+
+    if accept_ratio == 0.0:
+        Nmcmc_exact = (1. + 1. / tau) * minmcmc
+    else:
+        Nmcmc_exact = (
+            (1. - 1. / tau) * minmcmc +
+            (safety / tau) * (2. / accept_ratio - 1.)
+        )
+
+    Nmcmc_exact = float(min(Nmcmc_exact, maxmcmc))
+    Nmcmc = max(safety, int(Nmcmc_exact))
+
+    return Nmcmc
+
+
+class DynestySetupError(Exception):
+    pass