def test_choose_pool(self):
import schwimmbad
pool = choose_pool(mpi=False, processes=1, use_dill=True)
assert pool.is_master() is True
assert isinstance(pool, schwimmbad.serial.SerialPool)
pool = choose_pool(mpi=False, processes=2, use_dill=True)
assert pool.is_master() is True
assert isinstance(pool, MultiPool)
def optimize(self, n_particles=50, n_iterations=250, verbose=False, threadCount=1):
"""
:param n_particles: number of PSO particles, will be ignored if self._particle_swarm is False
:param n_iterations: number of PSO iterations, will be ignored if self._particle_swarm is False
:param verbose: whether to print stuff
:param threadCount: integer; number of threads in multi-threading mode
:return: keyword arguments that map (x_image, y_image) to the same source coordinate (source_x, source_y)
"""
if self._particle_swarm:
if threadCount > 1:
pool = choose_pool(mpi=False, processes=threadCount)
else:
pool = None
kwargs = self._fit_pso(n_particles, n_iterations, pool, verbose)
else:
kwargs = self._param_class.kwargs_lens
kwargs_lens_final, source_penalty = self._fit_amoeba(kwargs, verbose)
args_lens_final = self._param_class.kwargs_to_args(kwargs_lens_final)
source_x_array, source_y_array = self.fast_rayshooting.ray_shooting_fast(args_lens_final)
source_x, source_y = np.mean(source_x_array), np.mean(source_y_array)
if verbose:
print('optimization done.')
print('Recovered source position: ', (source_x_array, source_y_array))
return kwargs_lens_final, [source_x, source_y]
def pso(self, n_particles, n_iterations, lower_start=None, upper_start=None,
threadCount=1, init_pos=None, mpi=False, print_key='PSO'):
"""
Return the best fit for the lens model on catalogue basis with
particle swarm optimizer.
:param n_particles: number of particles in the sampling process
:param n_iterations: number of iterations of the swarm
:param lower_start: numpy array, lower end parameter of the values of the starting particles
:param upper_start: numpy array, upper end parameter of the values of the starting particles
:param threadCount: number of threads in the computation (only applied if mpi=False)
:param init_pos: numpy array, position of the initial best guess model
:param mpi: bool, if True, makes instance of MPIPool to allow for MPI execution
:param print_key: string, prints the process name in the progress bar (optional)
:return: kwargs_result (of best fit), [lnlikelihood of samples, positions of samples, velocity of sampels)
"""
if lower_start is None or upper_start is None:
lower_start, upper_start = np.array(self.lower_limit), np.array(self.upper_limit)
print("PSO initialises its particles with default values")
else:
lower_start = np.maximum(lower_start, self.lower_limit)
upper_start = np.minimum(upper_start, self.upper_limit)
pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)
if mpi is True and pool.is_master():
print('MPI option chosen for PSO.')
pso = ParticleSwarmOptimizer(self.chain.logL,
lower_start, upper_start, n_particles,
pool=pool)
if init_pos is None:
init_pos = (upper_start - lower_start) / 2 + lower_start
pso.set_global_best(init_pos, [0]*len(init_pos),
self.chain.logL(init_pos))
if pool.is_master():
print('Computing the %s ...' % print_key)
time_start = time.time()
result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations)
if pool.is_master():
kwargs_return = self.chain.param.args2kwargs(result)
print(pso.global_best.fitness * 2 / (max(
self.chain.effective_num_data_points(**kwargs_return), 1)), 'reduced X^2 of best position')
print(pso.global_best.fitness, 'logL')
print(self.chain.effective_num_data_points(**kwargs_return), 'effective number of data points')
print(kwargs_return.get('kwargs_lens', None), 'lens result')
print(kwargs_return.get('kwargs_source', None), 'source result')
print(kwargs_return.get('kwargs_lens_light', None), 'lens light result')
print(kwargs_return.get('kwargs_ps', None), 'point source result')
print(kwargs_return.get('kwargs_special', None), 'special param result')
time_end = time.time()
print(time_end - time_start, 'time used for ', print_key)
print('===================')
return result, [chi2_list, pos_list, vel_list]
def pso(self, n_particles=10, n_iterations=10, lowerLimit=-0.2, upperLimit=0.2, threadCount=1, mpi=False,
print_key='default'):
"""
returns the best fit for the lense model on catalogue basis with particle swarm optimizer
"""
init_pos = self.chain.get_args(self.chain.kwargs_data_init)
num_param = self.chain.num_param
lowerLimit = [lowerLimit] * num_param
upperLimit = [upperLimit] * num_param
pool = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)
pso = ParticleSwarmOptimizer(self.chain, lowerLimit, upperLimit,
n_particles, pool=pool)
if init_pos is not None:
pso.set_global_best(init_pos, [0]*len(init_pos),
self.chain.likelihood(init_pos))
if pool.is_master():
print('Computing the %s ...' % print_key)
time_start = time.time()
result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations)
kwargs_data = self.chain.update_data(result)
if pool.is_master():
time_end = time.time()
print("Shifts found: ", result)
print(time_end - time_start, 'time used for ', print_key)
return kwargs_data, [chi2_list, pos_list, vel_list]
def nautilus_sampling(self,
prior_type='uniform',
mpi=False,
thread_count=1,
verbose=True,
one_step=False,
**kwargs_nautilus):
"""
:param prior_type: string; prior type. Currently only 'uniform' supported
(in addition to Prior class in Likelihood module)
:param mpi: MPI option (currently not supported)
:param thread_count: integer; multi-threading option (currently not supported)
:param verbose: verbose statements of Nautilus
:param one_step: boolean, if True, only runs one iteration of filling the sampler and re-training.
This is meant for test purposes of the sampler to operate with little computational effort
:param kwargs_nautilus: additional keyword arguments for Nautilus
:return: points, log_w, log_l, log_z
"""
from nautilus import Prior, Sampler
prior = Prior()
# TODO better prior integration with Nautilus
if prior_type == 'uniform':
for i in range(self._num_param):
prior.add_parameter(dist=(self._lower_limit[i],
self._upper_limit[i]))
# assert self._num_param == prior.dimensionality()
# print(self._num_param, prior.dimensionality(), 'number of param, dimensionality')
else:
raise ValueError(
'prior_type %s is not supported for Nautilus wrapper.' %
prior_type)
# loop through prior
pool = choose_pool(mpi=mpi, processes=thread_count, use_dill=True)
sampler = Sampler(prior,
likelihood=self.likelihood,
pool=pool,
**kwargs_nautilus)
time_start = time.time()
if one_step is True:
sampler.add_bound()
sampler.fill_bound()
else:
sampler.run(verbose=verbose)
points, log_w, log_l = sampler.posterior()
log_z = sampler.evidence()
time_end = time.time()
if pool.is_master():
print(time_end - time_start, 'time taken for MCMC sampling')
return points, log_w, log_l, log_z
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(mean_start, sigma_start, n_walkers, dist='normal')
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
