def check_convergence_and_learn_proposal(self):
"""
Checks the convergence of the sampling process, and, if requested,
learns a new covariance matrix for the proposal distribution from the covariance
of the last samples.
"""
# Compute Rminus1 of means
self.been_waiting = 0
if more_than_one_process():
# Compute and gather means and covs
use_first = int(self.n() / 2)
mean = self.collection.mean(first=use_first)
cov = self.collection.cov(first=use_first)
acceptance_rate = self.get_acceptance_rate(use_first)
Ns, means, covs, acceptance_rates = mpi.array_gather(
[self.n(), mean, cov, acceptance_rate])
else:
# Compute and gather means, covs and CL intervals of last m-1 chain fractions
m = 1 + self.Rminus1_single_split
cut = int(len(self.collection) / m)
try:
acceptance_rate = self.get_acceptance_rate(cut)
Ns = np.ones(m - 1) * cut
means = np.array([
self.collection.mean(first=i * cut, last=(i + 1) * cut - 1)
for i in range(1, m)
])
covs = np.array([
self.collection.cov(first=i * cut, last=(i + 1) * cut - 1)
for i in range(1, m)
])
except always_stop_exceptions:
raise
except Exception:
self.log.info(
"Not enough points in chain to check convergence. "
"Waiting for next checkpoint.")
return
acceptance_rates = None
if is_main_process():
self.progress.at[self.i_learn, "N"] = sum(Ns)
self.progress.at[self.i_learn, "timestamp"] = \
datetime.datetime.now().isoformat()
acceptance_rate = (np.average(acceptance_rates, weights=Ns)
if acceptance_rates is not None else
acceptance_rate)
self.log.info(
" - Acceptance rate: %.3f" +
(" = avg(%r)" % list(acceptance_rates)
if acceptance_rates is not None else ""), acceptance_rate)
self.progress.at[self.i_learn, "acceptance_rate"] = acceptance_rate
# "Within" or "W" term -- our "units" for assessing convergence
# and our prospective new covariance matrix
mean_of_covs = np.average(covs, weights=Ns, axis=0)
# "Between" or "B" term
# We don't weight with the number of samples in the chains here:
# shorter chains will likely be outliers, and we want to notice them
cov_of_means = np.atleast_2d(np.cov(means.T)) # , fweights=Ns)
# For numerical stability, we turn mean_of_covs into correlation matrix:
# rho = (diag(Sigma))^(-1/2) * Sigma * (diag(Sigma))^(-1/2)
# and apply the same transformation to the mean of covs (same eigenvals!)
d = np.sqrt(np.diag(cov_of_means))
corr_of_means = (cov_of_means / d).T / d
norm_mean_of_covs = (mean_of_covs / d).T / d
success_means = False
converged_means = False
# Cholesky of (normalized) mean of covs and eigvals of Linv*cov_of_means*L
try:
L = np.linalg.cholesky(norm_mean_of_covs)
except np.linalg.LinAlgError:
self.log.warning(
"Negative covariance eigenvectors. "
"This may mean that the covariance of the samples does not "
"contain enough information at this point. "
"Skipping learning a new covmat for now.")
else:
Linv = np.linalg.inv(L)
try:
eigvals = np.linalg.eigvalsh(
Linv.dot(corr_of_means).dot(Linv.T))
success_means = True
except np.linalg.LinAlgError:
self.log.warning("Could not compute eigenvalues. "
"Skipping learning a new covmat for now.")
else:
Rminus1 = max(np.abs(eigvals))
self.progress.at[self.i_learn, "Rminus1"] = Rminus1
# For real square matrices, a possible def of the cond number is:
condition_number = Rminus1 / min(np.abs(eigvals))
self.log.debug(" - Condition number = %g",
condition_number)
self.log.debug(" - Eigenvalues = %r", eigvals)
self.log.info(
" - Convergence of means: R-1 = %f after %d accepted steps"
% (Rminus1, sum(Ns)) +
(" = sum(%r)" %
list(Ns) if more_than_one_process() else ""))
# Have we converged in means?
# (criterion must be fulfilled twice in a row)
converged_means = max(
Rminus1, self.Rminus1_last) < self.Rminus1_stop
else:
mean_of_covs = None
success_means = None
converged_means = False
Rminus1 = None
success_means, converged_means = mpi.share(
(success_means, converged_means))
# Check the convergence of the bounds of the confidence intervals
# Same as R-1, but with the rms deviation from the mean bound
# in units of the mean standard deviation of the chains
if converged_means:
if more_than_one_process():
mcsamples = self.collection.sampled_to_getdist_mcsamples(
first=use_first)
try:
bound = np.array([[
mcsamples.confidence(i,
limfrac=self.Rminus1_cl_level /
2.,
upper=which)
for i in range(self.model.prior.d())
] for which in [False, True]]).T
success_bounds = True
except:
bound = None
success_bounds = False
bounds = np.array(mpi.gather(bound))
else:
try:
mcsamples_list = [
self.collection.sampled_to_getdist_mcsamples(
first=i * cut, last=(i + 1) * cut - 1)
for i in range(1, m)
]
except always_stop_exceptions:
raise
except:
self.log.info(
"Not enough points in chain to check c.l. convergence. "
"Waiting for next checkpoint.")
return
try:
bounds = [
np.array([[
mcs.confidence(i,
limfrac=self.Rminus1_cl_level / 2.,
upper=which)
for i in range(self.model.prior.d())
] for which in [False, True]]).T
for mcs in mcsamples_list
]
success_bounds = True
except:
bounds = None
success_bounds = False
if is_main_process():
if success_bounds:
Rminus1_cl = (np.std(bounds, axis=0).T /
np.sqrt(np.diag(mean_of_covs)))
self.log.debug(" - normalized std's of bounds = %r",
Rminus1_cl)
Rminus1_cl = np.max(Rminus1_cl)
self.progress.at[self.i_learn, "Rminus1_cl"] = Rminus1_cl
self.log.info(
" - Convergence of bounds: R-1 = %f after %d " %
(Rminus1_cl,
(sum(Ns) if more_than_one_process() else self.n())) +
"accepted steps" +
(" = sum(%r)" %
list(Ns) if more_than_one_process() else ""))
if Rminus1_cl < self.Rminus1_cl_stop:
self.converged = True
self.log.info("The run has converged!")
self._Ns = Ns
else:
self.log.info(
"Computation of the bounds was not possible. "
"Waiting until the next converge check.")
# Broadcast and save the convergence status and the last R-1 of means
if success_means:
self.Rminus1_last, self.converged = mpi.share((
Rminus1, self.converged) if is_main_process() else None)
# Do we want to learn a better proposal pdf?
if self.learn_proposal and not self.converged:
good_Rminus1 = (self.learn_proposal_Rminus1_max >
self.Rminus1_last >
self.learn_proposal_Rminus1_min)
if not good_Rminus1:
self.mpi_info(
"Convergence less than requested for updates: "
"waiting until the next convergence check.")
return
mean_of_covs = mpi.share(mean_of_covs)
try:
self.proposer.set_covariance(mean_of_covs)
self.mpi_info(
" - Updated covariance matrix of proposal pdf.")
self.mpi_debug("%r", mean_of_covs)
except:
self.mpi_debug(
"Updating covariance matrix failed unexpectedly. "
"waiting until next covmat learning attempt.")
# Save checkpoint info
self.write_checkpoint()
def test_post_likelihood():
"""
Swaps likelihood "gaussian" for "target".
It also tests aggregated chi2's by removing and adding a likelihood to an existing
type.
"""
# Generate original chain
orig_interval = OutputOptions.output_inteveral_s
try:
OutputOptions.output_inteveral_s = 0
info_params_local = deepcopy(info_params)
info_params_local["dummy"] = 0
dummy_loglike_add = 0.1
dummy_loglike_remove = 0.01
info = {
"output": None,
"force": True,
"params": info_params_local,
"sampler": info_sampler,
"likelihood": {
"gaussian": {
"external": sampled_pdf,
"type": "A"
},
"dummy": {
"external": lambda dummy: 1,
"type": "BB"
},
"dummy_remove": {
"external": lambda dummy: dummy_loglike_add,
"type": "BB"
}
}
}
info_out, sampler = run(info)
samples_in = mpi.gather(sampler.products()["sample"])
if mpi.is_main_process():
mcsamples_in = MCSamplesFromCobaya(info_out, samples_in)
else:
mcsamples_in = None
info_out.update({
"post": {
"suffix": "foo",
"remove": {
"likelihood": {
"gaussian": None,
"dummy_remove": None
}
},
"add": {
"likelihood": {
"target": {
"external": target_pdf,
"type": "A",
"output_params": ["cprime"]
},
"dummy_add": {
"external": lambda dummy: dummy_loglike_remove,
"type": "BB"
}
}
}
}
})
info_post_out, products_post = post(info_out,
sampler.products()["sample"])
samples = mpi.gather(products_post["sample"])
# Load with GetDist and compare
if mcsamples_in:
target_mean, target_cov = mpi.share(_get_targets(mcsamples_in))
mcsamples = MCSamplesFromCobaya(info_post_out,
samples,
name_tag="sample")
new_mean = mcsamples.mean(["a", "b"])
new_cov = mcsamples.getCovMat().matrix
mpi.share((new_mean, new_cov))
else:
target_mean, target_cov = mpi.share()
new_mean, new_cov = mpi.share()
assert np.allclose(new_mean, target_mean)
assert np.allclose(new_cov, target_cov)
assert allclose(products_post["sample"]["chi2__A"],
products_post["sample"]["chi2__target"])
assert allclose(
products_post["sample"]["chi2__BB"],
products_post["sample"]["chi2__dummy"] +
products_post["sample"]["chi2__dummy_add"])
finally:
OutputOptions.output_inteveral_s = orig_interval
def run(self):
"""
Runs the sampler.
"""
self.mpi_info("Sampling!" + (
" (NB: no accepted step will be saved until %d burn-in samples " %
self.burn_in.value +
"have been obtained)" if self.burn_in.value else ""))
self.n_steps_raw = 0
last_output: float = 0
last_n = self.n()
state_check_every = 1
with mpi.ProcessState(self) as state:
while last_n < self.max_samples and not self.converged:
self.get_new_sample()
self.n_steps_raw += 1
if self.output_every.unit:
# if output_every in sec, print some info
# and dump at fixed time intervals
now = datetime.datetime.now()
now_sec = now.timestamp()
if now_sec >= last_output + self.output_every.value:
self.do_output(now)
last_output = now_sec
state.check_error()
if self.current_point.weight == 1:
# have added new point
# Callback function
n = self.n()
if n != last_n:
# and actually added
last_n = n
if (self.callback_function
and not (max(n, 1) % self.callback_every.value)
and self.current_point.weight == 1):
self.callback_function_callable(self)
self.last_point_callback = len(self.collection)
if more_than_one_process():
# Checking convergence and (optionally) learning
# the covmat of the proposal
if self.check_ready() and state.set(
mpi.State.READY):
self.log.info(self._msg_ready +
" (waiting for the rest...)")
if state.all_ready():
self.mpi_info("All chains are r%s",
self._msg_ready[1:])
self.check_convergence_and_learn_proposal()
self.i_learn += 1
else:
if self.check_ready():
self.log.debug(self._msg_ready)
self.check_convergence_and_learn_proposal()
self.i_learn += 1
elif self.current_point.weight % state_check_every == 0:
state.check_error()
# more frequent checks near beginning
state_check_every = min(10, state_check_every + 1)
if last_n == self.max_samples:
self.log.info(
"Reached maximum number of accepted steps allowed (%s). "
"Stopping.", self.max_samples)
# Write the last batch of samples ( < output_every (not in sec))
self.collection.out_update()
ns = mpi.gather(self.n())
self.mpi_info("Sampling complete after %d accepted steps.", sum(ns))
def body_of_sampler_test(info_sampler: SamplersDict,
dimension=1,
n_modes=1,
tmpdir="",
packages_path=None,
skip_not_installed=False,
fixed=False,
random_state=None):
# Info of likelihood and prior
ranges = np.array([[-1, 1] for _ in range(dimension)])
if fixed:
info = fixed_info.copy()
else:
info = generate_random_info(n_modes, ranges, random_state=random_state)
if mpi.is_main_process():
print("Original mean of the gaussian mode:")
print(info["likelihood"]["gaussian_mixture"]["means"])
print("Original covmat of the gaussian mode:")
print(info["likelihood"]["gaussian_mixture"]["covs"])
info["sampler"] = info_sampler
sampler_name = list(info_sampler)[0]
if random_state is not None:
info_sampler[sampler_name]["seed"] = random_state.integers(0, 2**31)
if sampler_name == "mcmc":
if "covmat" in info_sampler["mcmc"]:
info["sampler"]["mcmc"]["covmat_params"] = (list(
info["params"])[:dimension])
info["debug"] = False
info["debug_file"] = None
info["output"] = os.path.join(tmpdir, 'out_chain')
if packages_path:
info["packages_path"] = process_packages_path(packages_path)
updated_info, sampler = install_test_wrapper(skip_not_installed, run, info)
products = sampler.products()
products["sample"] = mpi.gather(products["sample"])
# Done! --> Tests
if mpi.is_main_process():
if sampler_name == "mcmc":
ignore_rows = 0.5
else:
ignore_rows = 0
results = MCSamplesFromCobaya(updated_info,
products["sample"],
ignore_rows=ignore_rows,
name_tag="sample")
clusters = None
if "clusters" in products:
clusters = [
MCSamplesFromCobaya(updated_info,
products["clusters"][i]["sample"],
name_tag="cluster %d" % (i + 1))
for i in products["clusters"]
]
# Plots!
if not is_travis():
try:
import getdist.plots as gdplots
from getdist.gaussian_mixtures import MixtureND
sampled_params = [
p for p, v in info["params"].items() if "prior" not in v
]
mixture = MixtureND(
info["likelihood"]["gaussian_mixture"]["means"],
info["likelihood"]["gaussian_mixture"]["covs"],
names=sampled_params,
label="truth")
g = gdplots.getSubplotPlotter()
to_plot = [mixture, results]
if clusters:
to_plot += clusters
g.triangle_plot(to_plot, params=sampled_params)
g.export("test.png")
except:
print("Plotting failed!")
# 1st test: KL divergence
if n_modes == 1:
cov_sample, mean_sample = results.getCov(
dimension), results.getMeans()
KL_final = KL_norm(
m1=info["likelihood"]["gaussian_mixture"]["means"][0],
S1=info["likelihood"]["gaussian_mixture"]["covs"][0],
m2=mean_sample[:dimension],
S2=cov_sample)
print("Final KL: ", KL_final)
assert KL_final <= KL_tolerance
# 2nd test: clusters
else:
if "clusters" in products:
assert len(products["clusters"]) >= n_modes, (
"Not all clusters detected!")
for i, c2 in enumerate(clusters):
cov_c2, mean_c2 = c2.getCov(), c2.getMeans()
KLs = [
KL_norm(m1=info["likelihood"]["gaussian_mixture"]
["means"][i_c1],
S1=info["likelihood"]["gaussian_mixture"]
["covs"][i_c1],
m2=mean_c2[:dimension],
S2=cov_c2[:dimension, :dimension])
for i_c1 in range(n_modes)
]
extra_tol = 4 * n_modes if n_modes > 1 else 1
print("Final KL for cluster %d: %g", i, min(KLs))
assert min(KLs) <= KL_tolerance * extra_tol
else:
assert 0, "Could not check sample convergence: multimodal but no clusters"
# 3rd test: Evidence
if "logZ" in products:
logZprior = sum(np.log(ranges[:, 1] - ranges[:, 0]))
assert (products["logZ"] - logZ_nsigmas * products["logZstd"] <
-logZprior <
products["logZ"] + logZ_nsigmas * products["logZstd"])