Modified GRAD_WINDOW_LENGTH calculation for sooner calculation of valid...

This MR includes two fixes and a little extra documentation for the ptemcee module.

• Fixes #638 (closed) by setting all numerical negative infinities (< -10^100) in the mean_log_posterior array to np.nan
• Modified the GRAD_WINDOW_LENGTH over which the gradient of the mean log posterior is calculated to be one of two cases:
  • 2 * ((ndim + 1) // 2) + 1 or
  • 2 * (nwalkers // 2) + 1, the original calculation of this variable in bilby, when ndim <= 3, as otherwise, the gradient calculation fails (because GRAD_WINDOW_LENGTH becomes too small)
• Added docstrings + clear variable names around the calculation of the Gelman-Rubin statistic

bilby/core/sampler/ptemcee.py

@@ -734,7 +734,7 @@ def check_iteration(
     mean_log_posterior,
     verbose=True,
 ):
-    """ Per-iteration logic to calculate the convergence check.
+    """Per-iteration logic to calculate the convergence check.
 
     To check convergence, this function does the following:
     1. Calculate the autocorrelation time (tau) for each dimension for each walker,
@@ -844,8 +844,16 @@ def check_iteration(
     if np.isnan(tau) or np.isinf(tau):
         if verbose:
             print_progress(
-                iteration, sampler, time_per_check, np.nan, np.nan,
-                np.nan, np.nan, np.nan, False, convergence_inputs,
+                iteration,
+                sampler,
+                time_per_check,
+                np.nan,
+                np.nan,
+                np.nan,
+                np.nan,
+                np.nan,
+                False,
+                convergence_inputs,
                 gelman_rubin_statistic,
             )
         return False, np.nan, np.nan, np.nan, np.nan
@@ -861,8 +869,11 @@ def check_iteration(
 
     # Calculate convergence boolean
     converged = gelman_rubin_statistic < ci.Q_tol and ci.nsamples < nsamples_effective
-    logger.debug("Convergence: Q<Q_tol={}, nsamples<nsamples_effective={}"
-                 .format(gelman_rubin_statistic < ci.Q_tol, ci.nsamples < nsamples_effective))
+    logger.debug(
+        "Convergence: Q<Q_tol={}, nsamples<nsamples_effective={}".format(
+            gelman_rubin_statistic < ci.Q_tol, ci.nsamples < nsamples_effective
+        )
+    )
 
     GRAD_WINDOW_LENGTH = 2 * ((ndim + 1) // 2) + 1
 
@@ -938,7 +949,7 @@ def check_iteration(
             gradient_mean_log_posterior,
             tau_usable,
             convergence_inputs,
-            gelman_rubin_statistic
+            gelman_rubin_statistic,
         )
 
     stop = converged and tau_usable
@@ -946,7 +957,7 @@ def check_iteration(
 
 
 def get_max_gradient(x, axis=0, window_length=11, polyorder=2, smooth=False):
-    """ Calculate the maximum value of the gradient in the input data.
+    """Calculate the maximum value of the gradient in the input data.
 
     Applies a Savitzky-Golay filter (`scipy.signal.savgol_filter`) to the input
     data x, along a particular axis. This filter smooths the data and, as configured
@@ -995,7 +1006,7 @@ def get_max_gradient(x, axis=0, window_length=11, polyorder=2, smooth=False):
 
 
 def get_Q_convergence(samples):
-    """ Calculate the Gelman-Rubin statistic as an estimate of convergence for
+    """Calculate the Gelman-Rubin statistic as an estimate of convergence for
     an ensemble of MCMC walkers.
 
     Calculates the Gelman-Rubin statistic, from Gelman and Rubin (1992).
@@ -1045,7 +1056,7 @@ def get_Q_convergence(samples):
         # grand mean (per dimension)
         mean = np.mean(per_walker_mean, axis=0)
         # variance b/t walkers
-        B = nsteps / (nwalkers - 1.) * np.sum((per_walker_mean - mean)**2, axis=0)
+        B = nsteps / (nwalkers - 1.0) * np.sum((per_walker_mean - mean) ** 2, axis=0)
 
         Vhat = (nsteps - 1) / nsteps * W + (nwalkers + 1) / (nwalkers * nsteps) * B
         Q_per_dim = np.sqrt(Vhat / W)
@@ -1123,7 +1134,7 @@ def print_progress(
 
 
 def calculate_tau_array(samples, search_parameter_keys, ci):
-    """ Calculate the autocorrelation time for zero-temperature chains.
+    """Calculate the autocorrelation time for zero-temperature chains.
 
     Calculates the autocorrelation time for each chain, for those parameters/
     dimensions that are not explicitly excluded in ci.ignore_keys_for_tau.
@@ -1311,7 +1322,7 @@ def plot_tau(
 def plot_mean_log_posterior(mean_log_posterior, outdir, label):
     import matplotlib.pyplot as plt
 
-    mean_log_posterior[mean_log_posterior < -1E100] = np.nan
+    mean_log_posterior[mean_log_posterior < -1e100] = np.nan
 
     ntemps, nsteps = mean_log_posterior.shape
     ymax = np.nanmax(mean_log_posterior)