fix: subcell_rsd and perturb_field

src/py21cmfast/src/BrightnessTemperatureBox.c

@@ -46,8 +46,7 @@ int ComputeBrightnessTemp(float redshift, struct UserParams *user_params, struct
         }
     }
 
-    float d1_low, d1_high, d2_low, d2_high, gradient_component, min_gradient_component, subcell_width, x_val1, x_val2, subcell_displacement;
-    float RSD_pos_new, RSD_pos_new_boundary_low,RSD_pos_new_boundary_high, fraction_within, fraction_outside, cell_distance;
+    float gradient_component, min_gradient_component;
 
     double dvdx, max_v_deriv;
     float const_factor, T_rad, pixel_Ts_factor, pixel_x_HI, pixel_deltax, H;
@@ -103,13 +102,8 @@ int ComputeBrightnessTemp(float redshift, struct UserParams *user_params, struct
 
     ave /= (float)HII_TOT_NUM_PIXELS;
 
-    x_val1 = 0.;
-    x_val2 = 1.;
-
     // Get RSDs
     if (flag_options->APPLY_RSDS){
-        subcell_width = (user_params->BOX_LEN/(float)user_params->HII_DIM)/(float)(astro_params->N_RSD_STEPS);
-
         ave = 0.;
         get_velocity_gradient(user_params, v, vel_gradient);
 
@@ -128,10 +122,12 @@ int ComputeBrightnessTemp(float redshift, struct UserParams *user_params, struct
                         for (k=0; k<HII_D_PARA; k++){
                             dvdx = clip(vel_gradient[HII_R_FFT_INDEX(i,j,k)], -max_v_deriv, max_v_deriv);
                             box->brightness_temp[HII_R_INDEX(i,j,k)] /= (dvdx/H + 1.0);
+                            ave += box->brightness_temp[HII_R_INDEX(i,j,k)];
                         }
                     }
                 }
             }
+            ave /= (float)HII_TOT_NUM_PIXELS;
         } else {
             // This is if we're doing both TS_FLUCT and SUBCELL_RSD
             min_gradient_component = 1.0;
@@ -165,7 +161,10 @@ int ComputeBrightnessTemp(float redshift, struct UserParams *user_params, struct
         }
 
         if(flag_options->SUBCELL_RSD) {
-            apply_subcell_rsds(user_params, cosmo_params, flag_options, astro_params, ionized_box, box, subcell_width, redshift, spin_temp, T_rad, v);
+            ave = apply_subcell_rsds(
+                user_params, cosmo_params, flag_options, astro_params, ionized_box, box,
+                redshift, spin_temp, T_rad, v, H
+            );
         }
     }
 

src/py21cmfast/src/GenerateICs.c

@@ -34,10 +34,10 @@
 #include "recombinations.c"
 #include "IonisationBox.c"
 #include "SpinTemperatureBox.c"
+#include "subcell_rsds.c"
 #include "BrightnessTemperatureBox.c"
 #include "FindHaloes.c"
 #include "PerturbHaloField.c"
-#include "subcell_rsds.c"
 
 
 void adj_complex_conj(fftwf_complex *HIRES_box, struct UserParams *user_params, struct CosmoParams *cosmo_params){

src/py21cmfast/src/PerturbField.c

@@ -35,11 +35,12 @@ void compute_perturbed_velocities(
             LOWRES_density_perturb_saved,
             sizeof(fftwf_complex)*HII_KSPACE_NUM_PIXELS
         );
+        LOG_SUPER_DEBUG("dDdt_over_D=%.6e, dimension=%d, switch_mid=%d, f_pixel_factor=%f", dDdt_over_D, dimension, switch_mid, f_pixel_factor);
     }
 
     #pragma omp parallel \
         shared(LOWRES_density_perturb,HIRES_density_perturb,dDdt_over_D,dimension,switch_mid) \
-        private(n_x,n_y,n_z,k_x,k_y,k_z,k_sq) \
+        private(n_x,n_y,n_z,k_x,k_y,k_z,k_sq, kvec) \
         num_threads(user_params->N_THREADS)
     {
         #pragma omp for
@@ -86,6 +87,8 @@ void compute_perturbed_velocities(
         }
     }
 
+    LOG_SUPER_DEBUG("density_perturb after modification by dDdt: ");
+    debugSummarizeBox(LOWRES_density_perturb, user_params->HII_DIM, user_params->NON_CUBIC_FACTOR, "  ");
 
     if(user_params->PERTURB_ON_HIGH_RES) {
 
@@ -128,6 +131,9 @@ void compute_perturbed_velocities(
             }
         }
     }
+    LOG_SUPER_DEBUG("velocity: ");
+    debugSummarizeBox(velocity, user_params->HII_DIM, user_params->NON_CUBIC_FACTOR, "  ");
+
 }
 
 int ComputePerturbField(
@@ -331,10 +337,12 @@ int ComputePerturbField(
 
         // go through the high-res box, mapping the mass onto the low-res (updated) box
         LOG_DEBUG("Perturb the density field");
-#pragma omp parallel shared(init_growth_factor,boxes,f_pixel_factor,resampled_box,dimension) \
-                        private(i,j,k,xi,xf,yi,yf,zi,zf,HII_i,HII_j,HII_k,d_x,d_y,d_z,t_x,t_y,t_z,xp1,yp1,zp1) num_threads(user_params->N_THREADS)
+        #pragma omp parallel \
+            shared(init_growth_factor,boxes,f_pixel_factor,resampled_box,dimension) \
+            private(i,j,k,xi,xf,yi,yf,zi,zf,HII_i,HII_j,HII_k,d_x,d_y,d_z,t_x,t_y,t_z,xp1,yp1,zp1) \
+            num_threads(user_params->N_THREADS)
         {
-#pragma omp for
+            #pragma omp for
             for (i=0; i<user_params->DIM;i++){
                 for (j=0; j<user_params->DIM;j++){
                     for (k=0; k<D_PARA;k++){
@@ -475,10 +483,12 @@ int ComputePerturbField(
         debugSummarizeBoxDouble(resampled_box, dimension, user_params->NON_CUBIC_FACTOR, "  ");
 
         // Resample back to a float for remaining algorithm
-#pragma omp parallel shared(LOWRES_density_perturb,HIRES_density_perturb,resampled_box,dimension) \
-                        private(i,j,k) num_threads(user_params->N_THREADS)
+        #pragma omp parallel \
+            shared(LOWRES_density_perturb,HIRES_density_perturb,resampled_box,dimension) \
+            private(i,j,k) \
+            num_threads(user_params->N_THREADS)
         {
-#pragma omp for
+            #pragma omp for
             for (i=0; i<dimension; i++){
                 for (j=0; j<dimension; j++){
                     for (k=0; k<(unsigned long long)(user_params->NON_CUBIC_FACTOR*dimension); k++){

src/py21cmfast/src/UsefulFunctions.c

@@ -613,7 +613,7 @@ void debugSummarizeBox(float *box, int size, float ncf, char *indent){
             }
         }
 
-        LOG_SUPER_DEBUG("%sCorners: %f %f %f %f %f %f %f %f",
+        LOG_SUPER_DEBUG("%sCorners: %.4e %.4e %.4e %.4e %.4e %.4e %.4e %.4e",
             indent,
             corners[0], corners[1], corners[2], corners[3],
             corners[4], corners[5], corners[6], corners[7]
@@ -631,7 +631,7 @@ void debugSummarizeBox(float *box, int size, float ncf, char *indent){
         }
         mean=sum/(size*size*((int)(size*ncf)));
 
-        LOG_SUPER_DEBUG("%sSum/Mean/Min/Max: %f, %f, %f, %f", indent, sum, mean, mn, mx);
+        LOG_SUPER_DEBUG("%sSum/Mean/Min/Max: %.4e, %.4e, %.4e, %.4e", indent, sum, mean, mn, mx);
     }
 }
 
@@ -654,7 +654,7 @@ void debugSummarizeBoxDouble(double *box, int size, float ncf, char *indent){
             }
         }
 
-        LOG_SUPER_DEBUG("%sCorners: %lf %lf %lf %lf %lf %lf %lf %lf",
+        LOG_SUPER_DEBUG("%sCorners: %.4e %.4e %.4e %.4e %.4e %.4e %.4e %.4e",
             indent,
             corners[0], corners[1], corners[2], corners[3],
             corners[4], corners[5], corners[6], corners[7]
@@ -672,7 +672,7 @@ void debugSummarizeBoxDouble(double *box, int size, float ncf, char *indent){
         }
         mean=sum/(size*size*((int)(size*ncf)));
 
-        LOG_SUPER_DEBUG("%sSum/Mean/Min/Max: %lf, %lf, %lf, %lf", indent, sum, mean, mn, mx);
+        LOG_SUPER_DEBUG("%sSum/Mean/Min/Max: %.4e, %.4e, %.4e, %.4e", indent, sum, mean, mn, mx);
     }
 }
 

src/py21cmfast/src/subcell_rsds.c

@@ -1,15 +1,15 @@
-void apply_subcell_rsds(
+double apply_subcell_rsds(
     struct UserParams *user_params,
     struct CosmoParams *cosmo_params,
     struct FlagOptions *flag_options,
     struct AstroParams *astro_params,
     struct IonizedBox *ionized_box,
     struct BrightnessTemp *box,
-    float subcell_width,
     float redshift,
     struct SpinTemperature *spin_temp,
     float T_rad,
-    float *v
+    float *v,
+    float H
 ) {
 
     char wisdom_filename[500];
@@ -28,20 +28,14 @@ void apply_subcell_rsds(
         delta_T_RSD_LOS[i] = (float *)calloc(user_params->HII_DIM,sizeof(float));
     }
 
-    float d1_low, d1_high, d2_low, d2_high, gradient_component, min_gradient_component;
+    float gradient_component, min_gradient_component;
+    float d1_low, d1_high, d2_low, d2_high;
     float x_val1, x_val2, subcell_displacement;
     float RSD_pos_new, RSD_pos_new_boundary_low,RSD_pos_new_boundary_high;
     float fraction_within, fraction_outside, cell_distance;
 
-    double dvdx, max_v_deriv;
-    float const_factor, pixel_Ts_factor, pixel_x_HI, pixel_deltax, H;
     float cellsize = user_params->BOX_LEN/(float)user_params->HII_DIM;
-
-    init_ps();
-
-    H = hubble(redshift);
-    const_factor = 27 * (cosmo_params->OMb*cosmo_params->hlittle*cosmo_params->hlittle/0.023) *
-    sqrt( (0.15/(cosmo_params->OMm)/(cosmo_params->hlittle)/(cosmo_params->hlittle)) * (1.+redshift)/10.0 );
+    float subcell_width = cellsize/(float)(astro_params->N_RSD_STEPS);
 
     // normalised units of cell length. 0 equals beginning of cell, 1 equals end of cell
     // These are the sub-cell central positions (x_pos_offset), and the corresponding normalised value (x_pos) between 0 and 1
@@ -50,6 +44,9 @@ void apply_subcell_rsds(
         x_pos[ii] = x_pos_offset[ii]/cellsize;
     }
 
+    x_val1 = 0.0;
+    x_val2 = 1.0;
+
     /*
     Note to convert the velocity v, to a displacement in redshift space, convert
     from s -> r + (1+z)*v/H(z). To convert the velocity within the array v to km/s, it
@@ -276,4 +273,5 @@ void apply_subcell_rsds(
     }
 
     ave /= (float)HII_TOT_NUM_PIXELS;
+    return(ave);
 }

tests/conftest.py

@@ -89,6 +89,7 @@ def setup_and_teardown_package(tmpdirec, request):
 
     log_level = request.config.getoption("--log-level-21") or logging.INFO
     logging.getLogger("py21cmfast").setLevel(log_level)
+    logging.getLogger("21cmFAST").setLevel(log_level)
 
     yield
 

tests/produce_integration_test_data.py

@@ -52,22 +52,21 @@
 LIGHTCONE_FIELDS = [
     "density",
     "velocity",
-    "xH_box",
     "Ts_box",
-    "z_re_box",
     "Gamma12_box",
     "dNrec_box",
     "x_e_box",
     "Tk_box",
     "J_21_LW_box",
+    "xH_box",
+    "z_re_box",
     "brightness_temp",
 ]
 
-COEVAL_FIELDS = [
-    "lowres_density",
-    "lowres_vx_2LPT",
-    "lowres_vx",
-] + LIGHTCONE_FIELDS
+COEVAL_FIELDS = LIGHTCONE_FIELDS.copy()
+COEVAL_FIELDS.insert(COEVAL_FIELDS.index("Ts_box"), "lowres_density")
+COEVAL_FIELDS.insert(COEVAL_FIELDS.index("Ts_box"), "lowres_vx_2LPT")
+COEVAL_FIELDS.insert(COEVAL_FIELDS.index("Ts_box"), "lowres_vx")
 
 OPTIONS = {
     "simple": [12, {}],

tests/test_integration_features.py

@@ -67,7 +67,10 @@ def test_power_spectra_coeval(name, module_direc, plt):
     if plt == mpl.pyplot:
         make_coeval_comparison_plot(test_k, true_powers, test_powers, plt)
 
-    for key, value in true_powers.items():
+    for key in prd.COEVAL_FIELDS:
+        if key not in true_powers:
+            continue
+        value = true_powers[key]
         print(f"Testing {key}")
         assert np.sum(~np.isclose(value, test_powers[key], atol=0, rtol=1e-2)) < 10
         np.testing.assert_allclose(value, test_powers[key], atol=0, rtol=1e-1)
@@ -102,7 +105,12 @@ def test_power_spectra_lightcone(name, module_direc, plt):
             test_k, lc.node_redshifts, true_powers, true_global, test_powers, lc, plt
         )
 
-    for key, value in true_powers.items():
+    for key in prd.LIGHTCONE_FIELDS:
+        if key not in true_powers:
+            continue
+
+        value = true_powers[key]
+
         if value[0] > 0:
             print(f"Testing {key}")
             # Ensure all but 10 of the values is within 1%, and none of the values