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

@@ -6,6 +6,7 @@ from collections import OrderedDict, namedtuple
 from copy import copy
 from distutils.version import LooseVersion
 from itertools import product
+from tqdm import tqdm
 
 import corner
 import json
@@ -120,9 +121,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)
@@ -151,7 +152,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):
@@ -175,7 +176,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, conversion_function=None, npool=1):
     """ Reweight a result to a new likelihood/prior using rejection sampling
 
     Parameters
@@ -206,22 +207,39 @@ 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)
+    # Saving important quantities to the reweighting procedure
+    result.weights = weights
+    result.efficiency = (np.sum(weights))**2 / np.sum(weights**2) / nposterior
+    result.old_log_likelihood_array = old_log_likelihood_array
+    result.new_log_likelihood_array = new_log_likelihood_array
+    result.old_log_prior_array = old_log_prior_array
+    result.new_log_prior_array = new_log_prior_array
 
     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 conversion_function is not None:
+        data_frame = result.posterior
+        if "npool" in inspect.getargspec(conversion_function).args:
+            data_frame = conversion_function(data_frame, new_likelihood, new_prior, npool=npool)
+        else:
+            data_frame = conversion_function(data_frame, new_likelihood, new_prior)
+        result.posterior = data_frame
 
     if label:
         result.label = label
@@ -244,7 +262,10 @@ class Result(object):
                  num_likelihood_evaluations=None, walkers=None,
                  max_autocorrelation_time=None, use_ratio=None,
                  parameter_labels=None, parameter_labels_with_unit=None,
-                 gzip=False, version=None):
+                 gzip=False, version=None,
+                 efficiency=None, weights=None, old_log_likelihood_array=None,
+                 new_log_likelihood_array=None, old_log_prior_array=None,
+                 new_log_prior_array=None):
         """ A class to store the results of the sampling run
 
         Parameters
@@ -330,6 +351,13 @@ class Result(object):
         self.version = version
         self.max_autocorrelation_time = max_autocorrelation_time
 
+        self.weights = weights
+        self.efficiency = efficiency
+        self.old_log_likelihood_array = old_log_likelihood_array
+        self.new_log_likelihood_array = new_log_likelihood_array
+        self.old_log_prior_array = old_log_prior_array
+        self.new_log_prior_array = new_log_prior_array
+
         self.prior_values = None
         self._kde = None
 
@@ -577,7 +605,9 @@ class Result(object):
             'log_likelihood_evaluations', 'log_prior_evaluations',
             'num_likelihood_evaluations', 'samples', 'nested_samples',
             'walkers', 'nburn', 'parameter_labels', 'parameter_labels_with_unit',
-            'version']
+            'version', 'efficiency', 'weights', 'old_log_likelihood_array',
+            'new_log_likelihood_array', 'old_log_prior_array',
+            'new_log_prior_array']
         dictionary = OrderedDict()
         for attr in save_attrs:
             try:
@@ -585,6 +615,9 @@ class Result(object):
             except ValueError as e:
                 logger.debug("Unable to save {}, message: {}".format(attr, e))
                 pass
+            except AttributeError as e:
+                logger.debug("Unable to save {}, message: {}".format(attr, e))
+                pass
         return dictionary
 
     def save_to_file(self, filename=None, overwrite=False, outdir=None,
@@ -1311,7 +1344,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
 
@@ -1331,12 +1364,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
 
@@ -1355,9 +1388,15 @@ class Result(object):
                              "Cannot copmute credible levels.")
         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)
+            if weights is None:
+                credible_level =\
+                    sum(self.posterior[parameter].values <
+                        self.injection_parameters[parameter]) / len(self.posterior)
+            else:
+                credible_level =\
+                    sum(np.array(self.posterior[parameter].values <
+                        self.injection_parameters[parameter]) * weights) / (sum(weights))
+                print(credible_level)
             return credible_level
         else:
             return np.nan
@@ -1785,7 +1824,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, weights=None,
                  **kwargs):
     """
     Make a P-P plot for a set of runs with injected signals.
@@ -1823,9 +1862,15 @@ def make_pp_plot(results, filename=None, save=True, confidence_interval=[0.68, 0
         keys = results[0].search_parameter_keys
 
     credible_levels = pd.DataFrame()
-    for result in results:
-        credible_levels = credible_levels.append(
-            result.get_all_injection_credible_levels(keys), ignore_index=True)
+    for i, result in enumerate(results):
+        if weights is not None:
+            credible_levels = credible_levels.append(
+                result.get_all_injection_credible_levels(keys, weights=weights[i]),
+                ignore_index=True)
+        else:
+            credible_levels = credible_levels.append(
+                result.get_all_injection_credible_levels(keys),
+                ignore_index=True)
 
     if lines is None:
         colors = ["C{}".format(i) for i in range(8)]