feat: add construct_lightcone_form_filelist

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

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):