adding functions

bilby/core/sampler/dynesty.py

@@ -86,7 +86,7 @@ class Dynesty(NestedSampler):
 
     def __init__(self, likelihood, priors, outdir='outdir', label='label', use_ratio=False, plot=False,
                  skip_import_verification=False, check_point=True, n_check_point=None, check_point_delta_t=600,
-                 resume=True, **kwargs):
+                 resume=True, dynesty_res=False, **kwargs):
         NestedSampler.__init__(self, likelihood=likelihood, priors=priors, outdir=outdir, label=label,
                                use_ratio=use_ratio, plot=plot,
                                skip_import_verification=skip_import_verification,
@@ -94,6 +94,7 @@ class Dynesty(NestedSampler):
         self.n_check_point = n_check_point
         self.check_point = check_point
         self.resume = resume
+        self.dynesty_res = dynesty_res
         if self.n_check_point is None:
             # If the log_likelihood_eval_time is not calculable then
             # check_point is set to False.
@@ -200,8 +201,14 @@ class Dynesty(NestedSampler):
 
         if self.plot:
             self.generate_trace_plots(out)
+            self.generate_evidence_plots(out)
+            print('This code has been edited by Virginia')
 
-        return self.result
+        if self.dynesty_res:
+            self.res = out
+            return self.result, self.res
+        else:
+            return self.result
 
     def _run_external_sampler_without_checkpointing(self):
         logger.debug("Running sampler without checkpointing")
@@ -402,6 +409,15 @@ class Dynesty(NestedSampler):
         fig.tight_layout()
         fig.savefig(filename)
 
+    def generate_evidence_plots(self, dynesty_results):
+        check_directory_exists_and_if_not_mkdir(self.outdir)
+        filename = '{}/{}_evidence.png'.format(self.outdir, self.label)
+        logger.debug("Writing evidence plot to {}".format(filename))
+        from dynesty import plotting as dyplot
+        fig, axes = dyplot.runplot(dynesty_results)
+        fig.tight_layout()
+        fig.savefig(filename)
+
     def _run_test(self):
         import dynesty
         self.sampler = dynesty.NestedSampler(