Curved-sky kernels and bin-averaging for angular correlation functions

benchmarks/test_correlation.py

@@ -229,3 +229,47 @@ def test_xi(set_up, corr_method, t1, t2, bm, er, kind, pref):
 
     xi *= pref
     assert np.all(np.fabs(xi - bms[bm]) < ers[er] * errfac)
+
+
+def test_correlation_curvedsky():
+
+    data = np.load(os.path.dirname(__file__) +
+                   "/data/curved_sky_cocoa.npz", allow_pickle=True)
+    COSMO = ccl.Cosmology(**data['cosmo_setup'][()])
+
+    nsrcs = 5
+    nlens = 10
+
+    theta_b = data['theta_b']
+    theta_u = theta_b[1:]
+    theta_l = theta_b[:-1]
+    ell_hr = np.arange(60000)
+
+    inp = dict(ell=ell_hr, theta=theta_l, theta_max=theta_u,
+               method='Legendre')
+
+    xip, xim, gammat, wtheta = [], [], [], []
+
+    for i in range(nsrcs*(nsrcs+1)//2):
+        C_ss_i = np.interp(ell_hr, data['ell'], data['C_ss'][i])
+        xip.append(ccl.correlation(COSMO, C_ell=C_ss_i, type='GG+', **inp))
+        xim.append(ccl.correlation(COSMO, C_ell=C_ss_i, type='GG-', **inp))
+
+    for i in range(nsrcs*nlens):
+        C_gs_i = np.interp(ell_hr, data['ell'], data['C_gs'][i])
+        gammat.append(ccl.correlation(COSMO, C_ell=C_gs_i, type='NG', **inp))
+
+    for i in range(nlens):
+        C_gg_i = np.interp(ell_hr, data['ell'], data['C_gg'][i])
+        wtheta.append(ccl.correlation(COSMO, C_ell=C_gg_i, type='NN', **inp))
+
+    xip = np.concatenate(xip)
+    xim = np.concatenate(xim)
+    gammat = np.concatenate(gammat)
+    wtheta = np.concatenate(wtheta)
+    dv_ccl = np.concatenate((xip, xim, gammat, wtheta))
+    dv_cca = data['dv_cocoa']
+
+    diff = dv_ccl/dv_cca - 1
+
+    assert np.max(np.abs(diff)) < 0.01

include/ccl_correlation.h

@@ -41,6 +41,36 @@ void ccl_correlation(ccl_cosmology *cosmo,
 		     int corr_type,int do_taper_cl,double *taper_cl_limits,int flag_method,
 		     int *status);
 
+
+             /**
+ * Computes the correlation function (wrapper)
+ * @param cosmo :Cosmological parameters
+ * @param n_ell : number of multipoles in the input power spectrum
+ * @param ell : multipoles at which the power spectrum is evaluated
+ * @param cls : input power spectrum
+ * @param n_theta : number of output values of the separation angle (theta)
+ * @param theta_min : Min/left  edge of the separation bin in degrees.
+ * @param theta_max : Max/right edge of the separation bin in degrees.
+ * @param wtheta : the values of the correlation function at the angles above will be returned in this array, which should be pre-allocated
+ * @param do_taper_cl :
+ * @param taper_cl_limits
+ * @param flag_method : method to compute the correlation function. Choose between:
+ *  - CCL_CORR_FFTLOG : fast integration with FFTLog
+ *  - CCL_CORR_BESSEL : direct integration over the Bessel function
+ *  - CCL_CORR_LGNDRE : brute-force sum over legendre polynomials
+ * @param corr_type : type of correlation function. Choose between:
+ *  - CCL_CORR_GG : spin0-spin0
+ *  - CCL_CORR_GL : spin0-spin2
+ *  - CCL_CORR_LP : spin2-spin2 (xi+)
+ *  - CCL_CORR_LM : spin2-spin2 (xi-)
+ * Currently supported spin-0 fields are number counts and CMB lensing. The only spin-2 is currently shear.
+ */
+void ccl_correlation_binned(ccl_cosmology *cosmo,
+		     int n_ell,double *ell,double *cls,
+		     int n_theta,double *theta_min,double *theta_max,double *wtheta,
+		     int corr_type,int do_taper_cl,double *taper_cl_limits,int flag_method,
+		     int *status);
+
 /**
  * Computes the 3dcorrelation function (wrapper)
  * @param cosmo :Cosmological parameters

pyccl/ccl_correlation.i

@@ -11,6 +11,8 @@
     (double* larr, int nlarr),
     (double* clarr, int nclarr),
     (double* theta, int nt),
+    (double* theta_min, int nt),
+    (double* theta_max, int ntmax),
     (double* r, int nr),
     (double* s, int ns),
     (double* sig, int nsig)}
@@ -27,6 +29,17 @@
         raise CCLError("Input shape for `theta` must match `(nout,)`!")
 %}
 
+%feature("pythonprepend") correlation_vec_binned %{
+    if numpy.shape(larr) != numpy.shape(clarr):
+        raise CCLError("Input shape for `larr` must match `clarr`!")
+
+    if numpy.shape(theta_min) != (nout,):
+        raise CCLError("Input shape for `theta_min` must match `(nout,)`!")
+
+    if numpy.shape(theta_max) != (nout,):
+        raise CCLError("Input shape for `theta_max` must match `(nout,)`!")
+%}
+
 %feature("pythonprepend") correlation_3d_vec %{
     if numpy.shape(r) != (nxi,):
         raise CCLError("Input shape for `r` must match `(nxi,)`!")
@@ -63,6 +76,30 @@ void correlation_vec(ccl_cosmology *cosmo, double* larr, int nlarr,
         output, corr_type, 0, NULL, method, status);
 }
 
+void correlation_vec_binned(ccl_cosmology *cosmo,
+                            double* larr, int nlarr,
+                            double* clarr, int nclarr,
+                            double* theta_min, int nt,
+                            double* theta_max, int ntmax,
+                            int corr_type, int method,
+                            int nout,
+                            double* output,
+                            int *status) {
+    if (ntmax != nt) {
+        *status = CCL_ERROR_INCONSISTENT;
+        return;
+    }
+
+    if (nout != nt) {
+        *status = CCL_ERROR_INCONSISTENT;
+        return;
+    }
+
+    ccl_correlation_binned(
+        cosmo, nlarr, larr, clarr, nt, theta_min, theta_max,
+        output, corr_type, 0, NULL, method, status);
+}
+
 void correlation_3d_vec(ccl_cosmology *cosmo,ccl_f2d_t *psp,
                         double a, double* r, int nr,
                         int nxi, double* xi, int *status) {

pyccl/correlations.py

@@ -42,7 +42,8 @@ class CorrelationTypes(Enum):
 }
 
 
-def correlation(cosmo, *, ell, C_ell, theta, type='NN', method='fftlog'):
+def correlation(cosmo, *, ell, C_ell, theta,
+                type='NN', method='fftlog', theta_max=None):
     r"""Compute the angular correlation function.
 
     .. math::
@@ -71,6 +72,12 @@ def correlation(cosmo, *, ell, C_ell, theta, type='NN', method='fftlog'):
         * :math:`s_a=2`, :math:`s_b=0` e.g. galaxy-shear, and :math:`\kappa`-shear
         * :math:`s_a=s_b=2` e.g. shear-shear.
 
+    Bin-averaging, where we predict the average correlation within a given theta-bin,
+    is also implemented. This is only available for the Legendre method, as the 
+    bin-averaging can be done analytically in that case. The method we use for the
+    NN, NG, and GG+/- correlations follow Eq 65, B1, and B5, respectively, from 
+    arxiv:2012.08568
+
     .. note::
         For scales smaller than :math:`\sim 0.1^{\circ}`, the input power
         spectrum should be sampled to sufficienly high :math:`\ell` to ensure
@@ -87,7 +94,9 @@ def correlation(cosmo, *, ell, C_ell, theta, type='NN', method='fftlog'):
                           spectrum.
         C_ell (array): Input angular power spectrum.
         theta (:obj:`float` or `array`): Angular separation(s) at which to
-            calculate the angular correlation function (in degrees).
+            calculate the angular correlation function (in degrees). If theta_max
+            is passed then this is interpreted as the minimum/left edge in a given
+            separation bin.
         type (:obj:`str`): Type of correlation function. Choices: ``'NN'`` (0x0),
             ``'NG'`` (0x2), ``'GG+'`` (2x2, :math:`\xi_+`),
             ``'GG-'`` (2x2, :math:`\xi_-`), where numbers refer to the spins
@@ -99,6 +108,9 @@ def correlation(cosmo, *, ell, C_ell, theta, type='NN', method='fftlog'):
             Choices: ``'Bessel'`` (direct integration over Bessel function),
             ``'FFTLog'`` (fast integration with FFTLog), ``'Legendre'``
             (brute-force sum over Legendre polynomials).
+        theta_max (:obj:`float` or `array1): Maximum angular separation(s) for the given
+            separation bin (in degrees). If passed, then the `theta` input is interpreted
+            as a minimum/left edge of the bin.
 
     Returns:
         (:obj:`float` or `array`): Value(s) of the correlation function at the
@@ -119,15 +131,40 @@ def correlation(cosmo, *, ell, C_ell, theta, type='NN', method='fftlog'):
     if scalar := isinstance(theta, (int, float)):
         theta = np.array([theta, ])
 
+        if theta_max is not None:
+            if isinstance(theta_max, (int, float)):
+                theta_max = np.array([theta_max, ])
+
+            assert theta.size == theta_max.size, \
+                (f"theta and theta_max have different sizes"
+                 f"({theta.size} != {theta_max.size})")
+
+            assert not np.allclose(theta, theta_max), \
+                "theta_min cannot be the same as theta_max"
+
     if np.all(np.array(C_ell) == 0):
         # short-cut and also avoid integration errors
         wth = np.zeros_like(theta)
     else:
-        # Call correlation function
-        wth, status = lib.correlation_vec(cosmo, ell, C_ell, theta,
-                                          correlation_types[type],
-                                          correlation_methods[method],
-                                          len(theta), status)
+        if theta_max is None:
+            wth, status = lib.correlation_vec(
+                cosmo, ell, C_ell, theta,
+                correlation_types[type],
+                correlation_methods[method],
+                len(theta), status)
+        else:
+            wth, status = lib.correlation_vec_binned(
+                cosmo,
+                ell,
+                C_ell,
+                theta,
+                theta_max,
+                correlation_types[type],
+                correlation_methods[method],
+                len(theta),
+                status,
+            )
+
     check(status, cosmo_in)
     if scalar:
         return wth[0]

pyccl/tests/test_correlations.py

@@ -29,6 +29,20 @@ def test_correlation_smoke(method):
         assert np.all(np.isfinite(corr))
         assert np.shape(corr) == np.shape(tval)
 
+    if method == 'legendre':
+        for tval in [t_arr, t_lst, t_scl, t_int]:
+
+            if isinstance(tval, list):
+                tmax = [t*2 for t in tval]
+            else:
+                tmax = tval*2
+
+            corr = ccl.correlation(
+                COSMO, ell=ell, C_ell=cl, theta=tval, type='NN', method=method,
+                theta_max=tmax)
+            assert np.all(np.isfinite(corr))
+            assert np.shape(corr) == np.shape(tval)
+
 
 @pytest.mark.parametrize(
     'rval',