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,10 +6,12 @@ import sys
 import pickle
 import signal
 
+import emcee
 import tqdm
 import matplotlib.pyplot as plt
 import numpy as np
 from pandas import DataFrame
+from scipy import linalg as lalg
 
 from ..utils import logger, check_directory_exists_and_if_not_mkdir, reflect
 from .base_sampler import Sampler, NestedSampler
@@ -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,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 +642,153 @@ 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', 2000)  # maximum number of steps
+    nact = kwargs.get('nact', 10)
+    min_act_check = kwargs.get('min_act_check', 10)
+    accept = 0
+    local_accept = 0
+    reject = 0
+    local_reject = 0
+    nfail = 0
+    local_nfail = 0
+    ncall_sub = 0
+    scale_sub = scale
+    nadjust = 1
+
+    drhat, dr, du, u_prop, logl_prop = np.nan, np.nan, np.nan, np.nan, np.nan
+    act = np.inf
+    u_list = [u]
+
+    last_check = 0
+    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
+                local_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
+            local_accept += 1
+            u_list.append(u_prop)
+        else:
+            reject += 1
+            local_reject += 1
+            u_list.append(u_list[-1])
+        ncall_sub += 1
+
+        if accept > min_act_check:
+            act = np.max([5, autocorr_new(np.array(u_list).T)])
+        if ncall_sub > walks:
+            logger.warning("Hit maximum number of walks")
+            break
+
+        if last_check > 10 * act:
+            local_acceptance_ratio = local_accept / (local_accept + local_nfail + local_reject)
+            nadjust += 0.1
+            if act < np.inf and local_acceptance_ratio < 0.1:
+                adjust = math.exp(-1. / n / nadjust)
+                logger.debug("Local acceptance ratio = {}".format(local_acceptance_ratio))
+                logger.debug(
+                    "Adjusting sub-chain random walk proposals down from {} to {}"
+                    .format(scale_sub, adjust * scale_sub))
+                scale_sub *= adjust
+            elif act < np.inf and local_acceptance_ratio > 0.3:
+                adjust = math.exp(1. / n / nadjust)
+                logger.debug("Local acceptance ratio = {}".format(local_acceptance_ratio))
+                logger.debug(
+                    "Adjusting sub-chain random walk proposals up from {} to {}"
+                    .format(scale_sub, adjust * scale_sub))
+                scale_sub *= adjust
+
+            last_check = 0
+            local_accept = 0
+            local_nfail = 0
+            local_reject = 0
+
+            cov = np.cov(np.array(u_list).T)
+            axes = lalg.cholesky(cov, lower=True)
+
+        last_check += 1
+
+    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