Ability to inject initial conditions (only density field)

docs/faqs/misc.rst

@@ -117,3 +117,13 @@ using the ``InitialConditions`` as an example::
 You would use this ``.set()`` method on each of the fields you needed to set. Now this
 data should be properly shared with the backend C-code, and the object can be used
 in subsequent steps within ``21cmFAST``.
+
+Can I inject my own initial conditions for the density field?
+-------------------------------------------------------------
+Yes, if you have a high resolution realization of the initial linear density field (at z=0),
+you can run
+
+    ics = p21c.compute_initial_conditions(inputs=p21c.InputParameters(), initial_density=my_hires_density_field)
+
+This will output an ``InitialConditions`` instance where all of its fields (initial densities and velocities) are consistent
+with your input ``my_hires_density_field``. You can then pass ``ics`` as an input to higher level functions (``run_coeval`` or ``run_lightcone``).

src/py21cmfast/drivers/coeval.py

@@ -840,11 +840,15 @@ def _setup_ics_and_pfs_for_scrolling(
     inputs: InputParameters,
     write: CacheConfig,
     progressbar: bool,
+    overdensity_z0: float | None = None,
     **iokw,
 ) -> tuple[InitialConditions, list[PerturbedField], list[HaloCatalog], dict]:
     if initial_conditions is None:
         initial_conditions = sf.compute_initial_conditions(
-            inputs=inputs, write=write.initial_conditions, **iokw
+            inputs=inputs,
+            write=write.initial_conditions,
+            initial_density=overdensity_z0,
+            **iokw,
         )
 
     # We can go ahead and purge some of the stuff in the initial_conditions, but only if

src/py21cmfast/drivers/global_evolution.py

@@ -133,7 +133,7 @@ def get_fields(cls, inputs: InputParameters) -> tuple:
             possible_outputs.append(TsBox.new(inputs, redshift=0))
         if inputs.matter_options.lagrangian_source_grid:
             possible_outputs.append(HaloBox.new(inputs, redshift=0))
-        field_names = ("neutral_fraction", "ionisation_rate_G12")
+        field_names = ("density", "neutral_fraction", "ionisation_rate_G12")
         for output in possible_outputs:
             field_names += tuple(output.arrays.keys())
         return field_names
@@ -231,6 +231,7 @@ def run_global_evolution(
     inputs: InputParameters,
     source_model: str | None = None,
     progressbar: bool = False,
+    overdensity_z0: float | None = None,
 ):
     r"""
     Compute the global evolution of all the fields in the simulation.
@@ -253,6 +254,8 @@ def run_global_evolution(
             and then mapping properties to the Eulerian grid using 2LPT.
     progressbar: bool, optional
         If True, a progress bar will be displayed throughout the simulation. Defaults to False.
+    overdensity_z0: float, optional
+        The linear matter overdensity at z=0. Default is 0. This input argument could be useful for doing linear perturbation analysis with 21cmFAST.
 
     Returns
     -------
@@ -344,6 +347,7 @@ def run_global_evolution(
         initial_conditions=None,
         write=CacheConfig.off(),
         progressbar=progressbar,
+        overdensity_z0=overdensity_z0,
         **iokw,
     )
 

src/py21cmfast/drivers/single_field.py

@@ -34,14 +34,20 @@
 
 
 @single_field_func
-def compute_initial_conditions(*, inputs: InputParameters) -> InitialConditions:
+def compute_initial_conditions(
+    *, inputs: InputParameters, initial_density: np.ndarray | float | None = None
+) -> InitialConditions:
     r"""
     Compute initial conditions.
 
     Parameters
     ----------
     inputs
         The InputParameters instance defining the run.
+    initial_density: np.ndarray or float, optional
+        A realization of the density field on the high resolution grid.
+        This input can also be used to determine the global density field,
+        in case we have a single cell in the box.
     regenerate : bool, optional
         Whether to force regeneration of data, even if matching cached data is found.
     cache
@@ -62,16 +68,43 @@ def compute_initial_conditions(*, inputs: InputParameters) -> InitialConditions:
         required_arrays = PerturbedField.new(
             redshift=0, inputs=inputs
         ).get_required_input_arrays(ics)
+
+        # Set the arrays to zero, or according to initial_density if the arrays are density fields
         for array in required_arrays:
+            if (initial_density is not None) and (
+                array in ["hires_density", "lowres_density"]
+            ):
+                value = initial_density
+            else:
+                value = 0.0
             setattr(
                 ics,
                 array,
                 Array(shape=shape, dtype=np.float32)
                 .initialize()
-                .with_value(val=np.zeros(shape)),
+                .with_value(val=value * np.ones(shape)),
             )
         return ics
     else:
+        if initial_density is not None:
+            if np.abs(initial_density.mean() > 1e-3):
+                warnings.warn(
+                    f"Your initial_density has mean {initial_density.mean()}. "
+                    + "Make sure you know what you are doing.",
+                    stacklevel=2,
+                )
+            shape = ics.hires_density.shape
+            if initial_density.shape != shape:
+                raise ValueError(
+                    "The shape of your high resolution initial_density is not consistent with inputs!"
+                    + f" According to inputs, initial_density must be of shape {shape}, got {initial_density.shape}."
+                )
+
+            ics.hires_density = (
+                Array(shape=shape, dtype=np.float32)
+                .initialize()
+                ._with_value_not_computed(val=initial_density)
+            )
         return ics.compute()
 
 

src/py21cmfast/src/InitialConditions.c

@@ -566,6 +566,7 @@ int ComputeInitialConditions(unsigned long long random_seed, InitialConditions *
 #endif
 
         int i, j, k;
+        unsigned long long int box_ct;
 
         gsl_rng *r[simulation_options_global->N_THREADS];
         seed_rng_threads(r, random_seed);
@@ -617,36 +618,83 @@ int ComputeInitialConditions(unsigned long long random_seed, InitialConditions *
 
         init_ps();
 
-        sample_ic_modes(HIRES_box, hi_dim, box_len, r);
+        // Check if the input hires density box is all zero
+        bool non_zero_input = false;
+#pragma omp parallel shared(boxes, non_zero_input) \
+    num_threads(simulation_options_global -> N_THREADS)
+        {
+#pragma omp for
+            for (box_ct = 0; box_ct < TOT_NUM_PIXELS; box_ct++) {
+                if (!non_zero_input && boxes->hires_density[box_ct]) {
+#pragma omp atomic write
+                    non_zero_input = true;
+                }
+            }
+        }
 
-        memcpy(HIRES_box_saved, HIRES_box, sizeof(fftwf_complex) * KSPACE_NUM_PIXELS);
+        // If we have a non zero input, we reverse the order of operations,
+        // namely we assign the input to HIRES_box, inverse FFT it, and save the result in
+        // HIRES_box_saved
+        if (non_zero_input) {
+#pragma omp parallel shared(boxes, HIRES_box) private(i, j, k) \
+    num_threads(simulation_options_global -> N_THREADS)
+            {
+                unsigned long long int index_r, index_f;
+#pragma omp for
+                for (i = 0; i < hi_dim[0]; i++) {
+                    for (j = 0; j < hi_dim[1]; j++) {
+                        for (k = 0; k < hi_dim[2]; k++) {
+                            index_r = grid_index_general(i, j, k, hi_dim);
+                            index_f = grid_index_fftw_r(i, j, k, hi_dim);
+                            // remember to add the factor of VOLUME/TOT_NUM_PIXELS when converting
+                            // from real space to k-space
+                            *((float *)HIRES_box + index_f) =
+                                boxes->hires_density[index_r] * VOLUME / TOT_NUM_PIXELS;
+                        }
+                    }
+                }
+            }
+            LOG_SUPER_DEBUG("Saved struct to HIRES_box.");
 
-        /* ASSIGN HIRES DENSITY */
-        // FFT back to real space
-        int stat =
-            dft_c2r_cube(matter_options_global->USE_FFTW_WISDOM, simulation_options_global->DIM,
-                         D_PARA, simulation_options_global->N_THREADS, HIRES_box);
-        if (stat > 0) Throw(stat);
-        LOG_SUPER_DEBUG("FFT'd hires boxes.");
+            int stat =
+                dft_r2c_cube(matter_options_global->USE_FFTW_WISDOM, simulation_options_global->DIM,
+                             D_PARA, simulation_options_global->N_THREADS, HIRES_box);
+            if (stat > 0) Throw(stat);
+            LOG_SUPER_DEBUG("Inverse FFT'd hires boxes.");
+
+            memcpy(HIRES_box_saved, HIRES_box, sizeof(fftwf_complex) * KSPACE_NUM_PIXELS);
+        } else {
+            sample_ic_modes(HIRES_box, hi_dim, box_len, r);
+
+            memcpy(HIRES_box_saved, HIRES_box, sizeof(fftwf_complex) * KSPACE_NUM_PIXELS);
+
+            /* ASSIGN HIRES DENSITY */
+            // FFT back to real space
+            int stat =
+                dft_c2r_cube(matter_options_global->USE_FFTW_WISDOM, simulation_options_global->DIM,
+                             D_PARA, simulation_options_global->N_THREADS, HIRES_box);
+            if (stat > 0) Throw(stat);
+            LOG_SUPER_DEBUG("FFT'd hires boxes.");
 
 #pragma omp parallel shared(boxes, HIRES_box) private(i, j, k) \
     num_threads(simulation_options_global -> N_THREADS)
-        {
-            unsigned long long int index_r, index_f;
+            {
+                unsigned long long int index_r, index_f;
 #pragma omp for
-            for (i = 0; i < hi_dim[0]; i++) {
-                for (j = 0; j < hi_dim[1]; j++) {
-                    for (k = 0; k < hi_dim[2]; k++) {
-                        index_r = grid_index_general(i, j, k, hi_dim);
-                        index_f = grid_index_fftw_r(i, j, k, hi_dim);
-                        boxes->hires_density[index_r] = *((float *)HIRES_box + index_f) / VOLUME;
+                for (i = 0; i < hi_dim[0]; i++) {
+                    for (j = 0; j < hi_dim[1]; j++) {
+                        for (k = 0; k < hi_dim[2]; k++) {
+                            index_r = grid_index_general(i, j, k, hi_dim);
+                            index_f = grid_index_fftw_r(i, j, k, hi_dim);
+                            boxes->hires_density[index_r] =
+                                *((float *)HIRES_box + index_f) / VOLUME;
+                        }
                     }
                 }
             }
-        }
-
-        LOG_SUPER_DEBUG("Saved HIRES_box to struct.");
 
+            LOG_SUPER_DEBUG("Saved HIRES_box to struct.");
+        }
         /* FILTER AND ASSIGN LOWRES DENSITY */
         memcpy(HIRES_box, HIRES_box_saved, sizeof(fftwf_complex) * KSPACE_NUM_PIXELS);
 

src/py21cmfast/wrapper/arrays.py

@@ -162,6 +162,13 @@ def with_value(self, val: np.ndarray) -> Self:
             self, value=val.astype(self.dtype, copy=False), state=self.state.computed()
         )
 
+    def _with_value_not_computed(self, val: np.ndarray) -> Self:
+        """Set the array to a given value and return a new Array, but without indicating the array has been computed."""
+        return attrs.evolve(
+            self,
+            value=val.astype(self.dtype, copy=False),
+        )
+
     def computed(self) -> Self:
         """Set the array to a given value and return a new Array."""
         return attrs.evolve(self, state=self.state.computed())

tests/test_global_evolution.py

@@ -129,3 +129,58 @@ def test_compatability_with_database():
     fname = DATA_PATH / "global_evolution.h5"
     global_evolution = GlobalEvolution.from_file(fname)
     assert isinstance(global_evolution, GlobalEvolution)
+
+
+def test_linear_perturbation_theory(default_input_struct_ts):
+    """
+    Test that we can do linear perturbation theory with 21cmFAST.
+
+    What this test does is to confirm that the contrast of the fields we compute does not depend
+    on our chosen initial perturbation, at first order.
+    """
+    delta1 = 1e-7
+    delta2 = 1e-8
+
+    # Compute global evolution three times.
+    # We stop at relatively high redshift, where linear perturbation theory becomes non-valid below it
+    inputs = default_input_struct_ts.with_logspaced_redshifts(zmin=20)
+    global_evolution_0 = p21c.run_global_evolution(inputs=inputs)
+    global_evolution_delta1 = p21c.run_global_evolution(
+        inputs=inputs, overdensity_z0=delta1
+    )
+    global_evolution_delta2 = p21c.run_global_evolution(
+        inputs=inputs, overdensity_z0=delta2
+    )
+
+    for quantity in global_evolution_0.quantities:
+        f_0 = global_evolution_0.quantities[quantity]
+
+        f_delta1 = global_evolution_delta1.quantities[quantity]
+        contrast1 = (f_delta1 - f_0) / delta1
+
+        f_delta2 = global_evolution_delta2.quantities[quantity]
+        contrast2 = (f_delta2 - f_0) / delta2
+
+        np.testing.assert_allclose(contrast1, contrast2, atol=0, rtol=1e-5)
+
+
+def test_linear_density_field(default_input_struct):
+    """Test that the linear density field grows linearly with time."""
+    delta1 = 1e-7
+    delta2 = 1e-8
+
+    global_evolution1 = p21c.run_global_evolution(
+        inputs=default_input_struct, overdensity_z0=delta1
+    )
+    global_evolution2 = p21c.run_global_evolution(
+        inputs=default_input_struct, overdensity_z0=delta2
+    )
+
+    growth_factor1 = global_evolution1.quantities["density"] / delta1
+    growth_factor2 = global_evolution2.quantities["density"] / delta2
+
+    # Test that the linear density field is linear
+    np.testing.assert_allclose(growth_factor1, growth_factor2, atol=0, rtol=1e-5)
+
+    # Test that the linear density field grows with time
+    assert np.all(np.diff(growth_factor1) < 0)