diff --git a/src/tuesday/simulators/py21cmfast/lightcones.py b/src/tuesday/simulators/py21cmfast/lightcones.py index b23fbf91..b704a54e 100644 --- a/src/tuesday/simulators/py21cmfast/lightcones.py +++ b/src/tuesday/simulators/py21cmfast/lightcones.py @@ -1,15 +1,17 @@ """Generate lightcones from 21cmFAST output cache.""" from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace import numpy as np from py21cmfast.drivers.lightcone import ( - AngularLightcone, LightCone, setup_lightcone_instance, ) +from py21cmfast.io import read_output_struct from py21cmfast.io.caching import RunCache -from py21cmfast.lightcones import Lightconer +from py21cmfast.lightconers import Lightconer def construct_lightcone_from_cache( @@ -49,14 +51,17 @@ def construct_lightcone_from_cache( inputs = cache.inputs node_redshifts = sorted(cache.BrightnessTemp.keys(), reverse=True) - lightconer.validate_options(cache.inputs.matter_options, cache.inputs.astro_options) + lightconer.validate_options( + cache.inputs, include_dvdr_in_tau21=False, apply_rsds=False + ) # Create the LightCone instance, loading from file if needed lightcone = setup_lightcone_instance( lightconer=lightconer, inputs=inputs, scrollz=node_redshifts, - global_quantities=global_quantities, + include_dvdr_in_tau21=False, + apply_rsds=False, photon_nonconservation_data={}, ) @@ -100,12 +105,74 @@ def construct_lightcone_from_cache( prev_coeval = coeval - # last redshift things - if iz == len(node_redshifts) - 1 and ( - isinstance(lightcone, AngularLightcone) and lightconer.get_los_velocity - ): - lightcone.lightcones["brightness_temp_with_rsds"] = lightcone.compute_rsds( - n_subcells=inputs.astro_params.N_RSD_STEPS - ) + return lightcone + + +def construct_lightcone_from_filelist( + filelist: Sequence[Path], + lightconer: Lightconer, +) -> LightCone: + """Construct a Lightcone from a list of OutputStruct files. + + This function assumes that the list of files are 21cmFAST output structs, each + at a different redshift. These may be from a RunCache, but don't need to be. + + Parameters + ---------- + filelist + A list of paths to 21cmFAST output struct files. + lightconer + The object used to generate lightcone slices. + global_quantities + A list of global quantities to extract. These must exist in the files + given. + """ + boxes = [read_output_struct(fl) for fl in filelist] + inputs = boxes[0].inputs + boxes = sorted(boxes, key=lambda b: b.redshift, reverse=True) + node_redshifts = [b.redshift for b in boxes] + + lightconer.validate_options(inputs, apply_rsds=False, include_dvdr_in_tau21=False) + + # Create the LightCone instance, loading from file if needed + lightcone = setup_lightcone_instance( + lightconer=lightconer, + inputs=inputs, + scrollz=node_redshifts, + photon_nonconservation_data={}, + include_dvdr_in_tau21=False, + apply_rsds=False, + ) + + lightcone._last_completed_node = -1 + lightcone._last_completed_lcidx = ( + np.sum( + lightcone.lightcone_redshifts + >= node_redshifts[lightcone._last_completed_node] + ) + - 1 + ) + + prev_box = None + for box in boxes: + # This hacks an OutputStruct to look like a Coeval, which has all the + # boxes defined directly as attributes. Later versions + # of 21cmFAST will not require this. + _box = SimpleNamespace() + _box.__dict__.update(box.__dict__) + _box.simulation_options = box.inputs.simulation_options + _box.cosmo_params = box.inputs.cosmo_params + for q in lightconer.quantities: + setattr(_box, q, box.get(q)) + + # Get lightcone slices + if prev_box is not None: + for quantity, idx, this_lc in lightconer.make_lightcone_slices( + _box, prev_box + ): + if this_lc is not None: + lightcone.lightcones[quantity][..., idx] = this_lc + + prev_box = _box return lightcone diff --git a/tests/test_simulators/test_py21cmfast/test_lightcones.py b/tests/test_simulators/test_py21cmfast/test_lightcones.py index 6c1a7fd4..f5c0935e 100644 --- a/tests/test_simulators/test_py21cmfast/test_lightcones.py +++ b/tests/test_simulators/test_py21cmfast/test_lightcones.py @@ -11,7 +11,10 @@ from py21cmfast.io.caching import OutputCache, RunCache from py21cmfast.wrapper import outputs -from tuesday.simulators.py21cmfast.lightcones import construct_lightcone_from_cache +from tuesday.simulators.py21cmfast.lightcones import ( + construct_lightcone_from_cache, + construct_lightcone_from_filelist, +) def create_mock_cache_output(cachedir: Path, ics: bool = False) -> RunCache: @@ -23,9 +26,7 @@ def create_mock_cache_output(cachedir: Path, ics: bool = False) -> RunCache: cachedir.mkdir() inputs = InputParameters.from_template( "simple", random_seed=1, node_redshifts=[10, 9, 8, 7, 6] - ).evolve_input_structs( - BOX_LEN=100, HII_DIM=50, DIM=100, APPLY_RSDS=False, KEEP_3D_VELOCITIES=True - ) + ).evolve_input_structs(BOX_LEN=100, HII_DIM=50, DIM=100, KEEP_3D_VELOCITIES=True) cache = RunCache.from_inputs(inputs, cache=OutputCache(cachedir)) for fldname, fld in attrs.asdict(cache, recurse=False).items(): @@ -61,8 +62,8 @@ def test_construct_rect_lightcone_from_cache(tmp_path: Path): cache = create_mock_cache_output(cachedir) lightconer = RectilinearLightconer.with_equal_cdist_slices( - min_redshift=6, - max_redshift=10, + min_redshift=6.5, + max_redshift=9.5, resolution=cache.inputs.simulation_options.BOX_LEN * un.Mpc / cache.inputs.simulation_options.HII_DIM, @@ -78,8 +79,16 @@ def test_construct_rect_lightcone_from_cache(tmp_path: Path): assert isinstance(lightcone, LightCone) # Check that the data is stored in the lightcone object - assert lightcone.lightcones["brightness_temp"].shape == (50, 50, 607) - assert lightcone.lightcones["density"].shape == (50, 50, 607) + assert lightcone.lightcones["brightness_temp"].shape == ( + 50, + 50, + len(lightconer.lc_redshifts), + ) + assert lightcone.lightcones["density"].shape == ( + 50, + 50, + len(lightconer.lc_redshifts), + ) assert lightcone.global_quantities["log10_mturn_acg"].shape == ( len(cache.inputs.node_redshifts), @@ -99,18 +108,74 @@ def test_construct_ang_lightcone_from_cache(tmp_path: Path): lightconer = AngularLightconer.like_rectilinear( simulation_options=cache.inputs.simulation_options, - max_redshift=10, - match_at_z=6.0, + max_redshift=9.5, + match_at_z=6.5, quantities=("density", "brightness_temp"), - get_los_velocity=True, ) lightcone = construct_lightcone_from_cache(cache, lightconer) assert isinstance(lightcone, LightCone) - assert lightcone.lightcones["brightness_temp"].shape == (2500, 607) - assert lightcone.lightcones["density"].shape == (2500, 607) - assert lightcone.global_quantities == {} + assert lightcone.lightcones["brightness_temp"].shape == ( + 2500, + len(lightconer.lc_redshifts), + ) + assert lightcone.lightcones["density"].shape == (2500, len(lightconer.lc_redshifts)) + + +def test_construct_rect_lightcone_from_filelist(tmp_path: Path): + """Test the construction of a lightcone from a cache.""" + + cachedir = tmp_path / "cache" + cache = create_mock_cache_output(cachedir) + + lightconer = RectilinearLightconer.with_equal_cdist_slices( + min_redshift=6.5, + max_redshift=9.5, + resolution=cache.inputs.simulation_options.BOX_LEN + * un.Mpc + / cache.inputs.simulation_options.HII_DIM, + quantities=("density",), + ) + + filelist = [ + cache.PerturbedField[z] + for z in sorted(cache.inputs.node_redshifts, reverse=True) + ] + + lightcone = construct_lightcone_from_filelist(filelist, lightconer) + + assert isinstance(lightcone, LightCone) + + # Check that the data is stored in the lightcone object + assert lightcone.lightcones["density"].shape == ( + 50, + 50, + len(lightconer.lc_redshifts), + ) + + +def test_construct_ang_lightcone_from_filelist(tmp_path: Path): + """Test the construction of a lightcone from a cache.""" + cachedir = tmp_path / "cache" + cache = create_mock_cache_output(cachedir) + + lightconer = AngularLightconer.like_rectilinear( + simulation_options=cache.inputs.simulation_options, + max_redshift=9.5, + match_at_z=6.5, + quantities=("density",), + ) + + filelist = [ + cache.PerturbedField[z] + for z in sorted(cache.inputs.node_redshifts, reverse=True) + ] + + lightcone = construct_lightcone_from_filelist(filelist, lightconer) + + assert isinstance(lightcone, LightCone) + assert lightcone.lightcones["density"].shape == (2500, len(lightconer.lc_redshifts)) def test_exceptions(tmp_path: Path):