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

@@ -6,6 +6,7 @@ import sys
 import pickle
 import signal
 
+import emcee
 import tqdm
 import matplotlib.pyplot as plt
 import numpy as np
@@ -89,11 +90,12 @@ class Dynesty(NestedSampler):
                           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,
+                          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=2000, nact=2)
 
     def __init__(self, likelihood, priors, outdir='outdir', label='label',
                  use_ratio=False, plot=False, skip_import_verification=False,
@@ -198,6 +200,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 +234,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_with_act
+            dynesty.nestedsamplers._SAMPLING["rwalk"] = sample_rwalk_bilby_with_act
+            if self.kwargs.get("walks", 25) > self.kwargs.get("maxmcmc"):
+                raise DynestySetupError("You have maxmcmc > walks (minimum mcmc)")
+        elif self.kwargs.get("sample") == "rwalk_dynesty_fixed":
+            logger.info(
+                "Using the bilby-implemented rwalk sample method")
+            self._kwargs["sample"] = "rwalk"
+            dynesty.dynesty._SAMPLING["rwalk"] = sample_rwalk_bilby
+            dynesty.nestedsamplers._SAMPLING["rwalk"] = sample_rwalk_bilby
+        elif self.kwargs.get("sample") == "rwalk_dynesty":
+            self._kwargs["sample"] = "rwalk"
+            logger.info(
+                "Using the dynest-implemented rwalk sample method")
 
         self.sampler = dynesty.NestedSampler(
             loglikelihood=self.log_likelihood,
@@ -626,3 +644,130 @@ def sample_rwalk_bilby(args):
     blob = {'accept': accept, 'reject': reject, 'fail': nfail, 'scale': scale}
 
     return u, v, logl, ncall, blob
+
+
+def sample_rwalk_bilby_with_act(args):
+    """
+    Modified version of dynesty.sampling.sample_rwalk
+
+    """
+
+    # Unzipping.
+    (u, loglstar, axes, scale,
+     prior_transform, loglikelihood, kwargs) = args
+    rstate = np.random
+
+    # Bounds
+    nonbounded = kwargs.get('nonbounded', None)
+    periodic = kwargs.get('periodic', None)
+    reflective = kwargs.get('reflective', None)
+
+    # Setup.
+    n = len(u)
+    maxmcmc = kwargs.get('maxmcmc', 2000)  # Maximum number of steps
+    walks = kwargs.get('walks', 25)  # Minimum number of steps
+    nact = kwargs.get('nact', 2)  # Number of ACT
+
+    accept = 0
+    reject = 0
+    nfail = 0
+    scale_sub = scale
+    act = np.inf
+    u_list = [u]
+    v = prior_transform(u)
+    max_walk_warning = True
+
+    drhat, dr, du, u_prop, logl_prop = np.nan, np.nan, np.nan, np.nan, np.nan
+    while len(u_list) < nact * act or accept == 0 or len(u_list) < walks:
+        while True:
+            # Check scale-factor.
+            if scale_sub == 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_sub * 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:
+                nfail += 1
+                u_list.append(u_list[-1])
+
+        # Check proposed point.
+        v_prop = prior_transform(np.array(u_prop))
+        logl_prop = loglikelihood(np.array(v_prop))
+        if logl_prop >= loglstar:
+            u = u_prop
+            v = v_prop
+            logl = logl_prop
+            accept += 1
+            u_list.append(u_prop)
+        else:
+            reject += 1
+            u_list.append(u_list[-1])
+
+        if accept > walks:
+            act = np.max([2, autocorr_new(np.array(u_list).T)])
+
+        if len(u_list) > maxmcmc and accept > 0:
+            if max_walk_warning:
+                # Print a warning the first time we hit the boundary
+                logger.warning(
+                    "Hit maximum number of walks {}, try increasing maxmcmc"
+                    .format(maxmcmc))
+                max_walk_warning = False
+
+            if accept > 0:
+                # Break if we are above maxmcmc and have at least one accepted point
+                break
+
+    blob = {'accept': accept, 'reject': reject, 'fail': nfail, 'scale': scale,
+            'u_list': np.array(u_list), 'act': act}
+
+    ncall = accept + reject
+    return u, v, logl, ncall, blob
+
+
+def autocorr_new(y, c=10.0):
+    f = np.zeros(y.shape[1])
+    for yy in y:
+        f += emcee.autocorr.function_1d(yy)
+    f /= len(y)
+    taus = 2.0 * np.cumsum(f) - 1.0
+    window = emcee.autocorr.auto_window(taus, c)
+    act = taus[window]
+    if np.isnan(act):
+        return np.inf
+    return act
+
+
+class DynestySetupError(Exception):
+    pass