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
@@ -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,21 @@ 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
+        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 +641,125 @@ 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)
+    walks = kwargs.get('walks', 500)  # maximum number of steps
+    nact = kwargs.get('nact', 10)
+    min_act_check = kwargs.get('min_act_check', 10)
+    accept = 0
+    reject = 0
+    nfail = 0
+    ncall_sub = 0
+    scale_sub = scale
+
+    drhat, dr, du, u_prop, logl_prop = np.nan, np.nan, np.nan, np.nan, np.nan
+    act = np.inf
+    u_list = [u]
+    while len(u_list) < nact * act or accept == 0:
+        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])
+        ncall_sub += 1
+
+        if accept > min_act_check:
+            act = np.max([1, autocorr_new(np.array(u_list).T)])
+            #print(act, nact * act, len(u_list))
+        if ncall_sub > walks:
+            logger.warning("Hit maximum number of walks")
+            break
+
+        # Check if we're stuck generating bad points.
+        if accept < 0.1 * len(u_list):
+            adjust = math.exp(-1. / n)
+            logger.info(
+                "Adjusting sub-chain random walk proposals from {} to {}"
+                .format(scale_sub, adjust * scale_sub))
+            scale_sub *= adjust
+
+    blob = {'accept': accept, 'reject': reject, 'fail': nfail, 'scale': scale,
+            'u_list': np.array(u_list), 'act': act}
+
+    return u, v, logl, ncall_sub, 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