Some changes to the results file for reweighting and uses of reweighted results

I've made some changes to some of the output saved from reweighting as well as updating some functions to calculate confidence intervals using weights.

bilby/core/result.py

@@ -5,6 +5,7 @@ from copy import copy
 from distutils.version import LooseVersion
 from importlib import import_module
 from itertools import product
+from tqdm import tqdm
 
 import corner
 import h5py
@@ -113,10 +114,13 @@ def get_weights_for_reweighting(
     ln_weights: array
         An array of the natural-log weights
     new_log_likelihood_array: array
-        An array of the natural-log likelihoods
+        An array of the natural-log likelihoods from the new likelihood
     new_log_prior_array: array
-        An array of the natural-log priors
-
+        An array of the natural-log priors from the new likelihood
+    old_log_likelihood_array: array
+        An array of the natural-log likelihoods from the old likelihood
+    old_log_prior_array: array
+        An array of the natural-log priors from the old likelihood
     """
     nposterior = len(result.posterior)
     old_log_likelihood_array = np.zeros(nposterior)
@@ -124,9 +128,9 @@ def get_weights_for_reweighting(
     new_log_likelihood_array = np.zeros(nposterior)
     new_log_prior_array = np.zeros(nposterior)
 
-    for ii, sample in result.posterior.iterrows():
+    for ii, sample in tqdm(result.posterior.iterrows()):
         # Convert sample to dictionary
-        par_sample = {key: sample[key] for key in result.search_parameter_keys}
+        par_sample = {key: sample[key] for key in result.posterior}
 
         if old_likelihood is not None:
             old_likelihood.parameters.update(par_sample)
@@ -155,7 +159,7 @@ def get_weights_for_reweighting(
     ln_weights = (
         new_log_likelihood_array + new_log_prior_array - old_log_likelihood_array - old_log_prior_array)
 
-    return ln_weights, new_log_likelihood_array, new_log_prior_array
+    return ln_weights, new_log_likelihood_array, new_log_prior_array, old_log_likelihood_array, old_log_prior_array
 
 
 def rejection_sample(posterior, weights):
@@ -179,7 +183,7 @@ def rejection_sample(posterior, weights):
 
 
 def reweight(result, label=None, new_likelihood=None, new_prior=None,
-             old_likelihood=None, old_prior=None):
+             old_likelihood=None, old_prior=None, verbose_output=False):
     """ Reweight a result to a new likelihood/prior using rejection sampling
 
     Parameters
@@ -198,11 +202,22 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None,
     old_prior: bilby.core.prior.PriorDict, (optional)
         If given, calculate the old prior from this object. If not given,
         the values stored in the posterior are used.
+    verbose_output: bool, optional
+        Flag determining whether the weight array and associated prior and
+        likelihood evaluations are output as well as the result file
 
     Returns
     =======
     result: bilby.core.result.Result
         A copy of the result object with a reweighted posterior
+    new_log_likelihood_array: array, optional (if verbose_output=True)
+        An array of the natural-log likelihoods from the new likelihood
+    new_log_prior_array: array, optional (if verbose_output=True)
+        An array of the natural-log priors from the new likelihood
+    old_log_likelihood_array: array, optional (if verbose_output=True)
+        An array of the natural-log likelihoods from the old likelihood
+    old_log_prior_array: array, optional (if verbose_output=True)
+        An array of the natural-log priors from the old likelihood
 
     """
 
@@ -210,29 +225,38 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None,
     nposterior = len(result.posterior)
     logger.info("Reweighting posterior with {} samples".format(nposterior))
 
-    ln_weights, new_log_likelihood_array, new_log_prior_array = get_weights_for_reweighting(
-        result, new_likelihood=new_likelihood, new_prior=new_prior,
-        old_likelihood=old_likelihood, old_prior=old_prior)
+    ln_weights, new_log_likelihood_array, new_log_prior_array, old_log_likelihood_array, old_log_prior_array =\
+        get_weights_for_reweighting(
+            result, new_likelihood=new_likelihood, new_prior=new_prior,
+            old_likelihood=old_likelihood, old_prior=old_prior)
+
+    weights = np.exp(ln_weights)
 
     # Overwrite the likelihood and prior evaluations
     result.posterior["log_likelihood"] = new_log_likelihood_array
     result.posterior["log_prior"] = new_log_prior_array
 
-    weights = np.exp(ln_weights)
-
     result.posterior = rejection_sample(result.posterior, weights=weights)
+    result.posterior = result.posterior.reset_index(drop=True)
     logger.info("Rejection sampling resulted in {} samples".format(len(result.posterior)))
     result.meta_data["reweighted_using_rejection_sampling"] = True
 
     result.log_evidence += logsumexp(ln_weights) - np.log(nposterior)
-    result.priors = new_prior
+
+    if new_prior is not None:
+        for key, prior in new_prior.items():
+            result.priors[key] = prior
 
     if label:
         result.label = label
     else:
         result.label += "_reweighted"
 
-    return result
+    if verbose_output:
+        return result, weights, new_log_likelihood_array, \
+            new_log_prior_array, old_log_likelihood_array, old_log_prior_array
+    else:
+        return result
 
 
 class Result(object):
@@ -1373,7 +1397,7 @@ class Result(object):
                     self.prior_values[key]\
                         = priors[key].prob(self.posterior[key].values)
 
-    def get_all_injection_credible_levels(self, keys=None):
+    def get_all_injection_credible_levels(self, keys=None, weights=None):
         """
         Get credible levels for all parameters
 
@@ -1382,6 +1406,10 @@ class Result(object):
         keys: list, optional
             A list of keys for which return the credible levels, if None,
             defaults to search_parameter_keys
+        weights: array, optional
+            A list of weights for the posterior samples to calculate a set of
+            weighted credible intervals.
+            If None, assumes equal weights between samples.
 
         Returns
         =======
@@ -1393,12 +1421,12 @@ class Result(object):
         if self.injection_parameters is None:
             raise(TypeError, "Result object has no 'injection_parameters'. "
                              "Cannot compute credible levels.")
-        credible_levels = {key: self.get_injection_credible_level(key)
+        credible_levels = {key: self.get_injection_credible_level(key, weights=weights)
                            for key in keys
                            if isinstance(self.injection_parameters.get(key, None), float)}
         return credible_levels
 
-    def get_injection_credible_level(self, parameter):
+    def get_injection_credible_level(self, parameter, weights=None):
         """
         Get the credible level of the injected parameter
 
@@ -1408,6 +1436,11 @@ class Result(object):
         ==========
         parameter: str
             Parameter to get credible level for
+        weights: array, optional
+            A list of weights for the posterior samples to calculate a
+            weighted credible interval.
+            If None, assumes equal weights between samples.
+
         Returns
         =======
         float: credible level
@@ -1415,11 +1448,15 @@ class Result(object):
         if self.injection_parameters is None:
             raise(TypeError, "Result object has no 'injection_parameters'. "
                              "Cannot copmute credible levels.")
+
+        if weights is None:
+            weights = np.ones(len(self.posterior))
+
         if parameter in self.posterior and\
                 parameter in self.injection_parameters:
             credible_level =\
-                sum(self.posterior[parameter].values <
-                    self.injection_parameters[parameter]) / len(self.posterior)
+                sum(np.array(self.posterior[parameter].values <
+                    self.injection_parameters[parameter]) * weights) / (sum(weights))
             return credible_level
         else:
             return np.nan
@@ -1849,7 +1886,7 @@ def plot_multiple(results, filename=None, labels=None, colours=None,
 @latex_plot_format
 def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0.95, 0.997],
                  lines=None, legend_fontsize='x-small', keys=None, title=True,
-                 confidence_interval_alpha=0.1,
+                 confidence_interval_alpha=0.1, weight_list=None,
                  **kwargs):
     """
     Make a P-P plot for a set of runs with injected signals.
@@ -1873,6 +1910,8 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0
         A list of keys to use, if None defaults to search_parameter_keys
     confidence_interval_alpha: float, list, optional
         The transparency for the background condifence interval
+    weight_list: list, optional
+        List of the weight arrays for each set of posterior samples.
     kwargs:
         Additional kwargs to pass to matplotlib.pyplot.plot
 
@@ -1886,10 +1925,14 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0
     if keys is None:
         keys = results[0].search_parameter_keys
 
+    if weight_list is None:
+        weight_list = [None] * len(results)
+
     credible_levels = pd.DataFrame()
-    for result in results:
+    for i, result in enumerate(results):
         credible_levels = credible_levels.append(
-            result.get_all_injection_credible_levels(keys), ignore_index=True)
+            result.get_all_injection_credible_levels(keys, weights=weight_list[i]),
+            ignore_index=True)
 
     if lines is None:
         colors = ["C{}".format(i) for i in range(8)]