multinest_emcee_scans.py

###############################################################################
# run_11bin.py
###############################################################################
#
# Perform scan to obtain best-fit disk and bulge PS population parameters
#
# Here using the 4D efficiency function and the new 11 binning
# 
###############################################################################


import sys, os
import numpy as np
import pymultinest
from iminuit import Minuit
import corner
import emcee
from tqdm import *
from multiprocessing import Pool
import contextlib

# Import custom functions
sys.path.append("../likelihood/")
import likelihood_4d_11bin_lp as likelihood_4d
import likelihood_3d_11bin_lp as likelihood_3d
from minuit_functions import call_ll

class run_scan():
    def __init__(self, fixed_params, fixed_param_vals, floated_params, 
                 floated_param_priors, data_dir = '../data/', Lmin = 1.e31, 
                 Ns=200, Nang=1, smax_disk=40., theta_mask=0.,
                 share_lf=True, use_prior=False, Nang_prior=40, 
                 efficiency_file = "efficiency_4d_11bins.npy",
                 data_file = 'PSR_data_11bins.npy', skipcentral = True,
                 fourd=True, diskonly=False, maskplane=False): 
        """ Initialize scan class

            :param fixed_params: Array of parameters to be held fixed
            :param fixed_param_vals: Array of values for fixed parameters
            :param floated_params: Array of parameters to be floated
            :param floated_param_priors: Priors array for floated parameters
            :param data_dir: Directory containing the required maps
            :param Lmin: Minimum luminosity [erg/s]
            :param Lmax_disk: Maximum luminosity for disk [erg/s]
            :param Lmax_bulge: Maximum luminosity for bulge [erg/s]
            :param Ns: Number of integration point in z [kpc]
            :param Nang: Number of angular integration steps
            :param smax_disk: How far to integrate out to for disk  [kpc]
            :param theta_mask: How many inner degrees to mask
            :param share_betas: Whether to float a single beta for disk and bulge
            :param use_prior: Whether to use prior on total number of sources
            :param Nang_prior: Number of prior angular integrations steps
            :param efficiency_file: Filename of efficiency file in the data_dir
            :param data_file: Filename of data file in the data_dir
            :param fourd: Whether using 4D efficiency and likelihood
            :param skipcentral: Whether to skip central pixel
            :param diskonly: Don't count bulge counts
            :param maskplane: Mask the galactic plane
        """

        # Append variables to self
        self.fixed_params = np.array(fixed_params)
        self.fixed_param_vals = np.array(fixed_param_vals)
        self.floated_params = np.array(floated_params)
        self.floated_param_priors = np.array(floated_param_priors)
        self.data_dir = data_dir
        self.efficiency_file = efficiency_file
        self.data_file = data_file
        self.skipcentral = skipcentral
        self.fourd = fourd
        self.diskonly = diskonly
        self.maskplane = maskplane

        self.Lmin = Lmin
        self.Ns = Ns
        self.Nang = Nang
        self.smax_disk = smax_disk
        self.theta_mask = theta_mask
        self.share_lf = share_lf

        self.use_prior = use_prior
        self.Nang_prior = Nang_prior
                
        self.all_params = np.array(['N_bulge','N_disk','alpha','n','sigma','z0',
                                    'Lmax_disk','Lmax_bulge',
                                    'beta_disk','lambda_disk','Lb_disk',
                                    'beta_bulge','lambda_bulge','Lb_bulge'])
       
        # If using the same beta for bulge/disk adjust parameters
        if self.share_lf:  
            
            self.all_params = self.all_params[self.all_params != 'beta_bulge']
            self.floated_params = self.floated_params[self.floated_params != 'beta_bulge']
            self.fixed_params = self.fixed_params[self.fixed_params != 'beta_bulge']

            self.all_params[self.all_params == 'beta_disk'] = 'beta'
            self.floated_params[self.floated_params == 'beta_disk'] = 'beta'
            self.fixed_params[self.fixed_params == 'beta_disk'] = 'beta'

            self.all_params = self.all_params[self.all_params != 'lambda_bulge']
            self.floated_params = self.floated_params[self.floated_params != 'lambda_bulge']
            self.fixed_params = self.fixed_params[self.fixed_params != 'lambda_bulge']

            self.all_params[self.all_params == 'lambda_disk'] = 'lambda'
            self.floated_params[self.floated_params == 'lambda_disk'] = 'lambda'
            self.fixed_params[self.fixed_params == 'lambda_disk'] = 'lambda'

            self.all_params = self.all_params[self.all_params != 'Lb_bulge']
            self.floated_params = self.floated_params[self.floated_params != 'Lb_bulge']
            self.fixed_params = self.fixed_params[self.fixed_params != 'Lb_bulge']

            self.all_params[self.all_params == 'Lb_disk'] = 'Lb'
            self.floated_params[self.floated_params == 'Lb_disk'] = 'Lb'
            self.fixed_params[self.fixed_params == 'Lb_disk'] = 'Lb'

        # Load the data, set up fixed parameters and priors
        self.load_data()
        self.setup_fixed_params()
        self.setup_prior_thetas()

    def load_data(self):
        """ Load the binned efficiency and data files
        """
        self.PSR_data = np.load(self.data_dir + self.data_file)
        
        self.omega = np.load(self.data_dir + self.efficiency_file)

    def setup_fixed_params(self):
        """ Set up values and arrays for parameters to be held fixed
        """
        self.theta = np.zeros(len(self.all_params))
        for param in self.all_params:
            if param in self.fixed_params:
                param_pos = np.argwhere(np.array(self.fixed_params) == param)[0][0]
                param_all_pos = np.argwhere(np.array(self.all_params) == param)[0][0]
                self.theta[param_all_pos] = self.fixed_param_vals[param_pos]

    def setup_prior_thetas(self):
        """ Set up priors and arrays for parameters to be floated
        """ 
        self.param_all_pos_ary = []
        self.theta_min = []
        self.theta_max = []
        for param in self.all_params:
            if param in self.floated_params:
                param_pos = np.argwhere(np.array(self.floated_params) == param)[0][0]
                param_all_pos = np.argwhere(np.array(self.all_params) == param)[0][0]
                self.param_all_pos_ary.append(param_all_pos)
                self.theta_min += [self.floated_param_priors[param_pos][0]]
                self.theta_max += [self.floated_param_priors[param_pos][1]]

        self.theta_interval = list(np.array(self.theta_max) - np.array(self.theta_min))

    def prior_cube(self, cube, ndim=1, nparams=1):
        """ Cube of priors in the format required by MultiNest
        """
        for i in range(ndim):
            cube[i] = cube[i] * self.theta_interval[i] + self.theta_min[i]
        return cube

    def ll(self, theta, ndim = 1, nparams = 1):
        """ Log Likelihood in the form that can be used by Minuit or Multinest
        """
        theta_ll = np.array(self.theta)

        for float_idx, float_item in enumerate(self.param_all_pos_ary):
            theta_ll[float_item] = theta[float_idx]

        # if self.share_betas and ('beta' in self.floated_params):
        if self.share_lf:
            theta_ll = np.append(theta_ll, theta_ll[-3:])  

        if self.fourd:
            likelihood = likelihood_4d
        else:
            likelihood = likelihood_3d

        theta_ll[1-1] = 10**theta_ll[1-1]
        theta_ll[2-1] = 10**theta_ll[2-1]
        theta_ll[11-1] = 10**theta_ll[11-1]
        theta_ll[14-1] = 10**theta_ll[14-1]

        ll_val =  likelihood.ll(self.PSR_data, self.omega, *theta_ll, 
                                    Lmin=self.Lmin, Ns=self.Ns, Nang=self.Nang, 
                                    smax_disk=self.smax_disk, 
                                    theta_mask=self.theta_mask, 
                                    use_prior=self.use_prior, 
                                    Nang_prior=self.Nang_prior,
                                    skipcentral=self.skipcentral,
                                    diskonly=self.diskonly,
                                    maskplane=self.maskplane)
        return ll_val

    def perform_scan_multinest(self, chains_dir, nlive=100, pymultinest_options=None):
        """ Perform a scan with MultiNest
        """
        self.make_dirs([chains_dir])
        n_params = len(self.floated_params)
        if pymultinest_options is None:
            pymultinest_options_arg = {'importance_nested_sampling': False,
                                    'resume': False, 'verbose': True,
                                    'sampling_efficiency': 'model',
                                    'init_MPI': False, 'evidence_tolerance': 0.5,
                                    'const_efficiency_mode': False}
        else:
            pymultinest_options_arg = pymultinest_options

        pymultinest.run(self.ll, self.prior_cube, n_params, 
                        outputfiles_basename=chains_dir, 
                        n_live_points=nlive, **pymultinest_options_arg)

    def perform_scan_emcee(self, chains_dir, MPI=False):
        """ Perform a scan with emcee
        """
        self.make_dirs([chains_dir])
        ndim, nwalkers = len(self.floated_params), 200
        init_vals = np.array((self.floated_param_priors[:,0] + self.floated_param_priors[:,1])/2.)
        self.pos = [init_vals + 1e-4*np.random.randn(ndim) for i in range(nwalkers)]
        
        if MPI:
            pool = MPIPool()
            if not pool.is_master():
                pool.wait()
                sys.exit(0)
            self.sampler = emcee.EnsembleSampler(nwalkers, ndim, self.ll, pool=pool)
        else:
            self.sampler = emcee.EnsembleSampler(nwalkers, ndim, self.ll)
               
        self.chain = []
        for results in tqdm(self.sampler.sample(self.pos, iterations=1000)):
            self.chain.append(results[0])
        
        np.save(chains_dir + "/chains.npy", self.chain)
        pool.close()
        
    def perform_scan_minuit(self, verbose=1):
        """ Perform a scan with minuit
        """
        limit_dict = {}
        init_val_dict = {}
        step_size_dict = {}
        for iparam, param in enumerate(self.floated_params):
            limit_dict['limit_'+param] = (self.floated_param_priors[iparam][0],
                                          self.floated_param_priors[iparam][1])
            init_val_dict[param] = (self.floated_param_priors[iparam][0]
                                   +self.floated_param_priors[iparam][1])/2.
            step_size_dict['error_'+param] = .1
        other_kwargs = {'print_level': verbose, 'errordef': 1}
        z = limit_dict.copy()
        z.update(other_kwargs)
        z.update(limit_dict)
        z.update(init_val_dict)
        z.update(step_size_dict)
        f = call_ll(len(self.floated_params),self.ll,self.floated_params)
        self.m = Minuit(f,**z)
        self.m.migrad()
        self.max_LL = - self.m.fval

    def plot_corner(self, chains_dir, labels, truths=None, param_range=None):
        """ Make a corner plot
        """
        chain_file = chains_dir + 'post_equal_weights.dat'
        self.samples = np.array(np.loadtxt(chain_file)[:, :-1])

        levels =  1.0 - np.exp(-0.5 *np.array([1.0,2.0,3.0])**2)

        sams_log = np.zeros(np.shape(self.samples))
        sams_log[::,0] = self.samples[::,0]
        sams_log[::,1] = self.samples[::,1] 
        sams_log[::,2:] = self.samples[::,2:]

        if param_range is None:
            param_range = [1 for _ in range(len(self.floated_params))]
            
        corner.corner(sams_log,bins=40,smooth=True,labels=labels,quantiles=[0.16, 0.5, 0.84],show_titles=True,
                      range=param_range,color='firebrick',levels=levels,
                      use_math_text=True,truths=truths, truth_color='forestgreen')

    @staticmethod
    def make_dirs(dirs):
        """ Creates directories if they do not already exist 
        """
        for d in dirs:
            if not os.path.exists(d):
                try:
                    os.mkdir(d)
                except OSError as e:
                    if e.errno != 17:
                        raise