def test_std(logp=logp,seed=42):
np.random.seed(seed)
ndim = np.random.randint(2,5)
nwalkers = 2 * ndim
nsteps = np.random.randint(3000,5000)
sampler = zeus.EnsembleSampler(nwalkers,ndim,logp,verbose=False)
start = np.random.rand(nwalkers,ndim)
sampler.run_mcmc(start,nsteps)
assert np.all(np.abs(np.std(sampler.get_chain(flat=True),axis=0)-1.0) < 0.1)
def __init__(
self,
fit,
nwalkers=50,
processes=1,
initspread=0.01,
moves=None,
verbose=True,
sampler='emcee',
):
self.fit = fit
self.nwalkers = nwalkers
self.numpars = fit.pars.numFree()
self.initspread = initspread
self.verbose = verbose
self.sampler_mode = sampler
mcmcpars = fit.pars.copy()
def likefunc(parvals):
mcmcpars.setFree(parvals)
like = Likelihood(fit.images, fit.model, mcmcpars)
return like.total
pool = None if processes <= 1 else _MultiProcessPool(
likefunc, processes)
# for doing the mcmc sampling
if sampler == 'emcee':
if emcee is None:
raise RuntimeError('emcee module not installed')
self.sampler = emcee.EnsembleSampler(nwalkers,
self.numpars,
likefunc,
pool=pool,
moves=moves)
elif sampler == 'zeus':
if zeus is None:
raise RuntimeError('zeus module not installed')
self.sampler = zeus.EnsembleSampler(
nwalkers,
self.numpars,
likefunc,
pool=pool,
)
else:
raise RuntimeError('Unknown sampler')
# starting point
self.pos0 = None
# header items to write to output file
self.header = {
'burn': 0,
}
def __init__(self, name, nwalkers, ndim, logprob, pool):
self.name = name
self.nwalkers = nwalkers
self.ndim = ndim
if name == 'zeus':
import zeus
self.sampler = zeus.EnsembleSampler(nwalkers,
ndim,
logprob,
pool=pool,
verbose=False)
elif name == 'light':
import zeus
self.sampler = zeus.EnsembleSampler(nwalkers,
ndim,
logprob,
pool=pool,
verbose=False,
light_mode=True)
elif name == 'emcee':
import emcee
#self.state = emcee.State
self.sampler = emcee.EnsembleSampler(nwalkers,
ndim,
logprob,
pool=pool)
elif name == 'demc':
import emcee
self.sampler = emcee.EnsembleSampler(
nwalkers,
ndim,
logprob,
moves=[(emcee.moves.DEMove(), 0.9),
(emcee.moves.DESnookerMove(), 0.1)],
pool=pool)
def test_ncall(seed=42):
np.random.seed(seed)
def loglike(theta):
assert len(theta) == 5
a = theta[:-1]
b = theta[1:]
loglike.ncalls += 1
return -2 * (100 * (b - a**2)**2 + (1 - a)**2).sum()
loglike.ncalls = 0
ndim = 5
nsteps = 100
nwalkers = 2 * ndim
sampler = zeus.EnsembleSampler(nwalkers,ndim,loglike,verbose=False)
start = np.random.rand(nwalkers,ndim)
sampler.run_mcmc(start,nsteps)
assert loglike.ncalls == sampler.ncall + nwalkers
def fit(self, log_posterior, start, num_dim, prior_transform, save_dims=None, uid=None):
import zeus
filename = self.get_filename(uid)
if os.path.exists(filename):
self.logger.info("Not sampling, returning result from file.")
return self.load_file(filename)
if self.num_walkers is None:
self.num_walkers = num_dim * 4
self.logger.debug("Fitting framework with %d dimensions" % num_dim)
if save_dims is None:
save_dims = num_dim
self.logger.debug("Fitting framework with %d dimensions" % num_dim)
self.logger.info("Using Zeus Sampler")
callbacks = []
if self.autoconverge:
# Default convergence criteria from Zeus docos. Seem reasonable.
cb0 = zeus.callbacks.AutocorrelationCallback(ncheck=50, dact=0.01, nact=50, discard=0.5)
cb1 = zeus.callbacks.SplitRCallback(ncheck=50, epsilon=0.01, nsplits=2, discard=0.5)
cb2 = zeus.callbacks.MinIterCallback(nmin=50)
callbacks = [cb0, cb1, cb2]
pos = start(num_walkers=self.num_walkers)
self.logger.info("Sampling posterior now")
sampler = zeus.EnsembleSampler(self.num_walkers, num_dim, log_posterior)
sampler.run_mcmc(pos, self.num_steps, callbacks=callbacks)
self.logger.debug("Fit finished")
tau = zeus.AutoCorrTime(sampler.get_chain(discard=0.5))
burnin = int(2 * np.max(tau))
samples = sampler.get_chain(discard=burnin, flat=True).T
likelihood = sampler.get_log_prob(discard=burnin, flat=True)
self._save(samples, likelihood, filename, save_dims)
return {"chain": samples, "weights": np.ones(len(likelihood)), "posterior": likelihood}
def _fit(self,
model: AbstractPriorModel,
analysis,
log_likelihood_cap=None):
"""
Fit a model using Zeus and the Analysis class which contains the data and returns the log likelihood from
instances of the model, which the `NonLinearSearch` seeks to maximize.
Parameters
----------
model : ModelMapper
The model which generates instances for different points in parameter space.
analysis : Analysis
Contains the data and the log likelihood function which fits an instance of the model to the data, returning
the log likelihood the `NonLinearSearch` maximizes.
Returns
-------
A result object comprising the Samples object that inclues the maximum log likelihood instance and full
chains used by the fit.
"""
pool = self.make_pool()
fitness_function = self.fitness_function_from_model_and_analysis(
model=model, analysis=analysis)
if self.paths.is_object("zeus"):
zeus_sampler = self.zeus_pickled
zeus_state = zeus_sampler.get_last_sample()
initial_log_posterior_list = zeus_sampler.get_last_log_prob()
samples = self.samples_from(model=model)
total_iterations = zeus_sampler.iteration
if samples.converged:
iterations_remaining = 0
else:
iterations_remaining = self.config_dict_run[
"nsteps"] - total_iterations
logger.info(
"Existing Zeus samples found, resuming non-linear search.")
else:
zeus_sampler = zeus.EnsembleSampler(
nwalkers=self.config_dict_search["nwalkers"],
ndim=model.prior_count,
logprob_fn=fitness_function.__call__,
pool=pool,
)
zeus_sampler.ncall_total = 0
initial_unit_parameter_lists, initial_parameter_lists, initial_log_posterior_list = self.initializer.initial_samples_from_model(
total_points=zeus_sampler.nwalkers,
model=model,
fitness_function=fitness_function,
)
zeus_state = np.zeros(shape=(zeus_sampler.nwalkers,
model.prior_count))
logger.info(
"No Zeus samples found, beginning new non-linear search.")
for index, parameters in enumerate(initial_parameter_lists):
zeus_state[index, :] = np.asarray(parameters)
total_iterations = 0
iterations_remaining = self.config_dict_run["nsteps"]
while iterations_remaining > 0:
if self.iterations_per_update > iterations_remaining:
iterations = iterations_remaining
else:
iterations = self.iterations_per_update
for sample in zeus_sampler.sample(
start=zeus_state,
log_prob0=initial_log_posterior_list,
iterations=iterations,
progress=True,
):
pass
zeus_sampler.ncall_total += zeus_sampler.ncall
self.paths.save_object("zeus", zeus_sampler)
zeus_state = zeus_sampler.get_last_sample()
initial_log_posterior_list = zeus_sampler.get_last_log_prob()
total_iterations += iterations
iterations_remaining = self.config_dict_run[
"nsteps"] - total_iterations
samples = self.perform_update(model=model,
analysis=analysis,
during_analysis=True)
if self.auto_correlations_settings.check_for_convergence:
if zeus_sampler.iteration > self.auto_correlations_settings.check_size:
if samples.converged:
iterations_remaining = 0
auto_correlation_time = zeus.AutoCorrTime(
samples=zeus_sampler.get_chain())
discard = int(3.0 * np.max(auto_correlation_time))
thin = int(np.max(auto_correlation_time) / 2.0)
chain = zeus_sampler.get_chain(discard=discard,
thin=thin,
flat=True)
if "maxcall" in self.kwargs:
if zeus_sampler.ncall_total > self.kwargs["maxcall"]:
iterations_remaining = 0
logger.info("Zeus sampling complete.")
lg = Posterior(modelw)
for fNL in [-999, -500, -100, -10, 0, 10, 100, 500, 999]:
print(fNL, lg.logpost(fNL, y, invcov, x))
np.random.seed(42)
ndim = 1 # Number of parameters/dimensions (e.g. m and c)
nwalkers = 10 # Number of walkers to use. It should be at least twice the number of dimensions.
nsteps = 1000 # Number of steps/iterations.
start = 10. * np.random.randn(nwalkers,
ndim) # Initial positions of the walkers.
print(f'initial guess: {start}')
sampler = zeus.EnsembleSampler(nwalkers,
ndim,
lg.logpost,
args=[y, invcov, x],
maxiter=10000)
sampler.run_mcmc(start, nsteps) # Run sampling
sampler.summary # Print summary diagnostics
# flatten the chains, thin them by a factor of 15,
# and remove the burn-in (first half of the chain)
chain = sampler.get_chain(flat=True, discard=20, thin=5)
np.save(f'chains_{kind}_{case}.npy', chain) #
def multiprocessing_zeus():
'''
'''
#
fsps_emulator = Models.DESIspeculator()
# set prior
priors = Infer.load_priors([
Infer.UniformPrior(10., 10.5, label='sed'),
Infer.FlatDirichletPrior(4, label='sed'),
Infer.UniformPrior(np.array([6.9e-5, 6.9e-5, 0., 0., -2.2]), np.array([7.3e-3, 7.3e-3, 3., 4., 0.4]), label='sed')
])
random_theta = priors.sample()
wave, flux = fsps_emulator.sed(priors.transform(random_theta), 0.1)
desi_mcmc = Infer.desiMCMC(prior=priors)
t0 = time.time()
mcmc = desi_mcmc.run(
wave_obs=wave[0],
flux_obs=flux[0],
flux_ivar_obs=np.ones(flux.shape[1]),
zred=0.1,
sampler='zeus',
nwalkers=20,
burnin=10,
opt_maxiter=1000,
niter=100,
pool=None,
debug=True)
print()
print('running on series takes %.f' % (time.time() - t0))
print()
import zeus
import multiprocessing
ncpu = multiprocessing.cpu_count()
print('%i cpus' % ncpu)
t0 = time.time()
lnpost_args, lnpost_kwargs = desi_mcmc._lnPost_args_kwargs(
wave_obs=wave[0],
flux_obs=flux[0],
flux_ivar_obs=np.ones(flux.shape[1]),
zred=0.1)
start = desi_mcmc._initialize_walkers(lnpost_args, lnpost_kwargs, priors,
nwalkers=20, opt_maxiter=1000, debug=True)
print('--- burn-in ---')
pewl = Pool(processes=ncpu)
with pewl as pool:
zeus_sampler = zeus.EnsembleSampler(
desi_mcmc.nwalkers,
desi_mcmc.prior.ndim,
desi_mcmc.lnPost,
pool=pool,
args=lnpost_args,
kwargs=lnpost_kwargs)
zeus_sampler.run_mcmc(start, 10)
burnin = zeus_sampler.get_chain()
print('--- running main MCMC ---')
pewl = Pool(processes=ncpu)
with pewl as pool:
zeus_sampler = zeus.EnsembleSampler(
desi_mcmc.nwalkers,
desi_mcmc.prior.ndim,
desi_mcmc.lnPost,
pool=pool,
args=lnpost_args,
kwargs=lnpost_kwargs)
zeus_sampler.run_mcmc(burnin[-1], 100)
_chain = zeus_sampler.get_chain()
print()
print('running on parallel takes %.f' % (time.time() - t0))
print()
return None
def fit(self,
data,
s,
nsteps=100,
outbase="output/test_",
seed=42,
resume=True):
self.nsteps = nsteps
if self.backend == 'zeus':
param_init = []
for i in range(self.nparams):
if self.prior_types[i] == 'flat':
min, max = self.prior_params[i]
param_init.append((max - min) *
np.random.random(self.nwalkers) + min)
elif self.prior_types[i] == 'gauss':
mu, sigma = self.prior_params[i]
param_init.append(
np.random.normal(loc=mu,
scale=sigma,
size=self.nwalkers))
else:
raise NotImplementedError
start = np.array(param_init).T
if not os.path.isfile(outbase + "zeus_chains.txt") or ~resume:
np.random.seed(seed)
with Pool() as pool:
sampler = zeus.EnsembleSampler(
self.nwalkers,
self.nparams,
self.logpost,
args=[data, s],
vectorize=False) # Initialise the sampler
sampler.run_mcmc(start, nsteps) # Run sampling
sampler.summary # Print summary diagnostics
return sampler
else:
print(f"==> Chain file found, skipping the fit.")
elif self.backend == 'multinest':
print(
f"==> Using PyMultinest. Assuming params passed are flat priors."
)
def prior_pmn(cube, ndim, nparams):
for i in range(self.nparams):
min, max = self.prior_params[i]
cube[i] = cube[i] * (max - min) + min
def loglike_pmn(cube, ndim, nparams):
alpha, B, Snl = cube[0], cube[1], cube[2]
return self.loglike((alpha, B, Snl), data, s)
pmn.run(loglike_pmn, prior_pmn, self.nparams, outputfiles_basename=outbase, resume=resume, \
verbose=True, n_live_points=self.live_points, evidence_tolerance=self.tolerance)
else:
raise NotImplementedError
cmax = 10. # upper range of prior
cini = np.random.uniform(cmin, cmax, Nens) # initial c points
inisamples = np.array([mini, cini]).T # initial samples
ndims = inisamples.shape[1] # number of parameters/dimensions
Nburnin = 500 # number of burn-in samples
Nsamples = 500 # number of final posterior samples
# set additional args for the posterior (the data, the noise std. dev., and the abscissa)
argslist = (data, sigma, x)
# set up the sampler
sampler = zeus.EnsembleSampler(Nens, ndims, logposterior, args=argslist)
# pass the initial samples and total number of samples required
sampler.run_mcmc(inisamples, Nsamples + Nburnin)
# extract the samples (removing the burn-in)
postsamples = sampler.get_chain(flat=True, discrd=Nburnin)
# plot posterior samples (if corner.py is installed)
try:
import matplotlib as mpl
mpl.use("Agg") # force Matplotlib backend to Agg
import corner # import corner.py
except ImportError:
sys.exit(1)
'''The natural logarithm of the likelihood.'''
# unpack the model parameters
fnl = theta
# evaluate the model
md = model(x, fnl=fnl, has_fnl=True, has_rsd=True)
# return the log likelihood
return -0.5 * (y - md).dot(invcov.dot(y - md))
def logpost(theta, y, invcov, x):
'''The natural logarithm of the posterior.'''
return logprior(theta) + loglike(theta, y, invcov, x)
ndim = 1 # Number of parameters/dimensions (e.g. m and c)
nwalkers = 10 # Number of walkers to use. It should be at least twice the number of dimensions.
nsteps = 1000 # Number of steps/iterations.
start = 0.01 * np.random.randn(nwalkers,
ndim) # Initial positions of the walkers.
print(f'initial guess: {start}')
sampler = zeus.EnsembleSampler(nwalkers, ndim, logpost,
args=[y, invcov, x]) # Initialise the sampler
sampler.run_mcmc(start, nsteps) # Run sampling
#sampler.summary # Print summary diagnostics
# flatten the chains, thin them by a factor of 10, and remove the burn-in (first half of the chain)
chain = sampler.get_chain(flat=True, discard=100, thin=10)
np.save('chains.npy', chain)
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