Haloprop Correlation

src/py21cmfast/drivers/_param_config.py

@@ -224,8 +224,8 @@ def check_consistency(kwargs: dict[str, Any], outputs: dict[str, OutputStruct]):
     def _make_wisdoms(self, use_fftw_wisdom: bool):
         construct_fftw_wisdoms(use_fftw_wisdom=use_fftw_wisdom)
 
-    def _broadcast_inputs(self, inputs: InputParameters):
-        broadcast_input_struct(inputs=inputs)
+    def _broadcast_inputs(self, inputs: InputParameters, redshift: float | None = None):
+        broadcast_input_struct(inputs=inputs, redshift=redshift)
 
     def _free_cosmo_tables(self):
         free_cosmo_tables()
@@ -458,7 +458,7 @@ def __call__(self, **kwargs) -> OutputStruct:
             kwargs["inputs"] = inputs
 
         if out is None:
-            self._broadcast_inputs(inputs)
+            self._broadcast_inputs(inputs, redshift=current_redshift)
             self._make_wisdoms(inputs.matter_options.USE_FFTW_WISDOM)
             out = self._func(**kwargs)
             self._handle_write_to_cache(cache, write, out)

src/py21cmfast/drivers/single_field.py

@@ -288,6 +288,7 @@ def compute_halo_grid(
     initial_conditions: InitialConditions,
     inputs: InputParameters | None = None,
     halo_catalog: HaloCatalog | None = None,
+    previous_halo_catalog: HaloCatalog | None = None,
     previous_spin_temp: TsBox | None = None,
     previous_ionize_box: IonizedBox | None = None,
 ) -> HaloBox:
@@ -332,6 +333,17 @@ def compute_halo_grid(
         else:
             halo_catalog = HaloCatalog.dummy()
 
+    if previous_halo_catalog is None:
+        if (
+            inputs.matter_options.has_discrete_halos
+            and redshift < inputs.simulation_options.Z_HEAT_MAX
+        ):
+            raise ValueError(
+                "You must provide previous_halo_catalog for discrete halo models below Z_HEAT_MAX"
+            )
+        else:
+            previous_halo_catalog = HaloCatalog.dummy()
+
     # NOTE: due to the order, we use the previous spin temp here, like spin_temperature,
     #       but UNLIKE ionize_box, which uses the current box
     # TODO: think about the inconsistency here
@@ -361,6 +373,7 @@ def compute_halo_grid(
     return box.compute(
         initial_conditions=initial_conditions,
         halo_catalog=halo_catalog,
+        previous_halo_catalog=previous_halo_catalog,
         previous_ionize_box=previous_ionize_box,
         previous_spin_temp=previous_spin_temp,
     )

src/py21cmfast/src/BrightnessTemperatureBox.c

@@ -19,10 +19,11 @@
 #include "indexing.h"
 #include "logger.h"
 
-int ComputeBrightnessTemp(float redshift, TsBox *spin_temp, IonizedBox *ionized_box,
-                          PerturbedField *perturb_field, BrightnessTemp *box) {
+int ComputeBrightnessTemp(TsBox *spin_temp, IonizedBox *ionized_box, PerturbedField *perturb_field,
+                          BrightnessTemp *box) {
     int status;
     Try {  // Try block around whole function.
+        double redshift = get_current_redshift();
         LOG_DEBUG("Starting Brightness Temperature calculation for redshift %f", redshift);
         // Makes the parameter structs visible to a variety of functions/macros
         // Do each time to avoid Python garbage collection issues

src/py21cmfast/src/BrightnessTemperatureBox.h

@@ -5,7 +5,7 @@
 #include "InputParameters.h"
 #include "OutputStructs.h"
 
-int ComputeBrightnessTemp(float redshift, TsBox *spin_temp, IonizedBox *ionized_box,
-                          PerturbedField *perturb_field, BrightnessTemp *box);
+int ComputeBrightnessTemp(TsBox *spin_temp, IonizedBox *ionized_box, PerturbedField *perturb_field,
+                          BrightnessTemp *box);
 
 #endif

src/py21cmfast/src/HaloBox.h

@@ -17,9 +17,11 @@ typedef struct HaloProperties {
     double count;  // from integral
     double halo_mass;
     double stellar_mass;
-    double halo_sfr;
     double stellar_mass_mini;
-    double sfr_mini;
+    double sfr_10;
+    double sfr_100;
+    double sfr_10_mcg;
+    double sfr_100_mcg;
     double fescweighted_sfr;
     double n_ion;
     double halo_xray;
@@ -47,16 +49,17 @@ typedef struct IntegralCondition {
 void set_integral_constants(IntegralCondition *consts, double redshift, double M_min, double M_max,
                             double M_cell);
 
-int ComputeHaloBox(double redshift, InitialConditions *ini_boxes, HaloCatalog *halos,
+int ComputeHaloBox(InitialConditions *ini_boxes, HaloCatalog *halos, HaloCatalog *halos_prev,
                    TsBox *previous_spin_temp, IonizedBox *previous_ionize_box, HaloBox *grids);
 
 void get_cell_integrals(double dens, double l10_mturn_a, double l10_mturn_m,
                         ScalingConstants *consts, IntegralCondition *int_consts,
                         HaloProperties *properties);
-void set_halo_properties(double halo_mass, double M_turn_a, double M_turn_m,
-                         ScalingConstants *consts, double *input_rng, HaloProperties *output);
 
-int convert_halo_props(double redshift, InitialConditions *ics, TsBox *prev_ts,
-                       IonizedBox *prev_ion, HaloCatalog *halo_catalog,
-                       PerturbedHaloCatalog *halo_catalog_out);
+void set_halo_properties(double snapshot_time, double halo_mass, double M_turn_a, double M_turn_m,
+                         double prog_hm, double prog_sm[2], ScalingConstants *consts,
+                         double *input_rng, HaloProperties *output);
+
+int convert_halo_props(InitialConditions *ics, TsBox *prev_ts, IonizedBox *prev_ion,
+                       HaloCatalog *halo_catalog, PerturbedHaloCatalog *halo_catalog_out);
 #endif

src/py21cmfast/src/HaloCatalog.c

@@ -35,22 +35,29 @@ int pixel_in_halo(int grid_dim, int z_dim, int x, int x_index, int y, int y_inde
                   int z_index, float Rsq_curr_index);
 void free_halo_catalog(HaloCatalog *halos);
 
-int ComputeHaloCatalog(float redshift_desc, float redshift, InitialConditions *boxes,
-                       unsigned long long int random_seed, HaloCatalog *halos_desc,
-                       HaloCatalog *halos) {
+int ComputeHaloCatalog(InitialConditions *boxes, unsigned long long int random_seed,
+                       HaloCatalog *halos_desc, HaloCatalog *halos) {
     int status;
 
     Try {  // This Try brackets the whole function, so we don't indent.
-
+        double redshift = get_current_redshift();
+        bool from_catalog =
+            (matter_options_global->SOURCE_MODEL == 4 && get_descendant_redshift() > 0);
+        if (halos->sfh_computed || halos_desc->sfh_computed) {
+            LOG_ERROR(
+                "You have passed a halo catalog with SFH already computed to the stochastic "
+                "sampler. "
+                "This is not allowed.");
+            Throw(ValueError);
+        }
         // This happens if we are updating a halo field (no need to redo big halos)
-        if (matter_options_global->SOURCE_MODEL == 4 && redshift_desc > 0) {
+        if (from_catalog) {
             LOG_DEBUG("Halo sampling switched on, bypassing halo finder to update %llu halos...",
                       halos_desc->n_halos);
             // this would hold the two boxes used in the halo sampler, but here we are taking the
             // sample from a catalogue so we define a dummy here
             float *dummy_box = NULL;
-            stochastic_halofield(random_seed, redshift_desc, redshift, dummy_box, dummy_box,
-                                 halos_desc, halos);
+            stochastic_halofield(random_seed, dummy_box, dummy_box, halos_desc, halos);
             return 0;
         }
 
@@ -402,8 +409,8 @@ int ComputeHaloCatalog(float redshift_desc, float redshift, InitialConditions *b
                 }
             }
 
-            stochastic_halofield(random_seed, redshift_desc, redshift, boxes->lowres_density,
-                                 halo_overlap_box, halos_dexm, halos);
+            stochastic_halofield(random_seed, boxes->lowres_density, halo_overlap_box, halos_dexm,
+                                 halos);
 
             // Here, halos_dexm is allocated in the C, so free it
             free_halo_catalog(halos_dexm);
@@ -539,18 +546,18 @@ void init_halo_coords(HaloCatalog *halos, long long unsigned int n_halos) {
     halos->halo_masses = (float *)calloc(alloc_size, sizeof(float));
     halos->halo_coords = (float *)calloc(3 * alloc_size, sizeof(float));
 
-    halos->star_rng = (float *)calloc(alloc_size, sizeof(float));
-    halos->sfr_rng = (float *)calloc(alloc_size, sizeof(float));
-    halos->xray_rng = (float *)calloc(alloc_size, sizeof(float));
+    halos->sfr_10 = (float *)calloc(alloc_size, sizeof(float));
+    halos->sfr_100 = (float *)calloc(alloc_size, sizeof(float));
+    halos->stellar_mass = (float *)calloc(alloc_size, sizeof(float));
 }
 
 void free_halo_catalog(HaloCatalog *halos) {
     LOG_DEBUG("Freeing HaloCatalog instance.");
     free(halos->halo_masses);
     free(halos->halo_coords);
-    free(halos->star_rng);
-    free(halos->sfr_rng);
-    free(halos->xray_rng);
+    free(halos->sfr_10);
+    free(halos->sfr_100);
+    free(halos->stellar_mass);
     halos->n_halos = 0;
 }
 

src/py21cmfast/src/HaloCatalog.h

@@ -5,8 +5,7 @@
 #include "InputParameters.h"
 #include "OutputStructs.h"
 
-int ComputeHaloCatalog(float redshift_desc, float redshift, InitialConditions *boxes,
-                       unsigned long long int random_seed, HaloCatalog *halos_desc,
-                       HaloCatalog *halos);
+int ComputeHaloCatalog(InitialConditions *boxes, unsigned long long int random_seed,
+                       HaloCatalog *halos_desc, HaloCatalog *halos);
 
 #endif

src/py21cmfast/src/InputParameters.c

@@ -66,6 +66,26 @@ void Broadcast_struct_global_noastro(SimulationOptions *simulation_options,
     cosmo_params_global = cosmo_params;
 }
 
+void Broadcast_snapshot_info(int n_nodes, double *node_redshifts, int curr_node) {
+    node_redshifts_global.n_nodes = n_nodes;
+    node_redshifts_global.node_redshifts = node_redshifts;
+    node_redshifts_global.curr_node = curr_node;
+}
+
+double get_redshift_relative(int offset) {
+    int target_node = node_redshifts_global.curr_node + offset;
+    if (target_node >= 0 && target_node < node_redshifts_global.n_nodes) {
+        return node_redshifts_global.node_redshifts[target_node];
+    } else {
+        return -1.0;  // or some other sentinel value indicating out of bounds
+    }
+}
+
+// some useful aliases
+double get_current_redshift() { return get_redshift_relative(0); }
+double get_previous_redshift() { return get_redshift_relative(-1); }
+double get_descendant_redshift() { return get_redshift_relative(1); }
+
 void Free_cosmo_tables_global() {
     if (allocated_cosmo_tables) {
         if (matter_options_global->POWER_SPECTRUM == 5) {
@@ -88,6 +108,7 @@ MatterOptions *matter_options_global;
 CosmoParams *cosmo_params_global;
 AstroParams *astro_params_global;
 AstroOptions *astro_options_global;
+NodeRedshifts node_redshifts_global;
 CosmoTables *cosmo_tables_global;
 
 // data paths, wisdoms, etc

src/py21cmfast/src/InputParameters.h

@@ -13,4 +13,11 @@ void Broadcast_struct_global_all(SimulationOptions *simulation_options,
 void Broadcast_struct_global_noastro(SimulationOptions *simulation_options,
                                      MatterOptions *matter_options, CosmoParams *cosmo_params);
 
+void Broadcast_snapshot_info(int n_nodes, double *node_redshifts, int curr_node);
+
+double get_redshift_relative(int offset);
+double get_current_redshift();
+double get_previous_redshift();
+double get_descendant_redshift();
+
 #endif

src/py21cmfast/src/IonisationBox.c

@@ -1310,13 +1310,14 @@ void set_recombination_rates(IonizedBox *box, IonizedBox *previous_ionize_box,
     }
 }
 
-int ComputeIonizedBox(float redshift, float prev_redshift, PerturbedField *perturbed_field,
-                      PerturbedField *previous_perturbed_field, IonizedBox *previous_ionize_box,
-                      TsBox *spin_temp, HaloBox *halos, InitialConditions *ini_boxes,
-                      IonizedBox *box) {
+int ComputeIonizedBox(PerturbedField *perturbed_field, PerturbedField *previous_perturbed_field,
+                      IonizedBox *previous_ionize_box, TsBox *spin_temp, HaloBox *halos,
+                      InitialConditions *ini_boxes, IonizedBox *box) {
     int status;
 
     Try {  // This Try brackets the whole function, so we don't indent.
+        double redshift = get_current_redshift();
+        double prev_redshift = get_previous_redshift();
         LOG_DEBUG("input values:");
         LOG_DEBUG("redshift=%f, prev_redshift=%f", redshift, prev_redshift);
 #if LOG_LEVEL >= DEBUG_LEVEL

src/py21cmfast/src/IonisationBox.h

@@ -4,9 +4,8 @@
 #include "InputParameters.h"
 #include "OutputStructs.h"
 
-int ComputeIonizedBox(float redshift, float prev_redshift, PerturbedField *perturbed_field,
-                      PerturbedField *previous_perturbed_field, IonizedBox *previous_ionize_box,
-                      TsBox *spin_temp, HaloBox *halos, InitialConditions *ini_boxes,
-                      IonizedBox *box);
+int ComputeIonizedBox(PerturbedField *perturbed_field, PerturbedField *previous_perturbed_field,
+                      IonizedBox *previous_ionize_box, TsBox *spin_temp, HaloBox *halos,
+                      InitialConditions *ini_boxes, IonizedBox *box);
 
 #endif

src/py21cmfast/src/PerturbedField.c

@@ -382,8 +382,7 @@ void compute_perturbed_velocities(unsigned short axis, double redshift,
                       simulation_options_global->HII_DIM, HII_D_PARA, "  ");
 }
 
-int ComputePerturbedField(float redshift, InitialConditions *boxes,
-                          PerturbedField *perturbed_field) {
+int ComputePerturbedField(InitialConditions *boxes, PerturbedField *perturbed_field) {
     /*
      ComputePerturbedField uses the first-order Langragian displacement field to move the
      masses in the cells of the density field. The high-res density field is extrapolated
@@ -395,6 +394,7 @@ int ComputePerturbedField(float redshift, InitialConditions *boxes,
     int status;
     Try {  // This Try{} wraps the whole function, so we don't indent.
 
+        double redshift = get_current_redshift();
         // Makes the parameter structs visible to a variety of functions/macros
         // Do each time to avoid Python garbage collection issues
 

src/py21cmfast/src/PerturbedField.h

@@ -4,7 +4,6 @@
 #include "InputParameters.h"
 #include "OutputStructs.h"
 
-int ComputePerturbedField(float redshift, InitialConditions *boxes,
-                          PerturbedField *perturbed_field);
+int ComputePerturbedField(InitialConditions *boxes, PerturbedField *perturbed_field);
 
 #endif

src/py21cmfast/src/PerturbedHaloCatalog.c

@@ -22,13 +22,12 @@
 #include "indexing.h"
 #include "logger.h"
 
-int ComputePerturbedHaloCatalog(float redshift, InitialConditions *boxes, TsBox *prev_ts,
-                                IonizedBox *prev_ion, HaloCatalog *halos,
-                                PerturbedHaloCatalog *halos_perturbed) {
+int ComputePerturbedHaloCatalog(InitialConditions *boxes, TsBox *prev_ts, IonizedBox *prev_ion,
+                                HaloCatalog *halos, PerturbedHaloCatalog *halos_perturbed) {
     int status;
 
     Try {  // This Try brackets the whole function, so we don't indent.
-
+        double redshift = get_current_redshift(boxes);
         LOG_DEBUG("input value:");
         LOG_DEBUG("redshift=%f", redshift);
 #if LOG_LEVEL >= SUPER_DEBUG_LEVEL
@@ -134,8 +133,7 @@ int ComputePerturbedHaloCatalog(float redshift, InitialConditions *boxes, TsBox
         }
 
         LOG_DEBUG("starting haloprops");
-        convert_halo_props(redshift, boxes, prev_ts, prev_ion, halos, halos_perturbed);
-        // Divide out multiplicative factor to return to pristine state
+        convert_halo_props(boxes, prev_ts, prev_ion, halos, halos_perturbed);
         LOG_SUPER_DEBUG("Number of halos exactly on the box edge = %llu of %llu", n_exact_dim,
                         halos->n_halos);
         if (error_in_parallel) {

src/py21cmfast/src/PerturbedHaloCatalog.h

@@ -4,8 +4,7 @@
 #include "InputParameters.h"
 #include "OutputStructs.h"
 
-int ComputePerturbedHaloCatalog(float redshift, InitialConditions *boxes, TsBox *prev_ts,
-                                IonizedBox *prev_ion, HaloCatalog *halos,
-                                PerturbedHaloCatalog *halos_perturbed);
+int ComputePerturbedHaloCatalog(InitialConditions *boxes, TsBox *prev_ts, IonizedBox *prev_ion,
+                                HaloCatalog *halos, PerturbedHaloCatalog *halos_perturbed);
 
 #endif