Merge branch 'add_some_post_processing_tools' into 'master'

Add some post processing tools

See merge request lscsoft/bilby!258

CHANGELOG.md

@@ -11,6 +11,8 @@ Changes currently on master, but not under a tag.
 - Renamed `BBHPriorSet` to `BBHPriorDict`
 - Renamed `BNSPriorSet` to `BNSPriorDict`
 - Renamed `CalibrationPriorSet` to `CalibrationPriorDict`
+- Added method to result to get injection recovery credible levels
+- Added function to generate a pp-plot from many results to core/result.py
 - Fixed a bug which caused `Interferometer.detector_tensor` not to update when `latitude`, `longitude`, `xarm_azimuth`, `yarm_azimuth`, `xarm_tilt`, `yarm_tilt` were updated.
 
 ### Changes

bilby/core/result.py

+from __future__ import division
+
 import os
 from distutils.version import LooseVersion
 import numpy as np
@@ -9,8 +11,9 @@ import matplotlib.pyplot as plt
 from collections import OrderedDict, namedtuple
 
 from . import utils
-from .utils import logger, infer_parameters_from_function
-from .prior import PriorDict, DeltaFunction
+from .utils import (logger, infer_parameters_from_function,
+                    check_directory_exists_and_if_not_mkdir)
+from .prior import Prior, PriorDict, DeltaFunction
 
 
 def result_file_name(outdir, label):
@@ -136,6 +139,7 @@ class Result(object):
         self.log_bayes_factor = log_bayes_factor
         self.log_likelihood_evaluations = log_likelihood_evaluations
         self.sampling_time = sampling_time
+        self.max_autocorrelation_time = max_autocorrelation_time
 
     def __str__(self):
         """Print a summary """
@@ -327,7 +331,7 @@ class Result(object):
             elif k in self.parameter_labels:
                 latex_labels.append(k)
             else:
-                logger.info(
+                logger.debug(
                     'key {} not a parameter label or latex label'.format(k))
                 latex_labels.append(' '.join(k.split('_')))
         return latex_labels
@@ -399,6 +403,80 @@ class Result(object):
             fmt(summary.median), fmt(summary.minus), fmt(summary.plus))
         return summary
 
+    def plot_single_density(self, key, prior=None, cumulative=False,
+                            title=None, truth=None, save=True,
+                            file_base_name=None, bins=50, label_fontsize=16,
+                            title_fontsize=16, quantiles=[0.16, 0.84], dpi=300):
+        """ Plot a 1D marginal density, either probablility or cumulative.
+
+        Parameters
+        ----------
+        key: str
+            Name of the parameter to plot
+        prior: {bool (True), bilby.core.prior.Prior}
+            If true, add the stored prior probability density function to the
+            one-dimensional marginal distributions. If instead a Prior
+            is provided, this will be plotted.
+        cumulative: bool
+            If true plot the CDF
+        title: bool
+            If true, add 1D title of the median and (by default 1-sigma)
+            error bars. To change the error bars, pass in the quantiles kwarg.
+            See method `get_one_dimensional_median_and_error_bar` for further
+            details). If `quantiles=None` is passed in, no title is added.
+        truth: {bool, float}
+            If true, plot self.injection_parameters[parameter].
+            If float, plot this value.
+        save: bool:
+            If true, save plot to disk.
+        file_base_name: str, optional
+            If given, the base file name to use (by default `outdir/label_` is
+            used)
+        bins: int
+            The number of histogram bins
+        label_fontsize, title_fontsize: int
+            The fontsizes for the labels and titles
+        quantiles: list
+            A length-2 list of the lower and upper-quantiles to calculate
+            the errors bars for.
+        dpi: int
+            Dots per inch resolution of the plot
+
+        Returns
+        -------
+        figure: matplotlib.pyplot.figure
+            A matplotlib figure object
+        """
+        logger.info('Plotting {} marginal distribution'.format(key))
+        label = self.get_latex_labels_from_parameter_keys([key])[0]
+        fig, ax = plt.subplots()
+        ax.hist(self.posterior[key].values, bins=bins, density=True,
+                histtype='step', cumulative=cumulative)
+        ax.set_xlabel(label, fontsize=label_fontsize)
+        if truth is not None:
+            ax.axvline(truth, ls='-', color='orange')
+
+        summary = self.get_one_dimensional_median_and_error_bar(
+            key, quantiles=quantiles)
+        ax.axvline(summary.median - summary.minus, ls='--', color='C0')
+        ax.axvline(summary.median + summary.plus, ls='--', color='C0')
+        if title:
+            ax.set_title(summary.string, fontsize=title_fontsize)
+
+        if isinstance(prior, Prior):
+            theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300)
+            ax.plot(theta, Prior.prob(theta), color='C2')
+
+        if save:
+            fig.tight_layout()
+            if cumulative:
+                file_name = file_base_name + key + '_cdf'
+            else:
+                file_name = file_base_name + key + '_pdf'
+            fig.savefig(file_name, dpi=dpi)
+
+        return fig
+
     def plot_marginals(self, parameters=None, priors=None, titles=True,
                        file_base_name=None, bins=50, label_fontsize=16,
                        title_fontsize=16, quantiles=[0.16, 0.84], dpi=300):
@@ -433,57 +511,39 @@ class Result(object):
 
         Returns
         -------
-        figures: dictionary
-            A dictionary of the matplotlib figures
-
         """
-
         if isinstance(parameters, dict):
             plot_parameter_keys = list(parameters.keys())
             truths = list(parameters.values())
         elif parameters is None:
-            plot_parameter_keys = self.search_parameter_keys
-            truths = None
+            plot_parameter_keys = self.injection_parameters.keys()
+            truths = [self.injection_parameters.get(key, None) for key
+                      in plot_parameter_keys]
         else:
             plot_parameter_keys = list(parameters)
-            truths = None
+            truths = [self.injection_parameters.get(key, None) for key
+                      in plot_parameter_keys]
 
-        labels = self.get_latex_labels_from_parameter_keys(plot_parameter_keys)
         if file_base_name is None:
-            file_base_name = '{}/{}_'.format(self.outdir, self.label)
+            file_base_name = '{}/{}_1d/'.format(self.outdir, self.label)
+            check_directory_exists_and_if_not_mkdir(file_base_name)
 
         if priors is True:
             priors = getattr(self, 'priors', False)
-        elif isinstance(priors, (dict)) or priors in [False, None]:
+        elif isinstance(priors, dict) or priors in [False, None]:
             pass
         else:
             raise ValueError('Input priors={} not understood'.format(priors))
 
-        figures = dict()
         for i, key in enumerate(plot_parameter_keys):
-            fig, ax = plt.subplots()
-            ax.hist(self.posterior[key].values, bins=bins, density=True,
-                    histtype='step')
-            ax.set_xlabel(labels[i], fontsize=label_fontsize)
-            if truths is not None:
-                ax.axvline(truths[i], ls='--', color='orange')
-
-            summary = self.get_one_dimensional_median_and_error_bar(
-                key, quantiles=quantiles)
-            ax.axvline(summary.median - summary.minus, ls='--', color='C0')
-            ax.axvline(summary.median + summary.plus, ls='--', color='C0')
-            if titles:
-                ax.set_title(summary.string, fontsize=title_fontsize)
-
-            if isinstance(priors, dict):
-                theta = np.linspace(ax.get_xlim()[0], ax.get_xlim()[1], 300)
-                ax.plot(theta, priors[key].prob(theta), color='C2')
-
-            fig.tight_layout()
-            fig.savefig(file_base_name + key)
-            figures[key] = fig
-
-        return figures
+            if not isinstance(self.posterior[key].values[0], float):
+                continue
+            for cumulative in [False, True]:
+                self.plot_single_density(
+                    key, prior=priors[i], cumulative=cumulative, title=titles,
+                    truth=truths[i], save=True, file_base_name=file_base_name,
+                    bins=bins, label_fontsize=label_fontsize, dpi=dpi,
+                    title_fontsize=title_fontsize, quantiles=quantiles)
 
     def plot_corner(self, parameters=None, priors=None, titles=True, save=True,
                     filename=None, dpi=300, **kwargs):
@@ -702,7 +762,7 @@ class Result(object):
         try:
             if all(~np.isnan(self.posterior.log_likelihood)):
                 logger.info('Plotting maximum likelihood')
-                s = model_posterior.ix[self.posterior.log_likelihood.idxmax()]
+                s = model_posterior.iloc[self.posterior.log_likelihood.idxmax()]
                 ax.plot(xsmooth, model(xsmooth, **s), lw=1, color='k',
                         label=maxl_label)
         except AttributeError:
@@ -777,6 +837,49 @@ class Result(object):
                     self.prior_values[key]\
                         = priors[key].prob(self.posterior[key].values)
 
+    def get_all_injection_credible_levels(self):
+        """
+        Get credible levels for all parameters in self.injection_parameters
+
+        Returns
+        -------
+        credible_levels: dict
+            The credible levels at which the injected parameters are found.
+        """
+        if self.injection_parameters is None:
+            raise(TypeError, "Result object has no 'injection_parameters'. "
+                             "Cannot copmute credible levels.")
+        credible_levels = {key: self.get_injection_credible_level(key)
+                           for key in self.search_parameter_keys
+                           if isinstance(self.injection_parameters[key], float)}
+        return credible_levels
+
+    def get_injection_credible_level(self, parameter):
+        """
+        Get the credible level of the injected parameter
+
+        Calculated as CDF(injection value)
+
+        Parameters
+        ----------
+        parameter: str
+            Parameter to get credible level for
+        Returns
+        -------
+        float: credible level
+        """
+        if self.injection_parameters is None:
+            raise(TypeError, "Result object has no 'injection_parameters'. "
+                             "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)
+            return credible_level
+        else:
+            return np.nan
+
     def _check_attribute_match_to_other_object(self, name, other_object):
         """ Check attribute name exists in other_object and is the same
 
@@ -890,3 +993,45 @@ def plot_multiple(results, filename=None, labels=None, colours=None,
     if save:
         fig.savefig(filename)
     return fig
+
+
+def make_pp_plot(results, filename=None, save=True, **kwargs):
+    """
+    Make a P-P plot for a set of runs with injected signals.
+
+    Parameters
+    ----------
+    results: list
+        A list of Result objects, each of these should have injected_parameters
+    filename: str, optional
+        The name of the file to save, the default is "outdir/pp.png"
+    save: bool, optional
+        Whether to save the file, default=True
+    kwargs:
+        Additional kwargs to pass to matplotlib.pyplot.plot
+
+    Returns
+    -------
+    fig:
+        Matplotlib figure
+    """
+    fig = plt.figure()
+    credible_levels = pd.DataFrame()
+    for result in results:
+        credible_levels = credible_levels.append(
+            result.get_all_injection_credible_levels(), ignore_index=True)
+    n_parameters = len(credible_levels.keys())
+    x_values = np.linspace(0, 1, 101)
+    for key in credible_levels:
+        plt.plot(x_values, [sum(credible_levels[key].values < xx) /
+                            len(credible_levels) for xx in x_values],
+                 color='k', alpha=min([1, 4 / n_parameters]), **kwargs)
+    plt.plot([0, 1], [0, 1], linestyle='--', color='r')
+    plt.xlim(0, 1)
+    plt.ylim(0, 1)
+    plt.tight_layout()
+    if save:
+        if filename is None:
+            filename = 'outdir/pp.png'
+        plt.savefig(filename)
+    return fig

examples/other_examples/hyper_parameter_example.py

@@ -8,6 +8,7 @@ import matplotlib.pyplot as plt
 from bilby.core.likelihood import GaussianLikelihood
 from bilby.core.prior import Uniform
 from bilby.core.sampler import run_sampler
+from bilby.core.result import make_pp_plot
 from bilby.hyper.likelihood import HyperparameterLikelihood
 
 outdir = 'outdir'
@@ -31,8 +32,7 @@ fig1, ax1 = plt.subplots()
 fig2, ax2 = plt.subplots()
 
 # Make the sample sets
-samples = list()
-evidences = list()
+results = list()
 for i in range(Nevents):
     c0 = np.random.normal(true_mu_c0, true_sigma_c0)
     c1 = np.random.uniform(-1, 1)
@@ -47,11 +47,10 @@ for i in range(Nevents):
     result = run_sampler(
         likelihood=likelihood, priors=priors, sampler='nestle', nlive=1000,
         outdir=outdir, verbose=False, label='individual_{}'.format(i),
-        save=False)
+        save=False, injection_parameters=injection_parameters)
     ax2.hist(result.posterior.c0, color=line[0].get_color(), normed=True,
              alpha=0.5, label=labels[i])
-    samples.append(result.posterior)
-    evidences.append(result.log_evidence)
+    results.append(result)
 
 ax1.set_xlabel('x')
 ax1.set_ylabel('y(x)')
@@ -72,6 +71,8 @@ def run_prior(data):
     return 1 / 20
 
 
+samples = [result.posterior for result in results]
+evidences = [result.log_evidence for result in results]
 hp_likelihood = HyperparameterLikelihood(
     posteriors=samples, hyper_prior=hyper_prior,
     sampling_prior=run_prior, log_evidences=evidences, max_samples=500)
@@ -85,3 +86,4 @@ result = run_sampler(
     use_ratio=False, outdir=outdir, label='hyper_parameter',
     verbose=True, clean=True)
 result.plot_corner(truth=dict(mu=true_mu_c0, sigma=true_sigma_c0))
+make_pp_plot(results, filename=outdir + '/hyper_parameter_pp.png')

test/result_test.py

@@ -11,8 +11,9 @@ import os
 class TestResult(unittest.TestCase):
 
     def setUp(self):
+        np.random.seed(7)
         bilby.utils.command_line_args.test = False
-        priors = bilby.prior.PriorSet(dict(
+        priors = bilby.prior.PriorDict(dict(
             x=bilby.prior.Uniform(0, 1, 'x', latex_label='$x$', unit='s'),
             y=bilby.prior.Uniform(0, 1, 'y', latex_label='$y$', unit='m'),
             c=1,
@@ -230,5 +231,15 @@ class TestResult(unittest.TestCase):
         with self.assertRaises(ValueError):
             self.result.plot_corner(priors='test')
 
+    def test_get_credible_levels(self):
+        levels = self.result.get_all_injection_credible_levels()
+        self.assertDictEqual(levels, dict(x=0.68, y=0.72))
+
+    def test_get_credible_levels_raises_error_if_no_injection_parameters(self):
+        self.result.injection_parameters = None
+        with self.assertRaises(TypeError):
+            self.result.get_all_injection_credible_levels()
+
+
 if __name__ == '__main__':
     unittest.main()