class Sampler:
"""
wrapper of emcee.EnsembleSampler.
"""
def __init__(self, lnpost, p0, nwalkers=120, blobs_dtype=float):
"""
init
"""
self.lnpost = lnpost
blobs_dtype = blobs_dtype # Note: Here dtype must be specified, otherwise an error happens. #[("lnlike",float),]
self.sampler = EnsembleSampler(
nwalkers, p0.shape[1], lnpost, blobs_dtype=blobs_dtype
) # NOTE: dtype must be list of tuple (not tuple of tuple)
self.p0 = p0
self.p_last = p0
self.ndim = p0.shape[1]
def reset_sampler(self):
self.sampler.reset()
def sample(self, n_sample, burnin=False, use_pool=False):
"""
execute mcmc for given iteration steps.
"""
desc = "burnin" if burnin else "sample"
with Pool() as pool:
self.sampler.pool = pool if use_pool else None
iteration = tqdm(self.sampler.sample(self.p_last,
iterations=n_sample),
total=n_sample,
desc=desc)
for _ret in iteration:
self.p_last = _ret.coords # if uses_emcee3 else _ret[0] # for emcee2
lnposts = _ret.log_prob # if uses_emcee3 else _ret[1] # for emcee2
iteration.set_postfix(lnpost_min=np.min(lnposts),
lnpost_max=np.max(lnposts),
lnpost_mean=np.mean(lnposts))
if burnin:
self.reset_sampler()
def get_chain(self, **kwargs):
return self.sampler.get_chain(**kwargs)
def get_log_prob(self, **kwargs):
return self.sampler.get_log_prob(**kwargs)
def get_blobs(self, **kwargs):
return self.sampler.get_blobs(**kwargs)
def get_last_sample(self, **kwargs):
return self.sampler.get_last_sample(**kwargs)
def _save(self, fname_base):
np.save(fname_base + "_chain.npy", self.get_chain())
np.save(fname_base + "_lnprob.npy", self.get_log_prob())
np.save(fname_base + "_lnlike.npy", self.get_blobs())
def save(self, fname_base):
'''
Save MCMC results into "<fname_base>_chain/lnprob/lnlike.npy".
If fname_base is like "your_directory/your_prefix", create "your_directory" before saving.
'''
dirname = os.path.dirname(fname_base)
if dirname == "":
self._save(fname_base)
else:
if not os.path.isdir(dirname): os.mkdir(dirname)
self._save(fname_base)
def save_pickle(self, fname_base, overwrite=False):
fname = fname_base + '_.gz'
if os.path.exists(fname):
if overwrite:
warn(f"{fname} exsits already. It will be overwritten.")
else:
raise RuntimeError(
f"{fname} exsits already. If you want to overwrite it, set \"overwrite=True\"."
)
data = pickle.dumps(self)
with gzip.open(fname, mode='wb') as fp:
fp.write(data)
