def auto_correlations(self):
times = zeus.AutoCorrTime(samples=self.zeus_sampler.get_chain())
try:
previous_auto_correlation_times = zeus.AutoCorrTime(
samples=self.zeus_sampler.get_chain()
[:-self.auto_correlation_settings.check_size, :, :], )
except IndexError:
previous_auto_correlation_times = None
return AutoCorrelations(
check_size=self.auto_correlation_settings.check_size,
required_length=self.auto_correlation_settings.required_length,
change_threshold=self.auto_correlation_settings.change_threshold,
times=times,
previous_times=previous_auto_correlation_times,
)
def auto_correlations(self):
times = zeus.AutoCorrTime(samples=self.results_internal.get_chain())
try:
previous_auto_correlation_times = zeus.AutoCorrTime(
samples=self.results_internal.get_chain()
[:-self.auto_correlation_settings.check_size, :, :], )
except IndexError:
logger.debug("Unable to compute previous auto correlation times.")
previous_auto_correlation_times = None
return AutoCorrelations(
check_size=self.auto_correlation_settings.check_size,
required_length=self.auto_correlation_settings.required_length,
change_threshold=self.auto_correlation_settings.change_threshold,
times=times,
previous_times=previous_auto_correlation_times,
)
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.")