def test_power_spectra_lightcone(redshift, kwargs, module_direc):
print("Options used for the test: ", kwargs)
# First get pre-made data
with h5py.File(prd.get_filename(redshift, **kwargs), "r") as f:
power = f["power_lc"][...]
xHI = f["xHI"][...]
Tb = f["Tb"][...]
with config.use(direc=module_direc, regenerate=False, write=True):
with global_params.use(
zprime_step_factor=prd.DEFAULT_ZPRIME_STEP_FACTOR):
# Note that if zprime_step_factor is set in kwargs, it will over-ride this.
k, p, lc = prd.produce_lc_power_spectra(redshift, **kwargs)
assert np.allclose(power[:len(power) // 2],
p[:len(power) // 2],
atol=0,
rtol=1e-2)
assert np.allclose(power[(len(power) // 2):],
p[(len(power) // 2):],
atol=0,
rtol=5e-2)
assert np.allclose(xHI, lc.global_xH, atol=1e-5, rtol=1e-3)
assert np.allclose(Tb, lc.global_brightness_temp, atol=1e-5, rtol=1e-3)
def test_power_spectra_coeval(name, module_direc, plt):
redshift, kwargs = prd.OPTIONS[name]
print(f"Options used for the test at z={redshift}: ", kwargs)
# First get pre-made data
with h5py.File(prd.get_filename("power_spectra", name), "r") as fl:
true_powers = {
"_".join(key.split("_")[1:]): value[...]
for key, value in fl["coeval"].items() if key.startswith("power_")
}
# Now compute the Coeval object
with config.use(direc=module_direc, regenerate=False, write=True):
with global_params.use(
zprime_step_factor=prd.DEFAULT_ZPRIME_STEP_FACTOR):
# Note that if zprime_step_factor is set in kwargs, it will over-ride this.
test_k, test_powers, _ = prd.produce_coeval_power_spectra(
redshift, **kwargs)
if plt == mpl.pyplot:
make_coeval_comparison_plot(test_k, true_powers, test_powers, plt)
for key, value in true_powers.items():
print(f"Testing {key}")
assert np.sum(
~np.isclose(value, test_powers[key], atol=0, rtol=1e-2)) < 10
np.testing.assert_allclose(value, test_powers[key], atol=0, rtol=1e-1)
def produce_halo_field_data(redshift, **kwargs):
options_halo = get_all_options_halo(redshift, **kwargs)
with config.use(regenerate=True, write=False):
pt_halos = perturb_halo_list(**options_halo)
return pt_halos
def test_power_spectra_coeval(redshift, kwargs, module_direc):
print("Options used for the test: ", kwargs)
# First get pre-made data
with h5py.File(prd.get_filename(redshift, **kwargs), "r") as f:
power = f["power_coeval"][...]
with config.use(direc=module_direc, regenerate=False, write=True):
with global_params.use(
zprime_step_factor=prd.DEFAULT_ZPRIME_STEP_FACTOR):
# Note that if zprime_step_factor is set in kwargs, it will over-ride this.
k, p, bt = prd.produce_coeval_power_spectra(redshift, **kwargs)
assert np.allclose(power, p, atol=1e-3, rtol=1e-2)
def produce_perturb_field_data(redshift, **kwargs):
options = get_all_options(redshift, **kwargs)
options_ics = get_all_options_ics(**kwargs)
out = {
key: kwargs[key]
for key in kwargs
if key.upper() in (k.upper() for k in global_params.keys())
}
velocity_normalisation = 1e16
with config.use(regenerate=True, write=False):
init_box = initial_conditions(**options_ics)
pt_box = perturb_field(redshift=redshift, init_boxes=init_box, **out)
p_dens, k_dens = get_power(
pt_box.density,
boxlength=options["user_params"]["BOX_LEN"],
)
p_vel, k_vel = get_power(
pt_box.velocity * velocity_normalisation,
boxlength=options["user_params"]["BOX_LEN"],
)
def hist(kind, xmin, xmax, nbins):
data = getattr(pt_box, kind)
if kind == "velocity":
data = velocity_normalisation * data
bins, edges = np.histogram(
data,
bins=np.linspace(xmin, xmax, nbins),
range=[xmin, xmax],
density=True,
)
left, right = edges[:-1], edges[1:]
X = np.array([left, right]).T.flatten()
Y = np.array([bins, bins]).T.flatten()
return X, Y
X_dens, Y_dens = hist("density", -0.8, 2.0, 50)
X_vel, Y_vel = hist("velocity", -2, 2, 50)
return k_dens, p_dens, k_vel, p_vel, X_dens, Y_dens, X_vel, Y_vel, init_box
def produce_power_spectra_for_tests(name, redshift, force, direc, **kwargs):
fname = get_filename("power_spectra", name)
# Need to manually remove it, otherwise h5py tries to add to it
if fname.exists():
if force:
fname.unlink()
else:
print(f"Skipping {fname} because it already exists.")
return fname
# For tests, we *don't* want to use cached boxes, but we also want to use the
# cache between the power spectra and lightcone. So we create a temporary
# directory in which to cache results.
with config.use(direc=direc):
k, p, coeval = produce_coeval_power_spectra(redshift, **kwargs)
k_l, p_l, lc = produce_lc_power_spectra(redshift, **kwargs)
with h5py.File(fname, "w") as fl:
for k, v in kwargs.items():
fl.attrs[k] = v
fl.attrs["HII_DIM"] = coeval.user_params.HII_DIM
fl.attrs["DIM"] = coeval.user_params.DIM
fl.attrs["BOX_LEN"] = coeval.user_params.BOX_LEN
coeval_grp = fl.create_group("coeval")
coeval_grp["k"] = k
for key, val in p.items():
coeval_grp[f"power_{key}"] = val
lc_grp = fl.create_group("lightcone")
lc_grp["k"] = k_l
for key, val in p_l.items():
lc_grp[f"power_{key}"] = val
lc_grp["global_xH"] = lc.global_xH
lc_grp["global_brightness_temp"] = lc.global_brightness_temp
print(f"Produced {fname} with {kwargs}")
return fname
def test_power_spectra_lightcone(name, module_direc, plt):
redshift, kwargs = prd.OPTIONS[name]
print(f"Options used for the test at z={redshift}: ", kwargs)
# First get pre-made data
with h5py.File(prd.get_filename("power_spectra", name), "r") as fl:
true_powers = {}
true_global = {}
for key in fl["lightcone"].keys():
if key.startswith("power_"):
true_powers["_".join(
key.split("_")[1:])] = fl["lightcone"][key][...]
elif key.startswith("global_"):
true_global[key] = fl["lightcone"][key][...]
# Now compute the lightcone
with config.use(direc=module_direc, regenerate=False, write=True):
with global_params.use(
zprime_step_factor=prd.DEFAULT_ZPRIME_STEP_FACTOR):
# Note that if zprime_step_factor is set in kwargs, it will over-ride this.
test_k, test_powers, lc = prd.produce_lc_power_spectra(
redshift, **kwargs)
if plt == mpl.pyplot:
make_lightcone_comparison_plot(test_k, lc.node_redshifts, true_powers,
true_global, test_powers, lc, plt)
for key, value in true_powers.items():
print(f"Testing {key}")
# Ensure all but 10 of the values is within 1%, and none of the values
# is outside 10%
assert np.sum(
~np.isclose(value, test_powers[key], atol=0, rtol=1e-2)) < 10
assert np.allclose(value, test_powers[key], atol=0, rtol=1e-1)
for key, value in true_global.items():
print(f"Testing Global {key}")
assert np.allclose(value, getattr(lc, key), atol=0, rtol=1e-3)
def produce_power_spectra_for_tests(redshift, force, direc, **kwargs):
fname = get_filename(redshift, **kwargs)
# Need to manually remove it, otherwise h5py tries to add to it
if os.path.exists(fname):
if force:
os.remove(fname)
else:
return fname
# For tests, we *don't* want to use cached boxes, but we also want to use the
# cache between the power spectra and lightcone. So we create a temporary
# directory in which to cache results.
with config.use(direc=direc):
k, p, coeval = produce_coeval_power_spectra(redshift, **kwargs)
k_l, p_l, lc = produce_lc_power_spectra(redshift, **kwargs)
with h5py.File(fname, "w") as fl:
for k, v in kwargs.items():
fl.attrs[k] = v
fl.attrs["HII_DIM"] = coeval.user_params.HII_DIM
fl.attrs["DIM"] = coeval.user_params.DIM
fl.attrs["BOX_LEN"] = coeval.user_params.BOX_LEN
fl["power_coeval"] = p
fl["k_coeval"] = k
fl["power_lc"] = p_l
fl["k_lc"] = k_l
fl["xHI"] = lc.global_xH
fl["Tb"] = lc.global_brightness_temp
print(f"Produced {fname} with {kwargs}")
return fname
