def test_edges_timing_only(lc_core, lc_core_ctx):
lk = mcmc.LikelihoodEDGES(simulate=True)
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
model = lk.reduce_data(lc_core_ctx)
assert "freq_tb_min" in model
def test_wrong_lf_redshift():
with pytest.raises(ValueError):
cores = [
mcmc.CoreLuminosityFunction(redshift=9, sigma=0, name="lf"),
]
lks = [
mcmc.LikelihoodLuminosityFunction(name="lf"),
]
mcmc.build_computation_chain(cores, lks, setup=True)
def test_global_signal(lc_core, lc_core_ctx):
lk = mcmc.LikelihoodGlobalSignal(simulate=True)
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
model = lk.reduce_data(lc_core_ctx)
assert "frequencies" in model
def test_planck(lc_core, lc_core_ctx):
lk = mcmc.LikelihoodPlanck()
with pytest.raises(mcmc.NotAChain):
assert lk._is_lightcone
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
model = lk.reduce_data(lc_core_ctx)
assert "tau" in model
def test_greig(lc_core, lc_core_ctx):
lk = mcmc.LikelihoodGreig()
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
assert lc_core in lk.lightcone_modules
assert len(lk.coeval_modules) == 0
assert lk._require_spline
model = lk.reduce_data(lc_core_ctx)
assert "xHI" in model
def test_lightcone_core(lc_core, lc_core_ctx):
lk = mcmc.Likelihood1DPowerLightcone(simulate=True)
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
assert lc_core_ctx.contains("lightcone")
assert isinstance(lc_core_ctx.get("lightcone"), LightCone)
model = lk.reduce_data(lc_core_ctx)
lk.store(model, lc_core_ctx.getData())
assert all(k + "_0" in lc_core_ctx.getData() for k in model[0])
def test_forest(lc_core_lowz, lc_core_lowz_ctx):
lk = mcmc.LikelihoodForest(name="z5pt4")
with pytest.raises(mcmc.NotAChain):
assert lk._is_lightcone
core = mcmc.CoreForest(
redshift=5.4,
name="z5pt4",
user_params=lc_core_lowz.user_params,
flag_options=lc_core_lowz.flag_options,
n_realization=10,
)
with pytest.raises(NotImplementedError):
coeval_core_lowz = mcmc.CoreCoevalModule(
redshift=5.4,
max_redshift=8.0,
user_params=lc_core_lowz.user_params,
flag_options=lc_core_lowz.flag_options,
)
chain = mcmc.build_computation_chain([coeval_core_lowz, core],
lk,
setup=True)
lk.setup()
model = lk.reduce_data(chain.build_model_data())
with pytest.raises(ValueError):
lk_wrongz = mcmc.LikelihoodForest(name="z5pt5")
core_wrongz = mcmc.CoreForest(
redshift=5.5,
name="z5pt5",
user_params=lc_core_lowz.user_params,
flag_options=lc_core_lowz.flag_options,
n_realization=10,
)
chain = mcmc.build_computation_chain([lc_core_lowz, core_wrongz],
lk_wrongz,
setup=True)
lk_wrongz.setup()
model = lk_wrongz.reduce_data(chain.build_model_data())
mcmc.build_computation_chain([lc_core_lowz, core], lk, setup=True)
lk.setup()
model = lk.reduce_data(lc_core_lowz_ctx)
assert not np.all(model == 0)
def test_multi_datafile(tmpdir):
direc = tmpdir.mkdir("test_single_datafile")
dfile7 = direc.join("test_datafile7.npz").strpath
dfile8 = direc.join("test_datafile8.npz").strpath
nfile7 = direc.join("test_noisefile7.npz").strpath
nfile8 = direc.join("test_noisefile8.npz").strpath
core7 = CoreLuminosityFunction(redshift=[7],
n_muv_bins=100,
sigma=np.ones((1, 1)),
name="7")
core8 = CoreLuminosityFunction(redshift=[8],
n_muv_bins=80,
sigma=np.ones((1, 1)),
name="8")
lk7 = LikelihoodLuminosityFunction(simulate=True,
datafile=dfile7,
noisefile=nfile7,
name="7")
lk8 = LikelihoodLuminosityFunction(simulate=True,
datafile=dfile8,
noisefile=nfile8,
name="8")
chain = build_computation_chain([core7, core8], [lk7, lk8], setup=True)
assert isinstance(lk7.data, dict)
assert len(lk7.data["Muv"]) == 1
assert len(lk8.data["Muv"]) == 1
assert len(lk8.data["Muv"][0]) != len(lk7.data["Muv"][0])
chain({})
lk7read = LikelihoodLuminosityFunction(datafile=dfile7,
noisefile=nfile7,
name="7")
lk8read = LikelihoodLuminosityFunction(datafile=dfile8,
noisefile=nfile8,
name="8")
chain = build_computation_chain([core7, core8], [lk7read, lk8read],
setup=True)
assert isinstance(lk7read.data, dict)
assert len(lk7read.data["Muv"]) == 1
assert len(lk8read.data["Muv"]) == 1
assert len(lk8read.data["Muv"][0]) != len(lk7read.data["Muv"][0])
chain({})
def test_wrong_lf_paring():
redshifts = [6, 7, 8, 10]
with pytest.raises(ValueError):
cores = [
mcmc.CoreLuminosityFunction(redshift=z, sigma=0, name="lf")
for z in redshifts
]
lks = [mcmc.LikelihoodLuminosityFunction(name="lf") for z in redshifts]
mcmc.build_computation_chain(cores, lks, setup=True)
cores = [
mcmc.CoreLuminosityFunction(redshift=z, sigma=0, name="lfz%d" % z)
for z in redshifts
]
lks = [
mcmc.LikelihoodLuminosityFunction(name="lfz%d" % z) for z in redshifts
]
mcmc.build_computation_chain(cores, lks, setup=True)
def test_core_coeval_setup(core, likelihood_coeval):
with pytest.raises(
ValueError): # If simulate is not true, and no datafile given...
lk = mcmc.Likelihood1DPowerCoeval()
mcmc.build_computation_chain(core, lk)
chain = mcmc.build_computation_chain(core, likelihood_coeval)
assert isinstance(core.chain, LikelihoodComputationChain)
assert core.initial_conditions_seed is not None
ctx = chain.build_model_data()
assert ctx.get("xH_box") is not None
assert ctx.get("brightness_temp") is not None
assert not np.all(ctx.get("xH_box") == 0)
assert not np.all(ctx.get("brightness_temp") == 0)
def test_core_lc_init_cache(core_lc, likelihood_lc, caplog):
"""Ensures that perturb_field boxes are read in, rather than re-computed."""
chain = mcmc.build_computation_chain(core_lc, likelihood_lc, setup=True)
with caplog.at_level(logging.INFO):
ctx = chain.build_model_data()
for z in ctx.get("lightcone").node_redshifts:
assert (f"Existing z={z} perturb_field boxes found and read in"
in caplog.text)
def lc_core_ctx(lc_core):
lk = mcmc.LikelihoodPlanck()
chain = mcmc.build_computation_chain(lc_core, lk)
assert lk._is_lightcone
ctx = chain.createChainContext()
lc_core.build_model_data(ctx)
return ctx
def test_core_coeval_init_cache(core, likelihood_coeval, caplog):
"""Ensures that perturb_field boxes are read in, rather than re-computed."""
with caplog.at_level(logging.INFO):
chain = mcmc.build_computation_chain(core,
likelihood_coeval,
setup=True)
print(caplog.text)
with caplog.at_level(logging.INFO):
chain.build_model_data()
assert "perturb_field boxes found and read in" in caplog.text
def test_neutral_fraction(lc_core, lc_core_ctx):
with pytest.raises(ValueError):
mcmc.build_computation_chain(
mcmc.CoreLuminosityFunction(
redshift=7, sigma=1), # Bad core for neutral fraction
mcmc.LikelihoodNeutralFraction(),
)
lk = mcmc.LikelihoodNeutralFraction()
mcmc.build_computation_chain(lc_core, lk, setup=False)
lk.setup()
assert lc_core in lk.lightcone_modules
assert len(lk.coeval_modules) == 0
assert lk._require_spline
model = lk.reduce_data(lc_core_ctx)
assert "xHI" in model
def test_wrong_ctx_variable():
core = mcmc.CoreCoevalModule(
redshift=6,
user_params={
"HII_DIM": 35,
"BOX_LEN": 70
},
ctx_variables=("bad_key", "good key"),
)
lk = mcmc.Likelihood1DPowerCoeval(use_data=False)
chain = mcmc.build_computation_chain(core, lk, setup=False)
with pytest.raises(ValueError):
chain.build_model_data()
def lc_core_lowz_ctx(lc_core_lowz):
lk = mcmc.LikelihoodForest(name="z5pt4")
core = mcmc.CoreForest(
redshift=5.4,
name="z5pt4",
user_params=lc_core_lowz.user_params,
flag_options=lc_core_lowz.flag_options,
n_realization=10,
)
chain = mcmc.build_computation_chain([lc_core_lowz, core], lk)
assert lk._is_lightcone
return chain.build_model_data()
def test_single_datafile(tmpdir):
direc = tmpdir.mkdir("test_single_datafile")
dfile = direc.join("test_datafile.npz").strpath
nfile = direc.join("test_noisefile.npz").strpath
core = CoreLuminosityFunction(redshift=[7, 8, 9], sigma=np.ones((3, 1)))
lk = LikelihoodLuminosityFunction(simulate=True,
datafile=dfile,
noisefile=nfile)
chain = build_computation_chain(core, lk, setup=True)
assert os.path.exists(dfile)
assert isinstance(lk.data, dict)
assert isinstance(lk.noise, dict)
model = lk.get_fiducial_model()
assert isinstance(model, dict)
assert len(lk.data["Muv"]) == 3
assert len(model["Muv"]) == 3
chain({})
lk = LikelihoodLuminosityFunction(datafile=dfile, noisefile=nfile)
chain = build_computation_chain(core, lk, setup=True)
assert isinstance(lk.data, dict)
model = lk.get_fiducial_model()
assert isinstance(model, dict)
assert len(lk.data["Muv"]) == 3
assert len(model["Muv"]) == 3
chain({})
def test_init_pos_generator_good(core, likelihood_coeval, tmpdirec):
params = Params(("HII_EFF_FACTOR", [30.0, 10.0, 50.0, 10.0]),
("ION_Tvir_MIN", [4.7, 2, 8, 2]))
chain = mcmc.build_computation_chain(core,
likelihood_coeval,
params=params)
sampler = CosmoHammerSampler(
continue_sampling=False,
likelihoodComputationChain=chain,
storageUtil=HDFStorageUtil(str(tmpdirec / "POSGENERATORTEST")),
filePrefix=str(tmpdirec / "POSGENERATORTEST"),
reuseBurnin=False,
burninIterations=0,
sampleIterations=1,
walkersRatio=50,
)
pos = sampler.createInitPos()
assert all(chain.isValid(p) for p in pos)
def test_init_pos_generator_bad(core, likelihood_coeval, tmpdirec):
params = Params(
("HII_EFF_FACTOR", [30.0, 29.0, 31.0, 100.0]),
("ION_Tvir_MIN", [4.7, 4.6, 4.8, 20]),
)
chain = mcmc.build_computation_chain(core,
likelihood_coeval,
params=params)
sampler = CosmoHammerSampler(
continue_sampling=False,
likelihoodComputationChain=chain,
storageUtil=HDFStorageUtil(str(tmpdirec / "POSGENERATORTEST")),
filePrefix=str(tmpdirec / "POSGENERATORTEST"),
reuseBurnin=False,
burninIterations=0,
sampleIterations=1,
walkersRatio=50,
)
with pytest.raises(ValueError):
sampler.createInitPos()
def imaging(chain=None, cores=None, lk=None, freq_ind=0):
"""
Create a plot of the imaging capability of the current setup.
Uses the loaded cores to create a simulated sky, then "observe" this with given baselines. Then uses an
Instrumental2D likelihood to grid those baselines and transform back to the image plane. Every step of
this process is output as a panel in a plot to be compared.
Parameters
----------
chain : :class:`~py21cmmc.mcmc.cosmoHammer.LikelihoodComputationChain.LikelihoodComputationChain` instance, optional
A computation chain which contains loaded likelihoods and cores.
cores : list of :class:`~py21cmmc.mcmc.core.CoreBase` instances, optional
A list of cores defining the sky and instrument. Only required if `chain` is not given.
lk : :class:`~likelihood.LikelihoodInstrumental2D` class, optional
An instrumental likelihood, required if `chain` not given.
freq_ind : int, optional
The index of the frequency to actually show plots for (default is the first frequency channel).
Returns
-------
fig :
A matplotlib figure object.
"""
if chain is None and (cores is None or lk is None):
raise ValueError("Either chain or both cores and likelihood must be given.")
# Create a likelihood computation chain.
if chain is None:
chain = build_computation_chain(cores, lk)
chain.setup()
else:
lk = chain.getLikelihoodModules()[0]
cores = chain.getCoreModules()
if not isinstance(lk, LikelihoodInstrumental2D):
raise ValueError("likelihood needs to be a Instrumental2D likelihood")
if not hasattr(lk, "LikelihoodComputationChain"):
chain.setup()
# Call all core simulators.
ctx = chain.build_model_data()
visgrid = lk.grid_visibilities(ctx.get("visibilities"))
# Do a direct FT there and back, rather than baselines.
print(ctx.get("new_sky"))
direct_vis, direct_u = fft(ctx.get("new_sky")[:, :, freq_ind], L=lk._instr_core.sky_size, a=0, b=2 * np.pi)
direct_img, direct_l = ifft(direct_vis, Lk=(lk.uvgrid[1] - lk.uvgrid[0]) * len(lk.uvgrid), a=0, b=2 * np.pi)
# Get the reconstructed image
image_plane, image_grid = ifft(visgrid[:, :, freq_ind], Lk=(lk.uvgrid[1] - lk.uvgrid[0]) * len(lk.uvgrid), a=0, b=2 * np.pi)
# Make a figure.
if len(cores) == 2:
fig, ax = plt.subplots(2, 4, figsize=(12, 6))
mid_row = 0
else:
fig, ax = plt.subplots(3, max((4, len(cores))), figsize=(3*max((4, len(cores))), 9))
mid_row = 1
# Show original sky(s) (before Beam)
i = 0
for core in cores:
if isinstance(core, ForegroundsBase):
# TODO: frequency plotted here does not necessarily match the frequency in other plots.
mp = ax[0, i].imshow(ctx.get("foregrounds")[i][:, :, freq_ind].T, origin='lower',
extent=(-core.sky_size / 2, core.sky_size / 2) * 2)
ax[0, i].set_title("Orig. %s FG" % core.__class__.__name__)
cbar = plt.colorbar(mp, ax=ax[0, i])
ax[0, i].set_xlabel("l")
ax[0, i].set_ylabel("m")
cbar.set_label("Brightness Temp. [Jy/sr]")
i += 1
# # TODO: add lightcone plot
# if isinstance(core, CoreLightConeModule):
# mp = ax[0, i].imshow(ctx.get("foregrounds")[i][:, :, -freq_ind].T, origin='lower',
# extent=(-core.sky_size / 2, core.sky_size / 2) * 2)
# ax[0, i].set_title("Original %s foregrounds" % core.__class__.__name__)
# cbar = plt.colorbar(mp, ax=ax[0, i])
# ax[0, i].set_xlabel("l")
# ax[0, i].set_ylabel("m")
# cbar.set_label("Brightness Temp. [K]")
# i += 1
# Show tiled (if applicable) and attenuated sky
mp = ax[mid_row, 1].imshow(
ctx.get("new_sky")[:, :, freq_ind].T, origin='lower',
extent=(-lk._instr_core.sky_size / 2, lk._instr_core.sky_size / 2) * 2
)
ax[mid_row, 1].set_title("Tiled+Beam FG")
cbar = plt.colorbar(mp, ax=ax[mid_row, 1])
ax[mid_row, 1].set_xlabel("l")
ax[mid_row, 1].set_ylabel("m")
cbar.set_label("Brightness Temp. [K]")
# Show UV weights
mp = ax[mid_row, 2].imshow(
lk.nbl_uvnu[:, :, freq_ind].T, origin='lower',
extent=(lk.uvgrid.min(), lk.uvgrid.max()) * 2
)
ax[mid_row, 2].set_title("UV weights")
cbar = plt.colorbar(mp, ax=ax[mid_row, 2])
ax[mid_row, 2].set_xlabel("u")
ax[mid_row, 2].set_ylabel("v")
cbar.set_label("Weight")
# Show raw visibilities
wvlength = 3e8 / ctx.get("frequencies")[freq_ind]
mp = ax[mid_row, 3].scatter(ctx.get("baselines")[:, 0] / wvlength, ctx.get("baselines")[:, 1] / wvlength,
c=np.real(ctx.get("visibilities")[:, freq_ind]))
ax[mid_row, 3].set_title("Raw Vis.")
cbar = plt.colorbar(mp, ax=ax[mid_row, 3])
ax[mid_row, 3].set_xlabel("u")
ax[mid_row, 3].set_xlabel("v")
cbar.set_label("Re[Vis] [Jy?]")
# Show Gridded Visibilities
mp = ax[mid_row+1, 3].imshow(
np.real(visgrid[:, :, freq_ind].T), origin='lower',
extent=(lk.uvgrid.min(), lk.uvgrid.max()) * 2
)
ax[mid_row+1, 3].set_title("Gridded Vis")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 3])
ax[mid_row+1, 3].set_xlabel("u")
ax[mid_row+1, 3].set_ylabel("v")
cbar.set_label("Jy")
# Show directly-calculated UV plane
mp = ax[mid_row+1, 2].imshow(
np.real(direct_vis), origin='lower',
extent=(direct_u[0].min(), direct_u[0].max()) * 2
)
ax[mid_row+1, 2].set_title("Direct Vis")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 2])
ax[mid_row+1, 2].set_xlabel("u")
ax[mid_row+1, 2].set_ylabel("v")
cbar.set_label("Jy")
# Show final "image"
mp = ax[mid_row+1, 1].imshow(np.abs(image_plane).T, origin='lower',
extent=(image_grid[0].min(), image_grid[0].max(),) * 2)
ax[mid_row+1, 1].set_title("Recon. FG")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 1])
ax[mid_row+1, 1].set_xlabel("l")
ax[mid_row+1, 1].set_ylabel("m")
cbar.set_label("Flux Density. [Jy]")
# Show direct reconstruction
mp = ax[mid_row+1, 0].imshow(np.abs(direct_img).T, origin='lower',
extent=(direct_l[0].min(), direct_l[0].max(),) * 2)
ax[mid_row+1, 0].set_title("Recon. direct FG")
cbar = plt.colorbar(mp, ax=ax[mid_row+1, 0])
ax[mid_row+1, 0].set_xlabel("l")
ax[mid_row+1, 0].set_ylabel("m")
cbar.set_label("Flux Density. [Jy]")
plt.tight_layout()
return fig
def test_create_mock():
core = CoreLuminosityFunction(redshift=[7])
lk = LikelihoodLuminosityFunction(simulate=True)
with pytest.raises(ValueError):
build_computation_chain(core, lk, setup=True)