Change the way idxs are recalculated for dynesty

bilby/core/sampler/base_sampler.py

@@ -399,6 +399,40 @@ class Sampler(object):
 class NestedSampler(Sampler):
     npoints_equiv_kwargs = ['nlive', 'nlives', 'n_live_points', 'npoints', 'npoint', 'Nlive']
 
+    def reorder_loglikelihoods(self, unsorted_loglikelihoods, unsorted_samples,
+                               sorted_samples):
+        """ Reorders the stored log-likelihood after they have been reweighted
+
+        This creates a sorting index by matching the reweights `result.samples`
+        against the raw samples, then uses this index to sort the
+        loglikelihoods
+
+        Parameters
+        ----------
+        sorted_samples, unsorted_samples: array
+            Sorted and unsorted values of the samples. These should be of the same
+            shape and contain the same sample values, but in different orders
+        unsorted_loglikelihoods: array
+            The loglikelihoods corresponding to the unsorted_samples
+
+        Returns
+        -------
+        sorted_loglikelihoods: array
+            The loglikelihoods reordered to match that of the sorted_samples
+
+
+        """
+
+        idxs = []
+        for ii in range(len(unsorted_loglikelihoods)):
+            idx = np.where(np.all(sorted_samples[ii] == unsorted_samples, axis=1))
+            if len(idx) > 1:
+                raise ValueError(
+                    "Multiple matches found between sorted and unsorted samples")
+            else:
+                idxs.append(idx[0])
+        return unsorted_loglikelihoods(idxs)
+
 
 class MCMCSampler(Sampler):
     nwalkers_equiv_kwargs = ['nwalker', 'nwalkers', 'draws']

bilby/core/sampler/dynesty.py

@@ -164,9 +164,9 @@ class Dynesty(NestedSampler):
             out.samples, columns=self.search_parameter_keys)
         self.result.nested_samples['weights'] = weights
         self.result.nested_samples['log_likelihood'] = out.logl
-        idxs = [np.unique(np.where(self.result.samples[ii] == out.samples)[0])
-                for ii in range(len(out.logl))]
-        self.result.log_likelihood_evaluations = out.logl[idxs]
+        self.result.log_likelihood_evaluations = self.reorder_loglikelihoods(
+            unsorted_loglikelihoods=out.logl, unsorted_samples=out.samples,
+            sorted_samples=self.result.samples)
         self.result.log_evidence = out.logz[-1]
         self.result.log_evidence_err = out.logzerr[-1]
 

bilby/core/sampler/nestle.py

@@ -64,9 +64,9 @@ class Nestle(NestedSampler):
             out.samples, columns=self.search_parameter_keys)
         self.result.nested_samples['weights'] = out.weights
         self.result.nested_samples['log_likelihood'] = out.logl
-        idxs = [np.unique(np.where(self.result.samples[ii] == out.samples)[0])
-                for ii in range(len(out.logl))]
-        self.result.log_likelihood_evaluations = out.logl[idxs]
+        self.result.log_likelihood_evaluations = self.reorder_loglikelihoods(
+            unsorted_loglikelihoods=out.logl, unsorted_samples=out.samples,
+            sorted_samples=self.result.samples)
         self.result.log_evidence = out.logz
         self.result.log_evidence_err = out.logzerr
         return self.result