Skip to content
Snippets Groups Projects
Commit 45e60bfd authored by Rhys Green's avatar Rhys Green
Browse files

adding gregs patch

parent f569efd3
No related branches found
No related tags found
1 merge request!322Adding PTMCMC sampler
Hide whitespace changes
Inline Side-by-side
bilby/core/sampler/ptmcmc.py
+ 69
− 96
View file @ 45e60bfd
from __future__ import absolute_import, print_function
import glob
import os
import shutil
import numpy as np
# from pandas import DataFrame
# from ..utils import logger, get_progress_bar
from .base_sampler import MCMCSampler, SamplerNotInstalledError
from ..utils import logger, check_directory_exists_and_if_not_mkdir
class PTMCMCSampler(MCMCSampler):
"""bilby wrapper PTMCMC (https://github.com/jellis18/PTMCMCSampler/)
"""bilby wrapper of PTMCMC (https://github.com/jellis18/PTMCMCSampler/)
All positional and keyword arguments (i.e., the args and kwargs) passed to
`run_sampler` will be propagated to `PTMCMCSampler.PTMCMCSampler`, see
......@@ -17,26 +20,25 @@ class PTMCMCSampler(MCMCSampler):
Other Parameters
----------------
Niter: int (2*10**4 +1)
Niter: int (2*10**4 + 1)
The number of mcmc steps
burn: int (5 * 10**3)
If given, the fixed number of steps to discard as burn-in
thin: int (1)
The number of steps before saving the sample to the chain
custom_proposals: dict (None)
this is to add any proposal to the array of proposal,
this must be in the form of a dictionary with the
name of the proposal, then a list containing the jump
function and the weight e.g {'name' : [function , weight]}
see (https://github.com/rgreen1995/PTMCMCSampler/blob/master/examples/simple.ipynb)
and (http://jellis18.github.io/PTMCMCSampler/PTMCMCSampler.html#ptmcmcsampler-ptmcmcsampler-module)
for examples and more info.
logl_grad: func (None)
Gradient of likelihood if known (default = None)
logp_grad: func (None)
Gradient of prior if known (default = None)
verbose: bool (True)
Update current run-status to the screen
Add dictionary of proposals to the array of proposals, this must be in
the form of a dictionary with the name of the proposal, then a list
containing the jump function and the weight e.g {'name' : [function ,
weight]} see
(https://github.com/rgreen1995/PTMCMCSampler/blob/master/examples/simple.ipynb)
and
(http://jellis18.github.io/PTMCMCSampler/PTMCMCSampler.html#ptmcmcsampler-ptmcmcsampler-module)
for examples and more info. logl_grad: func (None) Gradient of
likelihood if known (default = None) logp_grad: func (None) Gradient
of prior if known (default = None) verbose: bool (True) Update current
run-status to the screen
"""
default_kwargs = {'p0': None, 'Niter': 2 * 10**4 + 1, 'neff': 10**4,
......@@ -47,14 +49,17 @@ class PTMCMCSampler(MCMCSampler):
'HMCweight': 0, 'MALAweight': 0, 'NUTSweight': 0,
'HMCstepsize': 0.1, 'HMCsteps': 300,
'groups': None, 'custom_proposals': None,
'loglargs': {}, 'loglkwargs': {}, 'logpargs': {}, 'logpkwargs': {},
'logl_grad': None, 'logp_grad': None, 'outDir': './temp'}
'loglargs': {}, 'loglkwargs': {}, 'logpargs': {},
'logpkwargs': {}, 'logl_grad': None, 'logp_grad': None,
'outDir': None}
def __init__(self, likelihood, priors, outdir='outdir', label='label', use_ratio=False, plot=False,
skip_import_verification=False, pos0=None, nburn=None, burn_in_fraction=0.25, **kwargs):
def __init__(self, likelihood, priors, outdir='outdir', label='label',
use_ratio=False, plot=False, skip_import_verification=False,
pos0=None, burn_in_fraction=0.25, **kwargs):
MCMCSampler.__init__(self, likelihood=likelihood, priors=priors, outdir=outdir, label=label,
use_ratio=use_ratio, plot=plot,
MCMCSampler.__init__(self, likelihood=likelihood, priors=priors,
outdir=outdir, label=label, use_ratio=use_ratio,
plot=plot,
skip_import_verification=skip_import_verification,
**kwargs)
......@@ -62,8 +67,9 @@ class PTMCMCSampler(MCMCSampler):
self.likelihood = likelihood
self.priors = priors
# PTMCMC is imported with Caps so need to overwrite this.
def _verify_external_sampler(self):
# PTMCMC is imported with Caps so need to overwrite the parent function
# which forces `__name__.lower()
external_sampler_name = self.__class__.__name__
try:
self.external_sampler = __import__(external_sampler_name)
......@@ -71,31 +77,15 @@ class PTMCMCSampler(MCMCSampler):
raise SamplerNotInstalledError(
"Sampler {} is not installed on this system".format(external_sampler_name))
@property
def kwargs(self):
""" Ensures that proper keyword arguments are used for the PTMCMC sampler.
Returns
-------
dict: Keyword arguments used for the PTMCMC
"""
return self.__kwargs
@kwargs.setter
def kwargs(self, kwargs):
self.__kwargs = self.default_kwargs.copy()
self.__kwargs.update(kwargs)
self._verify_kwargs_against_default_kwargs()
# def _translate_kwargs(self, kwargs):
# if 'nwalkers' not in kwargs:
# for equiv in self.nwalkers_equiv_kwargs:
# if equiv in kwargs:
# kwargs['nwalkers'] = kwargs.pop(equiv)
# if 'iterations' not in kwargs:
# if 'nsteps' in kwargs:
# kwargs['iterations'] = kwargs.pop('nsteps')
def _translate_kwargs(self, kwargs):
if 'Niter' not in kwargs:
for equiv in self.nsteps_equiv_kwargs:
if equiv in kwargs:
kwargs['Niter'] = kwargs.pop(equiv)
if 'burn' not in kwargs:
for equiv in self.nburn_equiv_kwargs:
if equiv in kwargs:
kwargs['burn'] = kwargs.pop(equiv)
@property
def custom_proposals(self):
......@@ -113,6 +103,8 @@ class PTMCMCSampler(MCMCSampler):
'outDir',
'verbose']
init_kwargs = {key: self.kwargs[key] for key in keys}
if init_kwargs['outDir'] is None:
init_kwargs['outDir'] = '{}/ptmcmc_temp_{}/'.format(self.outdir, self.label)
return init_kwargs
@property
......@@ -138,52 +130,44 @@ class PTMCMCSampler(MCMCSampler):
sampler_kwargs = {key: self.kwargs[key] for key in keys}
return sampler_kwargs
@property
def nsteps(self):
return self.kwargs['Niter']
@property
def nburn(self):
return self.kwargs['burn']
@staticmethod
def _import_external_sampler():
from PTMCMCSampler import PTMCMCSampler
import glob
import os
# OPTIMIZE:
# import acor
# from mpi4py import MPI
# return MPI, PTMCMCSampler
return PTMCMCSampler, glob, os
return PTMCMCSampler
def run_sampler(self):
# MPI , PTMCMCSampler = self._import_external_sampler()
PTMCMCSampler, glob, os = self._import_external_sampler()
init_kwargs = self.sampler_init_kwargs
sampler_kwargs = self.sampler_function_kwargs
PTMCMCSampler = self._import_external_sampler()
sampler = PTMCMCSampler.PTSampler(ndim=self.ndim, logp=self.log_prior,
logl=self.log_likelihood, cov=np.eye(self.ndim),
**init_kwargs)
**self.sampler_init_kwargs)
if self.custom_proposals is not None:
for proposal in self.custom_proposals:
print('adding ' + str(proposal) + ' to proposals with weight:' +
str(self.custom_proposals[proposal][1]))
logger.info('Adding {} to proposals with weight {}'.format(
proposal, self.custom_proposals[proposal][1]))
sampler.addProposalToCycle(self.custom_proposals[proposal][0],
self.custom_proposals[proposal][1])
else:
pass
sampler.sample(p0=self.p0, **sampler_kwargs)
# The next bit is very hacky, the ptmcmc writes the samples and
# other info to file so here i read this info, write it to the result
# object then delete it
data = np.loadtxt('temp/chain_1.txt')
jumpfiles = glob.glob('temp/*jump.txt')
sampler.sample(p0=self.p0, **self.sampler_function_kwargs)
samples, meta, loglike = self.__read_in_data()
self.result.nburn = self.sampler_function_kwargs['burn']
self.result.samples = samples[self.result.nburn:]
self.meta_data['sampler_meta'] = meta
self.result.log_likelihood_evaluations = loglike[self.result.nburn:]
self.result.sampler_output = np.nan
self.result.walkers = np.nan
self.result.log_evidence = np.nan
self.result.log_evidence_err = np.nan
return self.result
def __read_in_data(self):
""" Read the data stored by PTMCMC to disk """
temp_outDir = self.sampler_init_kwargs['outDir']
data = np.loadtxt('{}/chain_1.txt'.format(temp_outDir))
jumpfiles = glob.glob('{}/*jump.txt'.format(temp_outDir))
jumps = map(np.loadtxt, jumpfiles)
samples = data[:, :-4]
loglike = data[:, -3]
meta = {}
jump_accept = {}
for ct, j in enumerate(jumps):
label = jumpfiles[ct].split('/')[-1].split('_jump.txt')[0]
......@@ -191,24 +175,13 @@ class PTMCMCSampler(MCMCSampler):
PT_swap = {'swap_accept': data[-1]}
tot_accept = {'tot_accept': data[-2]}
log_post = {'log_post': data[:, -4]}
meta = {}
meta['tot_accept'] = tot_accept
meta['PT_swap'] = PT_swap
meta['proposals'] = jump_accept
meta['log_post'] = log_post
samples = data[:, :-4]
for f in glob.glob('./temp/*'):
os.remove(f)
os.rmdir('temp')
shutil.rmtree(temp_outDir)
return samples, meta, loglike
self.result.nburn = self.nburn
self.result.samples = samples[self.nburn:]
self.meta_data['sampler_meta'] = meta
self.result.log_likelihood_evaluations = loglike[self.nburn:]
# self.calculate_autocorrelation(sampler.chain.reshape((-1, self.ndim)))
# self.print_nburn_logging_info()
self.result.sampler_output = np.nan
self.result.walkers = np.nan
self.result.log_evidence = np.nan
self.result.log_evidence_err = np.nan
return self.result
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment