result.py

import logging
import os
import numpy as np
import deepdish
import pandas as pd
import tupak

try:
    from chainconsumer import ChainConsumer
except ImportError:
    def ChainConsumer():
        raise ImportError(
            "You do not have the optional module chainconsumer installed")


def result_file_name(outdir, label):
    """ Returns the standard filename used for a result file """
    return '{}/{}_result.h5'.format(outdir, label)


def read_in_result(outdir=None, label=None, filename=None):
    """ Read in a saved .h5 data file

    Parameters
    ----------
    outdir, label: str
        If given, use the default naming convention for saved results file
    filename: str
        If given, try to load from this filename

    Returns:
    result: tupak.result.Result instance

    """
    if filename is None:
        filename = result_file_name(outdir, label)
    if os.path.isfile(filename):
        return Result(deepdish.io.load(filename))
    else:
        raise ValueError("No information given to load file")


class Result(dict):
    def __init__(self, dictionary=None):
        if type(dictionary) is dict:
            for key in dictionary:
                setattr(self, key, dictionary[key])

    def __getattr__(self, name):
        try:
            return self[name]
        except KeyError:
            raise AttributeError(name)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

    def __repr__(self):
        """Print a summary """
        if hasattr(self, 'samples'):
            return ("nsamples: {:d}\n"
                    "noise_logz: {:6.3f}\n"
                    "logz: {:6.3f} +/- {:6.3f}\n"
                    "log_bayes_factor: {:6.3f} +/- {:6.3f}\n"
                    .format(len(self.samples), self.noise_logz, self.logz,
                            self.logzerr, self.log_bayes_factor, self.logzerr))
        else:
            return ''

    def get_result_dictionary(self):
        return dict(self)

    def save_to_file(self, outdir, label):
        """ Writes the Result to a deepdish h5 file """
        file_name = result_file_name(outdir, label)
        if os.path.isdir(outdir) is False:
            os.makedirs(outdir)
        if os.path.isfile(file_name):
            logging.info(
                'Renaming existing file {} to {}.old'.format(file_name,
                                                             file_name))
            os.rename(file_name, file_name + '.old')

        logging.info("Saving result to {}".format(file_name))
        try:
            deepdish.io.save(file_name, self.get_result_dictionary())
        except Exception as e:
            logging.error(
                "\n\n Saving the data has failed with the following message:\n {} \n\n"
                .format(e))

    def get_latex_labels_from_parameter_keys(self, keys):
        return_list = []
        for k in keys:
            if k in self.search_parameter_keys:
                idx = self.search_parameter_keys.index(k)
                return_list.append(self.parameter_labels[idx])
            elif k in self.parameter_labels:
                return_list.append(k)
            else:
                raise ValueError('key {} not a parameter label or latex label'
                                 .format(k))
        return return_list

    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:
            filename = '{}/{}_corner.png'.format(self.outdir, self.label)
            kwargs['filename'] = kwargs.get('filename', filename)
            logging.info('Saving corner plot to {}'.format(kwargs['filename']))
        if getattr(self, 'injection_parameters', None) is not None:
            # If no truth argument given, set these to the injection params
            injection_parameters = [self.injection_parameters[key]
                                    for key in self.search_parameter_keys]
            kwargs['truth'] = kwargs.get('truth', injection_parameters)

        if type(kwargs.get('truth')) == dict:
            old_keys = kwargs['truth'].keys()
            new_keys = self.get_latex_labels_from_parameter_keys(old_keys)
            for old, new in zip(old_keys, new_keys):
                kwargs['truth'][new] = kwargs['truth'].pop(old)
        if 'parameters' in kwargs:
            kwargs['parameters'] = self.get_latex_labels_from_parameter_keys(
                kwargs['parameters'])

        # Check all parameter_labels are a valid string
        for i, label in enumerate(self.parameter_labels):
            if label is None:
                self.parameter_labels[i] = 'Unknown'
        c = ChainConsumer()
        c.add_chain(self.samples, parameters=self.parameter_labels,
                    name=self.label)
        fig = c.plotter.plot(**kwargs)
        return fig

    def plot_walks(self, save=True, **kwargs):
        """ Plot the chain walks 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)
        if save:
            kwargs['filename'] = '{}/{}_walks.png'.format(self.outdir, self.label)
            logging.info('Saving walker plot to {}'.format(kwargs['filename']))
        if self.injection_parameters is not None:
            kwargs['truth'] = [self.injection_parameters[key] for key in self.search_parameter_keys]
        c = ChainConsumer()
        c.add_chain(self.samples, parameters=self.parameter_labels)
        fig = c.plotter.plot_walks(**kwargs)
        return fig

    def plot_distributions(self, save=True, **kwargs):
        """ Plot the chain walks 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)
        if save:
            kwargs['filename'] = '{}/{}_distributions.png'.format(self.outdir, self.label)
            logging.info('Saving distributions plot to {}'.format(kwargs['filename']))
        if self.injection_parameters is not None:
            kwargs['truth'] = [self.injection_parameters[key] for key in self.search_parameter_keys]
        c = ChainConsumer()
        c.add_chain(self.samples, parameters=self.parameter_labels)
        fig = c.plotter.plot_distributions(**kwargs)
        return fig

    def write_prior_to_file(self, outdir):
        """
        Write the prior distribution to file.

        :return:
        """
        outfile = outdir + '.prior'
        with open(outfile, "w") as prior_file:
            for key in self.prior:
                prior_file.write(self.prior[key])

    def samples_to_data_frame(self, likelihood=None, priors=None, conversion_function=None):
        """
        Convert array of samples to data frame.

        Parameters
        ----------
        likelihood: tupak.likelihood.GravitationalWaveTransient
            GravitationalWaveTransient used for sampling.
        priors: dict
            Dictionary of prior object, used to fill in delta function priors.
        conversion_function: function
            Function which adds in extra parameters to the data frame,
            should take the data_frame, likelihood and prior as arguments.
        """
        data_frame = pd.DataFrame(self.samples, columns=self.search_parameter_keys)
        if conversion_function is not None:
            conversion_function(data_frame, likelihood, priors)
        self.posterior = data_frame

    def construct_cbc_derived_parameters(self):
        """
        Construct widely used derived parameters of CBCs

        :return:
        """
        self.posterior['mass_chirp'] = (self.posterior.mass_1 * self.posterior.mass_2) ** 0.6 / (
                self.posterior.mass_1 + self.posterior.mass_2) ** 0.2
        self.posterior['q'] = self.posterior.mass_2 / self.posterior.mass_1
        self.posterior['eta'] = (self.posterior.mass_1 * self.posterior.mass_2) / (
                self.posterior.mass_1 + self.posterior.mass_2) ** 2

        self.posterior['chi_eff'] = (self.posterior.a_1 * np.cos(self.posterior.tilt_1)
                                     + self.posterior.q * self.posterior.a_2 * np.cos(self.posterior.tilt_2)) / (
                                                1 + self.posterior.q)
        self.posterior['chi_p'] = max(self.posterior.a_1 * np.sin(self.posterior.tilt_1),
                                      (4 * self.posterior.q + 3) / (3 * self.posterior.q + 4) * self.posterior.q
                                      * self.posterior.a_2 * np.sin(self.posterior.tilt_2))

    def check_attribute_match_to_other_object(self, name, other_object):
        """ Check attribute name exists in other_object and is the same """
        A = getattr(self, name, False)
        B = getattr(other_object, name, False)
        logging.debug('Checking {} value: {}=={}'.format(name, A, B))
        if (A is not False) and (B is not False):
            typeA = type(A)
            typeB = type(B)
            if typeA == typeB:
                if typeA in [str, float, int, dict, list]:
                    return A == B
                elif typeA in [np.ndarray]:
                    return np.all(A == B)
        return False