diff --git a/bilby/core/sampler/ptemcee.py b/bilby/core/sampler/ptemcee.py index 15edd1edc0adf734996d2be08796c62d530efb03..f314563e5037b16563454e6c2918d4153cfe2955 100644 --- a/bilby/core/sampler/ptemcee.py +++ b/bilby/core/sampler/ptemcee.py @@ -49,7 +49,7 @@ class Ptemcee(MCMCSampler): resume=True, nsamples=5000, burn_in_nact=50, thin_by_nact=1, autocorr_c=5, safety=1, frac_threshold=0.01, autocorr_tol=50, min_tau=1, check_point_deltaT=600, - threads=1, **kwargs): + threads=1, exit_code=77, store_walkers=False, **kwargs): super(Ptemcee, self).__init__( likelihood=likelihood, priors=priors, outdir=outdir, label=label, use_ratio=use_ratio, plot=plot, @@ -71,8 +71,10 @@ class Ptemcee(MCMCSampler): self.check_point_deltaT = check_point_deltaT self.threads = threads + self.store_walkers = store_walkers self.resume_file = "{}/{}_checkpoint_resume.pickle".format(self.outdir, self.label) + self.exit_code = exit_code @property def sampler_function_kwargs(self): @@ -256,7 +258,6 @@ class Ptemcee(MCMCSampler): # Get 0-likelihood samples and store in the result samples = sampler.chain[0, :, :, :] # nwalkers, nsteps, ndim - self.result.walkers = samples[:, :sampler.time:, :] self.result.samples = ( samples[:, self.nburn: sampler.time:self.thin, :].reshape((-1, self.ndim))) loglikelihood = sampler.loglikelihood[ @@ -264,7 +265,8 @@ class Ptemcee(MCMCSampler): ] # nwalkers, nsteps self.result.log_likelihood_evaluations = loglikelihood.reshape((-1)) - self.result.walkers = self.sampler.chain + if self.store_walkers: + self.result.walkers = self.sampler.chain self.result.nburn = self.nburn log_evidence, log_evidence_err = compute_evidence( @@ -283,7 +285,7 @@ class Ptemcee(MCMCSampler): self.write_current_state(plot=False) logger.warning("Closing pool") self.pool.close() - sys.exit(77) + sys.exit(self.exit_code) def write_current_state(self, plot=True): checkpoint(self.outdir, self.label, self.nsamples_effective,