def test_parameterization():
updated_info, sampler = run(info)
products = sampler.products()
sample = products["sample"]
from getdist.mcsamples import MCSamplesFromCobaya
gdsample = MCSamplesFromCobaya(updated_info, products["sample"])
for i, point in sample:
a = info[_params]["a"]
b = get_external_function(info[_params]["b"])(a, point["bprime"])
c = get_external_function(info[_params]["c"])(a, point["cprime"])
e = get_external_function(e_func)(b)
f = get_external_function(f_func)(b)
g = get_external_function(info[_params]["g"]["derived"])(x_func(
point["c"]))
h = get_external_function(info[_params]["h"])(info[_params]["i"])
j = get_external_function(info[_params]["j"])(b)
k = get_external_function(info[_params]["k"]["derived"])(f)
assert np.allclose(point[["b", "c", "e", "f", "g", "h", "j", "k"]],
[b, c, e, f, g, h, j, k])
# Test for GetDist too (except fixed ones, ignored by GetDist)
bcefffg_getdist = [
gdsample.samples[i][gdsample.paramNames.list().index(p)]
for p in ["b", "c", "e", "f", "g", "j", "k"]
]
assert np.allclose(bcefffg_getdist, [b, c, e, f, g, j, k])
def test_mcmc_drag_results():
info: InputDict = yaml_load(yaml_drag)
info['likelihood'] = {
'g1': {
'external': GaussLike
},
'g2': {
'external': GaussLike2
}
}
updated_info, sampler = run(info)
products = sampler.products()
from getdist.mcsamples import MCSamplesFromCobaya
products["sample"] = mpi.allgather(products["sample"])
gdample = MCSamplesFromCobaya(updated_info,
products["sample"],
ignore_rows=0.2)
assert abs(gdample.mean('a') - 0.2) < 0.03
assert abs(gdample.mean('b')) < 0.03
assert abs(gdample.std('a') - 0.293) < 0.03
assert abs(gdample.std('b') - 0.4) < 0.03
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
# Export the results to GetDist
from getdist.mcsamples import MCSamplesFromCobaya
gd_sample = MCSamplesFromCobaya(updated_info, products["sample"])
# Analyze and plot
mean = gd_sample.getMeans()[:2]
covmat = gd_sample.getCovMat().matrix[:2, :2]
print("Mean:")
print(mean)
print("Covariance matrix:")
print(covmat)
# %matplotlib inline # uncomment if running from the Jupyter notebook
import getdist.plots as gdplt
gdplot = gdplt.getSubplotPlotter()
gdplot.triangle_plot(gd_sample, ["a", "b"], filled=True)
except LoggedError as err:
pass
# Di it work? (e.g. did not get stuck)
success = all(comm.allgather(success))
if not success and rank == 0:
print("Sampling failed!")
# Now gather chains of all MPI processes in rank 0
all_chains = comm.gather(mcmc.products()["sample"], root=0)
# Pass all of them to GetDist in rank = 0
if rank == 0:
from getdist.mcsamples import MCSamplesFromCobaya
gd_sample = MCSamplesFromCobaya(upd_info, all_chains)
# Manually concatenate them in rank = 0 for some custom manipulation,
# skipping 1st 3rd of each chain
copy_and_skip_1st_3rd = lambda chain: chain[int(len(chain) / 3):]
if rank == 0:
full_chain = copy_and_skip_1st_3rd(all_chains[0])
for chain in all_chains[1:]:
full_chain.append(copy_and_skip_1st_3rd(chain))
# The combined chain is now `full_chain`
import matplotlib.pyplot as plt
plt.figure()
plt.hist(full_chain["x"], weights=full_chain["weight"])
plt.show()
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"])
chi2 = compute_chi2_theory(Db, Db_th, inv_cov, alpha)
return -0.5 * chi2
info = {
"params": {"alpha": {"prior": {"min": -5, "max": 5}, "latex": r"\alpha"}},
"sampler": {
"mcmc": {
"max_tries": 10 ** 8,
"Rminus1_stop": 0.01,
"Rminus1_cl_stop": 0.08,
}
},
}
if chi2_with_respect_to_theory:
info["likelihood"] = {"chi2": compute_loglike_theory}
else:
info["likelihood"] = {"chi2": compute_loglike_data}
updated_info, sampler = run(info)
sample = MCSamplesFromCobaya(updated_info, sampler.products()["sample"])
runs[spec_name] = {
"progress": sampler.products()["progress"],
"sample": sample,
"std_fisher": compute_sigma_alpha(Db, inv_cov),
}
spec_id += 1
pickle.dump(runs, open("results.pkl", "wb"))
def test_cosmo_run_resume_post(tmpdir, packages_path=None):
# only vary As, so fast chain. Chain does not need to converge (tested elsewhere).
info['output'] = os.path.join(tmpdir, 'testchain')
if packages_path:
info["packages_path"] = process_packages_path(packages_path)
run(info, force=True)
# note that continuing from files leads to text-file precision at read in, so a mix of
# precision in the output SampleCollection returned from run
run(info, resume=True, override={'sampler': {'mcmc': {'Rminus1_stop': 0.2}}})
updated_info, sampler = run(info['output'] + '.updated' + Extension.dill,
resume=True,
override={'sampler': {'mcmc': {'Rminus1_stop': 0.05}}})
results = mpi.allgather(sampler.products()["sample"])
samp = MCSamplesFromCobaya(updated_info, results, ignore_rows=0.2)
assert np.isclose(samp.mean('As100'), 100 * samp.mean('As'))
# post-processing
info_post: PostDict = {'add': {'params': {'h': None},
"likelihood": {"test_likelihood2": likelihood2}},
'remove': {'likelihood': ["test_likelihood"]},
'suffix': 'testpost',
'skip': 0.2, 'thin': 4
}
output_info, products = run(updated_info, override={'post': info_post}, force=True)
results2 = mpi.allgather(products["sample"])
samp2 = MCSamplesFromCobaya(output_info, results2)
samp_test = samp.copy()
samp_test.weighted_thin(4)
sigma8 = samp_test.getParams().sigma8
samp_test.reweightAddingLogLikes(-(sigma8 - 0.7) ** 2 / 0.1 ** 2
+ (sigma8 - 0.75) ** 2 / 0.07 ** 2)
assert np.isclose(samp_test.mean('sigma8'), samp2.mean('sigma8'))
# from getdist-format chain files
root = os.path.join(tmpdir, 'getdist_format')
if mpi.is_main_process():
samp.saveChainsAsText(root)
mpi.sync_processes()
from_txt = dict(updated_info, output=root)
post_from_text = dict(info_post, skip=0) # getdist already skipped
output_info, products = run(from_txt, override={'post': post_from_text}, force=True)
samp_getdist = MCSamplesFromCobaya(output_info, mpi.allgather(products["sample"]))
assert not products["stats"]["points_removed"]
assert samp2.numrows == samp_getdist.numrows
assert np.isclose(samp2.mean('sigma8'), samp_getdist.mean('sigma8'))
# again with inferred-inputs for params
info_conv = cosmomc_root_to_cobaya_info_dict(root)
# have to manually add consistent likelihoods if re-computing
info_conv['likelihood'] = info['likelihood']
info_conv['theory'] = info['theory']
post_from_text = dict(info_post, skip=0, suffix='getdist2') # getdist already skipped
output_info, products = run(info_conv, override={'post': post_from_text},
output=False)
samp_getdist2 = MCSamplesFromCobaya(output_info, mpi.allgather(products["sample"]))
assert np.isclose(samp2.mean('sigma8'), samp_getdist2.mean('sigma8'))
# from save info, no output
info_post['output'] = None
output_info, products = run({'output': info['output'], 'post': info_post}, force=True)
results3 = mpi.allgather(products["sample"])
samp3 = MCSamplesFromCobaya(output_info, results3)
assert np.isclose(samp3.mean("sigma8"), samp2.mean("sigma8"))
assert np.isclose(samp3.mean("joint"), samp2.mean("joint"))
samps4 = loadMCSamples(info['output'] + '.post.testpost')
assert np.isclose(samp3.mean("joint"), samps4.mean("joint"))
# test recover original answer swapping likelihoods back
info_revert = {'add': {'likelihood': info['likelihood']},
'remove': {'likelihood': ["test_likelihood2"]},
'suffix': 'revert',
'skip': 0, 'thin': 1,
'output': None
}
output_info, products = run({'output': info['output'] + '.post.testpost',
'post': info_revert}, force=True)
results_revert = mpi.allgather(products["sample"])
samp_revert = MCSamplesFromCobaya(output_info, results_revert)
samp_thin = MCSamplesFromCobaya(updated_info, results, ignore_rows=0.2)
samp_thin.weighted_thin(4)
assert samp_thin.numrows == samp_revert.numrows + products["stats"]["points_removed"]
if not products["stats"]["points_removed"]:
assert np.isclose(samp_revert.mean("sigma8"), samp_thin.mean("sigma8"))
else:
assert abs(samp_revert.mean("sigma8") - samp_thin.mean("sigma8")) < 0.01
assert not products["stats"]["points_removed"]
# no remove
info_post = {
'add': {'params': {'h': None}, "likelihood": {"test_likelihood2": likelihood2}},
'suffix': 'test2', 'skip': 0.2, 'thin': 4}
output_info, products = run(updated_info, override={'post': info_post}, force=True)
results2 = mpi.allgather(products["sample"])
samp2 = MCSamplesFromCobaya(output_info, results2)
assert "chi2__type1" in samp2.paramNames.list()
# check what has been saved to disk is consistent
samps4 = loadMCSamples(updated_info['output'] + '.post.test2')
assert samp2.paramNames.list() == samps4.paramNames.list()
assert np.isclose(samp2.mean("sigma8"), samps4.mean("sigma8"))
# adding new theory derived
info_post['add']['theory'] = {'new_param_theory': BTheory}
output_info, products = run(updated_info, override={'post': info_post}, output=False)
results3 = mpi.allgather(products["sample"])
samp3 = MCSamplesFromCobaya(output_info, results3)
assert np.isclose(samp3.mean("sigma8"), samp2.mean("sigma8"))
assert np.isclose(samp3.mean("As1000"), samp2.mean("As") * 1000)
info_post['add']['theory'] = {'new_param_theory': CTheory}
with pytest.raises(LoggedError) as e, NoLogging(logging.ERROR):
run(updated_info, override={'post': info_post}, output=False)
assert 'Parameter AsX no known value' in str(e)