Commit d0842436 authored by Gregory Ashton's avatar Gregory Ashton

Move plot_corner to result

- Makes the plot_corner method an attribute of the result instance
- Removes the return of sampler from run_sampler
- Adds label and outdir to the result
- Fixes the GW150914 and BasicTutorial to run with the new interface
parent 9dcfa247
import logging
import os
import deepdish
from chainconsumer import ChainConsumer
class Result(dict):
......@@ -35,3 +36,29 @@ class Result(dict):
logging.info("Saving result to {}".format(file_name))
deepdish.io.save(file_name, self)
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.samples, parameters=self.parameter_labels)
fig = c.plotter.plot(**kwargs)
return fig
......@@ -4,9 +4,7 @@ import inspect
import logging
import os
import sys
import numpy as np
from chainconsumer import ChainConsumer
from .result import Result
from .prior import Prior
......@@ -68,7 +66,11 @@ class Sampler(object):
if result is None:
self.__result = Result()
self.__result.search_parameter_keys = self.__search_parameter_keys
self.__result.labels = [self.priors[k].latex_label for k in self.__search_parameter_keys]
self.__result.parameter_labels = [
self.priors[k].latex_label for k in
self.__search_parameter_keys]
self.__result.label = self.label
self.__result.outdir = self.outdir
elif type(result) is Result:
self.__result = result
else:
......@@ -152,30 +154,6 @@ 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):
......@@ -302,9 +280,8 @@ def run_sampler(likelihood, priors, label='label', outdir='outdir',
Returns
------
result, sampler
An object containing the results, and the sampler instance (useful
for creating plots etc)
result
An object containing the results
"""
utils.check_directory_exists_and_if_not_mkdir(outdir)
......@@ -320,7 +297,7 @@ def run_sampler(likelihood, priors, label='label', outdir='outdir',
result.log_bayes_factor = result.logz - result.noise_logz
result.injection_parameters = injection_parameters
result.save_to_file(outdir=outdir, label=label)
return result, sampler
return result
else:
raise ValueError(
"Sampler {} not yet implemented".format(sampler))
......
......@@ -77,13 +77,14 @@ fig.savefig('data')
likelihood = peyote.likelihood.Likelihood(IFOs, waveform_generator)
# New way way of doing it, still not perfect
sampling_parameters['mass_1'] = peyote.prior.Uniform(lower=35, upper=37, name='mass1')
sampling_parameters['mass_1'] = peyote.prior.Uniform(
lower=35, upper=37, name='mass_1', latex_label='$m_1$')
sampling_parameters['luminosity_distance'] = peyote.prior.Uniform(lower=30, upper=200, name='luminosity_distance')
#sampling_parameters["geocent_time"].prior = peyote.prior.Uniform(lower=injection_parameters["geocent_time"] - 0.1,
# upper=injection_parameters["geocent_time"]+0.1)
result, sampler = peyote.sampler.run_sampler(
result = peyote.sampler.run_sampler(
likelihood, priors=sampling_parameters, label='BasicTutorial',
sampler='nestle', verbose=True, injection_parameters=injection_parameters)
sampler.plot_corner()
result.plot_corner()
print(result)
......@@ -43,8 +43,8 @@ waveformgenerator = peyote.waveform_generator.WaveformGenerator(
likelihood = peyote.likelihood.Likelihood(IFOs, waveformgenerator)
# Run the sampler
result, sampler = peyote.sampler.run_sampler(
result = peyote.sampler.run_sampler(
likelihood, prior, label='GW150914', sampler='pymultinest',
npoints=1024, resume=False, outdir=outdir, use_ratio=True)
truth = [maximum_posterior_estimates[x] for x in result.search_parameter_keys]
sampler.plot_corner(truth=truth)
result.plot_corner(truth=truth)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment