def test_mcmc(tmpdir, modules=None):
dimension = 3
# Random initial proposal
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank == 0:
S0 = random_cov(dimension * [[0, 1]], n_modes=1, O_std_min=0.01, O_std_max=0.5)
else:
S0 = None
S0 = comm.bcast(S0, root=0)
info_sampler = {"mcmc": {
# Bad guess for covmat, so big burn in and max_tries
"max_tries": 1000 * dimension, "burn_in": 100 * dimension,
# Learn proposal
# "learn_proposal": True, # default now!
# Proposal
"covmat": S0, }}
def check_gaussian(sampler_instance):
KL_proposer = KL_norm(S1=sampler_instance.model.likelihood["gaussian_mixture"].covs[0],
S2=sampler_instance.proposer.get_covariance())
KL_sample = KL_norm(m1=sampler_instance.model.likelihood["gaussian_mixture"].means[0],
S1=sampler_instance.model.likelihood["gaussian_mixture"].covs[0],
m2=sampler_instance.collection.mean(
first=int(sampler_instance.n() / 2)),
S2=sampler_instance.collection.cov(
first=int(sampler_instance.n() / 2)))
print("KL proposer: %g ; KL sample: %g" % (KL_proposer, KL_sample))
if rank == 0:
info_sampler["mcmc"].update({
# Callback to check KL divergence -- disabled in the automatic test
"callback_function": check_gaussian, "callback_every": 100})
body_of_test(
dimension=dimension, n_modes=1, info_sampler=info_sampler, tmpdir=str(tmpdir))
def body_of_test(dim, tmpdir=None):
mindim = 4
assert dim > mindim, "Needs dimension>%d for the test." % mindim
initial_random_covmat = random_cov(dim * [[0, 1]])
i_s = list(range(dim))
shuffle(i_s)
n_altered = int(dim / 4)
i_proposal = i_s[:n_altered]
i_ref = i_s[n_altered:2 * n_altered]
i_prior = i_s[2 * n_altered:3 * n_altered]
removed = list(chain(*(i_proposal, i_ref, i_prior)))
i_covmat = [i for i in range(dim) if i not in removed]
for i in removed:
diag = initial_random_covmat[i, i]
initial_random_covmat[:, i] = 0
initial_random_covmat[i, :] = 0
initial_random_covmat[i, i] = diag
# Prepare info, including refs, priors and reduced covmat
prefix = "a_"
input_order = list(range(dim))
shuffle(input_order)
info = {_likelihood: {"one": None}, _params: odict()}
for i in input_order:
p = prefix + str(i)
info[_params][p] = {_prior: {_p_dist: "norm", "loc": 0, "scale": 1000}}
sigma = np.sqrt(initial_random_covmat[i, i])
if i in i_proposal:
info[_params][p][_p_proposal] = sigma
elif i in i_ref:
info[_params][prefix + str(i)][_p_ref] = {
_p_dist: "norm",
"scale": sigma
}
elif i in i_prior:
info[_params][prefix + str(i)][_prior]["scale"] = sigma
reduced_covmat = initial_random_covmat[np.ix_(i_covmat, i_covmat)]
reduced_covmat_params = [prefix + str(i) for i in i_covmat]
info[_sampler] = {"mcmc": {}}
if tmpdir:
filename = os.path.join(str(tmpdir), "mycovmat.dat")
header = " ".join(reduced_covmat_params)
np.savetxt(filename, reduced_covmat, header=header)
info[_sampler]["mcmc"]["covmat"] = str(filename)
else:
info[_sampler]["mcmc"]["covmat_params"] = reduced_covmat_params
info[_sampler]["mcmc"]["covmat"] = reduced_covmat
to_compare = initial_random_covmat[np.ix_(input_order, input_order)]
def callback(sampler):
assert np.allclose(to_compare, sampler.proposer.get_covariance())
info[_sampler]["mcmc"].update({
"callback_function": callback,
"callback_every": 1,
"max_samples": 1,
"burn_in": 0
})
updated_info, products = run(info)
def body_of_test(dim, tmpdir=None, random_state=None):
mindim = 4
assert dim > mindim, "Needs dimension>%d for the test." % mindim
if mpi.is_main_process():
random_state = np.random.default_rng(random_state)
i_s = list(range(dim))
random_state.shuffle(i_s)
initial_random_covmat = random_cov(dim * [[0, 1]],
random_state=random_state)
mpi.share((i_s, initial_random_covmat))
else:
i_s, initial_random_covmat = mpi.share()
n_altered = int(dim / 4)
i_proposal = i_s[:n_altered]
i_ref = i_s[n_altered:2 * n_altered]
i_prior = i_s[2 * n_altered:3 * n_altered]
removed = list(chain(*(i_proposal, i_ref, i_prior)))
i_covmat = [i for i in range(dim) if i not in removed]
for i in removed:
diag = initial_random_covmat[i, i]
initial_random_covmat[:, i] = 0
initial_random_covmat[i, :] = 0
initial_random_covmat[i, i] = diag
# Prepare info, including refs, priors and reduced covmat
prefix = "a_"
if mpi.is_main_process():
input_order = list(range(dim))
random_state.shuffle(input_order)
else:
input_order = None
input_order = mpi.share(input_order)
info: InputDict = {"likelihood": {"one": None}, "params": {}}
for i in input_order:
p = prefix + str(i)
info["params"][p] = {
"prior": {
"dist": "norm",
"loc": 0,
"scale": 1000
}
}
sigma = np.sqrt(initial_random_covmat[i, i])
if i in i_proposal:
info["params"][p]["proposal"] = sigma
elif i in i_ref:
info["params"][prefix + str(i)]["ref"] = {
"dist": "norm",
"scale": sigma
}
elif i in i_prior:
info["params"][prefix + str(i)]["prior"]["scale"] = sigma
reduced_covmat = initial_random_covmat[np.ix_(i_covmat, i_covmat)]
reduced_covmat_params = [prefix + str(i) for i in i_covmat]
info["sampler"] = {"mcmc": {}}
if tmpdir:
filename = os.path.join(str(tmpdir), "mycovmat.dat")
header = " ".join(reduced_covmat_params)
np.savetxt(filename, reduced_covmat, header=header)
info["sampler"]["mcmc"]["covmat"] = str(filename)
else:
info["sampler"]["mcmc"]["covmat_params"] = reduced_covmat_params
info["sampler"]["mcmc"]["covmat"] = reduced_covmat
to_compare = initial_random_covmat[np.ix_(input_order, input_order)]
def callback(sampler):
assert np.allclose(to_compare, sampler.proposer.get_covariance())
info["sampler"]["mcmc"].update({
"callback_function": callback,
"callback_every": 1,
"max_samples": 1,
"burn_in": 0
})
run(info)