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"])
def body_of_test(dimension=1, n_modes=1, info_sampler=empty_dict, tmpdir="",
packages_path=None, skip_not_installed=False):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
# Info of likelihood and prior
ranges = np.array([[-1, 1] for _ in range(dimension)])
while True:
info = info_random_gaussian_mixture(
ranges=ranges, n_modes=n_modes, input_params_prefix="a_",
O_std_min=O_std_min, O_std_max=O_std_max, derived=True)
if n_modes == 1:
break
means = info["likelihood"]["gaussian_mixture"]["means"]
distances = chain(*[[np.linalg.norm(m1 - m2) for m2 in means[i + 1:]]
for i, m1 in enumerate(means)])
if min(distances) >= distance_factor * O_std_max:
break
if rank == 0:
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[kinds.sampler] = info_sampler
sampler_name = list(info_sampler)[0]
if sampler_name == "mcmc":
if "covmat" in info_sampler["mcmc"]:
info[kinds.sampler]["mcmc"]["covmat_params"] = (
list(info["params"])[:dimension])
info[_debug] = False
info[_debug_file] = None
# TODO: this looks weird/bug:?
info[_output_prefix] = getattr(tmpdir, "realpath()", lambda: tmpdir)()
if packages_path:
info[_packages_path] = process_packages_path(packages_path)
# Delay to one chain to check that MPI communication of the sampler is non-blocking
# if rank == 1:
# info["likelihood"]["gaussian_mixture"]["delay"] = 0.1
updated_info, sampler = install_test_wrapper(skip_not_installed, run, info)
products = sampler.products()
# Done! --> Tests
if rank == 0:
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!
try:
if is_travis():
raise ValueError
import getdist.plots as gdplots
from getdist.gaussian_mixtures import MixtureND
sampled_params = [
p for p, v in info["params"].items() if partag.prior not in v]
mixture = MixtureND(
info[kinds.likelihood]["gaussian_mixture"]["means"],
info[kinds.likelihood]["gaussian_mixture"]["covs"],
names=sampled_params, label="truth")
g = gdplots.getSubplotPlotter()
to_plot = [mixture, results]
if clusters:
to_plot = 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(), results.getMeans()
KL_final = KL_norm(m1=info[kinds.likelihood]["gaussian_mixture"]["means"][0],
S1=info[kinds.likelihood]["gaussian_mixture"]["covs"][0],
m2=mean_sample[:dimension],
S2=cov_sample[:dimension, :dimension])
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[kinds.likelihood]["gaussian_mixture"]["means"][i_c1],
S1=info[kinds.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"])