detect_bin_type outside of setup_bins, clean up tests.py

Corrfunc/tests.py

@@ -65,38 +65,33 @@ def test_linear_binning_mocks(isa='fallback'):
     autocorr = 1
     numbins_to_print = 5
     cosmology = 1
+    bin_types = ['custom', 'lin']
 
-    def allclose(a, b):
-        return all(np.allclose(a[name], b[name]) for name in a.dtype.names)
+    def allclose(a, *others):
+        toret = True
+        for b in others:
+            toret &= all(np.allclose(a[name], b[name]) for name in ['npairs'])
+        return toret
+
+    def test_bin_type(pair_counter, *args, **kwargs):
+        res = (pair_counter(*args, bin_type=bin_type, weight_type=weight_type, **kwargs) for bin_type in bin_types for weight_type in [None, "pair_product"])
+        assert allclose(*res)
 
     print("\nRunning 2-D correlation function xi(rp,pi)")
-    results_DDrppi_ref = DDrppi_mocks(autocorr, cosmology, nthreads,
-                                      pimax, binfile,
-                                      ra, dec, cz,
-                                      weights1=w, weight_type='pair_product',
-                                      output_rpavg=True, verbose=True, isa=isa, bin_type='custom')
-    results_DDrppi = DDrppi_mocks(autocorr, cosmology, nthreads,
-                                  pimax, binfile,
-                                  ra, dec, cz,
-                                  weights1=w, weight_type='pair_product',
-                                  output_rpavg=True, verbose=False, isa=isa, bin_type='lin')
-    assert allclose(results_DDrppi, results_DDrppi_ref)
+    test_bin_type(DDrppi_mocks, autocorr, cosmology, nthreads,
+                  pimax, binfile,
+                  ra, dec, cz, weights1=w,
+                  output_rpavg=True, verbose=True,
+                  isa=isa)
 
     nmu_bins = 10
     mu_max = 1.0
-
     print("\nRunning 2-D correlation function xi(s,mu)")
-    results_DDsmu_ref = DDsmu_mocks(autocorr, cosmology, nthreads,
-                                    mu_max, nmu_bins, binfile,
-                                    ra, dec, cz, weights1=w,
-                                    output_savg=True, verbose=True,
-                                    weight_type='pair_product', isa=isa, bin_type='custom')
-    results_DDsmu = DDsmu_mocks(autocorr, cosmology, nthreads,
-                                mu_max, nmu_bins, binfile,
-                                ra, dec, cz, weights1=w,
-                                output_savg=True, verbose=False,
-                                weight_type='pair_product', isa=isa, bin_type='lin')
-    assert allclose(results_DDsmu, results_DDsmu_ref)
+    test_bin_type(DDsmu_mocks, autocorr, cosmology, nthreads,
+                  mu_max, nmu_bins, binfile,
+                  ra, dec, cz, weights1=w,
+                  output_savg=True, verbose=True,
+                  isa=isa)
 
 
 def test_linear_binning_theory(isa='fallback'):
@@ -113,7 +108,7 @@ def test_linear_binning_theory(isa='fallback'):
 
     tstart = time.time()
     t0 = tstart
-    x, y, z = read_catalog()
+    x, y, z = np.array(read_catalog())[:,:1000]
     w = np.ones((1, len(x)), dtype=x.dtype)
     boxsize = 420.0
     t1 = time.time()
@@ -133,80 +128,54 @@ def test_linear_binning_theory(isa='fallback'):
             file.write("{0} {1}\n".format(low, hi))
     autocorr = 1
     periodic = 1
+    bin_types = ['custom', 'lin']
+
+    def allclose(a, *others):
+        toret = True
+        for b in others:
+            toret &= all(np.allclose(a[name], b[name]) for name in ['npairs'])
+        return toret
 
-    def allclose(a, b):
-        return all(np.allclose(a[name], b[name]) for name in a.dtype.names)
+    def test_bin_type(pair_counter, *args, **kwargs):
+        res = (pair_counter(*args, bin_type=bin_type, weight_type=weight_type, **kwargs) for bin_type in bin_types for weight_type in [None, "pair_product"])
+        assert allclose(*res)
 
     print("Running 3-D correlation function DD(r)")
-    results_DD_ref = DD(autocorr, nthreads, binfile, x, y, z,
-                        weights1=w, weight_type='pair_product',
-                        verbose=True, periodic=periodic, boxsize=boxsize,
-                        isa=isa, bin_type='custom')
-    results_DD = DD(autocorr, nthreads, binfile, x, y, z,
-                    weights1=w, weight_type='pair_product',
-                    verbose=False, periodic=periodic, boxsize=boxsize,
-                    isa=isa, bin_type='auto')
-    allclose(results_DD, results_DD_ref)
+
+    test_bin_type(DD, autocorr, nthreads, binfile, x, y, z, weights1=w,
+                  verbose=True, periodic=periodic, boxsize=boxsize, isa=isa)
+
 
     print("\nRunning 2-D correlation function DD(rp,pi)")
-    results_DDrppi_ref = DDrppi(autocorr, nthreads, pimax,
-                                binfile, x, y, z,
-                                weights1=w, weight_type='pair_product',
-                                verbose=False, periodic=periodic,
-                                boxsize=boxsize, isa=isa, bin_type='custom')
-    results_DDrppi = DDrppi(autocorr, nthreads, pimax,
-                            binfile, x, y, z,
-                            weights1=w, weight_type='pair_product',
-                            verbose=False, periodic=periodic,
-                            boxsize=boxsize, isa=isa, bin_type='lin')
-    allclose(results_DDrppi, results_DDrppi_ref)
+    test_bin_type(DDrppi, autocorr, nthreads, pimax,
+                  binfile, x, y, z, weights1=w,
+                  verbose=True, periodic=periodic,
+                  boxsize=boxsize, isa=isa)
 
     mu_max = 0.5
     nmu_bins = 10
 
     print("\nRunning 2-D correlation function DD(s,mu)")
-    results_DDsmu_ref = DDsmu(autocorr, nthreads, binfile,
-                              mu_max, nmu_bins,
-                              x, y, z,
-                              weights1=w, weight_type='pair_product',
-                              verbose=True, periodic=periodic,
-                              boxsize=boxsize, output_savg=True,
-                              isa=isa, bin_type='custom')
-    results_DDsmu = DDsmu(autocorr, nthreads, binfile,
-                          mu_max, nmu_bins,
-                          x, y, z,
-                          weights1=w, weight_type='pair_product',
-                          verbose=False, periodic=periodic,
-                          boxsize=boxsize, output_savg=True,
-                          isa=isa, bin_type='lin')
-    allclose(results_DDsmu, results_DDsmu_ref)
+    test_bin_type(DDsmu, autocorr, nthreads, binfile,
+                  mu_max, nmu_bins,
+                  x, y, z, weights1=w,
+                  verbose=True, periodic=periodic,
+                  boxsize=boxsize, output_savg=True, isa=isa)
 
     print("\nRunning 2-D projected correlation function wp(rp)")
-    results_wp_ref = wp(boxsize, pimax, nthreads,
-                        binfile, x, y, z,
-                        weights=w, weight_type='pair_product',
-                        verbose=True, isa=isa, bin_type='custom')
-    results_wp = wp(boxsize, pimax, nthreads,
-                    binfile, x, y, z,
-                    weights=w, weight_type='pair_product',
-                    verbose=False, isa=isa, bin_type='custom')
-    allclose(results_wp, results_wp_ref)
+    test_bin_type(wp, boxsize, pimax, nthreads,
+                  binfile, x, y, z, weights=w,
+                  verbose=True, isa=isa)
 
     print("\nRunning 3-D auto-correlation function xi(r)")
-    results_xi_ref = xi(boxsize, nthreads, binfile,
-                        x, y, z,
-                        weights=w, weight_type='pair_product',
-                        verbose=True, isa=isa, bin_type='custom')
-    results_xi = xi(boxsize, nthreads, binfile,
-                    x, y, z,
-                    weights=w, weight_type='pair_product',
-                    verbose=False, isa=isa, bin_type='lin')
-    allclose(results_xi, results_xi_ref)
+    test_bin_type(xi, boxsize, nthreads, binfile,
+                  x, y, z, weights=w,
+                  verbose=True, isa=isa)
 
 
 if __name__ == '__main__':
 
-    #tests()
+    tests()
     for isa in ['fallback','sse42','avx','avx512f']:
         test_linear_binning_mocks(isa=isa)
         test_linear_binning_theory(isa=isa)

Corrfunc/utils.py

@@ -548,11 +548,11 @@ def translate_bin_type_string_to_enum(bin_type):
     except NameError:
         if not isinstance(bin_type, str):
             raise TypeError(msg)
-    valid_bin_type = ['AUTO', 'CUSTOM', 'LIN']
+    valid_bin_types = ['AUTO', 'CUSTOM', 'LIN']
     bin_type_upper = bin_type.upper()
-    if bin_type_upper not in valid_bin_type:
+    if bin_type_upper not in valid_bin_types:
         msg = "Desired bin type = {0} is not in the list of valid "\
-              "bin types = {1}".format(bin_type, valid_bin_type)
+              "bin types = {1}".format(bin_type, valid_bin_types)
         raise ValueError(msg)
 
     enums = {'AUTO': 0,

mocks/DDrppi_mocks/countpairs_rp_pi_mocks_impl.c.src

@@ -297,12 +297,13 @@ int countpairs_mocks_DOUBLE(const int64_t ND1, DOUBLE *ra1, DOUBLE *dec1, DOUBLE
     double *rupp;
     int nrpbin ;
     double rpmin,rpmax;
-    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp,&(options->bin_type));
+    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp);
     if( ! (rpmin > 0.0 && rpmax > 0.0 && rpmin < rpmax && nrpbin > 0)) {
         fprintf(stderr,"Error: Could not setup with R bins correctly. (rmin = %lf, rmax = %lf, with nbins = %d). Expected non-zero rmin/rmax with rmax > rmin and nbins >=1 \n",
                 rpmin, rpmax, nrpbin);
         return EXIT_FAILURE;
     }
+    detect_bin_type(rupp, nrpbin, &(options->bin_type));
     if (options->verbose) {
         if (options->bin_type == BIN_LIN) fprintf(stderr,"Linear R binning\n");
         else fprintf(stderr,"Custom R binning\n");

mocks/DDrppi_mocks/countpairs_rp_pi_mocks_kernels.c.src

@@ -53,12 +53,11 @@ static inline int countpairs_rp_pi_mocks_avx512_intrinsics_DOUBLE(const int64_t
     const int32_t need_rpavg = src_rpavg != NULL;
     const int32_t need_weightavg = src_weightavg != NULL;
     const DOUBLE sqr_rpmin=rpmin*rpmin, sqr_rpmax=rpmax*rpmax;
-    DOUBLE inv_rpstep=0., rpmin_invstep=0.;
     AVX512_FLOATS m_inv_rpstep = AVX512_SETZERO_FLOAT();
     AVX512_FLOATS m_rpmin_invstep = AVX512_SETZERO_FLOAT();
     if (bin_type == BIN_LIN) {
-        inv_rpstep = (nbin - 1)/(rpmax - rpmin);
-        rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
+        const DOUBLE inv_rpstep = (nbin - 1)/(rpmax - rpmin);
+        const DOUBLE rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
         m_inv_rpstep = AVX512_SET_FLOAT(inv_rpstep);
         m_rpmin_invstep = AVX512_SET_FLOAT(rpmin_invstep);
     }

mocks/DDsmu_mocks/countpairs_s_mu_mocks_impl.c.src

@@ -297,12 +297,13 @@ int countpairs_mocks_s_mu_DOUBLE(const int64_t ND1, DOUBLE *ra1, DOUBLE *dec1, D
     double *supp;
     int nsbin;
     double smin,smax;
-    setup_bins(sbinfile,&smin,&smax,&nsbin,&supp,&(options->bin_type));
+    setup_bins(sbinfile,&smin,&smax,&nsbin,&supp);
     if( ! (smin > 0.0 && smax > 0.0 && smin < smax && nsbin > 0)) {
         fprintf(stderr,"Error: Could not setup with S bins correctly. (smin = %lf, smax = %lf, with nbins = %d). Expected non-zero smin/smax with smax > smin and nbins >=1 \n",
                 smin, smax, nsbin);
         return EXIT_FAILURE;
     }
+    detect_bin_type(supp, nsbin, &(options->bin_type));
     if (options->verbose) {
         if (options->bin_type == BIN_LIN) fprintf(stderr,"Linear S binning\n");
         else fprintf(stderr,"Custom S binning\n");

mocks/DDsmu_mocks/countpairs_s_mu_mocks_kernels.c.src

@@ -48,12 +48,11 @@ static inline int countpairs_s_mu_mocks_avx512_intrinsics_DOUBLE(const int64_t N
     const int32_t need_savg = src_savg != NULL;
     const int32_t need_weightavg = src_weightavg != NULL;
     const DOUBLE sqr_smin=smin*smin, sqr_smax=smax*smax;
-    DOUBLE inv_sstep=0., smin_invstep=0.;
     AVX512_FLOATS m_inv_sstep = AVX512_SETZERO_FLOAT();
     AVX512_FLOATS m_smin_invstep = AVX512_SETZERO_FLOAT();
     if (bin_type == BIN_LIN) {
-        inv_sstep = (nsbin - 1)/(smax - smin);
-        smin_invstep = 1 - smin*inv_sstep; //trick to avoid adding one to (r - rmin)/rstep
+        const DOUBLE inv_sstep = (nsbin - 1)/(smax - smin);
+        const DOUBLE smin_invstep = 1 - smin*inv_sstep; //trick to avoid adding one to (r - rmin)/rstep
         m_inv_sstep = AVX512_SET_FLOAT(inv_sstep);
         m_smin_invstep = AVX512_SET_FLOAT(smin_invstep);
     }

theory/DD/countpairs_impl.c.src

@@ -198,13 +198,14 @@ int countpairs_DOUBLE(const int64_t ND1, DOUBLE *X1, DOUBLE *Y1, DOUBLE *Z1,
     double *rupp=NULL;
     int nrpbin ;
     double rpmin,rpmax;
-    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp,&(options->bin_type));
+    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp);
     if( ! (rpmin >=0.0 && rpmax > 0.0 && rpmin < rpmax && nrpbin > 0)) {
         fprintf(stderr,"Error: Could not setup with R bins correctly. (rmin = %lf, rmax = %lf, with nbins = %d). "
                 "Expected non-zero rmin/rmax with rmax > rmin and nbins >=1 \n",
                 rpmin, rpmax, nrpbin);
         return EXIT_FAILURE;
     }
+    detect_bin_type(rupp, nrpbin, &(options->bin_type));
     if (options->verbose) {
         if (options->bin_type == BIN_LIN) fprintf(stderr,"Linear R binning\n");
         else fprintf(stderr,"Custom R binning\n");

theory/DD/countpairs_kernels.c.src

@@ -51,12 +51,11 @@ static inline int countpairs_avx512_intrinsics_DOUBLE(const int64_t N0, DOUBLE *
     const int32_t need_rpavg = src_rpavg != NULL;
     const int32_t need_weightavg = src_weightavg != NULL;
     const DOUBLE sqr_rpmin=rpmin*rpmin, sqr_rpmax=rpmax*rpmax;
-    DOUBLE inv_rpstep=0., rpmin_invstep=0.;
     AVX512_FLOATS m_inv_rpstep = AVX512_SETZERO_FLOAT();
     AVX512_FLOATS m_rpmin_invstep = AVX512_SETZERO_FLOAT();
     if (bin_type == BIN_LIN) {
-        inv_rpstep = (nbin - 1)/(rpmax - rpmin);
-        rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
+        const DOUBLE inv_rpstep = (nbin - 1)/(rpmax - rpmin);
+        const DOUBLE rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
         m_inv_rpstep = AVX512_SET_FLOAT(inv_rpstep);
         m_rpmin_invstep = AVX512_SET_FLOAT(rpmin_invstep);
     }
@@ -249,7 +248,7 @@ static inline int countpairs_avx512_intrinsics_DOUBLE(const int64_t N0, DOUBLE *
                     }
                 }//backwards loop over the bins
             }
-            if(need_rpavg || need_weightavg) {
+            if(need_rpavg || need_weightavg || (bin_type == BIN_LIN)) {
                 //Do I need this step of going via the union? accessing int[] -> AVX* vector might
                 //cause alignment problems but accessing the ints from an AVX*
                 //register should always be fine
@@ -260,6 +259,9 @@ static inline int countpairs_avx512_intrinsics_DOUBLE(const int64_t N0, DOUBLE *
 #endif
                 for(int jj=0;jj<AVX512_NVEC;jj++) {
                     const int kbin = union_rpbin.ibin[jj];
+                    if(bin_type == BIN_LIN){
+                        npairs[kbin]++;
+                    }
                     if(need_rpavg){
                         const DOUBLE r = union_mDperp.Dperp[jj];
                         rpavg[kbin] += r;
@@ -525,7 +527,8 @@ static inline int countpairs_avx_intrinsics_DOUBLE(const int64_t N0, DOUBLE *x0,
                 if (!need_rpavg) {
                     union_mDperp.m_Dperp = AVX_SQRT_FLOAT(r2);
                 }
-                m_rpbin = AVX_FMA_ADD_TRUNCATE_FLOATS(union_mDperp.m_Dperp,m_inv_rpstep,m_rpmin_invstep);
+                m_rpbin = AVX_FMA_ADD_FLOATS(union_mDperp.m_Dperp,m_inv_rpstep,m_rpmin_invstep);
+                union_rpbin.m_ibin = AVX_TRUNCATE_FLOAT_TO_INT(m_rpbin);
             }
             else {
                 //Loop backwards through nbins. m_mask_left contains all the points that are less than rpmax
@@ -544,14 +547,17 @@ static inline int countpairs_avx_intrinsics_DOUBLE(const int64_t N0, DOUBLE *x0,
                     }
                 }
             }
-            if(need_rpavg || need_weightavg) {
+            if(need_rpavg || need_weightavg || (bin_type == BIN_LIN)) {
                 union_rpbin.m_ibin = AVX_TRUNCATE_FLOAT_TO_INT(m_rpbin);
                 //protect the unroll pragma in case compiler is not icc.
 #if  __INTEL_COMPILER
 #pragma unroll(AVX_NVEC)
 #endif
                 for(int jj=0;jj<AVX_NVEC;jj++) {
                     const int kbin = union_rpbin.ibin[jj];
+                    if(bin_type == BIN_LIN){
+                        npairs[kbin]++;
+                    }
                     if(need_rpavg){
                         const DOUBLE r = union_mDperp.Dperp[jj];
                         rpavg[kbin] += r;
@@ -839,7 +845,8 @@ static inline int countpairs_sse_intrinsics_DOUBLE(const int64_t N0, DOUBLE *x0,
               if (!need_rpavg) {
                   union_mDperp.m_Dperp = SSE_SQRT_FLOAT(r2);
               }
-              m_rpbin = SSE_FMA_ADD_TRUNCATE_FLOATS(union_mDperp.m_Dperp,m_inv_rpstep,m_rpmin_invstep);
+              m_rpbin = SSE_FMA_ADD_FLOATS(union_mDperp.m_Dperp,m_inv_rpstep,m_rpmin_invstep);
+              union_rpbin.m_ibin = SSE_TRUNCATE_FLOAT_TO_INT(m_rpbin);
           }
           else {
               for(int kbin=nbin-1;kbin>=1;kbin--) {
@@ -850,21 +857,24 @@ static inline int countpairs_sse_intrinsics_DOUBLE(const int64_t N0, DOUBLE *x0,
                   npairs[kbin] += SSE_BIT_COUNT_INT(test2);
                   if(need_rpavg || need_weightavg){
                       m_rpbin = SSE_BLEND_FLOATS_WITH_MASK(m_rpbin,m_kbin[kbin], m_bin_mask);
+                      union_rpbin.m_ibin = SSE_TRUNCATE_FLOAT_TO_INT(m_rpbin);
                   }
                   int test3 = SSE_TEST_COMPARISON(m_mask_left);
                   if(test3 == 0) {
                       break;
                   }
               }
           }
-          if(need_rpavg || need_weightavg) {
-              union_rpbin.m_ibin = SSE_TRUNCATE_FLOAT_TO_INT(m_rpbin);
+          if(need_rpavg || need_weightavg || (bin_type == BIN_LIN)) {
               //protect the unroll pragma in case compiler is not icc.
 #if  __INTEL_COMPILER
 #pragma unroll(SSE_NVEC)
 #endif
               for(int jj=0;jj<SSE_NVEC;jj++) {
                   const int kbin = union_rpbin.ibin[jj];
+                  if(bin_type == BIN_LIN){
+                      npairs[kbin]++;
+                  }
                   if(need_rpavg){
                       const DOUBLE r = union_mDperp.Dperp[jj];
                       rpavg[kbin] += r;

theory/DDrppi/countpairs_rp_pi_impl.c.src

@@ -197,12 +197,13 @@ int countpairs_rp_pi_DOUBLE(const int64_t ND1, DOUBLE *X1, DOUBLE *Y1, DOUBLE *Z
     double *rupp;
     int nrpbin ;
     double rpmin,rpmax;
-    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp,&(options->bin_type));
+    setup_bins(binfile,&rpmin,&rpmax,&nrpbin,&rupp);
     if( ! (rpmin >= 0.0 && rpmax > 0.0 && rpmin < rpmax && nrpbin > 0)) {
         fprintf(stderr,"Error: Could not setup with R bins correctly. (rmin = %lf, rmax = %lf, with nbins = %d). Expected non-zero rmin/rmax with rmax > rmin and nbins >=1 \n",
                 rpmin, rpmax, nrpbin);
         return EXIT_FAILURE;
     }
+    detect_bin_type(rupp, nrpbin, &(options->bin_type));
     if (options->verbose) {
         if (options->bin_type == BIN_LIN) fprintf(stderr,"Linear R binning\n");
         else fprintf(stderr,"Custom R binning\n");

theory/DDrppi/countpairs_rp_pi_kernels.c.src

@@ -44,12 +44,11 @@ static inline int countpairs_rp_pi_avx512_intrinsics_DOUBLE(const int64_t N0, DO
     const int32_t need_rpavg = src_rpavg != NULL;
     const int32_t need_weightavg = src_weightavg != NULL;
     const DOUBLE sqr_rpmin=rpmin*rpmin, sqr_rpmax=rpmax*rpmax;
-    DOUBLE inv_rpstep=0., rpmin_invstep=0.;
     AVX512_FLOATS m_inv_rpstep = AVX512_SETZERO_FLOAT();
     AVX512_FLOATS m_rpmin_invstep = AVX512_SETZERO_FLOAT();
     if (bin_type == BIN_LIN) {
-        inv_rpstep = (nbin - 1)/(rpmax - rpmin);
-        rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
+        const DOUBLE inv_rpstep = (nbin - 1)/(rpmax - rpmin);
+        const DOUBLE rpmin_invstep = 1 - rpmin*inv_rpstep; //trick to avoid adding one to (r - rmin)/rstep
         m_inv_rpstep = AVX512_SET_FLOAT(inv_rpstep);
         m_rpmin_invstep = AVX512_SET_FLOAT(rpmin_invstep);
     }