Add a mean-log-likelihood method to improve the ACT estimation

Suggestion from @colm.talbot. I've then added a number of other tweaks in addition.

bilby/core/sampler/ptemcee.py

@@ -25,6 +25,7 @@ ConvergenceInputs = namedtuple(
         "autocorr_tol",
         "autocorr_tau",
         "gradient_tau",
+        "gradient_mean_logl",
         "Q_tol",
         "safety",
         "burn_in_nact",
@@ -153,6 +154,7 @@ class Ptemcee(MCMCSampler):
         safety=1,
         autocorr_tau=1,
         gradient_tau=0.1,
+        gradient_mean_logl=0.1,
         Q_tol=1.02,
         min_tau=1,
         check_point_deltaT=600,
@@ -163,6 +165,7 @@ class Ptemcee(MCMCSampler):
         ignore_keys_for_tau=None,
         pos0="prior",
         niterations_per_check=5,
+        log10beta_min=None,
         **kwargs
     ):
         super(Ptemcee, self).__init__(
@@ -205,6 +208,7 @@ class Ptemcee(MCMCSampler):
             mean_logl_frac=mean_logl_frac,
             thin_by_nact=thin_by_nact,
             gradient_tau=gradient_tau,
+            gradient_mean_logl=gradient_mean_logl,
             Q_tol=Q_tol,
             nsamples=nsamples,
             ignore_keys_for_tau=ignore_keys_for_tau,
@@ -238,6 +242,12 @@ class Ptemcee(MCMCSampler):
             self.priors[key].maximum for key in self.search_parameter_keys
         ]) - self._minima
 
+        self.log10beta_min = log10beta_min
+        if self.log10beta_min is not None:
+            betas = np.logspace(0, self.log10beta_min, self.ntemps)
+            logger.warning("Using betas {}".format(betas))
+            self.kwargs["betas"] = betas
+
     @property
     def sampler_function_kwargs(self):
         """ Kwargs passed to samper.sampler() """
@@ -541,19 +551,21 @@ class Ptemcee(MCMCSampler):
 
         # Get 0-likelihood samples and store in the result
         self.result.samples = self.chain_array[
-            :, self.nburn : self.iteration : self.thin, :
+            :, self.discard + self.nburn : self.iteration : self.thin, :
         ].reshape((-1, self.ndim))
         loglikelihood = self.log_likelihood_array[
-            0, :, self.nburn : self.iteration : self.thin
+            0, :, self.discard + self.nburn : self.iteration : self.thin
         ]  # nwalkers, nsteps
         self.result.log_likelihood_evaluations = loglikelihood.reshape((-1))
 
         if self.store_walkers:
             self.result.walkers = self.sampler.chain
         self.result.nburn = self.nburn
+        self.result.discard = self.discard
 
         log_evidence, log_evidence_err = compute_evidence(
-            sampler, self.log_likelihood_array, self.outdir, self.label, self.nburn,
+            sampler, self.log_likelihood_array, self.outdir,
+            self.label, self.discard, self.nburn,
             self.thin, self.iteration,
         )
         self.result.log_evidence = log_evidence
@@ -586,6 +598,7 @@ class Ptemcee(MCMCSampler):
             self.label,
             self.nsamples_effective,
             self.sampler,
+            self.discard,
             self.nburn,
             self.thin,
             self.search_parameter_keys,
@@ -600,34 +613,43 @@ class Ptemcee(MCMCSampler):
             self.time_per_check,
         )
 
-        if plot and not np.isnan(self.nburn):
-            # Generate the walkers plot diagnostic
-            plot_walkers(
-                self.chain_array[:, : self.iteration, :],
-                self.nburn,
-                self.thin,
-                self.search_parameter_keys,
-                self.outdir,
-                self.label,
-                self.discard,
-            )
+        if plot:
+            try:
+                # Generate the walkers plot diagnostic
+                plot_walkers(
+                    self.chain_array[:, : self.iteration, :],
+                    self.nburn,
+                    self.thin,
+                    self.search_parameter_keys,
+                    self.outdir,
+                    self.label,
+                    self.discard,
+                )
+            except Exception as e:
+                logger.info("Walkers plot failed with exception {}".format(e))
 
-            # Generate the tau plot diagnostic
-            plot_tau(
-                self.tau_list_n,
-                self.tau_list,
-                self.search_parameter_keys,
-                self.outdir,
-                self.label,
-                self.tau_int,
-                self.convergence_inputs.autocorr_tau,
-            )
+            try:
+                # Generate the tau plot diagnostic
+                plot_tau(
+                    self.tau_list_n,
+                    self.tau_list,
+                    self.search_parameter_keys,
+                    self.outdir,
+                    self.label,
+                    self.tau_int,
+                    self.convergence_inputs.autocorr_tau,
+                )
+            except Exception as e:
+                logger.info("tau plot failed with exception {}".format(e))
 
-            plot_mean_log_likelihood(
-                self.mean_log_likelihood,
-                self.outdir,
-                self.label,
-            )
+            try:
+                plot_mean_log_likelihood(
+                    self.mean_log_likelihood,
+                    self.outdir,
+                    self.label,
+                )
+            except Exception as e:
+                logger.info("mean_logl plot failed with exception {}".format(e))
 
 
 def get_mean_logl_min_it(mean_log_likelihood, frac):
@@ -686,7 +708,7 @@ def check_iteration(
     ci = convergence_inputs
     # Note: nsteps is the number of steps in the samples while iterations is
     # the current iteration number. So iteration > nsteps by the number of
-    # od discards
+    # of discards
     nwalkers, nsteps, ndim = samples.shape
 
     tau_array = calculate_tau_array(samples, search_parameter_keys, ci)
@@ -735,12 +757,16 @@ def check_iteration(
             check_taus, axis=0, window_length=11)
 
         if gradient_tau < ci.gradient_tau:
-            logger.debug("tau usable as {} < gradient_tau={}"
-                         .format(gradient_tau, ci.gradient_tau))
+            logger.info(
+                "tau usable as {} < gradient_tau={}"
+                .format(gradient_tau, ci.gradient_tau)
+            )
             tau_usable = True
         else:
-            logger.debug("tau not usable as {} > gradient_tau={}"
-                         .format(gradient_tau, ci.gradient_tau))
+            logger.info(
+                "tau not usable as {} > gradient_tau={}"
+                .format(gradient_tau, ci.gradient_tau)
+            )
             tau_usable = False
 
         check_mean_log_likelihood = mean_log_likelihood[:, -nsteps_to_check:]
@@ -748,25 +774,29 @@ def check_iteration(
             check_mean_log_likelihood, axis=1, window_length=GRAD_WINDOW_LENGTH,
             smooth=True)
 
-        if gradient_mean_logl < ci.gradient_tau:
-            logger.debug("tau usable as {} < gradient_tau={}"
-                         .format(gradient_tau, ci.gradient_tau))
+        if gradient_mean_logl < ci.gradient_mean_logl:
+            logger.info(
+                "tau usable as {} < gradient_mean_logl={}"
+                .format(gradient_mean_logl, ci.gradient_mean_logl)
+            )
             tau_usable = True
         else:
-            logger.debug("tau not usable as {} > gradient_tau={}"
-                         .format(gradient_tau, ci.gradient_tau))
+            logger.info(
+                "tau not usable as {} > gradient_mean_logl={}"
+                .format(gradient_mean_logl, ci.gradient_mean_logl)
+            )
             tau_usable = False
 
     else:
-        logger.debug("ACT is nan")
+        logger.info("ACT is nan")
         gradient_tau = np.nan
         tau_usable = False
 
     if nsteps < tau_int * ci.autocorr_tol:
-        logger.debug("ACT less than autocorr_tol")
+        logger.info("ACT less than autocorr_tol")
         tau_usable = False
     elif tau_int < ci.min_tau:
-        logger.debug("ACT less than min_tau")
+        logger.info("ACT less than min_tau")
         tau_usable = False
 
     # Print an update on the progress
@@ -900,6 +930,7 @@ def checkpoint(
     label,
     nsamples_effective,
     sampler,
+    discard,
     nburn,
     thin,
     search_parameter_keys,
@@ -919,7 +950,7 @@ def checkpoint(
     # Store the samples if possible
     if nsamples_effective > 0:
         filename = "{}/{}_samples.txt".format(outdir, label)
-        samples = np.array(chain_array)[:, nburn : iteration : thin, :].reshape(
+        samples = np.array(chain_array)[:, discard + nburn : iteration : thin, :].reshape(
             (-1, ndim)
         )
         df = pd.DataFrame(samples, columns=search_parameter_keys)
@@ -952,9 +983,13 @@ def plot_walkers(walkers, nburn, thin, parameter_labels, outdir, label,
                  discard=0):
     """ Method to plot the trace of the walkers in an ensemble MCMC plot """
     nwalkers, nsteps, ndim = walkers.shape
+    if np.isnan(nburn):
+        nburn = nsteps
+    if np.isnan(thin):
+        thin = 1
     idxs = np.arange(nsteps)
     fig, axes = plt.subplots(nrows=ndim, ncols=2, figsize=(8, 3 * ndim))
-    scatter_kwargs = dict(lw=0, marker="o", markersize=1, alpha=0.05,)
+    scatter_kwargs = dict(lw=0, marker="o", markersize=1, alpha=0.1,)
 
     # Plot the fixed burn-in
     if discard > 0:
@@ -984,6 +1019,8 @@ def plot_walkers(walkers, nburn, thin, parameter_labels, outdir, label,
             **scatter_kwargs
         )
         axh.hist(walkers[:, discard + nburn::thin, i].reshape((-1)), bins=50, alpha=0.8)
+
+    for i, (ax, axh) in enumerate(axes):
         axh.set_xlabel(parameter_labels[i])
         ax.set_ylabel(parameter_labels[i])
 
@@ -1025,12 +1062,12 @@ def plot_mean_log_likelihood(mean_log_likelihood, outdir, label):
     plt.close(fig)
 
 
-def compute_evidence(sampler, log_likelihood_array, outdir, label, nburn, thin,
+def compute_evidence(sampler, log_likelihood_array, outdir, label, discard, nburn, thin,
                      iteration, make_plots=True):
     """ Computes the evidence using thermodynamic integration """
     betas = sampler.betas
     # We compute the evidence without the burnin samples, but we do not thin
-    lnlike = log_likelihood_array[:, :, nburn : iteration]
+    lnlike = log_likelihood_array[:, :, discard + nburn : iteration]
     mean_lnlikes = np.mean(np.mean(lnlike, axis=1), axis=1)
 
     mean_lnlikes = mean_lnlikes[::-1]