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"])