def test_minimize_gaussian():
# parameters
dimension = 3
n_modes = 1
# MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
# Info of likelihood and prior
ranges = np.array([[0, 1] for _ in range(dimension)])
prefix = "a_"
info = info_random_gaussian_mixture(ranges=ranges,
n_modes=n_modes,
input_params_prefix=prefix,
derived=True)
mean = info[kinds.likelihood]["gaussian_mixture"]["means"][0]
cov = info[kinds.likelihood]["gaussian_mixture"]["covs"][0]
maxloglik = multivariate_normal.logpdf(mean, mean=mean, cov=cov)
if rank == 0:
print("Maximum of the gaussian mode to be found:")
print(mean)
info[kinds.sampler] = {"minimize": {"ignore_prior": True}}
info["debug"] = False
info["debug_file"] = None
# info["output_prefix"] = "./minigauss/"
from cobaya.run import run
updated_info, sampler = run(info)
products = sampler.products()
# Done! --> Tests
if rank == 0:
rel_error = abs(maxloglik -
-products["minimum"]["minuslogpost"]) / abs(maxloglik)
assert rel_error < 0.001
def generate_random_info(n_modes, ranges, random_state):
while True:
inf = 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,
random_state=random_state)
if n_modes == 1:
break
means = inf["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
return inf
def body_of_test(dimension=1,
n_modes=1,
info_sampler={},
tmpdir="",
modules=None):
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,
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[_sampler] = info_sampler
if list(info_sampler.keys())[0] == "mcmc":
if "covmat" in info_sampler["mcmc"]:
info[_sampler]["mcmc"]["covmat_params"] = (list(
info["params"].keys())[:dimension])
info[_debug] = False
info[_debug_file] = None
info[_output_prefix] = getattr(tmpdir, "realpath()", lambda: tmpdir)()
if modules:
info[_path_install] = process_modules_path(modules)
# 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, products = run(info)
# Done! --> Tests
if rank == 0:
if list(info_sampler.keys())[0] == "mcmc":
ignore_rows = 0.5
else:
ignore_rows = 0
results = loadCobayaSamples(updated_info,
products["sample"],
ignore_rows=ignore_rows,
name_tag="sample")
clusters = None
if "clusters" in products:
clusters = [
loadCobayaSamples(updated_info,
products["clusters"][i]["sample"],
name_tag="cluster %d" % (i + 1))
for i in products["clusters"]
]
# Plots!
try:
import getdist.plots as gdplots
from getdist.gaussian_mixtures import MixtureND
mixture = MixtureND(
info[_likelihood]["gaussian_mixture"]["means"],
info[_likelihood]["gaussian_mixture"]["covs"],
names=[p for p in info[_params] if "deriv" not in p],
label="truth")
g = gdplots.getSubplotPlotter()
to_plot = [mixture, results]
if clusters:
to_plot = to_plot + clusters
g.triangle_plot(to_plot, )
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[_likelihood]["gaussian_mixture"]["means"][0],
S1=info[_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"].keys()) >= n_modes, (
"Not all clusters detected!")
for c2 in 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
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_speeds(info_sampler={}, modules=None):
# Generate 2 3-dimensional gaussians
dim = 3
speed1, speed2 = 5, 20
ranges = [[i, i + 1] for i in range(2 * dim)]
prefix = "a_"
mean1, cov1 = [
info_random_gaussian_mixture(
ranges=[ranges[i] for i in range(dim)],
n_modes=1,
prefix=prefix,
O_std_min=0.01,
O_std_max=0.2,
derived=True)[_likelihood]["gaussian_mixture"][p][0]
for p in ["means", "covs"]
]
mean2, cov2 = [
info_random_gaussian_mixture(
ranges=[ranges[i] for i in range(dim, 2 * dim)],
n_modes=1,
prefix=prefix,
O_std_min=0.01,
O_std_max=0.2,
derived=True)[_likelihood]["gaussian_mixture"][p][0]
for p in ["means", "covs"]
]
global n1, n2
n1, n2 = 0, 0
# PolyChord measures its own speeds, so we need to "sleep"
sleep_unit = 1 / 50
sampler = list(info_sampler.keys())[0]
def like1(a_0, a_1, a_2, _derived=["sum_like1"]):
if sampler == "polychord":
sleep(1 / speed1 * sleep_unit)
global n1
n1 += 1
if _derived is not None:
_derived["sum_like1"] = a_0 + a_1 + a_2
return multivariate_normal.logpdf([a_0, a_1, a_2],
mean=mean1,
cov=cov1)
def like2(a_3, a_4, a_5, _derived=["sum_like2"]):
if sampler == "polychord":
sleep(1 / speed2 * sleep_unit)
global n2
n2 += 1
if _derived is not None:
_derived["sum_like2"] = a_3 + a_4 + a_5
return multivariate_normal.logpdf([a_3, a_4, a_5],
mean=mean2,
cov=cov2)
# Rearrange parameter in arbitrary order
perm = list(range(2 * dim))
shuffle(perm)
# Create info
info = {
"params":
odict([[
prefix + "%d" % i, {
"prior": dict(zip(["min", "max"], ranges[i]))
}
] for i in perm] + [["sum_like1", None], ["sum_like2", None]]),
"likelihood": {
"like1": {
"external": like1,
"speed": speed1
},
"like2": {
"external": like2,
"speed": speed2
}
}
}
info["sampler"] = info_sampler
print("Parameter order:", list(info["params"]))
# info["debug"] = True
info["modules"] = modules
# Adjust number of samples
n_cycles_all_params = 10
if sampler == "mcmc":
info["sampler"][sampler]["burn_in"] = 0
info["sampler"][sampler][
"max_samples"] = n_cycles_all_params * 10 * dim
# Force mixing of blocks:
info["sampler"][sampler]["covmat_params"] = list(info["params"])
info["sampler"][sampler]["covmat"] = 1 / 10000 * np.eye(
len(info["params"]))
i_1st, i_2nd = map(
lambda x: info["sampler"][sampler]["covmat_params"].index(x),
[prefix + "0", prefix + "%d" % dim])
info["sampler"][sampler]["covmat"][i_1st, i_2nd] = 1 / 100000
info["sampler"][sampler]["covmat"][i_2nd, i_1st] = 1 / 100000
elif sampler == "polychord":
info["sampler"][sampler]["nlive"] = 2 * dim
info["sampler"][sampler]["max_ndead"] = n_cycles_all_params * dim
else:
assert 0, "Unknown sampler for this test."
updated_info, products = run(info)
# Done! --> Tests
if sampler == "polychord":
tolerance = 0.2
assert abs((n2 - n1) / (n1) / (speed2 / speed1) - 1) < tolerance, (
"#evaluations off: %g > %g" %
(abs((n2 - n1) / (n1) / (speed2 / speed1) - 1), tolerance))
# For MCMC tests, notice that there is a certain tolerance to be allowed for,
# since for every proposed step the BlockedProposer cycles once, but the likelihood
# may is not evaluated if the proposed point falls outside the prior bounds
elif sampler == "mcmc" and info["sampler"][sampler].get("drag"):
assert abs((n2 - n1) / (n1) / (speed2 / speed1) - 1) < 0.1
elif sampler == "mcmc" and info["sampler"][sampler].get("oversample"):
# Testing oversampling: number of evaluations per param * oversampling factor
assert abs((n2 - n1) * dim / (n1 * dim) / (speed2 / speed1) - 1) < 0.1
elif sampler == "mcmc":
# Testing just correct blocking: same number of evaluations per param
assert abs((n2 - n1) * dim / (n1 * dim) - 1) < 0.1
# Finally, test some points of the chain to reproduce the correct likes and derived
# These are not AssertionError's to override the flakyness of the test
for _ in range(10):
i = choice(list(range(products["sample"].n())))
chi2_1_chain = -0.5 * products["sample"]["chi2__like1"][i]
chi2_1_good = like1(
_derived=None,
**{p: products["sample"][p][i]
for p in ["a_0", "a_1", "a_2"]})
chi2_2_chain = -0.5 * products["sample"]["chi2__like2"][i]
chi2_2_good = like2(
_derived=None,
**{p: products["sample"][p][i]
for p in ["a_3", "a_4", "a_5"]})
if not np.allclose([chi2_1_chain, chi2_2_chain],
[chi2_1_good, chi2_2_good]):
raise ValueError(
"Likelihoods not reproduced correctly. "
"Chain has %r but should be %r. " %
([chi2_1_chain, chi2_2_chain], [chi2_1_good, chi2_2_good]) +
"Full chain point: %r" % products["sample"][i])
derived_chain = products["sample"][["sum_like1",
"sum_like2"]].values[i]
derived_good = np.array([
sum(products["sample"][["a_0", "a_1", "a_2"]].values[i]),
sum(products["sample"][["a_3", "a_4", "a_5"]].values[i])
])
if not np.allclose(derived_chain, derived_good):
raise ValueError("Derived params not reproduced correctly. "
"Chain has %r but should be %r. " %
(derived_chain, derived_good) +
"Full chain point:\n%r" % products["sample"][i])
print(products["sample"])
def body_of_test_speeds(info_sampler,
manual_blocking=False,
packages_path=None,
skip_not_installed=False,
random_state=None):
# #dimensions and speed ratio mutually prime (e.g. 2,3,5)
dim0, dim1 = 5, 2
speed0, speed1 = 1, 10
ranges = [[i, i + 1] for i in range(dim0 + dim1)]
prefix = "a_"
params0, params1 = (lambda x: (x[:dim0], x[dim0:]))(
[prefix + str(d) for d in range(dim0 + dim1)])
derived0, derived1 = "sum_like0", "sum_like1"
random_state = np.random.default_rng(random_state)
mean0, cov0 = [
info_random_gaussian_mixture(
ranges=[ranges[i] for i in range(dim0)],
n_modes=1,
input_params_prefix=prefix,
O_std_min=0.01,
O_std_max=0.2,
derived=True,
random_state=random_state)["likelihood"]["gaussian_mixture"][p][0]
for p in ["means", "covs"]
]
mean1, cov1 = [
info_random_gaussian_mixture(
ranges=[ranges[i] for i in range(dim0, dim0 + dim1)],
n_modes=1,
input_params_prefix=prefix,
O_std_min=0.01,
O_std_max=0.2,
derived=True,
random_state=random_state)["likelihood"]["gaussian_mixture"][p][0]
for p in ["means", "covs"]
]
n_evals = [0, 0]
def like0(**kwargs):
n_evals[0] += 1
input_params = [kwargs[p] for p in params0]
derived = {derived0: sum(input_params)}
return multivariate_normal.logpdf(input_params, mean=mean0,
cov=cov0), derived
def like1(**kwargs):
n_evals[1] += 1
input_params = [kwargs[p] for p in params1]
derived = {derived1: sum(input_params)}
return multivariate_normal.logpdf(input_params, mean=mean1,
cov=cov1), derived
# Rearrange parameter in arbitrary order
perm = list(range(dim0 + dim1))
random_state.shuffle(perm)
# Create info
info = {
"params":
dict(
{
prefix + "%d" % i: {
"prior": dict(zip(["min", "max"], ranges[i]))
}
for i in perm
},
sum_like0=None,
sum_like1=None),
"likelihood": {
"like0": {
"external": like0,
"speed": speed0,
"input_params": params0,
"output_params": derived0
},
"like1": {
"external": like1,
"speed": speed1,
"input_params": params1,
"output_params": derived1
}
},
"sampler":
info_sampler
}
sampler_name = list(info_sampler)[0]
info_sampler[sampler_name]["seed"] = random_state.integers(0, 2**31)
if manual_blocking:
over0, over1 = speed0, speed1
info["sampler"][sampler_name]["blocking"] = [[over0, params0],
[over1, params1]]
print("Parameter order:", list(info["params"]))
# info["debug"] = True
info["packages_path"] = packages_path
# Adjust number of samples
n_cycles_all_params = 10
info_sampler = info["sampler"][sampler_name]
if sampler_name == "mcmc":
info_sampler["measure_speeds"] = False
info_sampler["burn_in"] = 0
info_sampler["max_samples"] = \
info_sampler.get("max_samples", n_cycles_all_params * 10 * (dim0 + dim1))
# Force mixing of blocks:
info_sampler["covmat_params"] = list(info["params"])
info_sampler["covmat"] = 1 / 10000 * np.eye(len(info["params"]))
i_0th, i_1st = map(lambda x: info_sampler["covmat_params"].index(x),
[prefix + "0", prefix + "%d" % dim0])
info_sampler["covmat"][i_0th, i_1st] = 1 / 100000
info_sampler["covmat"][i_1st, i_0th] = 1 / 100000
# info_sampler["learn_proposal"] = False
elif sampler_name == "polychord":
info_sampler["nlive"] = dim0 + dim1
info_sampler["max_ndead"] = n_cycles_all_params * (dim0 + dim1)
else:
assert False, "Unknown sampler for this test."
updated_info, sampler = install_test_wrapper(skip_not_installed, run, info)
products = sampler.products()
# TEST: same (steps block i / speed i / dim i) (steps block 1 = evals[1] - evals[0])
def test_func(_n_evals, _dim0, _speed0, _dim1, _speed1):
return (abs((_n_evals[1] - _n_evals[0]) / _speed1 / _dim1 /
(_n_evals[0] / _speed0 / _dim0)) - 1)
# Tolerance accounting for random starts of the proposers (PolyChord and MCMC) and for
# steps outside prior bounds, where likelihoods are not evaluated (MCMC only)
tolerance = 0.1
if sampler_name == "polychord":
assert test_func(n_evals, dim0, speed0, dim1, speed1) <= tolerance, (
("%g > %g" %
(test_func(n_evals, dim0, speed0, dim1, speed1), tolerance)))
elif sampler_name == "mcmc" and info["sampler"][sampler_name].get("drag"):
assert test_func(
n_evals, dim0, speed0, dim1, 2 * speed1) <= tolerance, (
("%g > %g" %
(test_func(n_evals, dim0, speed0, dim1, speed1), tolerance)))
elif sampler_name == "mcmc" and (
info["sampler"][sampler_name].get("oversample")
or info["sampler"][sampler_name].get("oversample_power", 0) > 0):
assert test_func(n_evals, dim0, speed0, dim1, speed1) <= tolerance, (
("%g > %g" %
(test_func(n_evals, dim0, speed0, dim1, speed1), tolerance)))
elif sampler_name == "mcmc": # just blocking
assert test_func(n_evals, dim0, speed0, dim1, speed1) <= tolerance, (
("%g > %g" %
(test_func(n_evals, dim0, speed0, dim1, speed1), tolerance)))
else:
raise ValueError("This should not happen!")
# Finally, test some points of the chain to reproduce the correct likes and derived
# These are not AssertionError's to override the flakiness of the test
for _ in range(10):
i = random_state.choice(list(range(len(products["sample"]))))
chi2_0_chain = -0.5 * products["sample"]["chi2__like0"][i]
chi2_0_good = like0(**{p: products["sample"][p][i]
for p in params0})[0]
chi2_1_chain = -0.5 * products["sample"]["chi2__like1"][i]
chi2_1_good = like1(**{p: products["sample"][p][i]
for p in params1})[0]
if not np.allclose([chi2_0_chain, chi2_1_chain],
[chi2_0_good, chi2_1_good]):
raise ValueError(
"Likelihoods not reproduced correctly. "
"Chain has %r but should be %r. " %
([chi2_0_chain, chi2_1_chain], [chi2_0_good, chi2_1_good]) +
"Full chain point: %r" % products["sample"][i])
derived_chain = products["sample"][["sum_like0",
"sum_like1"]].values[i]
derived_good = np.array([
sum(products["sample"][params0].values[i]),
sum(products["sample"][params1].values[i])
])
if not np.allclose(derived_chain, derived_good):
raise ValueError("Derived params not reproduced correctly. "
"Chain has %r but should be %r. " %
(derived_chain, derived_good) +
"Full chain point:\n%r" % products["sample"][i])
print(products["sample"])
