Merge pull request #380 from 21cmfast/halobox-memory-improvements

Halobox memory improvements
21cmfast · Apr 22, 2024 · d257b08 · d257b08
2 parents 0fa768b + e1b6e4a
commit d257b08
Showing 1 changed file with 28 additions and 5 deletions.
diff --git a/src/py21cmfast/wrapper.py b/src/py21cmfast/wrapper.py
@@ -2826,7 +2826,7 @@ def run_coeval(
         ib_tracker = [0] * len(redshift)
         bt = [0] * len(redshift)
         # At first we don't have any "previous" st or ib.
-        st, ib, pf = None, None, None
+        st, ib, pf, hb = None, None, None, None
         # optional fields which remain None if their flags are off
         hb2, ph2 = None, None
 
@@ -2859,11 +2859,11 @@ def run_coeval(
                     previous_spin_temp=st,
                     **kw,
                 )
+
+            if flag_options.USE_TS_FLUCT:
                 # append the halo redshift array so we have all halo boxes [z,zmax]
                 z_halos += [z]
                 hbox_arr += [hb2]
-
-            if flag_options.USE_TS_FLUCT:
                 if flag_options.USE_HALO_FIELD:
                     xray_source_box = xray_source(
                         redshift=z,
@@ -2912,6 +2912,16 @@ def run_coeval(
                 except OSError:
                     pass
 
+            # we only need the SFR fields at previous redshifts for XraySourceBox
+            if hb is not None:
+                try:
+                    hb.prepare(
+                        keep=["halo_sfr", "halo_sfr_mini", "log10_Mcrit_LW_ave"],
+                        force=always_purge,
+                    )
+                except OSError:
+                    pass
+
             if z in redshift:
                 logger.debug(f"PID={os.getpid()} doing brightness temp for z={z}")
                 ib_tracker[redshift.index(z)] = ib2
@@ -2936,6 +2946,7 @@ def run_coeval(
                 ib = ib2
                 pf = pf2
                 _bt = None
+                hb = hb2
 
             perturb_files.append((z, os.path.join(direc, pf2.filename)))
             if flag_options.USE_HALO_FIELD:
@@ -3446,6 +3457,7 @@ def run_lightcone(
         prev_coeval = None
         st2 = None
         hbox2 = None
+        hbox = None
 
         if lightcone_filename and not Path(lightcone_filename).exists():
             lightcone.save(lightcone_filename)
@@ -3470,10 +3482,10 @@ def run_lightcone(
                     perturbed_field=pf2,
                     **kw,
                 )
-                z_halos.append(z)
-                hboxes.append(hbox2)
 
                 if flag_options.USE_TS_FLUCT:
+                    z_halos.append(z)
+                    hboxes.append(hbox2)
                     xray_source_box = xray_source(
                         redshift=z,
                         z_halos=z_halos,
@@ -3587,7 +3599,18 @@ def run_lightcone(
                 except OSError:
                     pass
 
+            # we only need the SFR fields at previous redshifts for XraySourceBox
+            if hbox is not None:
+                try:
+                    hbox.prepare(
+                        keep=["halo_sfr", "halo_sfr_mini", "log10_Mcrit_LW_ave"],
+                        force=always_purge,
+                    )
+                except OSError:
+                    pass
+
             pf = pf2
+            hbox = hbox2
 
             if z <= return_at_z:
                 # Optionally return when the lightcone is only partially filled