Source code for lenstools.statistics.samplers

"""

.. module:: samplers
	:platform: Unix
	:synopsis: This module contains wrappers for some pre implemented parameter samplers


.. moduleauthor:: Andrea Petri <apetri@phys.columbia.edu>


"""

from __future__ import division,print_function,with_statement

import numpy as np
import emcee

from .ensemble import Ensemble
from ..simulations.logs import logdriver
from ..utils.algorithms import precision_bias_correction

###############################
######Multiquadric kernel######
###############################

def multiquadric(x,e):
	return np.sqrt(1.+x/e)

##################################
##Log of the gaussian likelihood##
##################################

def lnprobgauss(p,emulator,data,icov,pslice_values,sample_indices):

	if pslice_values is not None:
		
		pc = np.empty(len(pslice_values)+len(sample_indices))
		for i in pslice_values:
			pc[i] = pslice_values[i]
		for n,i in enumerate(sample_indices):
			pc[i] = p[n]

	else:
		pc =p

	diff = data - emulator.predict(pc,raw=True)
	return -0.5*diff.dot(icov).dot(diff)

###############################
######emcee sampler############
###############################

[docs]def emcee_sampler(emulator,observed_feature,features_covariance,correct=None,pslice=None,nwalkers=16,nburn=100,nchain=1000,pool=None): """ Parameter posterior sampling based on the MCMC algorithm implemented by the emcee package :param emulator: feature emulator :type emulator: :py:class:`Emulator` :param observed_feature: observed feature to condition the parameter estimation :type observed_feature: array :param features_covariance: covariance matrix of the features :type features_covariance: array :param correct: if not None, correct for the bias in the inverse covariance estimator assuming the covariance was estimated by 'correct' simulations :type correct: int. :param pslice: specify slices of the parameter space in which some parameters are keps as constants :type pslice: dict. :param nwalkers: number of chains :type nwalkers: int. :param nburn: length of the burn-in chain :type nburn: int. :param nchain: length of the MCMC chain :type nchain: int. :param pool: MPI Pool for parallelization of computations :type pool: MPIPool :returns: ensemble of samples from the posterior probability distribution :rtype: :py:class:`Ensemble` """ #Train the emulator with fast interpolation emulator.train(method=multiquadric) parameters = emulator["parameters"] parameters_to_sample = emulator.parameter_names #Parameter space to sample if pslice is None: parameters = parameters.values pslice_values = None sample_indices = None else: pslice_values = dict((emulator.parameter_names.index(p),pslice[p]) for p in pslice) sample_indices = [ emulator.parameter_names.index(p) for p in emulator.parameter_names if p not in pslice ] assert len(pslice_values)+len(sample_indices)==len(emulator.parameter_names) parameters = parameters.values[:,sample_indices] #Extremes of the sampling space pmin,pmax = parameters.min(0),parameters.max(0) #Feature name feature_name = emulator.feature_names[0] #Compute the inverse covariance icov = np.linalg.inv(features_covariance) if correct is not None: icov *= precision_bias_correction(correct,len(icov)) #Initialize the walkers positions ndim = len(pmin) p0 = pmin + np.random.uniform(size=(nwalkers,ndim))*(pmax-pmin) #Initialize the sampler sampler = emcee.EnsembleSampler(nwalkers,ndim,lnprobgauss,args=[emulator,observed_feature,icov,pslice_values,sample_indices],pool=pool) #Burn-in logdriver.info("Running emcee burn-in: feature name={0}, feature dimension={1}, parameter dimension={2}, steps={3}".format(feature_name,len(observed_feature),ndim,nburn)) pos,prob,state = sampler.run_mcmc(p0,nburn) sampler.reset() #Sampling logdriver.info("Running emcee MCMC chain: feature name={0}, feature dimension={1}, parameter dimension={2}, steps={3}".format(feature_name,len(observed_feature),ndim,nchain)) sampler.run_mcmc(pos,nchain,rstate0=state) #Return sampled parameters if pslice is None: return Ensemble(sampler.flatchain,columns=emulator.parameter_names) else: return Ensemble(sampler.flatchain,columns=filter(lambda p:p not in pslice,emulator.parameter_names))