Adds basic corner plotting functionality

- Demonstrated in GW150914.py and BasicTutorial.py

peyote/sampler.py

@@ -6,6 +6,7 @@ import os
 import sys
 
 import numpy as np
+from chainconsumer import ChainConsumer
 
 from .result import Result
 from .prior import Prior
@@ -137,6 +138,30 @@ class Sampler(object):
         logging.info("Using sampler {} with kwargs {}".format(
             self.__class__.__name__, self.kwargs))
 
+    def plot_corner(self, save=True, **kwargs):
+        """ Plot a corner-plot using chain-consumer
+
+        Parameters
+        ----------
+        save: bool
+            If true, save the image using the given label and outdir
+
+        Returns
+        -------
+        fig:
+            A matplotlib figure instance
+        """
+
+        # Set some defaults (unless already set)
+        kwargs['figsize'] = kwargs.get('figsize', 'GROW')
+        if save:
+            kwargs['filename'] = '{}/{}_corner.png'.format(self.outdir, self.label)
+            logging.info('Saving corner plot to {}'.format(kwargs['filename']))
+        c = ChainConsumer()
+        c.add_chain(self.result.samples, parameters=self.result.labels)
+        fig = c.plotter.plot(**kwargs)
+        return fig
+
 
 class Nestle(Sampler):
 
@@ -250,7 +275,7 @@ def run_sampler(likelihood, priors, label='label', outdir='outdir',
         result.log_bayes_factor = result.logz - result.noise_logz
         print("")
         result.save_to_file(outdir=outdir, label=label)
-        return result
+        return result, sampler
     else:
         raise ValueError(
             "Sampler {} not yet implemented".format(sampler))

tutorials/BasicTutorial.py

 import numpy as np
 import pylab as plt
 
-import dynesty.plotting as dyplot
-import corner
 import peyote.prior
 
 peyote.utils.setup_logger()
@@ -84,13 +82,8 @@ sampling_parameters['luminosity_distance'] = peyote.prior.Uniform(lower=30, uppe
 #sampling_parameters["geocent_time"].prior = peyote.prior.Uniform(lower=injection_parameters["geocent_time"] - 0.1,
 #                                                                  upper=injection_parameters["geocent_time"]+0.1)
 
-result = peyote.sampler.run_sampler(likelihood, priors=sampling_parameters, sampler='nestle', verbose=True)
+result, sampler = peyote.sampler.run_sampler(
+    likelihood, priors=sampling_parameters, label='BasicTutorial',
+    sampler='nestle', verbose=True)
+sampler.plot_corner()
 print(result)
-truths = [injection_parameters[x] for x in result.search_parameter_keys]
-
-
-fig = corner.corner(result.samples, truths=truths, labels=result.search_parameter_keys)
-fig.savefig('corner')
-
-fig, axes = dyplot.traceplot(result['sampler_output'])
-fig.savefig('trace')

tutorials/GW150914.py

@@ -21,14 +21,16 @@ maximum_posterior_estimates = dict(
 # Define the prior
 prior = peyote.prior.parse_floats_to_fixed_priors(maximum_posterior_estimates)
 prior['ra'] = peyote.prior.create_default_prior(name='ra')
+prior['psi'] = peyote.prior.create_default_prior(name='psi')
 prior['phase'] = peyote.prior.create_default_prior(name='phase')
 prior['dec'] = peyote.prior.create_default_prior(name='dec')
 prior['iota'] = peyote.prior.create_default_prior(name='iota')
-prior['mass_2'] = peyote.prior.create_default_prior(name='mass_2')
+#prior['mass_1'] = peyote.prior.create_default_prior(name='mass_1')
+#prior['mass_2'] = peyote.prior.create_default_prior(name='mass_2')
 prior['geocent_time'] = peyote.prior.Uniform(
-    time_of_event-2, time_of_event+2, name='geocent_time')
-prior['luminosity_distance'] = peyote.prior.create_default_prior(
-    name='luminosity_distance')
+    time_of_event-1, time_of_event+1, name='geocent_time')
+#prior['luminosity_distance'] = peyote.prior.create_default_prior(
+#    name='luminosity_distance')
 
 # Create the waveformgenerator
 waveformgenerator = peyote.waveform_generator.WaveformGenerator(
@@ -39,12 +41,8 @@ waveformgenerator = peyote.waveform_generator.WaveformGenerator(
 likelihood = peyote.likelihood.Likelihood(IFOs, waveformgenerator)
 
 # Run the sampler
-result = peyote.sampler.run_sampler(
-    likelihood, prior, sampler='pymultinest', n_live_points=300, verbose=True,
-    resume=False, outdir=outdir, use_ratio=True)
-
-truths = [maximum_posterior_estimates[x] for x in result.search_parameter_keys]
-fig = corner.corner(result.samples, labels=result.search_parameter_keys,
-                    truths=truths)
-fig.savefig('{}/corner.png'.format(outdir))
-import IPython; IPython.embed()
+result, sampler = peyote.sampler.run_sampler(
+    likelihood, prior, sampler='pymultinest', n_live_points=400, verbose=True,
+    resume=True, outdir=outdir, use_ratio=True)
+truth = [maximum_posterior_estimates[x] for x in result.search_parameter_keys]
+sampler.plot_corner(truth=truth)