mcmc_based_searches.py

""" Searches using MCMC-based methods """
from __future__ import division, absolute_import, print_function

import sys
import os
import copy
import logging
from collections import OrderedDict
import subprocess

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import emcee
import corner
import dill as pickle

import pyfstat.core as core
from pyfstat.core import tqdm, args, read_par
import pyfstat.optimal_setup_functions as optimal_setup_functions
import pyfstat.helper_functions as helper_functions


class MCMCSearch(core.BaseSearchClass):
    """MCMC search using ComputeFstat

    Parameters
    ----------
    label, outdir: str
        A label and directory to read/write data from/to
    theta_prior: dict
        Dictionary of priors and fixed values for the search parameters.
        For each parameters (key of the dict), if it is to be held fixed
        the value should be the constant float, if it is be searched, the
        value should be a dictionary of the prior.
    tref, minStartTime, maxStartTime: int
        GPS seconds of the reference time, start time and end time
    sftfilepattern: str
        Pattern to match SFTs using wildcards (*?) and ranges [0-9];
        mutiple patterns can be given separated by colons.
    detectors: str
        Two character reference to the detectors to use, specify None for no
        contraint and comma separate for multiple references.
    nsteps: list (m,)
        List specifying the number of steps to take, the last two entries
        give the nburn and nprod of the 'production' run, all entries
        before are for iterative initialisation steps (usually just one)
        e.g. [1000, 1000, 500].
    nwalkers, ntemps: int,
        The number of walkers and temperates to use in the parallel
        tempered PTSampler.
    log10temperature_min float < 0
        The  log_10(tmin) value, the set of betas passed to PTSampler are
        generated from `np.logspace(0, log10temperature_min, ntemps)`.
    theta_initial: dict, array, (None)
        A dictionary of distribution about which to distribute the
        initial walkers about
    rhohatmax: float,
        Upper bound for the SNR scale parameter (required to normalise the
        Bayes factor) - this needs to be carefully set when using the
        evidence.
    binary: bool
        If true, search over binary parameters
    BSGL: bool
        If true, use the BSGL statistic
    SSBPrec: int
        SSBPrec (SSB precision) to use when calling ComputeFstat
    minCoverFreq, maxCoverFreq: float
        Minimum and maximum instantaneous frequency which will be covered
        over the SFT time span as passed to CreateFstatInput
    injectSources: dict
        If given, inject these properties into the SFT files before running
        the search
    assumeSqrtSX: float
        Don't estimate noise-floors, but assume (stationary) per-IFO sqrt{SX}

    Attributes
    ----------
    symbol_dictionary: dict
        Key, val pairs of the parameters (i.e. `F0`, `F1`), to Latex math
        symbols for plots
    unit_dictionary: dict
        Key, val pairs of the parameters (i.e. `F0`, `F1`), and the
        units (i.e. `Hz`)
    transform_dictionary: dict
        Key, val pairs of the parameters (i.e. `F0`, `F1`), where the key is
        itself a dictionary which can item `multiplier`, `subtractor`, or
        `unit` by which to transform by and update the units.

    """

    symbol_dictionary = dict(
        F0='$f$', F1='$\dot{f}$', F2='$\ddot{f}$', Alpha=r'$\alpha$',
        Delta='$\delta$', asini='asini', period='P', ecc='ecc', tp='tp',
        argp='argp')
    unit_dictionary = dict(
        F0='Hz', F1='Hz/s', F2='Hz/s$^2$', Alpha=r'rad', Delta='rad',
        asini='', period='s', ecc='', tp='', argp='')
    transform_dictionary = {}

    @helper_functions.initializer
    def __init__(self, label, outdir, theta_prior, tref, minStartTime,
                 maxStartTime, sftfilepattern=None, detectors=None,
                 nsteps=[100, 100], nwalkers=100, ntemps=1,
                 log10temperature_min=-5, theta_initial=None,
                 scatter_val=1e-10, rhohatmax=1000, binary=False, BSGL=False,
                 SSBprec=None, minCoverFreq=None, maxCoverFreq=None,
                 injectSources=None, assumeSqrtSX=None):

        if os.path.isdir(outdir) is False:
            os.mkdir(outdir)
        self._add_log_file()
        logging.info('Set-up MCMC search for model {}'.format(self.label))
        if sftfilepattern:
            logging.info('Using data {}'.format(self.sftfilepattern))
        else:
            logging.info('No sftfilepattern given')
        if injectSources:
            logging.info('Inject sources: {}'.format(injectSources))
        self.pickle_path = '{}/{}_saved_data.p'.format(self.outdir, self.label)
        self._unpack_input_theta()
        self.ndim = len(self.theta_keys)
        if self.log10temperature_min:
            self.betas = np.logspace(0, self.log10temperature_min, self.ntemps)
        else:
            self.betas = None

        if args.clean and os.path.isfile(self.pickle_path):
            os.rename(self.pickle_path, self.pickle_path+".old")

        self._set_likelihoodcoef()
        self._log_input()

    def _set_likelihoodcoef(self):
        self.likelihoodcoef = np.log(70./self.rhohatmax**4)

    def _log_input(self):
        logging.info('theta_prior = {}'.format(self.theta_prior))
        logging.info('nwalkers={}'.format(self.nwalkers))
        logging.info('scatter_val = {}'.format(self.scatter_val))
        logging.info('nsteps = {}'.format(self.nsteps))
        logging.info('ntemps = {}'.format(self.ntemps))
        logging.info('log10temperature_min = {}'.format(
            self.log10temperature_min))

    def _initiate_search_object(self):
        logging.info('Setting up search object')
        self.search = core.ComputeFstat(
            tref=self.tref, sftfilepattern=self.sftfilepattern,
            minCoverFreq=self.minCoverFreq, maxCoverFreq=self.maxCoverFreq,
            detectors=self.detectors, BSGL=self.BSGL, transient=False,
            minStartTime=self.minStartTime, maxStartTime=self.maxStartTime,
            binary=self.binary, injectSources=self.injectSources,
            assumeSqrtSX=self.assumeSqrtSX, SSBprec=self.SSBprec)

    def logp(self, theta_vals, theta_prior, theta_keys, search):
        H = [self._generic_lnprior(**theta_prior[key])(p) for p, key in
             zip(theta_vals, theta_keys)]
        return np.sum(H)

    def logl(self, theta, search):
        for j, theta_i in enumerate(self.theta_idxs):
            self.fixed_theta[theta_i] = theta[j]
        FS = search.compute_fullycoherent_det_stat_single_point(
            *self.fixed_theta)
        return FS + self.likelihoodcoef

    def _unpack_input_theta(self):
        full_theta_keys = ['F0', 'F1', 'F2', 'Alpha', 'Delta']
        if self.binary:
            full_theta_keys += [
                'asini', 'period', 'ecc', 'tp', 'argp']
        full_theta_keys_copy = copy.copy(full_theta_keys)

        full_theta_symbols = ['$f$', '$\dot{f}$', '$\ddot{f}$', r'$\alpha$',
                              r'$\delta$']
        if self.binary:
            full_theta_symbols += [
                'asini', 'period', 'ecc', 'tp', 'argp']

        self.theta_keys = []
        fixed_theta_dict = {}
        for key, val in self.theta_prior.iteritems():
            if type(val) is dict:
                fixed_theta_dict[key] = 0
                self.theta_keys.append(key)
            elif type(val) in [float, int, np.float64]:
                fixed_theta_dict[key] = val
            else:
                raise ValueError(
                    'Type {} of {} in theta not recognised'.format(
                        type(val), key))
            full_theta_keys_copy.pop(full_theta_keys_copy.index(key))

        if len(full_theta_keys_copy) > 0:
            raise ValueError(('Input dictionary `theta` is missing the'
                              'following keys: {}').format(
                                  full_theta_keys_copy))

        self.fixed_theta = [fixed_theta_dict[key] for key in full_theta_keys]
        self.theta_idxs = [full_theta_keys.index(k) for k in self.theta_keys]
        self.theta_symbols = [full_theta_symbols[i] for i in self.theta_idxs]

        idxs = np.argsort(self.theta_idxs)
        self.theta_idxs = [self.theta_idxs[i] for i in idxs]
        self.theta_symbols = [self.theta_symbols[i] for i in idxs]
        self.theta_keys = [self.theta_keys[i] for i in idxs]

    def _check_initial_points(self, p0):
        for nt in range(self.ntemps):
            logging.info('Checking temperature {} chains'.format(nt))
            initial_priors = np.array([
                self.logp(p, self.theta_prior, self.theta_keys, self.search)
                for p in p0[nt]])
            number_of_initial_out_of_bounds = sum(initial_priors == -np.inf)

            if number_of_initial_out_of_bounds > 0:
                logging.warning(
                    'Of {} initial values, {} are -np.inf due to the prior'
                    .format(len(initial_priors),
                            number_of_initial_out_of_bounds))

                p0 = self._generate_new_p0_to_fix_initial_points(
                    p0, nt, initial_priors)

    def _generate_new_p0_to_fix_initial_points(self, p0, nt, initial_priors):
        logging.info('Attempting to correct intial values')
        idxs = np.arange(self.nwalkers)[initial_priors == -np.inf]
        count = 0
        while sum(initial_priors == -np.inf) > 0 and count < 100:
            for j in idxs:
                p0[nt][j] = (p0[nt][np.random.randint(0, self.nwalkers)]*(
                             1+np.random.normal(0, 1e-10, self.ndim)))
            initial_priors = np.array([
                self.logp(p, self.theta_prior, self.theta_keys,
                          self.search)
                for p in p0[nt]])
            count += 1

        if sum(initial_priors == -np.inf) > 0:
            logging.info('Failed to fix initial priors')
        else:
            logging.info('Suceeded to fix initial priors')

        return p0

    def setup_burnin_convergence_testing(
            self, n=10, test_type='autocorr', windowed=False, **kwargs):
        """ Set up convergence testing during the MCMC simulation

        Parameters
        ----------
        n: int
            Number of steps after which to test convergence
        test_type: str ['autocorr', 'GR']
            If 'autocorr' use the exponential autocorrelation time (kwargs
            passed to `get_autocorr_convergence`). If 'GR' use the Gelman-Rubin
            statistic (kwargs passed to `get_GR_convergence`)
        windowed: bool
            If True, only calculate the convergence test in a window of length
            `n`
        **kwargs:
            Passed to either `_test_autocorr_convergence()` or
            `_test_GR_convergence()` depending on `test_type`.

        """
        logging.info('Setting up convergence testing')
        self.convergence_n = n
        self.convergence_windowed = windowed
        self.convergence_test_type = test_type
        self.convergence_kwargs = kwargs
        self.convergence_diagnostic = []
        self.convergence_diagnosticx = []
        if test_type in ['autocorr']:
            self._get_convergence_test = self._test_autocorr_convergence
        elif test_type in ['GR']:
            self._get_convergence_test = self._test_GR_convergence
        else:
            raise ValueError('test_type {} not understood'.format(test_type))

    def _test_autocorr_convergence(self, i, sampler, test=True, n_cut=5):
        try:
            acors = np.zeros((self.ntemps, self.ndim))
            for temp in range(self.ntemps):
                if self.convergence_windowed:
                    j = i-self.convergence_n
                else:
                    j = 0
                x = np.mean(sampler.chain[temp, :, j:i, :], axis=0)
                acors[temp, :] = emcee.autocorr.exponential_time(x)
            c = np.max(acors, axis=0)
        except emcee.autocorr.AutocorrError:
            logging.info('Failed to calculate exponential autocorrelation')
            c = np.zeros(self.ndim) + np.nan
        except AttributeError:
            logging.info('Unable to calculate exponential autocorrelation')
            c = np.zeros(self.ndim) + np.nan

        self.convergence_diagnosticx.append(i - self.convergence_n/2.)
        self.convergence_diagnostic.append(list(c))

        if test:
            return i > n_cut * np.max(c)

    def _test_GR_convergence(self, i, sampler, test=True, R=1.1):
        if self.convergence_windowed:
            s = sampler.chain[0, :, i-self.convergence_n+1:i+1, :]
        else:
            s = sampler.chain[0, :, :i+1, :]
        N = float(self.convergence_n)
        M = float(self.nwalkers)
        W = np.mean(np.var(s, axis=1), axis=0)
        per_walker_mean = np.mean(s, axis=1)
        mean = np.mean(per_walker_mean, axis=0)
        B = N / (M-1.) * np.sum((per_walker_mean-mean)**2, axis=0)
        Vhat = (N-1)/N * W + (M+1)/(M*N) * B
        c = np.sqrt(Vhat/W)
        self.convergence_diagnostic.append(c)
        self.convergence_diagnosticx.append(i - self.convergence_n/2.)

        if test and np.max(c) < R:
            return True
        else:
            return False

    def _test_convergence(self, i, sampler, **kwargs):
        if np.mod(i+1, self.convergence_n) == 0:
            return self._get_convergence_test(i, sampler, **kwargs)
        else:
            return False

    def _run_sampler_with_conv_test(self, sampler, p0, nprod=0, nburn=0):
        logging.info('Running {} burn-in steps with convergence testing'
                     .format(nburn))
        iterator = tqdm(sampler.sample(p0, iterations=nburn), total=nburn)
        for i, output in enumerate(iterator):
            if self._test_convergence(i, sampler, test=True,
                                      **self.convergence_kwargs):
                logging.info(
                    'Converged at {} before max number {} of steps reached'
                    .format(i, nburn))
                self.convergence_idx = i
                break
        iterator.close()
        logging.info('Running {} production steps'.format(nprod))
        j = nburn
        iterator = tqdm(sampler.sample(output[0], iterations=nprod),
                        total=nprod)
        for result in iterator:
            self._test_convergence(j, sampler, test=False,
                                   **self.convergence_kwargs)
            j += 1
        return sampler

    def _run_sampler(self, sampler, p0, nprod=0, nburn=0):
        if hasattr(self, 'convergence_n'):
            self._run_sampler_with_conv_test(sampler, p0, nprod, nburn)
        else:
            for result in tqdm(sampler.sample(p0, iterations=nburn+nprod),
                               total=nburn+nprod):
                pass

        self.mean_acceptance_fraction = np.mean(
            sampler.acceptance_fraction, axis=1)
        logging.info("Mean acceptance fraction: {}"
                     .format(self.mean_acceptance_fraction))
        if self.ntemps > 1:
            self.tswap_acceptance_fraction = sampler.tswap_acceptance_fraction
            logging.info("Tswap acceptance fraction: {}"
                         .format(sampler.tswap_acceptance_fraction))
        try:
            self.autocorr_time = sampler.get_autocorr_time(c=4)
            logging.info("Autocorrelation length: {}".format(
                self.autocorr_time))
        except emcee.autocorr.AutocorrError as e:
            self.autocorr_time = np.nan
            logging.warning(
                'Autocorrelation calculation failed with message {}'.format(e))

        return sampler

    def _estimate_run_time(self):
        """ Print the estimated run time

        Uses timing coefficients based on a Lenovo T460p Intel(R)
        Core(TM) i5-6300HQ CPU @ 2.30GHz.

        """
        # Todo: add option to time on a machine, and move coefficients to
        # ~/.pyfstat.conf
        if (type(self.theta_prior['Alpha']) == dict or
                type(self.theta_prior['Delta']) == dict):
            tau0S = 7.3e-5
            tau0LD = 4.2e-7
        else:
            tau0S = 5.0e-5
            tau0LD = 6.2e-8
        Nsfts = (self.maxStartTime - self.minStartTime) / 1800.
        numb_evals = np.sum(self.nsteps)*self.nwalkers*self.ntemps
        a = tau0S * numb_evals
        b = tau0LD * Nsfts * numb_evals
        logging.info('Estimated run-time = {} s = {:1.0f}:{:1.0f} m'.format(
            a+b, *divmod(a+b, 60)))

    def run(self, proposal_scale_factor=2, create_plots=True, c=5, **kwargs):
        """ Run the MCMC simulatation

        Parameters
        ----------
        proposal_scale_factor: float
            The proposal scale factor used by the sampler, see Goodman & Weare
            (2010). If the acceptance fraction is too low, you can raise it by
            decreasing the a parameter; and if it is too high, you can reduce
            it by increasing the a parameter [Foreman-Mackay (2013)].
        create_plots: bool
            If true, save trace plots of the walkers
        c: int
            The minimum number of autocorrelation times needed to trust the
            result when estimating the autocorrelation time (see
            emcee.autocorr.integrated_time for further details. Default is 5
        **kwargs:
            Passed to _plot_walkers to control the figures

        Returns
        -------
        sampler: emcee.ptsampler.PTSampler
            The emcee ptsampler object

        """

        self.old_data_is_okay_to_use = self._check_old_data_is_okay_to_use()
        if self.old_data_is_okay_to_use is True:
            logging.warning('Using saved data from {}'.format(
                self.pickle_path))
            d = self.get_saved_data_dictionary()
            self.samples = d['samples']
            self.lnprobs = d['lnprobs']
            self.lnlikes = d['lnlikes']
            self.all_lnlikelihood = d['all_lnlikelihood']
            return

        self._initiate_search_object()
        self._estimate_run_time()

        sampler = emcee.PTSampler(
            self.ntemps, self.nwalkers, self.ndim, self.logl, self.logp,
            logpargs=(self.theta_prior, self.theta_keys, self.search),
            loglargs=(self.search,), betas=self.betas, a=proposal_scale_factor)

        p0 = self._generate_initial_p0()
        p0 = self._apply_corrections_to_p0(p0)
        self._check_initial_points(p0)

        ninit_steps = len(self.nsteps) - 2
        for j, n in enumerate(self.nsteps[:-2]):
            logging.info('Running {}/{} initialisation with {} steps'.format(
                j, ninit_steps, n))
            sampler = self._run_sampler(sampler, p0, nburn=n)
            if create_plots:
                fig, axes = self._plot_walkers(sampler,
                                               symbols=self.theta_symbols,
                                               **kwargs)
                fig.tight_layout()
                fig.savefig('{}/{}_init_{}_walkers.png'.format(
                    self.outdir, self.label, j))

            p0 = self._get_new_p0(sampler)
            p0 = self._apply_corrections_to_p0(p0)
            self._check_initial_points(p0)
            sampler.reset()

        if len(self.nsteps) > 1:
            nburn = self.nsteps[-2]
        else:
            nburn = 0
        nprod = self.nsteps[-1]
        logging.info('Running final burn and prod with {} steps'.format(
            nburn+nprod))
        sampler = self._run_sampler(sampler, p0, nburn=nburn, nprod=nprod)
        if create_plots:
            fig, axes = self._plot_walkers(sampler, symbols=self.theta_symbols,
                                           nprod=nprod, **kwargs)
            fig.tight_layout()
            fig.savefig('{}/{}_walkers.png'.format(self.outdir, self.label),
                        )

        samples = sampler.chain[0, :, nburn:, :].reshape((-1, self.ndim))
        lnprobs = sampler.lnprobability[0, :, nburn:].reshape((-1))
        lnlikes = sampler.lnlikelihood[0, :, nburn:].reshape((-1))
        all_lnlikelihood = sampler.lnlikelihood[:, :, nburn:]
        self.samples = samples
        self.lnprobs = lnprobs
        self.lnlikes = lnlikes
        self.all_lnlikelihood = all_lnlikelihood
        self._save_data(sampler, samples, lnprobs, lnlikes, all_lnlikelihood)
        return sampler

    def _get_rescale_multiplier_for_key(self, key):
        """ Get the rescale multiplier from the transform_dictionary

        Can either be a float, a string (in which case it is interpretted as
        a attribute of the MCMCSearch class, e.g. minStartTime, or non-existent
        in which case 0 is returned
        """
        if key not in self.transform_dictionary:
            return 1

        if 'multiplier' in self.transform_dictionary[key]:
            val = self.transform_dictionary[key]['multiplier']
            if type(val) == str:
                if hasattr(self, val):
                    multiplier = getattr(
                        self, self.transform_dictionary[key]['multiplier'])
                else:
                    raise ValueError(
                        "multiplier {} not a class attribute".format(val))
            else:
                multiplier = val
        else:
            multiplier = 1
        return multiplier

    def _get_rescale_subtractor_for_key(self, key):
        """ Get the rescale subtractor from the transform_dictionary

        Can either be a float, a string (in which case it is interpretted as
        a attribute of the MCMCSearch class, e.g. minStartTime, or non-existent
        in which case 0 is returned
        """
        if key not in self.transform_dictionary:
            return 0

        if 'subtractor' in self.transform_dictionary[key]:
            val = self.transform_dictionary[key]['subtractor']
            if type(val) == str:
                if hasattr(self, val):
                    subtractor = getattr(
                        self, self.transform_dictionary[key]['subtractor'])
                else:
                    raise ValueError(
                        "subtractor {} not a class attribute".format(val))
            else:
                subtractor = val
        else:
            subtractor = 0
        return subtractor

    def _scale_samples(self, samples, theta_keys):
        """ Scale the samples using the transform_dictionary """
        for key in theta_keys:
            if key in self.transform_dictionary:
                idx = theta_keys.index(key)
                s = samples[:, idx]
                subtractor = self._get_rescale_subtractor_for_key(key)
                s = s - subtractor
                multiplier = self._get_rescale_multiplier_for_key(key)
                s *= multiplier
                samples[:, idx] = s

        return samples

    def _get_labels(self):
        """ Combine the units, symbols and rescaling to give labels """

        labels = []
        for key in self.theta_keys:
            label = None
            s = self.symbol_dictionary[key]
            s.replace('_{glitch}', r'_\textrm{glitch}')
            u = self.unit_dictionary[key]
            if key in self.transform_dictionary:
                if 'symbol' in self.transform_dictionary[key]:
                    s = self.transform_dictionary[key]['symbol']
                if 'label' in self.transform_dictionary[key]:
                    label = self.transform_dictionary[key]['label']
                if 'unit' in self.transform_dictionary[key]:
                    u = self.transform_dictionary[key]['unit']
            if label is None:
                label = '{} \n [{}]'.format(s, u)
            labels.append(label)
        return labels

    def plot_corner(self, figsize=(7, 7), add_prior=False, nstds=None,
                    label_offset=0.4, dpi=300, rc_context={},
                    tglitch_ratio=False, fig_and_axes=None, save_fig=True,
                    **kwargs):
        """ Generate a corner plot of the posterior

        Using the `corner` package (https://pypi.python.org/pypi/corner/),
        generate estimates of the posterior from the production samples.

        Parameters
        ----------
        figsize: tuple (7, 7)
            Figure size in inches (passed to plt.subplots)
        add_prior: bool, str
            If true, plot the prior as a red line. If 'full' then for uniform
            priors plot the full extent of the prior.
        nstds: float
            The number of standard deviations to plot centered on the mean
        label_offset: float
            Offset the labels from the plot: useful to precent overlapping the
            tick labels with the axis labels
        dpi: int
            Passed to plt.savefig
        rc_context: dict
            Dictionary of rc values to set while generating the figure (see
            matplotlib rc for more details)
        tglitch_ratio: bool
            If true, and tglitch is a parameter, plot posteriors as the
            fractional time at which the glitch occurs instead of the actual
            time
        fig_and_axes: tuple
            fig and axes to plot on, the axes must be of the right shape,
            namely (ndim, ndim)
        save_fig: bool
            If true, save the figure, else return the fig, axes
        **kwargs:
            Passed to corner.corner

        Returns
        -------
        fig, axes:
            The matplotlib figure and axes, only returned if save_fig = False

        """

        if 'truths' in kwargs and len(kwargs['truths']) != self.ndim:
            logging.warning('len(Truths) != ndim, Truths will be ignored')
            kwargs['truths'] = None

        if self.ndim < 2:
            with plt.rc_context(rc_context):
                if fig_and_axes is None:
                    fig, ax = plt.subplots(figsize=figsize)
                else:
                    fig, ax = fig_and_axes
                ax.hist(self.samples, bins=50, histtype='stepfilled')
                ax.set_xlabel(self.theta_symbols[0])

            fig.savefig('{}/{}_corner.png'.format(
                self.outdir, self.label), dpi=dpi)
            return

        with plt.rc_context(rc_context):
            if fig_and_axes is None:
                fig, axes = plt.subplots(self.ndim, self.ndim,
                                         figsize=figsize)
            else:
                fig, axes = fig_and_axes

            samples_plt = copy.copy(self.samples)
            labels = self._get_labels()

            samples_plt = self._scale_samples(samples_plt, self.theta_keys)

            if tglitch_ratio:
                for j, k in enumerate(self.theta_keys):
                    if k == 'tglitch':
                        s = samples_plt[:, j]
                        samples_plt[:, j] = (
                            s - self.minStartTime)/(
                                self.maxStartTime - self.minStartTime)
                        labels[j] = r'$R_{\textrm{glitch}}$'

            if type(nstds) is int and 'range' not in kwargs:
                _range = []
                for j, s in enumerate(samples_plt.T):
                    median = np.median(s)
                    std = np.std(s)
                    _range.append((median - nstds*std, median + nstds*std))
            elif 'range' in kwargs:
                _range = kwargs.pop('range')
            else:
                _range = None

            hist_kwargs = kwargs.pop('hist_kwargs', dict())
            if 'normed' not in hist_kwargs:
                hist_kwargs['normed'] = True

            fig_triangle = corner.corner(samples_plt,
                                         labels=labels,
                                         fig=fig,
                                         bins=50,
                                         max_n_ticks=4,
                                         plot_contours=True,
                                         plot_datapoints=True,
                                         label_kwargs={'fontsize': 8},
                                         data_kwargs={'alpha': 0.1,
                                                      'ms': 0.5},
                                         range=_range,
                                         hist_kwargs=hist_kwargs,
                                         **kwargs)

            axes_list = fig_triangle.get_axes()
            axes = np.array(axes_list).reshape(self.ndim, self.ndim)
            plt.draw()
            for ax in axes[:, 0]:
                ax.yaxis.set_label_coords(-label_offset, 0.5)
            for ax in axes[-1, :]:
                ax.xaxis.set_label_coords(0.5, -label_offset)
            for ax in axes_list:
                ax.set_rasterized(True)
                ax.set_rasterization_zorder(-10)
            plt.tight_layout(h_pad=0.0, w_pad=0.0)
            fig.subplots_adjust(hspace=0.05, wspace=0.05)

            if add_prior:
                self._add_prior_to_corner(axes, self.samples, add_prior)

            if save_fig:
                fig_triangle.savefig('{}/{}_corner.png'.format(
                    self.outdir, self.label), dpi=dpi)
            else:
                return fig, axes

    def _add_prior_to_corner(self, axes, samples, add_prior):
        for i, key in enumerate(self.theta_keys):
            ax = axes[i][i]
            s = samples[:, i]
            lnprior = self._generic_lnprior(**self.theta_prior[key])
            if add_prior == 'full' and self.theta_prior[key]['type'] == 'unif':
                lower = self.theta_prior[key]['lower']
                upper = self.theta_prior[key]['upper']
                r = upper-lower
                xlim = [lower-0.05*r, upper+0.05*r]
                x = np.linspace(xlim[0], xlim[1], 1000)
            else:
                xlim = ax.get_xlim()
                x = np.linspace(s.min(), s.max(), 1000)
            multiplier = self._get_rescale_multiplier_for_key(key)
            subtractor = self._get_rescale_subtractor_for_key(key)
            ax.plot((x-subtractor)*multiplier,
                    [np.exp(lnprior(xi)) for xi in x], '-C3',
                    label='prior')

            for j in range(i, self.ndim):
                axes[j][i].set_xlim(xlim[0], xlim[1])
            for k in range(0, i):
                axes[i][k].set_ylim(xlim[0], xlim[1])

    def plot_prior_posterior(self, normal_stds=2):
        """ Plot the posterior in the context of the prior """
        fig, axes = plt.subplots(nrows=self.ndim, figsize=(8, 4*self.ndim))
        N = 1000
        from scipy.stats import gaussian_kde

        for i, (ax, key) in enumerate(zip(axes, self.theta_keys)):
            prior_dict = self.theta_prior[key]
            prior_func = self._generic_lnprior(**prior_dict)
            if prior_dict['type'] == 'unif':
                x = np.linspace(prior_dict['lower'], prior_dict['upper'], N)
                prior = prior_func(x)
                prior[0] = 0
                prior[-1] = 0
            elif prior_dict['type'] == 'log10unif':
                upper = prior_dict['log10upper']
                lower = prior_dict['log10lower']
                x = np.linspace(lower, upper, N)
                prior = [prior_func(xi) for xi in x]
            elif prior_dict['type'] == 'norm':
                lower = prior_dict['loc'] - normal_stds * prior_dict['scale']
                upper = prior_dict['loc'] + normal_stds * prior_dict['scale']
                x = np.linspace(lower, upper, N)
                prior = prior_func(x)
            elif prior_dict['type'] == 'halfnorm':
                lower = prior_dict['loc']
                upper = prior_dict['loc'] + normal_stds * prior_dict['scale']
                x = np.linspace(lower, upper, N)
                prior = [prior_func(xi) for xi in x]
            elif prior_dict['type'] == 'neghalfnorm':
                upper = prior_dict['loc']
                lower = prior_dict['loc'] - normal_stds * prior_dict['scale']
                x = np.linspace(lower, upper, N)
                prior = [prior_func(xi) for xi in x]
            else:
                raise ValueError('Not implemented for prior type {}'.format(
                    prior_dict['type']))
            priorln = ax.plot(x, prior, 'C3', label='prior')
            ax.set_xlabel(self.theta_symbols[i])

            s = self.samples[:, i]
            while len(s) > 10**4:
                # random downsample to avoid slow calculation of kde
                s = np.random.choice(s, size=int(len(s)/2.))
            kde = gaussian_kde(s)
            ax2 = ax.twinx()
            postln = ax2.plot(x, kde.pdf(x), 'k', label='posterior')
            ax2.set_yticklabels([])
            ax.set_yticklabels([])

        lns = priorln + postln
        labs = [l.get_label() for l in lns]
        axes[0].legend(lns, labs, loc=1, framealpha=0.8)

        fig.savefig('{}/{}_prior_posterior.png'.format(
            self.outdir, self.label))

    def plot_cumulative_max(self, **kwargs):
        """ Plot the cumulative twoF for the maximum posterior estimate

        See the pyfstat.core.plot_twoF_cumulative function for further details
        """
        d, maxtwoF = self.get_max_twoF()
        for key, val in self.theta_prior.iteritems():
            if key not in d:
                d[key] = val

        if hasattr(self, 'search') is False:
            self._initiate_search_object()
        if self.binary is False:
            self.search.plot_twoF_cumulative(
                self.label, self.outdir, F0=d['F0'], F1=d['F1'], F2=d['F2'],
                Alpha=d['Alpha'], Delta=d['Delta'],
                tstart=self.minStartTime, tend=self.maxStartTime,
                **kwargs)
        else:
            self.search.plot_twoF_cumulative(
                self.label, self.outdir, F0=d['F0'], F1=d['F1'], F2=d['F2'],
                Alpha=d['Alpha'], Delta=d['Delta'], asini=d['asini'],
                period=d['period'], ecc=d['ecc'], argp=d['argp'], tp=d['argp'],
                tstart=self.minStartTime, tend=self.maxStartTime, **kwargs)

    def _generic_lnprior(self, **kwargs):
        """ Return a lambda function of the pdf

        Parameters
        ----------
        **kwargs:
            A dictionary containing 'type' of pdf and shape parameters

        """

        def log_of_unif(x, a, b):
            above = x < b
            below = x > a
            if type(above) is not np.ndarray:
                if above and below:
                    return -np.log(b-a)
                else:
                    return -np.inf
            else:
                idxs = np.array([all(tup) for tup in zip(above, below)])
                p = np.zeros(len(x)) - np.inf
                p[idxs] = -np.log(b-a)
                return p

        def log_of_log10unif(x, log10lower, log10upper):
            log10x = np.log10(x)
            above = log10x < log10upper
            below = log10x > log10lower
            if type(above) is not np.ndarray:
                if above and below:
                    return -np.log(x*np.log(10)*(log10upper-log10lower))
                else:
                    return -np.inf
            else:
                idxs = np.array([all(tup) for tup in zip(above, below)])
                p = np.zeros(len(x)) - np.inf
                p[idxs] = -np.log(x*np.log(10)*(log10upper-log10lower))
                return p

        def log_of_halfnorm(x, loc, scale):
            if x < loc:
                return -np.inf
            else:
                return -0.5*((x-loc)**2/scale**2+np.log(0.5*np.pi*scale**2))

        def cauchy(x, x0, gamma):
            return 1.0/(np.pi*gamma*(1+((x-x0)/gamma)**2))

        def exp(x, x0, gamma):
            if x > x0:
                return np.log(gamma) - gamma*(x - x0)
            else:
                return -np.inf

        if kwargs['type'] == 'unif':
            return lambda x: log_of_unif(x, kwargs['lower'], kwargs['upper'])
        if kwargs['type'] == 'log10unif':
            return lambda x: log_of_log10unif(
                x, kwargs['log10lower'], kwargs['log10upper'])
        elif kwargs['type'] == 'halfnorm':
            return lambda x: log_of_halfnorm(x, kwargs['loc'], kwargs['scale'])
        elif kwargs['type'] == 'neghalfnorm':
            return lambda x: log_of_halfnorm(
                -x, kwargs['loc'], kwargs['scale'])
        elif kwargs['type'] == 'norm':
            return lambda x: -0.5*((x - kwargs['loc'])**2/kwargs['scale']**2
                                   + np.log(2*np.pi*kwargs['scale']**2))
        else:
            logging.info("kwargs:", kwargs)
            raise ValueError("Print unrecognise distribution")

    def _generate_rv(self, **kwargs):
        dist_type = kwargs.pop('type')
        if dist_type == "unif":
            return np.random.uniform(low=kwargs['lower'], high=kwargs['upper'])
        if dist_type == "log10unif":
            return 10**(np.random.uniform(low=kwargs['log10lower'],
                                          high=kwargs['log10upper']))
        if dist_type == "norm":
            return np.random.normal(loc=kwargs['loc'], scale=kwargs['scale'])
        if dist_type == "halfnorm":
            return np.abs(np.random.normal(loc=kwargs['loc'],
                                           scale=kwargs['scale']))
        if dist_type == "neghalfnorm":
            return -1 * np.abs(np.random.normal(loc=kwargs['loc'],
                                                scale=kwargs['scale']))
        if dist_type == "lognorm":
            return np.random.lognormal(
                mean=kwargs['loc'], sigma=kwargs['scale'])
        else:
            raise ValueError("dist_type {} unknown".format(dist_type))

    def _plot_walkers(self, sampler, symbols=None, alpha=0.8, color="k",
                      temp=0, lw=0.1, nprod=0, add_det_stat_burnin=False,
                      fig=None, axes=None, xoffset=0, plot_det_stat=False,
                      context='ggplot', subtractions=None, labelpad=0.05):
        """ Plot all the chains from a sampler """

        if context not in plt.style.available:
            raise ValueError((
                'The requested context {} is not available; please select a'
                ' context from `plt.style.available`').format(context))

        if np.ndim(axes) > 1:
            axes = axes.flatten()

        shape = sampler.chain.shape
        if len(shape) == 3:
            nwalkers, nsteps, ndim = shape
            chain = sampler.chain[:, :, :]
        if len(shape) == 4:
            ntemps, nwalkers, nsteps, ndim = shape
            if temp < ntemps:
                logging.info("Plotting temperature {} chains".format(temp))
            else:
                raise ValueError(("Requested temperature {} outside of"
                                  "available range").format(temp))
            chain = sampler.chain[temp, :, :, :]

        if subtractions is None:
            subtractions = [0 for i in range(ndim)]
        else:
            if len(subtractions) != self.ndim:
                raise ValueError('subtractions must be of length ndim')

        if plot_det_stat:
            extra_subplots = 1
        else:
            extra_subplots = 0
        with plt.style.context((context)):
            plt.rcParams['text.usetex'] = True
            if fig is None and axes is None:
                fig = plt.figure(figsize=(4, 3.0*ndim))
                ax = fig.add_subplot(ndim+extra_subplots, 1, 1)
                axes = [ax] + [fig.add_subplot(ndim+extra_subplots, 1, i)
                               for i in range(2, ndim+1)]

            idxs = np.arange(chain.shape[1])
            burnin_idx = chain.shape[1] - nprod
            if hasattr(self, 'convergence_idx'):
                convergence_idx = self.convergence_idx
            else:
                convergence_idx = burnin_idx
            if ndim > 1:
                for i in range(ndim):
                    axes[i].ticklabel_format(useOffset=False, axis='y')
                    cs = chain[:, :, i].T
                    if burnin_idx > 0:
                        axes[i].plot(xoffset+idxs[:convergence_idx+1],
                                     cs[:convergence_idx+1]-subtractions[i],
                                     color="C3", alpha=alpha,
                                     lw=lw)
                        axes[i].axvline(xoffset+convergence_idx,
                                        color='k', ls='--', lw=0.25)
                    axes[i].plot(xoffset+idxs[burnin_idx:],
                                 cs[burnin_idx:]-subtractions[i],
                                 color="k", alpha=alpha, lw=lw)

                    axes[i].set_xlim(0, xoffset+idxs[-1])
                    if symbols:
                        if subtractions[i] == 0:
                            axes[i].set_ylabel(symbols[i], labelpad=labelpad)
                        else:
                            axes[i].set_ylabel(
                                symbols[i]+'$-$'+symbols[i]+'$_0$',
                                labelpad=labelpad)

                    if hasattr(self, 'convergence_diagnostic'):
                        ax = axes[i].twinx()
                        axes[i].set_zorder(ax.get_zorder()+1)
                        axes[i].patch.set_visible(False)
                        c_x = np.array(self.convergence_diagnosticx)
                        c_y = np.array(self.convergence_diagnostic)
                        break_idx = np.argmin(np.abs(c_x - burnin_idx))
                        ax.plot(c_x[:break_idx], c_y[:break_idx, i], '-C0',
                                zorder=-10)
                        ax.plot(c_x[break_idx:], c_y[break_idx:, i], '-C0',
                                zorder=-10)
                        if self.convergence_test_type == 'autocorr':
                            ax.set_ylabel(r'$\tau_\mathrm{exp}$')
                        elif self.convergence_test_type == 'GR':
                            ax.set_ylabel('PSRF')
                        ax.ticklabel_format(useOffset=False)
            else:
                axes[0].ticklabel_format(useOffset=False, axis='y')
                cs = chain[:, :, temp].T
                if burnin_idx:
                    axes[0].plot(idxs[:burnin_idx], cs[:burnin_idx],
                                 color="C3", alpha=alpha, lw=lw)
                axes[0].plot(idxs[burnin_idx:], cs[burnin_idx:], color="k",
                             alpha=alpha, lw=lw)
                if symbols:
                    axes[0].set_ylabel(symbols[0], labelpad=labelpad)

            axes[-1].set_xlabel(r'$\textrm{Number of steps}$', labelpad=0.2)

            if plot_det_stat:
                if len(axes) == ndim:
                    axes.append(fig.add_subplot(ndim+1, 1, ndim+1))

                lnl = sampler.lnlikelihood[temp, :, :]
                if burnin_idx and add_det_stat_burnin:
                    burn_in_vals = lnl[:, :burnin_idx].flatten()
                    try:
                        twoF_burnin = (burn_in_vals[~np.isnan(burn_in_vals)]
                                       - self.likelihoodcoef)
                        axes[-1].hist(twoF_burnin, bins=50, histtype='step',
                                      color='C3')
                    except ValueError:
                        logging.info('Det. Stat. hist failed, most likely all '
                                     'values where the same')
                        pass
                else:
                    twoF_burnin = []
                prod_vals = lnl[:, burnin_idx:].flatten()
                try:
                    twoF = prod_vals[~np.isnan(prod_vals)]-self.likelihoodcoef
                    axes[-1].hist(twoF, bins=50, histtype='step', color='k')