From 836c8ef3ce417089dc07670336e159fe4a59f5fd Mon Sep 17 00:00:00 2001 From: daviesje Date: Sun, 1 Sep 2024 22:23:35 +0200 Subject: [PATCH 1/7] WIP filter tests --- src/py21cmfast/inputs.py | 4 + src/py21cmfast/src/HaloField.c | 4 +- src/py21cmfast/src/InitialConditions.c | 8 +- src/py21cmfast/src/IonisationBox.c | 23 +- src/py21cmfast/src/PerturbField.c | 6 +- src/py21cmfast/src/SpinTemperatureBox.c | 5 +- .../src/_functionprototypes_wrapper.h | 2 + src/py21cmfast/src/filtering.c | 241 ++++++------------ src/py21cmfast/src/filtering.h | 9 +- tests/test_filtering.py | 148 +++++++++++ 10 files changed, 262 insertions(+), 188 deletions(-) create mode 100644 tests/test_filtering.py diff --git a/src/py21cmfast/inputs.py b/src/py21cmfast/inputs.py index a3d9f1b2..ddd67186 100644 --- a/src/py21cmfast/inputs.py +++ b/src/py21cmfast/inputs.py @@ -1304,6 +1304,10 @@ def validate_all_inputs( The input params may be modified in-place in this function, but if so, a warning should be emitted. """ + if global_params.HII_FILTER not in [0, 1, 2]: + raise ValueError( + "global_params.HII_FITLER is not within the available options, 0: real space top-hat, 1: sharp k-space, 2: Gaussian." + ) if astro_params is not None: if astro_params.R_BUBBLE_MAX > user_params.BOX_LEN: astro_params.update(R_BUBBLE_MAX=user_params.BOX_LEN) diff --git a/src/py21cmfast/src/HaloField.c b/src/py21cmfast/src/HaloField.c index 02b8854d..02a0afa4 100644 --- a/src/py21cmfast/src/HaloField.c +++ b/src/py21cmfast/src/HaloField.c @@ -213,7 +213,7 @@ int ComputeHaloField(float redshift_desc, float redshift, UserParams *user_param // now filter the box on scale R // 0 = top hat in real space, 1 = top hat in k space - filter_box(density_field, 0, global_params.HALO_FILTER, R); + filter_box(density_field, 0, global_params.HALO_FILTER, R, 0.); // do the FFT to get delta_m box dft_c2r_cube(user_params->USE_FFTW_WISDOM, grid_dim, z_dim, user_params->N_THREADS, density_field); @@ -341,7 +341,7 @@ int ComputeHaloField(float redshift_desc, float redshift, UserParams *user_param dft_r2c_cube(user_params->USE_FFTW_WISDOM, grid_dim, z_dim, user_params->N_THREADS, density_field); if (user_params->DIM != user_params->HII_DIM) { //the tophat filter here will smoothe the grid to HII_DIM - filter_box(density_field, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(density_field, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } dft_c2r_cube(user_params->USE_FFTW_WISDOM, user_params->DIM, D_PARA, user_params->N_THREADS, density_field); diff --git a/src/py21cmfast/src/InitialConditions.c b/src/py21cmfast/src/InitialConditions.c index 7b39a9d4..046d2f67 100644 --- a/src/py21cmfast/src/InitialConditions.c +++ b/src/py21cmfast/src/InitialConditions.c @@ -270,7 +270,7 @@ int ComputeInitialConditions( // Only filter if we are perturbing on the low-resolution grid if(!user_params->PERTURB_ON_HIGH_RES) { if (user_params->DIM != user_params->HII_DIM) { - filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } // FFT back to real space @@ -344,7 +344,7 @@ int ComputeInitialConditions( //we only care about the lowres vcb box, so we filter it directly. if (user_params->DIM != user_params->HII_DIM) { - filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } //fft each velocity component back to real space @@ -428,7 +428,7 @@ int ComputeInitialConditions( // Filter only if we require perturbing on the low-res grid if(!user_params->PERTURB_ON_HIGH_RES) { if (user_params->DIM != user_params->HII_DIM) { - filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } } @@ -783,7 +783,7 @@ int ComputeInitialConditions( // Filter only if we require perturbing on the low-res grid if(!user_params->PERTURB_ON_HIGH_RES) { if (user_params->DIM != user_params->HII_DIM) { - filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_box, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } } diff --git a/src/py21cmfast/src/IonisationBox.c b/src/py21cmfast/src/IonisationBox.c index 06adf42f..6c931e9f 100644 --- a/src/py21cmfast/src/IonisationBox.c +++ b/src/py21cmfast/src/IonisationBox.c @@ -790,26 +790,21 @@ LOG_SUPER_DEBUG("excursion set normalisation, mean_f_coll_MINI: %e", box->mean_f ((R - cell_length_factor * (user_params->BOX_LEN / (double) (user_params->HII_DIM))) > FRACT_FLOAT_ERR)) { if (flag_options->USE_TS_FLUCT) { - filter_box(xe_filtered, 1, global_params.HII_FILTER, R); + filter_box(xe_filtered, 1, global_params.HII_FILTER, R, 0.); } if (recomb_filter_flag) { - filter_box(N_rec_filtered, 1, global_params.HII_FILTER, R); + filter_box(N_rec_filtered, 1, global_params.HII_FILTER, R, 0.); } if (flag_options->USE_HALO_FIELD) { - if(flag_options->USE_EXP_FILTER){ - filter_box_mfp(stars_filtered, 1, R, exp_mfp); - filter_box_mfp(sfr_filtered, 1, R, exp_mfp); - } - else{ - filter_box(stars_filtered, 1, global_params.HII_FILTER, R); - filter_box(sfr_filtered, 1, global_params.HII_FILTER, R); - } + int filter_hf = flag_options->USE_EXP_FILTER ? 3 : global_params.HII_FILTER; + filter_box(stars_filtered, 1, filter_hf, R, exp_mfp); + filter_box(sfr_filtered, 1, filter_hf, R, exp_mfp); } - filter_box(deltax_filtered, 1, global_params.HII_FILTER, R); + filter_box(deltax_filtered, 1, global_params.HII_FILTER, R, 0.); if(flag_options->USE_MINI_HALOS){ - filter_box(prev_deltax_filtered, 1, global_params.HII_FILTER, R); - filter_box(log10_Mturnover_MINI_filtered, 1, global_params.HII_FILTER, R); - filter_box(log10_Mturnover_filtered, 1, global_params.HII_FILTER, R); + filter_box(prev_deltax_filtered, 1, global_params.HII_FILTER, R, 0.); + filter_box(log10_Mturnover_MINI_filtered, 1, global_params.HII_FILTER, R, 0.); + filter_box(log10_Mturnover_filtered, 1, global_params.HII_FILTER, R, 0.); } } diff --git a/src/py21cmfast/src/PerturbField.c b/src/py21cmfast/src/PerturbField.c index 1f02f0db..02f453a3 100644 --- a/src/py21cmfast/src/PerturbField.c +++ b/src/py21cmfast/src/PerturbField.c @@ -115,7 +115,7 @@ void compute_perturbed_velocities( // smooth the high resolution field ready for resampling if (user_params->DIM != user_params->HII_DIM) - filter_box(HIRES_density_perturb, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_density_perturb, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); dft_c2r_cube(user_params->USE_FFTW_WISDOM, user_params->DIM, D_PARA, user_params->N_THREADS, HIRES_density_perturb); @@ -592,7 +592,7 @@ int ComputePerturbField( // Now filter the box if (user_params->DIM != user_params->HII_DIM) { - filter_box(HIRES_density_perturb, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0)); + filter_box(HIRES_density_perturb, 0, 0, L_FACTOR*user_params->BOX_LEN/(user_params->HII_DIM+0.0), 0.); } // FFT back to real space @@ -647,7 +647,7 @@ int ComputePerturbField( //smooth the field if (!global_params.EVOLVE_DENSITY_LINEARLY && global_params.SMOOTH_EVOLVED_DENSITY_FIELD){ - filter_box(LOWRES_density_perturb, 1, 2, global_params.R_smooth_density*user_params->BOX_LEN/(float)user_params->HII_DIM); + filter_box(LOWRES_density_perturb, 1, 2, global_params.R_smooth_density*user_params->BOX_LEN/(float)user_params->HII_DIM, 0.); } LOG_SUPER_DEBUG("LOWRES_density_perturb after smoothing: "); diff --git a/src/py21cmfast/src/SpinTemperatureBox.c b/src/py21cmfast/src/SpinTemperatureBox.c index 11d66796..0affacb3 100644 --- a/src/py21cmfast/src/SpinTemperatureBox.c +++ b/src/py21cmfast/src/SpinTemperatureBox.c @@ -567,7 +567,7 @@ void fill_Rbox_table(float **result, fftwf_complex *unfiltered_box, double * R_a // don't filter on cell size if (R > L_FACTOR*(user_params_global->BOX_LEN / user_params_global->HII_DIM)){ - filter_box(box, 1, global_params.HEAT_FILTER, R); + filter_box(box, 1, global_params.HEAT_FILTER, R, 0.); } // now fft back to real space @@ -608,6 +608,7 @@ void fill_Rbox_table(float **result, fftwf_complex *unfiltered_box, double * R_a //NOTE: I've moved this to a function to help in simplicity, it is not clear whether it is faster // to do all of one radii at once (more clustered FFT and larger thread blocks) or all of one box (better memory locality) +//TODO: filter speed tests void one_annular_filter(float *input_box, float *output_box, double R_inner, double R_outer, double *u_avg, double *f_avg){ int i,j,k; unsigned long long int ct; @@ -652,7 +653,7 @@ void one_annular_filter(float *input_box, float *output_box, double R_inner, dou // Don't filter on the cell scale if(R_inner > 0){ - filter_box_annulus(filtered_box, 1, R_inner, R_outer); + filter_box(filtered_box, 1, 4, R_inner, R_outer); } // now fft back to real space diff --git a/src/py21cmfast/src/_functionprototypes_wrapper.h b/src/py21cmfast/src/_functionprototypes_wrapper.h index a7a5c1ad..0e4d901f 100644 --- a/src/py21cmfast/src/_functionprototypes_wrapper.h +++ b/src/py21cmfast/src/_functionprototypes_wrapper.h @@ -157,6 +157,8 @@ int single_test_sample( UserParams *user_params, CosmoParams *cosmo_params, As int test_halo_props(double redshift, UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions * flag_options, float * vcb_grid, float *J21_LW_grid, float *z_re_grid, float *Gamma12_ion_grid, PerturbHaloField *halos, float *halo_props_out); +int test_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions *flag_options + , float *input_box, double R, double R_param, int filter_flag, double *result); /* Miscellaneous exposed functions for testing */ double dicke(double z); diff --git a/src/py21cmfast/src/filtering.c b/src/py21cmfast/src/filtering.c index 00ba9186..cb89f021 100644 --- a/src/py21cmfast/src/filtering.c +++ b/src/py21cmfast/src/filtering.c @@ -17,64 +17,57 @@ #include "indexing.h" #include "dft.h" -void filter_box_annulus(fftwf_complex *box, int RES, float R_inner, float R_outer){ - int n_x, n_z, n_y, dimension, midpoint; - float k_x, k_y, k_z, k_mag, kRinner, kRouter; - float f_inner, f_outer; - - switch(RES) { - case 0: - dimension = user_params_global->DIM; - midpoint = MIDDLE; - break; - case 1: - dimension = user_params_global->HII_DIM; - midpoint = HII_MIDDLE; - break; - default: - LOG_ERROR("Resolution for filter functions must be 0(DIM) or 1(HII_DIM)"); - Throw(ValueError); - break; - } - // loop through k-box - -#pragma omp parallel shared(box) private(n_x,n_y,n_z,k_x,k_y,k_z,k_mag,kRinner,kRouter,f_inner,f_outer) num_threads(user_params_global->N_THREADS) - { -#pragma omp for - for (n_x=0; n_xmidpoint) {k_x =(n_x-dimension) * DELTA_K;} - else {k_x = n_x * DELTA_K;} +double real_tophat_filter(double kR){ + if (kR < 1e-4) + return 1; + return 3.0*pow(kR, -3) * (sin(kR) - cos(kR)*kR); +} - for (n_y=0; n_ymidpoint) {k_y =(n_y-dimension) * DELTA_K;} - else {k_y = n_y * DELTA_K;} +//TODO: TEST USING kR^2 INSTEAD FOR SPEED +// ALSO TEST ASSIGNMENT vs MULTIPLICATION +double sharp_k_filter(double kR){ + // equates integrated volume to the real space top-hat (9pi/2)^(-1/3) + if (kR*0.413566994 > 1) + return 0.; + return 1; +} - for (n_z=0; n_z<=(int)(user_params_global->NON_CUBIC_FACTOR*midpoint); n_z++){ - k_z = n_z * DELTA_K_PARA; +double gaussian_filter(double kR_squared){ + return exp(-0.643*0.643*kR_squared/2.); +} - k_mag = sqrt(k_x*k_x + k_y*k_y + k_z*k_z); +double exp_mfp_filter(double k, double R, double mfp, double R_constant, double limit){ + double kR,kR2,f; + + kR = k*R; + if(kR < 1e-4) + return limit; + + kR2 = k*mfp; + //Davies & Furlanetto MFP-eps(r) window function + //The filter no longer approaches 1 at k->0, so we can't use the limit + f = (kR2*kR2*R + 2*mfp + R)*kR2*cos(kR); + f += (-kR2*kR2*mfp + kR2*kR2*R + limit + R)*sin(kR); + f *= R_constant; + f -= 2*kR2*mfp; + f *= -3.0*mfp/(kR*R*R*(kR2*kR2+1)*(kR2*kR2+1)); + return f; +} - kRinner = k_mag*R_inner; - kRouter = k_mag*R_outer; +double spherical_shell_filter(double k, double R_outer, double R_inner){ + double kR_inner = k*R_inner; + double kR_outer = k*R_outer; - if (kRinner > 1e-4){ - f_inner = 3.0/(pow(kRouter, 3) - pow(kRinner, 3)) * (sin(kRinner) - cos(kRinner)*kRinner); - f_outer = 3.0/(pow(kRouter, 3) - pow(kRinner, 3)) * (sin(kRouter) - cos(kRouter)*kRouter); - if(RES==1) { box[HII_C_INDEX(n_x, n_y, n_z)] *= (f_outer - f_inner); } - if(RES==0) { box[C_INDEX(n_x, n_y, n_z)] *= (f_outer - f_inner); } - } + if (kR_outer < 1e-4) + return 1.; - } - } - } // end looping through k box - } - return; + return 3.0/(pow(kR_outer, 3) - pow(kR_inner, 3)) \ + * (sin(kR_outer) - cos(kR_outer)*kR_outer \ + - sin(kR_inner) + cos(kR_inner)*kR_inner); } -void filter_box(fftwf_complex *box, int RES, int filter_type, float R){ - int n_x, n_z, n_y, dimension, midpoint; //TODO: figure out why defining as ULL breaks this - float k_x, k_y, k_z, k_mag, kR; - +void filter_box(fftwf_complex *box, int RES, int filter_type, float R, float R_param){ + int dimension, midpoint; //TODO: figure out why defining as ULL breaks this switch(RES) { case 0: dimension = user_params_global->DIM; @@ -90,11 +83,21 @@ void filter_box(fftwf_complex *box, int RES, int filter_type, float R){ break; } - // loop through k-box + //setup constants if needed + float R_constant_1, R_constant_2; + if(filter_type == 3){ + R_constant_1 = exp(-R/R_param); //independent of k + //k->0 limit of the mfp filter + R_constant_2 = (2*R_param*R_param - R_constant_1*(2*R_param*R_param + 2*R_param*R + R*R))*3*R_param/(R*R*R); + } -#pragma omp parallel shared(box) private(n_x,n_y,n_z,k_x,k_y,k_z,k_mag,kR) num_threads(user_params_global->N_THREADS) + // loop through k-box + #pragma omp parallel num_threads(user_params_global->N_THREADS) { -#pragma omp for + int n_x, n_z, n_y; + float k_x, k_y, k_z, k_mag_sq, kR; + unsigned long long grid_index; + #pragma omp for for (n_x=0; n_xmidpoint) {k_x =(n_x-dimension) * DELTA_K;} else {k_x = n_x * DELTA_K;} @@ -102,42 +105,36 @@ void filter_box(fftwf_complex *box, int RES, int filter_type, float R){ for (n_y=0; n_ymidpoint) {k_y =(n_y-dimension) * DELTA_K;} else {k_y = n_y * DELTA_K;} + for (n_z=0; n_z<=(int)(user_params_global->NON_CUBIC_FACTOR*midpoint); n_z++){ k_z = n_z * DELTA_K_PARA; + k_mag_sq = k_x*k_x + k_y*k_y + k_z*k_z; - if (filter_type == 0){ // real space top-hat - - k_mag = sqrt(k_x*k_x + k_y*k_y + k_z*k_z); - - kR = k_mag*R; // real space top-hat + grid_index = RES==1 ? HII_C_INDEX(n_x, n_y, n_z) : C_INDEX(n_x, n_y, n_z); - if (kR > 1e-4){ - if(RES==1) { box[HII_C_INDEX(n_x, n_y, n_z)] *= 3.0*pow(kR, -3) * (sin(kR) - cos(kR)*kR); } - if(RES==0) { box[C_INDEX(n_x, n_y, n_z)] *= 3.0*pow(kR, -3) * (sin(kR) - cos(kR)*kR); } - } + if (filter_type == 0){ // real space top-hat + kR = sqrt(k_mag_sq)*R; + box[grid_index] *= real_tophat_filter(kR); } else if (filter_type == 1){ // k-space top hat - + //NOTE: why was this commented???? // This is actually (kR^2) but since we zero the value and find kR > 1 this is more computationally efficient - // as we don't need to evaluate the slower sqrt function -// kR = 0.17103765852*( k_x*k_x + k_y*k_y + k_z*k_z )*R*R; - - k_mag = sqrt(k_x*k_x + k_y*k_y + k_z*k_z); - kR = k_mag*R; // real space top-hat - - kR *= 0.413566994; // equates integrated volume to the real space top-hat (9pi/2)^(-1/3) - if (kR > 1){ - if(RES==1) { box[HII_C_INDEX(n_x, n_y, n_z)] = 0; } - if(RES==0) { box[C_INDEX(n_x, n_y, n_z)] = 0; } - } + // kR = 0.17103765852*( k_x*k_x + k_y*k_y + k_z*k_z )*R*R; + kR = sqrt(k_mag_sq)*R; + box[grid_index] *= sharp_k_filter(kR); } else if (filter_type == 2){ // gaussian // This is actually (kR^2) but since we zero the value and find kR > 1 this is more computationally efficient - // as we don't need to evaluate the slower sqrt function - kR = 0.643*0.643*( k_x*k_x + k_y*k_y + k_z*k_z )*R*R; -// kR *= 0.643; // equates integrated volume to the real space top-hat - if(RES==1) { box[HII_C_INDEX(n_x, n_y, n_z)] *= pow(E, -kR/2.0); } - if(RES==0) { box[C_INDEX(n_x, n_y, n_z)] *= pow(E, -kR/2.0); } + kR = k_mag_sq*R*R; + box[grid_index] *= gaussian_filter(kR); + } + //The next two filters are not given by the HII_FILTER global, but used for specific grids + else if (filter_type == 3){ // exponentially decaying tophat, param == scale of decay (MFP) + //NOTE: This should be optimized, I havne't looked at it in a while + box[grid_index] *= exp_mfp_filter(sqrt(k_mag_sq),R,R_param,R_constant_1,R_constant_2); + } + else if (filter_type == 4){ //spherical shell, R_param == inner radius + box[grid_index] *= spherical_shell_filter(sqrt(k_mag_sq),R,R_param); } else{ if ( (n_x==0) && (n_y==0) && (n_z==0) ) @@ -151,79 +148,15 @@ void filter_box(fftwf_complex *box, int RES, int filter_type, float R){ return; } -void filter_box_mfp(fftwf_complex *box, int RES, float R, float mfp){ - int n_x, n_z, n_y, dimension, midpoint; - float k_x, k_y, k_z, k_mag, f, kR, kl; - float const1; - const1 = exp(-R/mfp); //independent of k, move it out of the loop - // LOG_DEBUG("Filtering box with R=%.2e, L=%.2e",R,mfp); - - switch(RES) { - case 0: - dimension = user_params_global->DIM; - midpoint = MIDDLE; - break; - case 1: - dimension = user_params_global->HII_DIM; - midpoint = HII_MIDDLE; - break; - default: - LOG_ERROR("Resolution for filter functions must be 0(DIM) or 1(HII_DIM)"); - Throw(ValueError); - break; - } - // loop through k-box - -#pragma omp parallel shared(box) private(n_x,n_y,n_z,k_x,k_y,k_z,k_mag,kR,kl,f) num_threads(user_params_global->N_THREADS) - { -#pragma omp for - for (n_x=0; n_xmidpoint) {k_x =(n_x-dimension) * DELTA_K;} - else {k_x = n_x * DELTA_K;} - - for (n_y=0; n_ymidpoint) {k_y =(n_y-dimension) * DELTA_K;} - else {k_y = n_y * DELTA_K;} - for (n_z=0; n_z<=(int)(user_params_global->NON_CUBIC_FACTOR*midpoint); n_z++){ - k_z = n_z * DELTA_K_PARA; - - k_mag = sqrt(k_x*k_x + k_y*k_y + k_z*k_z); - - kR = k_mag*R; - kl = k_mag*mfp; - - //Davies & Furlanetto MFP-eps(r) window function - //The filter no longer approaches 1 at k->0, so we can't use the limit - if (kR > 1e-4){ - //build the filter - f = (kl*kl*R + 2*mfp + R)*kl*cos(kR); - f += (-kl*kl*mfp + kl*kl*R + mfp + R)*sin(kR); - f *= const1; - f -= 2*kl*mfp; - f *= -3.0*mfp/(kR*R*R*(kl*kl+1)*(kl*kl+1)); - } - else{ - // k-> 0 limit - f = 2*mfp*mfp + 2*mfp*R + R*R; - f *= -const1; - f += 2*mfp*mfp; - f *= 3*mfp/(R*R*R); - } - if(RES==1) { box[HII_C_INDEX(n_x, n_y, n_z)] *= f; } - if(RES==0) { box[C_INDEX(n_x, n_y, n_z)] *= f; } - } - } - } // end looping through k box - } - return; -} - -int test_mfp_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions *flag_options - , float *input_box, double R, double mfp, double *result){ +//Test function to filter a box without computing a whole output box +int test_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions *flag_options + , float *input_box, double R, double R_param, int filter_flag, double *result){ int i,j,k; unsigned long long int ii; - //setup the box + Broadcast_struct_global_all(user_params,cosmo_params,astro_params,flag_options); + + //setup the box fftwf_complex *box_unfiltered = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*HII_KSPACE_NUM_PIXELS); fftwf_complex *box_filtered = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*HII_KSPACE_NUM_PIXELS); @@ -232,28 +165,22 @@ int test_mfp_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroPar for (k=0; kUSE_FFTW_WISDOM, user_params->HII_DIM, HII_D_PARA, user_params->N_THREADS, box_unfiltered); - //QUESTION: why do this here instead of at the end? for(ii=0;iiUSE_EXP_FILTER) - filter_box_mfp(box_filtered, 1, R, mfp); - else - filter_box(box_filtered,1,global_params.HII_FILTER,R); - + filter_box(box_filtered,1,filter_flag,R,R_param); dft_c2r_cube(user_params->USE_FFTW_WISDOM, user_params->HII_DIM, HII_D_PARA, user_params->N_THREADS, box_filtered); for (i=0; iHII_DIM; i++) for (j=0; jHII_DIM; j++) for (k=0; k -void filter_box(fftwf_complex *box, int RES, int filter_type, float R); -void filter_box_annulus(fftwf_complex *box, int RES, float R_inner, float R_outer); -void filter_box_mfp(fftwf_complex *box, int RES, float R, float mfp); - -int test_mfp_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions *flag_options - , float *input_box, double R, double mfp, double *result); +void filter_box(fftwf_complex *box, int RES, int filter_type, float R, float R_param); +int test_filter(UserParams *user_params, CosmoParams *cosmo_params, AstroParams *astro_params, FlagOptions *flag_options + , float *input_box, double R, double R_param, int filter_flag, double *result); #endif diff --git a/tests/test_filtering.py b/tests/test_filtering.py new file mode 100644 index 00000000..6c183b08 --- /dev/null +++ b/tests/test_filtering.py @@ -0,0 +1,148 @@ +import pytest + +import matplotlib as mpl +import numpy as np + +from py21cmfast import ( + AstroParams, + CosmoParams, + FlagOptions, + PerturbHaloField, + UserParams, + global_params, +) +from py21cmfast.c_21cmfast import ffi, lib + +from . import produce_integration_test_data as prd + +RELATIVE_TOLERANCE = 1e-4 + +options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from +R_PARAM_LIST = [1.5, 5, 10, 30, 60] + + +@pytest.mark.parametrize("R", R_PARAM_LIST) +@pytest.mark.parametrize("filter_flag", options_filter) +def test_filters(filter_flag, R, plt): + opts = prd.get_all_options(redshift=10.0) + + up = UserParams(opts["user_params"]) + cp = CosmoParams(opts["cosmo_params"]) + ap = AstroParams(opts["astro_params"]) + fo = FlagOptions(opts["flag_options"]) + + # testing a single pixel source + input_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f4") + input_box_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = 1.0 + + # testing a uniform grid + input_box_uniform = np.full((up.HII_DIM,) * 3, 1.0, dtype="f4") + + output_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f8") + output_box_uniform = np.zeros((up.HII_DIM,) * 3, dtype="f8") + + # use MFP=20 for the exp filter, use a 3 cell shell for the annular filter + R_param_list = [0.0, 0.0, 0.0, 20, max(R - 3 * (up.BOX_LEN / up.HII_DIM), 0)] + + lib.test_filter( + up(), + cp(), + ap(), + fo(), + ffi.cast("float *", input_box_centre.ctypes.data), + R, + R_param_list[filter_flag], + filter_flag, + ffi.cast("double *", output_box_centre.ctypes.data), + ) + + lib.test_filter( + up(), + cp(), + ap(), + fo(), + ffi.cast("float *", input_box_uniform.ctypes.data), + R, + R_param_list[filter_flag], + filter_flag, + ffi.cast("double *", output_box_uniform.ctypes.data), + ) + + r_from_centre = np.linalg.norm( + np.mgrid[0 : up.HII_DIM, 0 : up.HII_DIM, 0 : up.HII_DIM] + - np.array( + [ + up.HII_DIM // 2, + up.HII_DIM // 2, + up.HII_DIM // 2, + ] + )[:, None, None, None], + axis=0, + ) + # single pixel boxes have a specific shape based on the filter + R_ratio = r_from_centre / R + if filter_flag == 0.0: + # output is uniform sphere around the centre point + exp_vol = 4 / 3 * np.pi * R**3 + expected_output = 1 / exp_vol * (R_ratio < 1) + elif filter_flag == 1.0: + # output is sinc function + expected_output = np.sin(R_ratio) / R_ratio + elif filter_flag == 2.0: + # output is Gaussian + expected_output = 1 / np.sqrt(2 * np.pi * R) * np.exp(-(R_ratio**2) / 2) + elif filter_flag == 3.0: + exp_vol = 4 / 3 * np.pi * R**3 + expected_output = 1 / exp_vol * (R_ratio < 1) * np.exp(-R_ratio) + elif filter_flag == 4.0: + R_i = R_param_list[4] + exp_vol = 4 / 3 * np.pi * (R**3 - R_i**3) + expected_output = 1 / exp_vol * (R_ratio < 1) * (R_i / R > 1) + + # uniform boxes should come out uniform aside from normalisation + norm_factor = ( + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ) + if plt == mpl.pyplot: + filter_plot_symmetric( + [input_box_centre, input_box_uniform], + [output_box_centre, output_box_uniform], + [expected_output, input_box_uniform / norm_factor[filter_flag]], + up.HII_DIM, + plt, + ) + + np.testing.assert_allclose( + input_box_centre, expected_output, rtol=RELATIVE_TOLERANCE + ) + np.testing.assert_allclose( + input_box_uniform, + norm_factor[filter_flag] * output_box_uniform, + rtol=RELATIVE_TOLERANCE, + ) + + +# since the filters are symmetric I'm doing an R vs value plot instead of imshowing slices +def filter_plot_symmetric(inputs, outputs, truths, dimension, plt): + if not (len(inputs) == len(truths) == len(outputs)): + raise ValueError( + f"inputs {len(inputs)}, outputs {len(outputs)} and truths {len(truths)}" + "must have the same length." + ) + + n_plots = len(inputs) + fig, axs = plt.subplots( + n_plots, 3, figsize=(16.0, 9.0 * n_plots / 3.0), layout="constrained" + ) + # fig.get_layout_engine().set(w_pad=2 / 72, h_pad=2 / 72, hspace=0.0, + # wspace=0.0) + + for idx, (i, o, t) in enumerate(zip(inputs, outputs, truths)): + r = np.linalg.norm(np.mgrid[0:dimension, 0:dimension, 0:dimension], axis=0) + axs[idx, 0].scatter(r, i, s=1) + axs[idx, 1].scatter(r, o, s=1) + axs[idx, 2].scatter(r, t, s=1) From f3435ca624abe25e9e87ae6d0872dddebad2bacf Mon Sep 17 00:00:00 2001 From: James Davies Date: Mon, 2 Sep 2024 17:22:31 +0200 Subject: [PATCH 2/7] WIP filters 0,3,4 test as expected --- tests/test_filtering.py | 201 ++++++++++++++++++++++++++++++---------- 1 file changed, 150 insertions(+), 51 deletions(-) diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 6c183b08..62c9598b 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -2,6 +2,8 @@ import matplotlib as mpl import numpy as np +from matplotlib.colors import LogNorm, Normalize +from scipy.stats import binned_statistic as binstat from py21cmfast import ( AstroParams, @@ -15,12 +17,78 @@ from . import produce_integration_test_data as prd -RELATIVE_TOLERANCE = 1e-4 +# tolerance for aliasing errors +RELATIVE_TOLERANCE = 1e-1 options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from R_PARAM_LIST = [1.5, 5, 10, 30, 60] +def get_expected_output_centre(r_in, R_filter, R_param, filter_flag): + # single pixel boxes have a specific shape based on the filter + R_ratio = r_in / R_filter + if filter_flag == 0: + # output is uniform sphere around the centre point + exp_vol = 4 / 3 * np.pi * R_filter**3 + return (R_ratio < 1) / exp_vol + elif filter_flag == 1: + # output is sinc function + R_ratio /= 0.413566994 + R_filter /= 0.413566994 + exp_vol = R_filter**3 / 3.0 + return (np.sin(R_ratio) / R_ratio) / exp_vol + elif filter_flag == 2: + # output is Gaussian + R_ratio /= 0.643 + R_filter /= 0.643 + exp_vol = (2 * np.pi) ** (3 / 2) * R_filter**3 + return np.exp(-(R_ratio**2) / 2) / exp_vol + elif filter_flag == 3: + # output is sphere with exponential damping + exp_vol = 4 / 3 * np.pi * R_filter**3 + return (R_ratio < 1) * np.exp(-r_in / R_param) / exp_vol + elif filter_flag == 4: + # output is spherical shell + exp_vol = 4 / 3 * np.pi * (R_filter**3 - R_param**3) + return (R_ratio < 1) * (R_param <= r_in) / exp_vol + + +def get_expected_output_uniform(in_box, R_filter, R_param, filter_flag): + # uniform boxes should come out uniform, the exp filter will be + if filter_flag == 3: + norm_factor = ( + R_param - R_param * np.exp(-R_filter / R_param) + ) / R_filter # TODO: this is wrong, change it + else: + norm_factor = 1 + return in_box * norm_factor + + +# return binned mean & 1-2 sigma quantiles +def get_binned_stats(x_arr, y_arr, bins, stats): + x_in = x_arr.flatten() + y_in = y_arr.flatten() + result = {} + + statistic_dict = { + "pc1u": lambda x: np.percentile(x, 84), + "pc1l": lambda x: np.percentile(x, 16), + "pc2u": lambda x: np.percentile(x, 97.5), + "pc2l": lambda x: np.percentile(x, 2.5), + # used to mark percentiles in errorbar plots, since these cannot be negative + "err1u": lambda x: np.maximum(np.percentile(x, 84) - np.mean(x), 0), + "err2u": lambda x: np.maximum(np.percentile(x, 97.5) - np.mean(x), 0), + "err1l": lambda x: np.maximum(np.mean(x) - np.percentile(x, 16), 0), + "err2l": lambda x: np.maximum(np.mean(x) - np.percentile(x, 2.5), 0), + } + + for stat in stats: + spstatkey = statistic_dict[stat] if stat in statistic_dict.keys() else stat + result[stat], _, _ = binstat(x_in, y_in, bins=bins, statistic=spstatkey) + + return result + + @pytest.mark.parametrize("R", R_PARAM_LIST) @pytest.mark.parametrize("filter_flag", options_filter) def test_filters(filter_flag, R, plt): @@ -42,7 +110,12 @@ def test_filters(filter_flag, R, plt): output_box_uniform = np.zeros((up.HII_DIM,) * 3, dtype="f8") # use MFP=20 for the exp filter, use a 3 cell shell for the annular filter - R_param_list = [0.0, 0.0, 0.0, 20, max(R - 3 * (up.BOX_LEN / up.HII_DIM), 0)] + if filter_flag == 3: + R_param = 20 + elif filter_flag == 4: + R_param = max(R - 3 * (up.BOX_LEN / up.HII_DIM), 0) + else: + R_param = 0 lib.test_filter( up(), @@ -51,7 +124,7 @@ def test_filters(filter_flag, R, plt): fo(), ffi.cast("float *", input_box_centre.ctypes.data), R, - R_param_list[filter_flag], + R_param, filter_flag, ffi.cast("double *", output_box_centre.ctypes.data), ) @@ -63,11 +136,13 @@ def test_filters(filter_flag, R, plt): fo(), ffi.cast("float *", input_box_uniform.ctypes.data), R, - R_param_list[filter_flag], + R_param, filter_flag, ffi.cast("double *", output_box_uniform.ctypes.data), ) + R_cells = R / up.BOX_LEN * up.HII_DIM + Rp_cells = R_param / up.BOX_LEN * up.HII_DIM r_from_centre = np.linalg.norm( np.mgrid[0 : up.HII_DIM, 0 : up.HII_DIM, 0 : up.HII_DIM] - np.array( @@ -79,70 +154,94 @@ def test_filters(filter_flag, R, plt): )[:, None, None, None], axis=0, ) - # single pixel boxes have a specific shape based on the filter - R_ratio = r_from_centre / R - if filter_flag == 0.0: - # output is uniform sphere around the centre point - exp_vol = 4 / 3 * np.pi * R**3 - expected_output = 1 / exp_vol * (R_ratio < 1) - elif filter_flag == 1.0: - # output is sinc function - expected_output = np.sin(R_ratio) / R_ratio - elif filter_flag == 2.0: - # output is Gaussian - expected_output = 1 / np.sqrt(2 * np.pi * R) * np.exp(-(R_ratio**2) / 2) - elif filter_flag == 3.0: - exp_vol = 4 / 3 * np.pi * R**3 - expected_output = 1 / exp_vol * (R_ratio < 1) * np.exp(-R_ratio) - elif filter_flag == 4.0: - R_i = R_param_list[4] - exp_vol = 4 / 3 * np.pi * (R**3 - R_i**3) - expected_output = 1 / exp_vol * (R_ratio < 1) * (R_i / R > 1) - - # uniform boxes should come out uniform aside from normalisation - norm_factor = ( - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, + exp_output_centre = get_expected_output_centre( + r_from_centre, R_cells, Rp_cells, filter_flag ) + exp_output_uniform = get_expected_output_uniform( + input_box_uniform, R_cells, Rp_cells, filter_flag + ) + if plt == mpl.pyplot: - filter_plot_symmetric( - [input_box_centre, input_box_uniform], - [output_box_centre, output_box_uniform], - [expected_output, input_box_uniform / norm_factor[filter_flag]], - up.HII_DIM, - plt, + r_bins = np.linspace(0, (up.HII_DIM / 2 * np.sqrt(3)), num=32) + uniform_bin = np.ones_like(r_bins) + binned_truth_centre = get_expected_output_centre( + r_bins, R_cells, Rp_cells, filter_flag + ) + binned_truth_uniform = get_expected_output_uniform( + uniform_bin, R_cells, Rp_cells, filter_flag + ) + filter_plot( + inputs=[input_box_centre, input_box_uniform], + outputs=[output_box_centre, output_box_uniform], + binned_truths=[binned_truth_centre, binned_truth_uniform], + truths=[exp_output_centre, exp_output_uniform], + r_bins=r_bins, + r_grid=r_from_centre, + slice_index=up.HII_DIM // 2, + slice_axis=0, + plt=plt, ) np.testing.assert_allclose( - input_box_centre, expected_output, rtol=RELATIVE_TOLERANCE + input_box_centre, exp_output_centre, rtol=RELATIVE_TOLERANCE ) np.testing.assert_allclose( input_box_uniform, - norm_factor[filter_flag] * output_box_uniform, + exp_output_uniform, rtol=RELATIVE_TOLERANCE, ) # since the filters are symmetric I'm doing an R vs value plot instead of imshowing slices -def filter_plot_symmetric(inputs, outputs, truths, dimension, plt): - if not (len(inputs) == len(truths) == len(outputs)): +def filter_plot( + inputs, outputs, binned_truths, truths, r_bins, r_grid, slice_index, slice_axis, plt +): + if not (len(inputs) == len(binned_truths) == len(outputs)): raise ValueError( - f"inputs {len(inputs)}, outputs {len(outputs)} and truths {len(truths)}" + f"inputs {len(inputs)}, outputs {len(outputs)} and truths {len(binned_truths)}" "must have the same length." ) n_plots = len(inputs) fig, axs = plt.subplots( - n_plots, 3, figsize=(16.0, 9.0 * n_plots / 3.0), layout="constrained" + n_plots, 4, figsize=(16.0, 12.0 * n_plots / 4.0), layout="constrained" ) - # fig.get_layout_engine().set(w_pad=2 / 72, h_pad=2 / 72, hspace=0.0, - # wspace=0.0) - - for idx, (i, o, t) in enumerate(zip(inputs, outputs, truths)): - r = np.linalg.norm(np.mgrid[0:dimension, 0:dimension, 0:dimension], axis=0) - axs[idx, 0].scatter(r, i, s=1) - axs[idx, 1].scatter(r, o, s=1) - axs[idx, 2].scatter(r, t, s=1) + fig.get_layout_engine().set(w_pad=2 / 72, h_pad=2 / 72, hspace=0.0, wspace=0.0) + + r_cen = (r_bins[1:] + r_bins[:-1]) / 2 + + axs[0, 0].set_title("Input") + axs[0, 1].set_title("output") + axs[0, 2].set_title("Expected") + axs[0, 3].set_title("Radii") + for idx, (i, o, t, tt) in enumerate(zip(inputs, outputs, binned_truths, truths)): + axs[idx, 0].pcolormesh( + i.take(indices=slice_index, axis=slice_axis), + norm=Normalize(vmin=0, vmax=o.max()), + ) + axs[idx, 1].pcolormesh( + o.take(indices=slice_index, axis=slice_axis), + norm=Normalize(vmin=0, vmax=o.max()), + ) + axs[idx, 2].pcolormesh( + tt.take(indices=slice_index, axis=slice_axis), + norm=Normalize(vmin=0, vmax=o.max()), + ) + + stats_o = get_binned_stats(r_grid, o, r_bins, stats=["mean", "err1u", "err1l"]) + axs[idx, 3].errorbar( + r_cen, + stats_o["mean"], + markerfacecolor="b", + elinewidth=1, + capsize=3, + markersize=5, + marker="o", + color="b", + yerr=[stats_o["err1l"], stats_o["err1u"]], + label="filtered grid", + ) + axs[idx, 3].plot(r_bins, t, "m:", label="Expected") + axs[idx, 3].grid() + axs[idx, 3].set_xlabel("dist from centre") + axs[idx, 3].set_ylabel("cell value") From 452f935eaeb0a34c2a03029a8fb4962f9796d8cd Mon Sep 17 00:00:00 2001 From: James Davies Date: Tue, 3 Sep 2024 16:21:45 +0200 Subject: [PATCH 3/7] fix test comparison cases --- tests/test_filtering.py | 121 +++++++++++++++++++++------------------- 1 file changed, 64 insertions(+), 57 deletions(-) diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 62c9598b..72bd1391 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -21,9 +21,13 @@ RELATIVE_TOLERANCE = 1e-1 options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from -R_PARAM_LIST = [1.5, 5, 10, 30, 60] +R_PARAM_LIST = [1.5, 5, 10, 25] # default test HII_DIM = 50 +# NOTE: These don't directly test against the expected FFT of these filters applied +# to a central cell, but the continuous FT filters applied to a delta function. +# this makes it a test of both the filter construction, as well as the aliasing. +# NOTE: These do not include the periodic boundary conditions, so face issues with R ~> HII_DIM/2. def get_expected_output_centre(r_in, R_filter, R_param, filter_flag): # single pixel boxes have a specific shape based on the filter R_ratio = r_in / R_filter @@ -32,17 +36,20 @@ def get_expected_output_centre(r_in, R_filter, R_param, filter_flag): exp_vol = 4 / 3 * np.pi * R_filter**3 return (R_ratio < 1) / exp_vol elif filter_flag == 1: - # output is sinc function - R_ratio /= 0.413566994 - R_filter /= 0.413566994 - exp_vol = R_filter**3 / 3.0 - return (np.sin(R_ratio) / R_ratio) / exp_vol + # output is the tophat FFT + R_ratio *= 1 / 0.413566994 # == this is the 2*pi*k factor for equating volumes + result = (np.sin(R_ratio) - R_ratio * np.cos(R_ratio)) / ( + 2 * np.pi**2 * r_in**3 + ) + result[r_in == 0] = ( + 1 / 6 / np.pi**2 * (0.413566994 * R_filter) ** 3 + ) # r->0 limit + return result elif filter_flag == 2: # output is Gaussian - R_ratio /= 0.643 - R_filter /= 0.643 - exp_vol = (2 * np.pi) ** (3 / 2) * R_filter**3 - return np.exp(-(R_ratio**2) / 2) / exp_vol + const = (0.643 * R_filter) ** 2 + exp_vol = (2 * np.pi * const) ** (3.0 / 2.0) + return np.exp(-((r_in) ** 2 / const / 2)) / exp_vol elif filter_flag == 3: # output is sphere with exponential damping exp_vol = 4 / 3 * np.pi * R_filter**3 @@ -53,17 +60,6 @@ def get_expected_output_centre(r_in, R_filter, R_param, filter_flag): return (R_ratio < 1) * (R_param <= r_in) / exp_vol -def get_expected_output_uniform(in_box, R_filter, R_param, filter_flag): - # uniform boxes should come out uniform, the exp filter will be - if filter_flag == 3: - norm_factor = ( - R_param - R_param * np.exp(-R_filter / R_param) - ) / R_filter # TODO: this is wrong, change it - else: - norm_factor = 1 - return in_box * norm_factor - - # return binned mean & 1-2 sigma quantiles def get_binned_stats(x_arr, y_arr, bins, stats): x_in = x_arr.flatten() @@ -102,13 +98,7 @@ def test_filters(filter_flag, R, plt): # testing a single pixel source input_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f4") input_box_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = 1.0 - - # testing a uniform grid - input_box_uniform = np.full((up.HII_DIM,) * 3, 1.0, dtype="f4") - output_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f8") - output_box_uniform = np.zeros((up.HII_DIM,) * 3, dtype="f8") - # use MFP=20 for the exp filter, use a 3 cell shell for the annular filter if filter_flag == 3: R_param = 20 @@ -129,18 +119,7 @@ def test_filters(filter_flag, R, plt): ffi.cast("double *", output_box_centre.ctypes.data), ) - lib.test_filter( - up(), - cp(), - ap(), - fo(), - ffi.cast("float *", input_box_uniform.ctypes.data), - R, - R_param, - filter_flag, - ffi.cast("double *", output_box_uniform.ctypes.data), - ) - + # expected outputs given in cell units R_cells = R / up.BOX_LEN * up.HII_DIM Rp_cells = R_param / up.BOX_LEN * up.HII_DIM r_from_centre = np.linalg.norm( @@ -154,27 +133,33 @@ def test_filters(filter_flag, R, plt): )[:, None, None, None], axis=0, ) + # average value of R in a sphere with same volume of the cell + r_from_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = ( + np.pi * (0.62035) ** 4 + ) exp_output_centre = get_expected_output_centre( r_from_centre, R_cells, Rp_cells, filter_flag ) - exp_output_uniform = get_expected_output_uniform( - input_box_uniform, R_cells, Rp_cells, filter_flag - ) if plt == mpl.pyplot: r_bins = np.linspace(0, (up.HII_DIM / 2 * np.sqrt(3)), num=32) - uniform_bin = np.ones_like(r_bins) + r_cen = (r_bins[1:] + r_bins[:-1]) / 2 binned_truth_centre = get_expected_output_centre( - r_bins, R_cells, Rp_cells, filter_flag - ) - binned_truth_uniform = get_expected_output_uniform( - uniform_bin, R_cells, Rp_cells, filter_flag + r_cen, R_cells, Rp_cells, filter_flag ) filter_plot( - inputs=[input_box_centre, input_box_uniform], - outputs=[output_box_centre, output_box_uniform], - binned_truths=[binned_truth_centre, binned_truth_uniform], - truths=[exp_output_centre, exp_output_uniform], + inputs=[ + input_box_centre, + ], + outputs=[ + output_box_centre, + ], + binned_truths=[ + binned_truth_centre, + ], + truths=[ + exp_output_centre, + ], r_bins=r_bins, r_grid=r_from_centre, slice_index=up.HII_DIM // 2, @@ -182,13 +167,25 @@ def test_filters(filter_flag, R, plt): plt=plt, ) + # All filters should be normalised aside from the exp filter + if filter_flag == 3: + norm_factor = ( + R_param - R_param * np.exp(-R / R_param) + ) / R # TODO: this is wrong, change it + else: + norm_factor = 1 + # firstly, make sure the filters are normalised np.testing.assert_allclose( - input_box_centre, exp_output_centre, rtol=RELATIVE_TOLERANCE + input_box_centre.sum() * norm_factor, output_box_centre.sum(), rtol=1e-4 ) + print("SUMS:") + print(f"INPUT: {input_box_centre.sum()}") + print(f"OUTPUT: {output_box_centre.sum()}") + print(f"EXPECTED: {exp_output_centre.sum()}") + + # then make sure we get the right shapes (more lenient for aliasing) np.testing.assert_allclose( - input_box_uniform, - exp_output_uniform, - rtol=RELATIVE_TOLERANCE, + output_box_centre, exp_output_centre, rtol=RELATIVE_TOLERANCE ) @@ -204,7 +201,11 @@ def filter_plot( n_plots = len(inputs) fig, axs = plt.subplots( - n_plots, 4, figsize=(16.0, 12.0 * n_plots / 4.0), layout="constrained" + n_plots, + 4, + figsize=(16.0, 12.0 * n_plots / 4.0), + layout="constrained", + squeeze=False, ) fig.get_layout_engine().set(w_pad=2 / 72, h_pad=2 / 72, hspace=0.0, wspace=0.0) @@ -241,7 +242,13 @@ def filter_plot( yerr=[stats_o["err1l"], stats_o["err1u"]], label="filtered grid", ) - axs[idx, 3].plot(r_bins, t, "m:", label="Expected") + axs[idx, 3].plot(r_cen, t, "m:", label="Expected") axs[idx, 3].grid() axs[idx, 3].set_xlabel("dist from centre") axs[idx, 3].set_ylabel("cell value") + + axst = axs[idx, 3].twinx() + axst.plot(r_cen, stats_o["mean"] / t - 1, "r-") + axst.set_ylim(-10, 10) + axst.set_ylabel("out-truth/truth") + # axst.set_yscale('log') From cdf8c3e3f7e8a725235d82a68c838282cc3ee274 Mon Sep 17 00:00:00 2001 From: James Davies Date: Tue, 3 Sep 2024 17:22:54 +0200 Subject: [PATCH 4/7] fix bug in mfp filter, set expected mean --- src/py21cmfast/src/filtering.c | 2 +- tests/test_filtering.py | 11 ++++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/src/py21cmfast/src/filtering.c b/src/py21cmfast/src/filtering.c index cb89f021..a3580af8 100644 --- a/src/py21cmfast/src/filtering.c +++ b/src/py21cmfast/src/filtering.c @@ -47,7 +47,7 @@ double exp_mfp_filter(double k, double R, double mfp, double R_constant, double //Davies & Furlanetto MFP-eps(r) window function //The filter no longer approaches 1 at k->0, so we can't use the limit f = (kR2*kR2*R + 2*mfp + R)*kR2*cos(kR); - f += (-kR2*kR2*mfp + kR2*kR2*R + limit + R)*sin(kR); + f += (-kR2*kR2*mfp + kR2*kR2*R + mfp + R)*sin(kR); f *= R_constant; f -= 2*kR2*mfp; f *= -3.0*mfp/(kR*R*R*(kR2*kR2+1)*(kR2*kR2+1)); diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 72bd1391..05925c85 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -169,19 +169,17 @@ def test_filters(filter_flag, R, plt): # All filters should be normalised aside from the exp filter if filter_flag == 3: + # ratio of exponential and sphere volume integrals norm_factor = ( - R_param - R_param * np.exp(-R / R_param) - ) / R # TODO: this is wrong, change it + 2 * R_param**3 + - R_param * np.exp(-R / R_param) * (2 * R_param**2 + 2 * R_param * R + R**2) + ) / (3 * R**3) else: norm_factor = 1 # firstly, make sure the filters are normalised np.testing.assert_allclose( input_box_centre.sum() * norm_factor, output_box_centre.sum(), rtol=1e-4 ) - print("SUMS:") - print(f"INPUT: {input_box_centre.sum()}") - print(f"OUTPUT: {output_box_centre.sum()}") - print(f"EXPECTED: {exp_output_centre.sum()}") # then make sure we get the right shapes (more lenient for aliasing) np.testing.assert_allclose( @@ -251,4 +249,3 @@ def filter_plot( axst.plot(r_cen, stats_o["mean"] / t - 1, "r-") axst.set_ylim(-10, 10) axst.set_ylabel("out-truth/truth") - # axst.set_yscale('log') From 1b59bd1e81ee01e7d0bc0a043b25f8f12cf1c15a Mon Sep 17 00:00:00 2001 From: James Davies Date: Mon, 7 Oct 2024 17:45:09 +0200 Subject: [PATCH 5/7] add taylor series limits to filters, fix expfilter tests --- src/py21cmfast/src/filtering.c | 48 ++++++++++++++++++++-------------- tests/test_filtering.py | 47 ++++++++++++++------------------- 2 files changed, 47 insertions(+), 48 deletions(-) diff --git a/src/py21cmfast/src/filtering.c b/src/py21cmfast/src/filtering.c index a3580af8..65c48912 100644 --- a/src/py21cmfast/src/filtering.c +++ b/src/py21cmfast/src/filtering.c @@ -18,8 +18,9 @@ #include "dft.h" double real_tophat_filter(double kR){ + //Second order taylor expansion around kR==0 if (kR < 1e-4) - return 1; + return 1 - kR*kR/10; return 3.0*pow(kR, -3) * (sin(kR) - cos(kR)*kR); } @@ -36,21 +37,27 @@ double gaussian_filter(double kR_squared){ return exp(-0.643*0.643*kR_squared/2.); } -double exp_mfp_filter(double k, double R, double mfp, double R_constant, double limit){ - double kR,kR2,f; - - kR = k*R; - if(kR < 1e-4) - return limit; +double exp_mfp_filter(double k, double R, double mfp, double exp_term){ + double f; + + double kR = k*R; + double ratio = mfp/R; + //Second order taylor expansion around kR==0 + //NOTE: the taylor coefficients could be stored and passed in + // but there aren't any super expensive operations here + // assuming the integer pow calls are optimized by the compiler + // test with the profiler + if(kR < 1e-4){ + double ts_0 = 6*pow(ratio,3) - exp_term*(6*pow(ratio,3) + 6*pow(ratio,2) + 3*ratio); + return ts_0 + (exp_term*(2*pow(ratio,2) + 0.5*ratio) - 2*ts_0*pow(ratio,2))*kR*kR; + } - kR2 = k*mfp; //Davies & Furlanetto MFP-eps(r) window function - //The filter no longer approaches 1 at k->0, so we can't use the limit - f = (kR2*kR2*R + 2*mfp + R)*kR2*cos(kR); - f += (-kR2*kR2*mfp + kR2*kR2*R + mfp + R)*sin(kR); - f *= R_constant; - f -= 2*kR2*mfp; - f *= -3.0*mfp/(kR*R*R*(kR2*kR2+1)*(kR2*kR2+1)); + f = (kR*kR*pow(ratio,2) + 2*ratio + 1)*ratio*cos(kR); + f += (kR*kR*(pow(ratio,2)-pow(ratio,3)) + ratio + 1)*sin(kR)/kR; + f *= exp_term; + f -= 2*pow(ratio,2); + f *= -3*ratio/pow(pow(kR*ratio,2) + 1,2); return f; } @@ -58,8 +65,11 @@ double spherical_shell_filter(double k, double R_outer, double R_inner){ double kR_inner = k*R_inner; double kR_outer = k*R_outer; + //Second order taylor expansion around kR_outer==0 if (kR_outer < 1e-4) - return 1.; + return 1. - kR_outer*kR_outer/10 * \ + (pow(R_inner/R_outer,5) - 1) / \ + (pow(R_inner/R_outer,3) - 1); return 3.0/(pow(kR_outer, 3) - pow(kR_inner, 3)) \ * (sin(kR_outer) - cos(kR_outer)*kR_outer \ @@ -84,11 +94,9 @@ void filter_box(fftwf_complex *box, int RES, int filter_type, float R, float R_p } //setup constants if needed - float R_constant_1, R_constant_2; + double R_const; if(filter_type == 3){ - R_constant_1 = exp(-R/R_param); //independent of k - //k->0 limit of the mfp filter - R_constant_2 = (2*R_param*R_param - R_constant_1*(2*R_param*R_param + 2*R_param*R + R*R))*3*R_param/(R*R*R); + R_const = exp(-R/R_param); } // loop through k-box @@ -131,7 +139,7 @@ void filter_box(fftwf_complex *box, int RES, int filter_type, float R, float R_p //The next two filters are not given by the HII_FILTER global, but used for specific grids else if (filter_type == 3){ // exponentially decaying tophat, param == scale of decay (MFP) //NOTE: This should be optimized, I havne't looked at it in a while - box[grid_index] *= exp_mfp_filter(sqrt(k_mag_sq),R,R_param,R_constant_1,R_constant_2); + box[grid_index] *= exp_mfp_filter(sqrt(k_mag_sq),R,R_param,R_const); } else if (filter_type == 4){ //spherical shell, R_param == inner radius box[grid_index] *= spherical_shell_filter(sqrt(k_mag_sq),R,R_param); diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 05925c85..de1205e8 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -17,9 +17,6 @@ from . import produce_integration_test_data as prd -# tolerance for aliasing errors -RELATIVE_TOLERANCE = 1e-1 - options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from R_PARAM_LIST = [1.5, 5, 10, 25] # default test HII_DIM = 50 @@ -60,7 +57,7 @@ def get_expected_output_centre(r_in, R_filter, R_param, filter_flag): return (R_ratio < 1) * (R_param <= r_in) / exp_vol -# return binned mean & 1-2 sigma quantiles +# return binned quantities def get_binned_stats(x_arr, y_arr, bins, stats): x_in = x_arr.flatten() y_in = y_arr.flatten() @@ -76,6 +73,8 @@ def get_binned_stats(x_arr, y_arr, bins, stats): "err2u": lambda x: np.maximum(np.percentile(x, 97.5) - np.mean(x), 0), "err1l": lambda x: np.maximum(np.mean(x) - np.percentile(x, 16), 0), "err2l": lambda x: np.maximum(np.mean(x) - np.percentile(x, 2.5), 0), + "errmin": lambda x: np.mean(x) - np.amin(x), + "errmax": lambda x: np.amax(x) - np.mean(x), } for stat in stats: @@ -148,18 +147,10 @@ def test_filters(filter_flag, R, plt): r_cen, R_cells, Rp_cells, filter_flag ) filter_plot( - inputs=[ - input_box_centre, - ], - outputs=[ - output_box_centre, - ], - binned_truths=[ - binned_truth_centre, - ], - truths=[ - exp_output_centre, - ], + inputs=[input_box_centre], + outputs=[output_box_centre], + binned_truths=[binned_truth_centre], + truths=[exp_output_centre], r_bins=r_bins, r_grid=r_from_centre, slice_index=up.HII_DIM // 2, @@ -170,21 +161,19 @@ def test_filters(filter_flag, R, plt): # All filters should be normalised aside from the exp filter if filter_flag == 3: # ratio of exponential and sphere volume integrals - norm_factor = ( - 2 * R_param**3 - - R_param * np.exp(-R / R_param) * (2 * R_param**2 + 2 * R_param * R + R**2) - ) / (3 * R**3) + R_q = R_param / R + norm_factor = 6 * R_q**3 - np.exp(-1 / R_q) * ( + 6 * R_q**3 + 6 * R_q**2 + 3 * R_q + ) else: norm_factor = 1 # firstly, make sure the filters are normalised np.testing.assert_allclose( - input_box_centre.sum() * norm_factor, output_box_centre.sum(), rtol=1e-4 + input_box_centre.sum() * norm_factor, output_box_centre.sum(), atol=1e-4 ) # then make sure we get the right shapes (more lenient for aliasing) - np.testing.assert_allclose( - output_box_centre, exp_output_centre, rtol=RELATIVE_TOLERANCE - ) + np.testing.assert_allclose(output_box_centre, exp_output_centre, atol=1e-4) # since the filters are symmetric I'm doing an R vs value plot instead of imshowing slices @@ -227,7 +216,9 @@ def filter_plot( norm=Normalize(vmin=0, vmax=o.max()), ) - stats_o = get_binned_stats(r_grid, o, r_bins, stats=["mean", "err1u", "err1l"]) + stats_o = get_binned_stats( + r_grid, o, r_bins, stats=["mean", "errmin", "errmax"] + ) axs[idx, 3].errorbar( r_cen, stats_o["mean"], @@ -237,15 +228,15 @@ def filter_plot( markersize=5, marker="o", color="b", - yerr=[stats_o["err1l"], stats_o["err1u"]], + yerr=[stats_o["errmin"], stats_o["errmax"]], label="filtered grid", ) - axs[idx, 3].plot(r_cen, t, "m:", label="Expected") + axs[idx, 3].plot(r_cen, t, "m:", linewidth=2, label="Expected") axs[idx, 3].grid() axs[idx, 3].set_xlabel("dist from centre") axs[idx, 3].set_ylabel("cell value") axst = axs[idx, 3].twinx() - axst.plot(r_cen, stats_o["mean"] / t - 1, "r-") + axst.plot(r_cen, stats_o["mean"] / (t + 1e-16) - 1, "r-") axst.set_ylim(-10, 10) axst.set_ylabel("out-truth/truth") From a1d5d54a4d155d6baa19d1626d4984b494018a9d Mon Sep 17 00:00:00 2001 From: James Davies Date: Mon, 7 Oct 2024 18:57:34 +0200 Subject: [PATCH 6/7] improve filter test plotting --- tests/test_c_interpolation_tables.py | 4 +- tests/test_filtering.py | 115 ++++++++++++++++++--------- 2 files changed, 79 insertions(+), 40 deletions(-) diff --git a/tests/test_c_interpolation_tables.py b/tests/test_c_interpolation_tables.py index ac6c588e..ae085ffb 100644 --- a/tests/test_c_interpolation_tables.py +++ b/tests/test_c_interpolation_tables.py @@ -1453,7 +1453,9 @@ def print_failure_stats(test, truth, input_arr, abs_tol, rel_tol, name): f"max abs diff of failures {np.fabs(truth - test)[sel_failed].max():.4e} relative {(np.fabs(truth - test) / truth)[sel_failed].max():.4e}" ) - print(f"first 10 = {truth.flatten()[:10]} {test.flatten()[:10]}") + print( + f"first 10 = {truth[sel_failed].flatten()[:10]} {test[sel_failed].flatten()[:10]}" + ) def massfunc_table_comparison_plot( diff --git a/tests/test_filtering.py b/tests/test_filtering.py index de1205e8..9b0115b3 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -16,6 +16,7 @@ from py21cmfast.c_21cmfast import ffi, lib from . import produce_integration_test_data as prd +from .test_c_interpolation_tables import print_failure_stats options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from R_PARAM_LIST = [1.5, 5, 10, 25] # default test HII_DIM = 50 @@ -123,13 +124,9 @@ def test_filters(filter_flag, R, plt): Rp_cells = R_param / up.BOX_LEN * up.HII_DIM r_from_centre = np.linalg.norm( np.mgrid[0 : up.HII_DIM, 0 : up.HII_DIM, 0 : up.HII_DIM] - - np.array( - [ - up.HII_DIM // 2, - up.HII_DIM // 2, - up.HII_DIM // 2, - ] - )[:, None, None, None], + - np.array([up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2])[ + :, None, None, None + ], axis=0, ) # average value of R in a sphere with same volume of the cell @@ -140,6 +137,8 @@ def test_filters(filter_flag, R, plt): r_from_centre, R_cells, Rp_cells, filter_flag ) + abs_tol_pixel = exp_output_centre.max() * 1e-2 + rel_tol_pixel = 2e-1 if plt == mpl.pyplot: r_bins = np.linspace(0, (up.HII_DIM / 2 * np.sqrt(3)), num=32) r_cen = (r_bins[1:] + r_bins[:-1]) / 2 @@ -155,6 +154,8 @@ def test_filters(filter_flag, R, plt): r_grid=r_from_centre, slice_index=up.HII_DIM // 2, slice_axis=0, + abs_tol=abs_tol_pixel, + rel_tol=rel_tol_pixel, plt=plt, ) @@ -169,16 +170,37 @@ def test_filters(filter_flag, R, plt): norm_factor = 1 # firstly, make sure the filters are normalised np.testing.assert_allclose( - input_box_centre.sum() * norm_factor, output_box_centre.sum(), atol=1e-4 + input_box_centre.sum() * norm_factor, output_box_centre.sum(), atol=1e-6 ) # then make sure we get the right shapes (more lenient for aliasing) - np.testing.assert_allclose(output_box_centre, exp_output_centre, atol=1e-4) + # absolute tolerance is set by the maximum, and some relative tolerance is added for aliasing + print_failure_stats( + output_box_centre, + exp_output_centre, + [r_from_centre], + abs_tol=abs_tol_pixel, + rel_tol=rel_tol_pixel, + name="pixel", + ) + np.testing.assert_allclose( + output_box_centre, exp_output_centre, rtol=rel_tol_pixel, atol=abs_tol_pixel + ) # since the filters are symmetric I'm doing an R vs value plot instead of imshowing slices def filter_plot( - inputs, outputs, binned_truths, truths, r_bins, r_grid, slice_index, slice_axis, plt + inputs, + outputs, + binned_truths, + truths, + r_bins, + r_grid, + slice_index, + slice_axis, + abs_tol, + rel_tol, + plt, ): if not (len(inputs) == len(binned_truths) == len(outputs)): raise ValueError( @@ -194,24 +216,20 @@ def filter_plot( layout="constrained", squeeze=False, ) - fig.get_layout_engine().set(w_pad=2 / 72, h_pad=2 / 72, hspace=0.0, wspace=0.0) + fig.get_layout_engine().set(w_pad=0 / 72, h_pad=0 / 72, hspace=0.0, wspace=0.0) r_cen = (r_bins[1:] + r_bins[:-1]) / 2 - axs[0, 0].set_title("Input") - axs[0, 1].set_title("output") - axs[0, 2].set_title("Expected") - axs[0, 3].set_title("Radii") + axs[0, 0].set_title("Output") + axs[0, 1].set_title("Expected") + axs[0, 2].set_title("Radii") + axs[0, 3].set_title("Pixels") for idx, (i, o, t, tt) in enumerate(zip(inputs, outputs, binned_truths, truths)): axs[idx, 0].pcolormesh( - i.take(indices=slice_index, axis=slice_axis), - norm=Normalize(vmin=0, vmax=o.max()), - ) - axs[idx, 1].pcolormesh( o.take(indices=slice_index, axis=slice_axis), norm=Normalize(vmin=0, vmax=o.max()), ) - axs[idx, 2].pcolormesh( + axs[idx, 1].pcolormesh( tt.take(indices=slice_index, axis=slice_axis), norm=Normalize(vmin=0, vmax=o.max()), ) @@ -219,24 +237,43 @@ def filter_plot( stats_o = get_binned_stats( r_grid, o, r_bins, stats=["mean", "errmin", "errmax"] ) - axs[idx, 3].errorbar( - r_cen, - stats_o["mean"], - markerfacecolor="b", - elinewidth=1, - capsize=3, - markersize=5, - marker="o", - color="b", - yerr=[stats_o["errmin"], stats_o["errmax"]], - label="filtered grid", + + lns = [] + lns.append( + axs[idx, 2].errorbar( + r_cen, + stats_o["mean"], + markerfacecolor="b", + elinewidth=1, + capsize=3, + markersize=5, + marker="o", + color="b", + yerr=[stats_o["errmin"], stats_o["errmax"]], + label="filtered grid", + zorder=2, + ) ) - axs[idx, 3].plot(r_cen, t, "m:", linewidth=2, label="Expected") - axs[idx, 3].grid() - axs[idx, 3].set_xlabel("dist from centre") - axs[idx, 3].set_ylabel("cell value") - axst = axs[idx, 3].twinx() - axst.plot(r_cen, stats_o["mean"] / (t + 1e-16) - 1, "r-") - axst.set_ylim(-10, 10) - axst.set_ylabel("out-truth/truth") + lns.append( + axs[idx, 2].plot(r_cen, t, "m:", linewidth=2, label="Expected", zorder=3)[0] + ) + axs[idx, 2].grid() + axs[idx, 2].set_xlabel("dist from centre") + axs[idx, 2].set_ylabel("cell value") + axs[idx, 2].legend() + + err_base = np.linspace(tt.min(), tt.max(), num=100) + err_max = err_base + (abs_tol + np.fabs(err_base) * rel_tol) + err_min = err_base - (abs_tol + np.fabs(err_base) * rel_tol) + axs[idx, 3].plot(err_base, err_max, "k:") + axs[idx, 3].plot(err_base, err_min, "k:") + axs[idx, 3].plot(err_base, err_base, "k--") + print(err_base, err_max, err_min) + + axs[idx, 3].scatter(tt, o, s=1, alpha=0.5, rasterized=True) + # axs[idx, 3].set_yscale('symlog') + # axs[idx, 3].set_xscale('symlog') + axs[idx, 3].grid() + axs[idx, 3].set_xlabel("expected cell value") + axs[idx, 3].set_ylabel("filtered cell value") From 2b6717589d0d43ab3b053d5eccc572bc516ef228 Mon Sep 17 00:00:00 2001 From: James Davies Date: Tue, 8 Oct 2024 16:22:19 +0200 Subject: [PATCH 7/7] widen filter testing tolerance --- tests/test_filtering.py | 86 +++++++++++++++++++++++++++++------------ 1 file changed, 61 insertions(+), 25 deletions(-) diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 9b0115b3..b38aed85 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -19,7 +19,12 @@ from .test_c_interpolation_tables import print_failure_stats options_filter = [0, 1, 2, 3, 4] # cell densities to draw samples from -R_PARAM_LIST = [1.5, 5, 10, 25] # default test HII_DIM = 50 +R_PARAM_LIST = [ + 1.5, + 5, + 10, + 20, +] # default test HII_DIM = 50, we want max R < BOX_LEN*HII_DIM/3 # NOTE: These don't directly test against the expected FFT of these filters applied @@ -99,11 +104,11 @@ def test_filters(filter_flag, R, plt): input_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f4") input_box_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = 1.0 output_box_centre = np.zeros((up.HII_DIM,) * 3, dtype="f8") - # use MFP=20 for the exp filter, use a 3 cell shell for the annular filter + # use MFP=20 for the exp filter, use a 4 cell shell for the annular filter if filter_flag == 3: R_param = 20 elif filter_flag == 4: - R_param = max(R - 3 * (up.BOX_LEN / up.HII_DIM), 0) + R_param = max(R - 4 * (up.BOX_LEN / up.HII_DIM), 0) else: R_param = 0 @@ -129,26 +134,44 @@ def test_filters(filter_flag, R, plt): ], axis=0, ) - # average value of R in a sphere with same volume of the cell - r_from_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = ( - np.pi * (0.62035) ** 4 - ) + # prevent divide by zero in the central cell + r_from_centre[up.HII_DIM // 2, up.HII_DIM // 2, up.HII_DIM // 2] = 1e-6 + exp_output_centre = get_expected_output_centre( r_from_centre, R_cells, Rp_cells, filter_flag ) - abs_tol_pixel = exp_output_centre.max() * 1e-2 - rel_tol_pixel = 2e-1 + # these are very wide tolerances, just to make sure there aren't + # cells 100x the expected values + abs_tol_pixel = exp_output_centre.max() * 0.8 + rel_tol_pixel = 0 + # we take bins of 2 pixels to smooth over sharp edged filters + r_bins = np.arange(0, int(up.HII_DIM / 2 * np.sqrt(3)), 2) + r_cen = (r_bins[1:] + r_bins[:-1]) / 2 + # binned_truth_centre = get_expected_output_centre( + # r_cen, R_cells, Rp_cells, filter_flag + # ) + binned_truth_centre = get_binned_stats( + r_from_centre, + exp_output_centre, + r_bins, + stats=["mean", "errmin", "errmax", "err1l", "err1u"], + ) + binned_truth_centre = binned_truth_centre["mean"] + + stats_o = get_binned_stats( + r_from_centre, + output_box_centre, + r_bins, + stats=["mean", "errmin", "errmax", "err1l", "err1u"], + ) + if plt == mpl.pyplot: - r_bins = np.linspace(0, (up.HII_DIM / 2 * np.sqrt(3)), num=32) - r_cen = (r_bins[1:] + r_bins[:-1]) / 2 - binned_truth_centre = get_expected_output_centre( - r_cen, R_cells, Rp_cells, filter_flag - ) filter_plot( inputs=[input_box_centre], outputs=[output_box_centre], binned_truths=[binned_truth_centre], + binned_stats=[stats_o], truths=[exp_output_centre], r_bins=r_bins, r_grid=r_from_centre, @@ -170,11 +193,26 @@ def test_filters(filter_flag, R, plt): norm_factor = 1 # firstly, make sure the filters are normalised np.testing.assert_allclose( - input_box_centre.sum() * norm_factor, output_box_centre.sum(), atol=1e-6 + input_box_centre.sum() * norm_factor, output_box_centre.sum(), atol=1e-4 + ) + + # secondly, make sure binned results are reasonable + rel_tol_bins = 1e-1 + abs_tol_bins = exp_output_centre.max() * 1e-1 + print_failure_stats( + stats_o["mean"], + binned_truth_centre, + [r_cen], + abs_tol=abs_tol_bins, + rel_tol=rel_tol_bins, + name="bins", + ) + np.testing.assert_allclose( + binned_truth_centre, stats_o["mean"], atol=abs_tol_bins, rtol=rel_tol_bins ) - # then make sure we get the right shapes (more lenient for aliasing) - # absolute tolerance is set by the maximum, and some relative tolerance is added for aliasing + # lastly, make sure no pixels are way out of line. + # this has a wide tolerance due to significant aliasing print_failure_stats( output_box_centre, exp_output_centre, @@ -193,6 +231,7 @@ def filter_plot( inputs, outputs, binned_truths, + binned_stats, truths, r_bins, r_grid, @@ -224,7 +263,9 @@ def filter_plot( axs[0, 1].set_title("Expected") axs[0, 2].set_title("Radii") axs[0, 3].set_title("Pixels") - for idx, (i, o, t, tt) in enumerate(zip(inputs, outputs, binned_truths, truths)): + for idx, (i, o, bo, t, tt) in enumerate( + zip(inputs, outputs, binned_stats, binned_truths, truths) + ): axs[idx, 0].pcolormesh( o.take(indices=slice_index, axis=slice_axis), norm=Normalize(vmin=0, vmax=o.max()), @@ -234,22 +275,18 @@ def filter_plot( norm=Normalize(vmin=0, vmax=o.max()), ) - stats_o = get_binned_stats( - r_grid, o, r_bins, stats=["mean", "errmin", "errmax"] - ) - lns = [] lns.append( axs[idx, 2].errorbar( r_cen, - stats_o["mean"], + bo["mean"], markerfacecolor="b", elinewidth=1, capsize=3, markersize=5, marker="o", color="b", - yerr=[stats_o["errmin"], stats_o["errmax"]], + yerr=[bo["errmin"], bo["errmax"]], label="filtered grid", zorder=2, ) @@ -269,7 +306,6 @@ def filter_plot( axs[idx, 3].plot(err_base, err_max, "k:") axs[idx, 3].plot(err_base, err_min, "k:") axs[idx, 3].plot(err_base, err_base, "k--") - print(err_base, err_max, err_min) axs[idx, 3].scatter(tt, o, s=1, alpha=0.5, rasterized=True) # axs[idx, 3].set_yscale('symlog')