Fexible non-Limber parameters (#1211)

* added flexibility to FKEM

* tested

pyccl/_nonlimber_FKEM.py

@@ -35,8 +35,8 @@ def _get_general_params(b):
 
 
 def _nonlimber_FKEM(
-    cosmo, clt1, clt2, p_of_k_a, p_of_k_a_lin, ls, l_limber, limber_max_error
-):
+        cosmo, clt1, clt2, p_of_k_a,
+        ls, l_limber, **params):
     """clt1, clt2 are lists of tracer in a tracer object
     cosmo (:class:`~pyccl.core.Cosmology`): A Cosmology object.
     psp non-linear power spectrum
@@ -55,6 +55,10 @@ def _nonlimber_FKEM(
     fll_t2 = clt2.get_f_ell(ls)
     status = 0
 
+    p_of_k_a_lin = params['pk_linear']
+    limber_max_error = params['limber_max_error']
+    Nchi = params['Nchi']
+    chi_min = params['chi_min']
     if (not (isinstance(p_of_k_a, str) and isinstance(p_of_k_a_lin, str)) and
        not (isinstance(p_of_k_a, ccl.Pk2D)
             and isinstance(p_of_k_a_lin, ccl.Pk2D)
@@ -86,9 +90,12 @@ def _nonlimber_FKEM(
     min_chis_t2 = np.min([np.min(i) for i in chis_t2])
     max_chis_t1 = np.max([np.max(i) for i in chis_t1])
     max_chis_t2 = np.max([np.max(i) for i in chis_t2])
-    chi_min = np.min([min_chis_t1, min_chis_t2])
+    if chi_min is None:
+        chi_min = np.min([min_chis_t1, min_chis_t2])
     chi_max = np.max([max_chis_t1, max_chis_t2])
-    Nchi = min(min(len(i) for i in chis_t1), min(len(i) for i in chis_t2))
+    if Nchi is None:
+        Nchi = min(min(len(i) for i in chis_t1),
+                   min(len(i) for i in chis_t2))
     """zero chi_min will result in a divide-by-zero error.
     If it is zero, we set it to something very small
     """

pyccl/cells.py

@@ -20,6 +20,8 @@ def angular_cl(
     limber_max_error=0.01,
     limber_integration_method="qag_quad",
     non_limber_integration_method="FKEM",
+    fkem_chi_min=None,
+    fkem_Nchi=None,
     p_of_k_a_lin=DEFAULT_POWER_SPECTRUM,
     return_meta=False
 ):
@@ -49,6 +51,20 @@ def angular_cl(
             for the non-Limber integrals. Currently the only method implemented
             is ``'FKEM'`` (see the `N5K paper <https://arxiv.org/abs/2212.04291>`_
             for details).
+        fkem_chi_min: Minimum comoving distance used by `FKEM` to sample the
+            tracer radial kernels. If ``None``, the minimum distance over which
+            the kernels are defined will be used (capped to 1E-6 Mpc if this
+            value is zero). Users are encouraged to experiment with this parameter
+            and ``fkem_Nchi`` to ensure the robustness of the output
+            :math:`C_\\ell`s.
+        fkem_Nchi: Number of values of the comoving distance over which `FKEM`
+            will interpolate the radial kernels. If ``None`` the smallest number
+            over which the kernels are currently sampled will be used. Note that
+            `FKEM` will use a logarithmic sampling for distances between
+            ``fkem_chi_min`` and the maximum distance over which the tracers
+            are defined.  Users are encouraged to experiment with this parameter
+            and ``fkem_chi_min`` to ensure the robustness of the output
+            :math:`C_\\ell`s.
         p_of_k_a_lin (:class:`~pyccl.pk2d.Pk2D`, :obj:`str` or :obj:`None`): 
             3D linear Power spectrum to project, for special use in
             PT calculations using the FKEM non-limber integration technique.
@@ -110,17 +126,20 @@ def angular_cl(
     if not (np.diff(ell_use) > 0).all():
         raise ValueError("ell values must be monotonically increasing")
 
+    fkem_params = {'pk_linear': p_of_k_a_lin,
+                   'limber_max_error': limber_max_error,
+                   'Nchi': fkem_Nchi,
+                   'chi_min': fkem_chi_min}
     if auto_limber or (type(l_limber) is not str and ell_use[0] < l_limber):
         if non_limber_integration_method == "FKEM":
             l_limber, cl_non_limber, status = _nonlimber_FKEM(
                 cosmo,
                 tracer1,
                 tracer2,
                 p_of_k_a,
-                p_of_k_a_lin,
                 ell_use,
                 l_limber,
-                limber_max_error,
+                **fkem_params
             )
         check(status, cosmo=cosmo)
     else:

pyccl/tests/test_cells.py

@@ -125,6 +125,36 @@ def test_cells_raise_weird_pk():
         ccl.angular_cl(COSMO, LENS, LENS, ells, p_of_k_a=lambda k, a: 10)
 
 
+def test_fkem_chi_params():
+    # Redshift distribution
+    z = np.linspace(0, 4.72, 60)
+    nz = z**2*np.exp(-0.5*((z-1.5)/0.7)**2)
+
+    # Bias
+    bz = np.ones_like(z)
+
+    # Power spectra
+    ls = np.unique(np.geomspace(2, 2000, 128).astype(int)).astype(float)
+    cosmo = ccl.CosmologyVanillaLCDM()
+    tracer_gal = ccl.NumberCountsTracer(cosmo, has_rsd=False,
+                                        dndz=(z, nz), bias=(z, bz))
+    cl_gg = ccl.angular_cl(cosmo, tracer_gal, tracer_gal, ls,
+                           l_limber=-1)
+    cl_ggn = ccl.angular_cl(cosmo, tracer_gal, tracer_gal, ls,
+                            l_limber=1000)
+    cl_ggn_b = ccl.angular_cl(cosmo, tracer_gal, tracer_gal, ls,
+                              l_limber=1000, fkem_chi_min=1.0,
+                              fkem_Nchi=100)
+
+    ell_good = ls > 100
+
+    # Check that, above ell, the non-Limber calculation does not
+    # agree with Limber when using the default FKEM sampling params
+    assert not np.all(np.fabs(cl_ggn/cl_gg-1)[ell_good] < 0.01)
+    # Check that it works with custom ones.
+    assert np.all(np.fabs(cl_ggn_b/cl_gg-1)[ell_good] < 0.01)
+
+
 def test_cells_mg():
     # Check that if we feed the non-linear matter power spectrum from a MG
     # cosmology into a Calculator and get Cells using MG tracers, we get the