lenstronomy.Sampling.Samplers.polychord_sampler — lenstronomy 1.8.2 documentation

# Source code for lenstronomy.Sampling.Samplers.polychord_sampler

    __author__ = 'aymgal'

    import os
    import shutil
    import numpy as np
    import sys
    import copy

    from lenstronomy.Sampling.Samplers.base_nested_sampler import NestedSampler
    import lenstronomy.Util.sampling_util as utils

    __all__ = ['DyPolyChordSampler']

    class DyPolyChordSampler(NestedSampler):
        """
        Wrapper for dynamical nested sampling algorithm DyPolyChord
        by E. Higson, M. Hobson, W. Handley, A. Lasenby

        papers : arXiv:1704.03459, arXiv:1804.06406
        doc : https://dypolychord.readthedocs.io
        """

        def __init__(self, likelihood_module, prior_type='uniform',
                     prior_means=None, prior_sigmas=None, width_scale=1, sigma_scale=1,
                     output_dir=None, output_basename='-',
                     resume_dyn_run=False,
                     polychord_settings=None,
                     remove_output_dir=False, use_mpi=False): #, num_mpi_procs=1):
            """
            :param likelihood_module: likelihood_module like in likelihood.py (should be callable)
            :param prior_type: 'uniform' of 'gaussian', for converting the unit hypercube to param cube
            :param prior_means: if prior_type is 'gaussian', mean for each param
            :param prior_sigmas: if prior_type is 'gaussian', std dev for each param
            :param width_scale: scale the widths of the parameters space by this factor
            :param sigma_scale: if prior_type is 'gaussian', scale the gaussian sigma by this factor
            :param output_dir: name of the folder that will contain output files
            :param output_basename: prefix for output files
            :param seed_increment: seed increment for dypolychord RNG with MPI. Check dypolychord documentation for details.
            :param resume_dyn_run: if True, previous resume files will not be deleted so that previous run can be resumed
            :param polychord_settings: settings dictionary to send to pypolychord. Check dypolychord documentation for details.
            :param remove_output_dir: remove the output_dir folder after completion
            :param use_mpi: Use MPI computing if `True`
            """
            self._check_install()
            super(DyPolyChordSampler, self).__init__(likelihood_module, prior_type,
                                                     prior_means, prior_sigmas,
                                                     width_scale, sigma_scale)

            # if use_mpi:
            #     mpi_str = 'mpirun -np {}'.format(num_mpi_procs)
            # else:
            #     mpi_str = None
            if polychord_settings == None:
                polychord_settings = {}
            self._use_mpi = use_mpi

            self._output_dir= output_dir
            self._is_master = True

            if self._use_mpi:
                from mpi4py import MPI
                self._comm = MPI.COMM_WORLD

                if self._comm.Get_rank() != 0:
                    self._is_master = False
            else:
                self._comm = None

            if not resume_dyn_run:
                if self._is_master:
                    if os.path.exists(self._output_dir):
                        shutil.rmtree(self._output_dir, ignore_errors=True)
                    os.mkdir(self._output_dir)

            self._output_basename = output_basename
            self._settings = copy.deepcopy(polychord_settings)
            self._settings['file_root'] = self._output_basename
            self._settings['base_dir'] = self._output_dir

            if self._all_installed:
                # create the dyPolyChord callable object
                self._sampler = self._RunPyPolyChord(self.log_likelihood,
                                                     self.prior, self.n_dims)
            else:
                self._sampler = None

            self._rm_output = remove_output_dir
            self._has_warned = False

        def prior(self, cube):
            """
            compute the mapping between the unit cube and parameter cube

            'copy=True' below because cube can not be modified in-place (read-only)

            :param cube: unit hypercube, sampled by the algorithm
            :return: hypercube in parameter space
            """
            if self.prior_type == 'gaussian':
                p = utils.cube2args_gaussian(cube, self.lowers, self.uppers,
                                             self.means, self.sigmas, self.n_dims,
                                             copy=True)
            elif self.prior_type == 'uniform':
                p = utils.cube2args_uniform(cube, self.lowers, self.uppers,
                                            self.n_dims, copy=True)
            return p

        def log_likelihood(self, args):
            """
            compute the log-likelihood given list of parameters

            :param args: parameter values
            :return: log-likelihood (from the likelihood module)
            """
            phi = []
            logL = self._ll.likelihood(args)
            if not np.isfinite(logL):
                if not self._has_warned:
                    print("WARNING : logL is not finite : return very low value instead")
                logL = -1e15
                self._has_warned = True
            return float(logL), phi

        def run(self, dynamic_goal, kwargs_run):
            """
            run the DyPolyChord dynamical nested sampler

            see https://dypolychord.readthedocs.io for content of kwargs_run

            :param dynamic_goal: 0 for evidence computation, 1 for posterior computation
            :param kwargs_run: kwargs directly passed to dyPolyChord.run_dypolychord
            :return: samples, means, logZ, logZ_err, logL, ns_run
            """
            print("prior type :", self.prior_type)
            print("parameter names :", self.param_names)

            if self._all_installed:
                # TODO : put a default dynamic_goal ?
                # dynamic_goal = 0 for evidence-only, 1 for posterior-only

                self._dyPolyChord.run_dypolychord(self._sampler, dynamic_goal,
                                                  settings_dict_in=self._settings,
                                                  comm=self._comm, **kwargs_run)

                if self._is_master:
                    ns_run = self._ns_process_run(self._settings['file_root'],
                                               self._settings['base_dir'])

            else:
                # in case DyPolyChord or NestCheck was not compiled properly, for unit tests
                ns_run = {
                    'theta': np.zeros((1, self.n_dims)),
                    'logl': np.zeros(self.n_dims),
                    'output': {
                        'logZ': np.zeros(self.n_dims),
                        'logZerr': np.zeros(self.n_dims),
                        'param_means': np.zeros(self.n_dims)
                    }
                }
                self._write_equal_weights(ns_run['theta'], ns_run['logl'])

            if self._is_master:
                samples, logL = self._get_equal_weight_samples()
                # logL     = ns_run['logl']
                # samples_w = ns_run['theta']
                logZ     = ns_run['output']['logZ']
                logZ_err = ns_run['output']['logZerr']
                means    = ns_run['output']['param_means']

                print('The log evidence estimate using the first run is {}'
                      .format(logZ))
                print('The estimated mean of the first parameter is {}'
                      .format(means[0]))

                if self._rm_output:
                    shutil.rmtree(self._output_dir, ignore_errors=True)

                return samples, means, logZ, logZ_err, logL, ns_run
            else:
                sys.exit(0)

        def _get_equal_weight_samples(self):
            """
            Inspired by pymultinest's Analyzer,
            because DyPolyChord has more or less the same output conventions as MultiNest
            """
            file_name = '{}_equal_weights.txt'.format(self._output_basename)
            file_path = os.path.join(self._output_dir, file_name)
            data = np.loadtxt(file_path, ndmin=2)
            logL = -0.5 * data[:, 0]
            samples = data[:, 1:]
            return samples, logL

        def _write_equal_weights(self, samples, logL):
            # write fake output file for unit tests
            file_name = '{}_equal_weights.txt'.format(self._output_basename)
            file_path = os.path.join(self._output_dir, file_name)
            data = np.zeros((samples.shape[0], 1+samples.shape[1]))
            data[:, 0]  = -2. * logL
            data[:, 1:] = samples
            np.savetxt(file_path, data, fmt='% .14E')

        def _check_install(self):
            try:
                import dyPolyChord
                from dyPolyChord import pypolychord_utils
            except:
                print("Warning : dyPolyChord not properly installed. \
    You can get it from : https://github.com/ejhigson/dyPolyChord")
                dypolychord_installed = False
            else:
                dypolychord_installed = True
                self._dyPolyChord = dyPolyChord
                self._RunPyPolyChord = pypolychord_utils.RunPyPolyChord

            try:
                from nestcheck import data_processing
            except:
                print("Warning : nestcheck not properly installed (results might be unexpected). \
    You can get it from : https://github.com/ejhigson/nestcheck")
                nestcheck_installed = False
            else:
                nestcheck_installed = True
                self._ns_process_run = data_processing.process_polychord_run

            self._all_installed = dypolychord_installed and nestcheck_installed