test: many more tests of lightconing

.flake8

@@ -15,7 +15,7 @@ ignore =
     # Docstring in imperative mood. This should *not* be the case for @property's, but can't ignore them atm.
     D401
 max-line-length = 88
-max-complexity = 50
+max-complexity = 55
 docstring-convention=numpy
 inline-quotes="
 ignore-decorators=click.option

setup.cfg

@@ -25,3 +25,26 @@ skip = migrations
 
 [pydocstyle]
 convention = numpy
+
+[coverage:run]
+branch = True
+
+[coverage:report]
+; Regexes for lines to exclude from consideration
+exclude_also =
+    ; Don't complain about missing debug-only code:
+    def __repr__
+    if self\.debug
+
+    ; Don't complain if tests don't hit defensive assertion code:
+    raise AssertionError
+    raise NotImplementedError
+
+    ; Don't complain if non-runnable code isn't run:
+    if 0:
+    if __name__ == .__main__.:
+
+    ; Don't complain about abstract methods, they aren't run:
+    @(abc\.)?abstractmethod
+
+ignore_errors = True

src/py21cmfast/cache_tools.py

@@ -212,9 +212,9 @@ def get_boxes_at_redshift(
                 logger.debug(
                     f"{file} {paramtype} don't match: {getattr(obj, paramtype)} vs. {paramobj}"
                 )
-                continue
-
-        out[obj.__class__.__name__].append(obj)
+                break
+        else:
+            out[obj.__class__.__name__].append(obj)
 
     return out
 

src/py21cmfast/inputs.py

@@ -561,7 +561,7 @@ def POWER_SPECTRUM(self):
                 isinstance(self._POWER_SPECTRUM, str)
                 and self._POWER_SPECTRUM.upper() != "CLASS"
             ):
-                logger.warning(
+                logger.warn(
                     "Automatically setting POWER_SPECTRUM to 5 (CLASS) as you are using "
                     "relative velocities"
                 )
@@ -618,7 +618,7 @@ def FAST_FCOLL_TABLES(self):
         """Check that USE_INTERPOLATION_TABLES is True."""
         if not self._FAST_FCOLL_TABLES or self.USE_INTERPOLATION_TABLES:
             return self._FAST_FCOLL_TABLES
-        logger.warning(
+        logger.warn(
             "You cannot turn on FAST_FCOLL_TABLES without USE_INTERPOLATION_TABLES."
         )
         return False
@@ -708,7 +708,7 @@ def SUBCELL_RSD(self):
     def USE_HALO_FIELD(self):
         """Automatically setting USE_HALO_FIELD to False if USE_MINI_HALOS."""
         if self._USE_HALO_FIELD and self.USE_MINI_HALOS:
-            logger.warning(
+            logger.warn(
                 "You have set USE_MINI_HALOS to True but USE_HALO_FIELD is also True! "
                 "Automatically setting USE_HALO_FIELD to False."
             )
@@ -725,7 +725,7 @@ def M_MIN_in_Mass(self):
     def USE_MASS_DEPENDENT_ZETA(self):
         """Automatically setting USE_MASS_DEPENDENT_ZETA to True if USE_MINI_HALOS."""
         if self.USE_MINI_HALOS and not self._USE_MASS_DEPENDENT_ZETA:
-            logger.warning(
+            logger.warn(
                 "You have set USE_MINI_HALOS to True but USE_MASS_DEPENDENT_ZETA is False! "
                 "Automatically setting USE_MASS_DEPENDENT_ZETA to True."
             )
@@ -736,7 +736,7 @@ def USE_MASS_DEPENDENT_ZETA(self):
     def INHOMO_RECO(self):
         """Automatically setting INHOMO_RECO to True if USE_MINI_HALOS."""
         if self.USE_MINI_HALOS and not self._INHOMO_RECO:
-            logger.warning(
+            warnings.warn(
                 "You have set USE_MINI_HALOS to True but INHOMO_RECO to False! "
                 "Automatically setting INHOMO_RECO to True."
             )
@@ -747,7 +747,7 @@ def INHOMO_RECO(self):
     def USE_TS_FLUCT(self):
         """Automatically setting USE_TS_FLUCT to True if USE_MINI_HALOS."""
         if self.USE_MINI_HALOS and not self._USE_TS_FLUCT:
-            logger.warning(
+            logger.warn(
                 "You have set USE_MINI_HALOS to True but USE_TS_FLUCT to False! "
                 "Automatically setting USE_TS_FLUCT to True."
             )
@@ -758,7 +758,7 @@ def USE_TS_FLUCT(self):
     def PHOTON_CONS(self):
         """Automatically setting PHOTON_CONS to False if USE_MINI_HALOS."""
         if self.USE_MINI_HALOS and self._PHOTON_CONS:
-            logger.warning(
+            logger.warn(
                 "USE_MINI_HALOS is not compatible with PHOTON_CONS! "
                 "Automatically setting PHOTON_CONS to False."
             )
@@ -920,7 +920,7 @@ def R_BUBBLE_MAX(self):
         if not self._R_BUBBLE_MAX:
             return 50.0 if self.INHOMO_RECO else 15.0
         if self.INHOMO_RECO and self._R_BUBBLE_MAX != 50:
-            logger.warning(
+            logger.warn(
                 "You are setting R_BUBBLE_MAX != 50 when INHOMO_RECO=True. "
                 "This is non-standard (but allowed), and usually occurs upon manual "
                 "update of INHOMO_RECO"
@@ -1004,7 +1004,7 @@ def validate_all_inputs(
             )
 
             if config["ignore_R_BUBBLE_MAX_error"]:
-                warnings.warn(msg)
+                logger.warn(msg)
             else:
                 raise ValueError(msg)
 
@@ -1013,6 +1013,6 @@ def validate_all_inputs(
         and not user_params.USE_RELATIVE_VELOCITIES
         and not flag_options.FIX_VCB_AVG
     ):
-        logger.warning(
+        logger.warn(
             "USE_MINI_HALOS needs USE_RELATIVE_VELOCITIES to get the right evolution!"
         )

src/py21cmfast/wrapper.py

@@ -2488,6 +2488,7 @@ def run_lightcone(
     hooks=None,
     always_purge: bool = False,
     lightcone_filename: str | Path = None,
+    return_at_z: float = 0.0,
     **global_kwargs,
 ):
     r"""
@@ -2564,6 +2565,11 @@ def run_lightcone(
         The filename to which to save the lightcone. The lightcone is returned in
         memory, and can be saved manually later, but including this filename will
         save the lightcone on each iteration, which can be helpful for checkpointing.
+    return_at_z
+        If given, evaluation of the lightcone will be stopped at the given redshift,
+        and the partial lightcone object will be returned. Lightcone evaluation can
+        continue if the returned lightcone is saved to file, and this file is passed
+        as `lightcone_filename`.
     \*\*global_kwargs :
         Any attributes for :class:`~py21cmfast.inputs.GlobalParams`. This will
         *temporarily* set global attributes for the duration of the function. Note that
@@ -2618,8 +2624,8 @@ def run_lightcone(
             redshift = perturb.redshift
         elif redshift is not None:
             warnings.warn(
-                DeprecationWarning,
                 "passing redshift directly is deprecated, please use the Lightconer interface instead",
+                category=DeprecationWarning,
             )
 
         if user_params.MINIMIZE_MEMORY and not write:
@@ -2629,7 +2635,11 @@ def run_lightcone(
 
         max_redshift = (
             global_params.Z_HEAT_MAX
-            if (flag_options.INHOMO_RECO or flag_options.USE_TS_FLUCT)
+            if (
+                flag_options.INHOMO_RECO
+                or flag_options.USE_TS_FLUCT
+                or (max_redshift is None and lightconer is None)
+            )
             else (
                 max_redshift
                 if max_redshift is not None
@@ -2649,8 +2659,10 @@ def run_lightcone(
         lightconer.validate_options(user_params, flag_options)
 
         # Get the redshift through which we scroll and evaluate the ionization field.
-        scrollz = _logscroll_redshifts(
-            redshift, global_params.ZPRIME_STEP_FACTOR, max_redshift
+        scrollz = np.array(
+            _logscroll_redshifts(
+                redshift, global_params.ZPRIME_STEP_FACTOR, max_redshift
+            )
         )
 
         lcz = lightconer.lc_redshifts
@@ -2703,7 +2715,13 @@ def run_lightcone(
 
         if lightcone_filename and Path(lightcone_filename).exists():
             lightcone = LightCone.read(lightcone_filename)
-            scrollz = scrollz[np.array(scrollz) < lightcone._current_redshift]
+            scrollz = scrollz[scrollz < lightcone._current_redshift]
+            if len(scrollz) == 0:
+                # The entire lightcone is already full!
+                logger.info(
+                    f"Lightcone already full at z={lightcone._current_redshift}. Returning."
+                )
+                return lightcone
             lc = lightcone.lightcones
         else:
             lcn_cls = (
@@ -2781,6 +2799,12 @@ def run_lightcone(
         if flag_options.PHOTON_CONS:
             calibrate_photon_cons(**kw)
 
+        if return_at_z > lightcone.redshift and not write:
+            raise ValueError(
+                "Returning before the final redshift requires caching in order to "
+                "continue the simulation later. Set write=True!"
+            )
+
         # Iterate through redshift from top to bottom
         if lightcone.redshift != lightcone._current_redshift:
             logger.info(
@@ -2795,10 +2819,16 @@ def run_lightcone(
                 flag_options=flag_options,
                 astro_params=astro_params,
             )
-            st = cached_boxes["TsBox"][0]
-            prev_perturb = cached_boxes["PerturbedField"][0]
-            ib = cached_boxes["IonizedBox"][0]
-            lc_index = lightcone._current_index
+            try:
+                st = cached_boxes["TsBox"][0] if flag_options.USE_TS_FLUCT else None
+                prev_perturb = cached_boxes["PerturbedField"][0]
+                ib = cached_boxes["IonizedBox"][0]
+            except (KeyError, IndexError):
+                raise OSError(
+                    f"No component boxes found at z={lightcone._current_redshift} with "
+                    f"seed {lightcone.random_seed} and direc={direc}. You need to have "
+                    "run with write=True to continue from a checkpoint."
+                )
             pf = prev_perturb
         else:
             st, ib, prev_perturb = None, None, None
@@ -2814,6 +2844,8 @@ def run_lightcone(
         log10_mturnovers_mini = np.zeros(len(scrollz))
         coeval = None
         prev_coeval = None
+        st2 = None
+        pt_halos = None
 
         if lightcone_filename and not Path(lightcone_filename).exists():
             lightcone.save(lightcone_filename)
@@ -2846,8 +2878,8 @@ def run_lightcone(
                 previous_ionize_box=ib,
                 perturbed_field=pf2,
                 previous_perturbed_field=prev_perturb,
-                spin_temp=st2 if flag_options.USE_TS_FLUCT else None,
-                pt_halos=pt_halos if flag_options.USE_HALO_FIELD else None,
+                spin_temp=st2,
+                pt_halos=pt_halos,
                 cleanup=(cleanup and iz == (len(scrollz) - 1)),
                 **kw,
             )
@@ -2908,6 +2940,12 @@ def run_lightcone(
                 ):
                     if this_lc is not None:
                         lightcone.lightcones[quantity][..., idx] = this_lc
+                        lc_index = idx
+
+                if lightcone_filename:
+                    lightcone.make_checkpoint(
+                        lightcone_filename, redshift=z, index=lc_index
+                    )
 
             # Save current ones as old ones.
             if flag_options.USE_TS_FLUCT:
@@ -2925,10 +2963,9 @@ def run_lightcone(
 
             pf = pf2
 
-            if lightcone_filename:
-                lightcone.make_checkpoint(
-                    lightcone_filename, redshift=z, index=lc_index
-                )
+            if z <= return_at_z:
+                # Optionally return when the lightcone is only partially filled
+                break
 
         if flag_options.PHOTON_CONS:
             photon_nonconservation_data = _get_photon_nonconservation_data()

tests/conftest.py

@@ -135,21 +135,21 @@ def low_redshift():
 
 
 @pytest.fixture(scope="session")
-def perturb_field(ic, redshift):
+def perturbed_field(ic, redshift):
     """A default perturb_field"""
     return wrapper.perturb_field(redshift=redshift, init_boxes=ic, write=True)
 
 
 @pytest.fixture(scope="session")
-def rectlcn(perturb_field, max_redshift) -> RectilinearLightconer:
+def rectlcn(perturbed_field, max_redshift) -> RectilinearLightconer:
     return RectilinearLightconer.with_equal_cdist_slices(
-        min_redshift=perturb_field.redshift,
+        min_redshift=perturbed_field.redshift,
         max_redshift=max_redshift,
-        resolution=perturb_field.user_params.cell_size,
-        cosmo=perturb_field.cosmo_params.cosmo,
+        resolution=perturbed_field.user_params.cell_size,
+        cosmo=perturbed_field.cosmo_params.cosmo,
     )
 
 
 @pytest.fixture(scope="session")
-def lc(perturb_field, rectlcn):
-    return run_lightcone(lightconer=rectlcn, perturb=perturb_field)
+def lc(perturbed_field, rectlcn):
+    return run_lightcone(lightconer=rectlcn, perturb=perturbed_field)