def test_sample_ball_truncated():
mean = np.array([0, 1])
sigma = np.array([0.1, 1])
lower_limit = np.array([0, 0])
upper_limit = np.array([3, 4])
samples = sampling_util.sample_ball_truncated(mean, sigma, lower_limit, upper_limit, size=1000)
assert len(samples) == 1000
var0 = samples[:, 0]
var1 = samples[:, 1]
assert np.max(var1) <= upper_limit[1]
assert np.max(var0) <= upper_limit[0]
assert np.min(var1) >= lower_limit[1]
assert np.min(var0) >= lower_limit[0]
def mcmc_emcee(self,
n_walkers,
n_run,
n_burn,
mean_start,
sigma_start,
mpi=False,
progress=False,
threadCount=1,
initpos=None,
backup_filename=None,
start_from_backup=False):
"""
Run MCMC with emcee.
For details, please have a look at the documentation of the emcee packager.
:param n_walkers: number of walkers in the emcee process
:type n_walkers: integer
:param n_run: number of sampling (after burn-in) of the emcee
:type n_run: integer
:param n_burn: number of burn-in iterations (those will not be saved in the output sample)
:type n_burn: integer
:param mean_start: mean of the parameter position of the initialising sample
:type mean_start: numpy array of length the number of parameters
:param sigma_start: spread of the parameter values (uncorrelated in each dimension) of the initialising sample
:type sigma_start: numpy array of length the number of parameters
:param mpi: if True, initializes an MPIPool to allow for MPI execution of the sampler
:type mpi: bool
:param progress: if True, prints the progress bar
:type progress: bool
:param threadCount: number of threats in multi-processing (not applicable for MPI)
:type threadCount: integer
:param initpos: initial walker position to start sampling (optional)
:type initpos: numpy array of size num param x num walkser
:param backup_filename: name of the HDF5 file where sampling state is saved (through emcee backend engine)
:type backup_filename: string
:param start_from_backup: if True, start from the state saved in `backup_filename`.
Otherwise, create a new backup file with name `backup_filename` (any already existing file is overwritten!).
:type start_from_backup: bool
:return: samples, ln likelihood value of samples
:rtype: numpy 2d array, numpy 1d array
"""
num_param, _ = self.chain.param.num_param()
if initpos is None:
initpos = sampling_util.sample_ball_truncated(mean_start,
sigma_start,
self.lower_limit,
self.upper_limit,
size=n_walkers)
pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)
if backup_filename is not None:
backend = emcee.backends.HDFBackend(backup_filename,
name="lenstronomy_mcmc_emcee")
if pool.is_master():
print(
"Warning: All samples (including burn-in) will be saved in backup file '{}'."
.format(backup_filename))
if start_from_backup:
initpos = None
n_run_eff = n_run
else:
n_run_eff = n_burn + n_run
backend.reset(n_walkers, num_param)
if pool.is_master():
print("Warning: backup file '{}' has been reset!".format(
backup_filename))
else:
backend = None
n_run_eff = n_burn + n_run
time_start = time.time()
sampler = emcee.EnsembleSampler(n_walkers,
num_param,
self.chain.logL,
pool=pool,
backend=backend)
sampler.run_mcmc(initpos, n_run_eff, progress=progress)
flat_samples = sampler.get_chain(discard=n_burn, thin=1, flat=True)
dist = sampler.get_log_prob(flat=True, discard=n_burn, thin=1)
if pool.is_master():
print('Computing the MCMC...')
print('Number of walkers = ', n_walkers)
print('Burn-in iterations: ', n_burn)
print('Sampling iterations (in current run):', n_run_eff)
time_end = time.time()
print(time_end - time_start, 'time taken for MCMC sampling')
return flat_samples, dist
def mcmc_zeus(self,
n_walkers,
n_run,
n_burn,
mean_start,
sigma_start,
mpi=False,
threadCount=1,
progress=False,
initpos=None,
backend_filename=None,
moves=None,
tune=True,
tolerance=0.05,
patience=5,
maxsteps=10000,
mu=1.0,
maxiter=10000,
pool=None,
vectorize=False,
blobs_dtype=None,
verbose=True,
check_walkers=True,
shuffle_ensemble=True,
light_mode=False):
"""
Lightning fast MCMC with zeus: https://github.com/minaskar/zeus
For the full list of arguments for the EnsembleSampler, see see `the zeus docs <https://zeus-mcmc.readthedocs.io/en/latest/api/sampler.html>`_.
If you use the zeus sampler, you should cite the following papers: 2105.03468, 2002.06212.
:param n_walkers: number of walkers per parameter
:type n_walkers: integer
:param n_run: number of sampling steps
:type n_run: integer
:param n_burn: number of burn-in steps
:type n_burn: integer
:param mean_start: mean of the parameter position of the initialising sample
:type mean_start: numpy array of length the number of parameters
:param sigma_start: spread of the parameter values (uncorrelated in each dimension) of the initialising sample
:type sigma_start: numpy array of length the number of parameters
:param progress:
:type progress: bool
:param initpos: initial walker position to start sampling (optional)
:type initpos: numpy array of size num param x num walkser
:param backup_filename: name of the HDF5 file where sampling state is saved (through zeus callback function)
:type backup_filename: string
:return: samples, ln likelihood value of samples
:rtype: numpy 2d array, numpy 1d array
"""
import zeus
print('Using zeus to perform the MCMC.')
num_param, _ = self.chain.param.num_param()
if initpos is None:
initpos = sampling_util.sample_ball_truncated(mean_start,
sigma_start,
self.lower_limit,
self.upper_limit,
size=n_walkers)
if backend_filename is not None:
backend = zeus.callbacks.SaveProgressCallback(
filename=backend_filename, ncheck=1)
n_run_eff = n_burn + n_run
else:
backend = None
n_run_eff = n_burn + n_run
pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)
sampler = zeus.EnsembleSampler(nwalkers=n_walkers,
ndim=num_param,
logprob_fn=self.chain.logL,
moves=moves,
tune=tune,
tolerance=tolerance,
patience=patience,
maxsteps=maxsteps,
mu=mu,
maxiter=maxiter,
pool=pool,
vectorize=vectorize,
blobs_dtype=blobs_dtype,
verbose=verbose,
check_walkers=check_walkers,
shuffle_ensemble=shuffle_ensemble,
light_mode=light_mode)
sampler.run_mcmc(initpos,
n_run_eff,
progress=progress,
callbacks=backend)
flat_samples = sampler.get_chain(flat=True, thin=1, discard=n_burn)
dist = sampler.get_log_prob(flat=True, thin=1, discard=n_burn)
return flat_samples, dist