Make tracer creation more robust (#928)

* Add option for spline integration. Fail loudly when integration does not converge.

* Specify number of chi samples in tracers

* Rename n_chi to n_sample. Add docstring.

* Default to have warnings shown.

* Add spline integration fall-back for lensing kernel.

* Fixes for openmp

* Fix typo.

* Added flags for integration option. Warn if status message is being overwritten.

* Switch between integration options. Check if lensing kernel integration error is sufficiently small compared to the peak of the kernel.

* Allow get_kernel to return the internal spline representation.

* Added tests for tracer integration options, delta n(z)s.

* Remove left-over #define

* Tests for n_samples argument

* Add test for n_samples=None

* Update docstrings

* Remove superfluous warning.

* Explicit abolute test tolerances.

* Disable openmp in debug mode by default.

* Correction for low-res n(z).

* Account for case where z[0] != 0

* Make sure test covers edge case if test n(z) gets changed in the future.

* Add warning if n(z) is poorly sampled.

CMakeLists.txt

@@ -60,18 +60,22 @@ find_package(OpenMP)
 
 # Compilation flags
 set(CMAKE_C_FLAGS_RELEASE "-O3 -fomit-frame-pointer -fno-common -fPIC -std=gnu99 -DHAVE_ANGPOW")
-set(CMAKE_C_FLAGS_DEBUG   "-Og -g -fomit-frame-pointer -fno-common -fPIC -std=gnu99 -DHAVE_ANGPOW")
+set(CMAKE_C_FLAGS_DEBUG   "-O0 -g -fomit-frame-pointer -fno-common -fPIC -std=gnu99 -DHAVE_ANGPOW")
 
 if ("${CMAKE_C_COMPILER_ID}" MATCHES "^(Apple)?Clang$") 
-    if(FORCE_OPENMP OR OpenMP_C_FOUND)
+    if(OpenMP_C_FOUND)
         # OpenMP flags for macOS clang
         set(CMAKE_C_FLAGS_RELEASE "${CMAKE_C_FLAGS_RELEASE} -Xpreprocessor -fopenmp")
-        set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -Xpreprocessor -fopenmp")
+        if(FORCE_OPENMP)
+            set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -Xpreprocessor -fopenmp")
+        endif()
     endif()
 else()
     # OpenMP flags for all other compilers clang
     set(CMAKE_C_FLAGS_RELEASE "${CMAKE_C_FLAGS_RELEASE} -fopenmp")
-    set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -fopenmp")
+    if(FORCE_OPENMP)
+        set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -fopenmp")
+    endif()
 endif()
 
 if(DEFINED ENV{CONDA_PREFIX})

include/ccl_core.h

@@ -190,6 +190,9 @@ typedef struct ccl_gsl_params {
   size_t HM_LIMIT;
   int HM_INT_METHOD;
 
+  // Flags for using spline integration
+  bool NZ_NORM_SPLINE_INTEGRATION;
+  bool LENSING_KERNEL_SPLINE_INTEGRATION;
 } ccl_gsl_params;
 
 extern const ccl_gsl_params default_gsl_params;

include/ccl_error.h

@@ -46,6 +46,7 @@ CCL_BEGIN_DECLS
 #define CCL_ERROR_LINEAR_POWER_INIT 1061
 #define CCL_ERROR_SIGMA_INIT 1062
 #define CCL_ERROR_HMF_INIT 1063
+#define CCL_ERROR_OVERWRITE 1064
 
 typedef enum {
   CCL_DEBUG_MODE_OFF = 0,

pyccl/ccl.i

@@ -21,7 +21,7 @@ from .errors import CCLError
 %init %{
     import_array();
     // Tell CCL to not print to stdout/stderr for debugging.
-    ccl_set_debug_policy(CCL_DEBUG_MODE_OFF);
+    ccl_set_debug_policy(CCL_DEBUG_MODE_ON);
 %}
 
 // Automatically document arguments and output types of all functions

pyccl/tests/test_tracers.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import pyccl as ccl
+from pyccl.errors import CCLWarning
 
 COSMO = ccl.Cosmology(
     Omega_c=0.27, Omega_b=0.045, h=0.67, sigma8=0.8, n_s=0.96,
@@ -11,28 +12,32 @@ def dndz(z):
     return np.exp(-((z-0.5)/0.1)**2)
 
 
-def get_tracer(tracer_type):
+def get_tracer(tracer_type, cosmo=None, **tracer_kwargs):
+    if cosmo is None:
+        cosmo = COSMO
     z = np.linspace(0., 1., 2000)
     n = dndz(z)
     b = np.sqrt(1. + z)
 
     if tracer_type == 'nc':
         ntr = 3
-        tr = ccl.NumberCountsTracer(COSMO, True,
+        tr = ccl.NumberCountsTracer(cosmo, True,
                                     dndz=(z, n),
                                     bias=(z, b),
-                                    mag_bias=(z, b))
+                                    mag_bias=(z, b),
+                                    **tracer_kwargs)
     elif tracer_type == 'wl':
         ntr = 2
-        tr = ccl.WeakLensingTracer(COSMO,
+        tr = ccl.WeakLensingTracer(cosmo,
                                    dndz=(z, n),
-                                   ia_bias=(z, b))
+                                   ia_bias=(z, b),
+                                   **tracer_kwargs)
     elif tracer_type == 'cl':
         ntr = 1
-        tr = ccl.CMBLensingTracer(COSMO, 1100.)
+        tr = ccl.CMBLensingTracer(cosmo, 1100., **tracer_kwargs)
     else:
         ntr = 0
-        tr = ccl.Tracer()
+        tr = ccl.Tracer(**tracer_kwargs)
     return tr, ntr
 
 
@@ -151,3 +156,115 @@ def test_tracer_nz_support():
 
     with pytest.raises(ValueError):
         _ = ccl.CMBLensingTracer(calculator_cosmo, z_source=2.0)
+
+
+def test_tracer_nz_norm_spline_vs_gsl_intergation():
+    # Create a new Cosmology object so that we're not messing with the other
+    # tests
+    cosmo = ccl.Cosmology(Omega_c=0.27, Omega_b=0.045, h=0.67,
+                          sigma8=0.8, n_s=0.96,
+                          transfer_function='bbks',
+                          matter_power_spectrum='linear')
+    cosmo.cosmo.gsl_params.NZ_NORM_SPLINE_INTEGRATION = True
+    tr_wl, _ = get_tracer("wl", cosmo)
+    tr_nc, _ = get_tracer("nc", cosmo)
+
+    w_wl_spline, _ = tr_wl.get_kernel(chi=None)
+    w_nc_spline, _ = tr_nc.get_kernel(chi=None)
+
+    cosmo.cosmo.gsl_params.NZ_NORM_SPLINE_INTEGRATION = False
+    tr_wl, _ = get_tracer("wl", cosmo)
+    tr_nc, _ = get_tracer("nc", cosmo)
+
+    w_wl_gsl, _ = tr_wl.get_kernel(chi=None)
+    w_nc_gsl, _ = tr_nc.get_kernel(chi=None)
+
+    for w_spline, w_gsl in zip(w_wl_spline, w_wl_gsl):
+        assert np.allclose(w_spline, w_gsl, atol=0, rtol=1e-8)
+    for w_spline, w_gsl in zip(w_nc_spline, w_nc_gsl):
+        assert np.allclose(w_spline, w_gsl, atol=0, rtol=1e-8)
+
+
+@pytest.mark.parametrize('z_min, z_max, n_z_samples', [(0.0, 1.0, 2000),
+                                                       (0.0, 1.0, 1000),
+                                                       (0.0, 1.0, 500),
+                                                       (0.0, 1.0, 100),
+                                                       (0.3, 1.0, 1000)])
+def test_tracer_lensing_kernel_spline_vs_gsl_intergation(z_min, z_max,
+                                                         n_z_samples):
+    # Create a new Cosmology object so that we're not messing with the other
+    # tests
+    cosmo = ccl.Cosmology(Omega_c=0.27, Omega_b=0.045, h=0.67,
+                          sigma8=0.8, n_s=0.96,
+                          transfer_function='bbks',
+                          matter_power_spectrum='linear')
+    z = np.linspace(z_min, z_max, n_z_samples)
+    n = dndz(z)
+
+    # Make sure case where z[0] > 0 and n[0] > 0 is tested for
+    if z_min > 0:
+        assert n[0] > 0
+
+    cosmo.cosmo.gsl_params.LENSING_KERNEL_SPLINE_INTEGRATION = True
+    tr_wl = ccl.WeakLensingTracer(cosmo, dndz=(z, n))
+    w_wl_spline, _ = tr_wl.get_kernel(chi=None)
+
+    cosmo.cosmo.gsl_params.LENSING_KERNEL_SPLINE_INTEGRATION = False
+    tr_wl = ccl.WeakLensingTracer(cosmo, dndz=(z, n))
+    w_wl_gsl, chi = tr_wl.get_kernel(chi=None)
+
+    # Peak of kernel is ~1e-5
+    if n_z_samples >= 1000:
+        assert np.allclose(w_wl_spline[0], w_wl_gsl[0], atol=1e-10, rtol=1e-9)
+    else:
+        assert np.allclose(w_wl_spline[0], w_wl_gsl[0], atol=5e-9, rtol=1e-5)
+
+
+def test_tracer_delta_function_nz():
+    z = np.linspace(0., 1., 2000)
+    z_s_idx = int(z.size*0.8)
+    z_s = z[z_s_idx]
+    n = np.zeros_like(z)
+    n[z_s_idx] = 2.0
+
+    tr_wl = ccl.WeakLensingTracer(COSMO, dndz=(z, n))
+
+    # Single source plane tracer to compare against
+    chi_kappa, w_kappa = ccl.tracers.get_kappa_kernel(COSMO, z_source=z_s,
+                                                      nsamples=100)
+
+    # Use the same comoving distances
+    w = tr_wl.get_kernel(chi=chi_kappa)
+
+    assert np.allclose(w[0], w_kappa, atol=1e-8, rtol=1e-6)
+    # at z=z_source, interpolation becomes apparent, so for this test we
+    # ignore these data points.
+    assert np.allclose(w[0][:-2], w_kappa[:-2], atol=1e-11, rtol=1e-11)
+
+
+@pytest.mark.parametrize('tracer_type', ['nc', 'wl', 'cl'])
+def test_tracer_n_sample_smoke(tracer_type):
+    tr, ntr = get_tracer(tracer_type, n_samples=50)
+    if tracer_type != "cl":
+        # n_samples=None should fall back to using the samples from the n(z).
+        tr, ntr = get_tracer(tracer_type, n_samples=None)
+
+
+def test_tracer_n_sample_wl():
+    z = np.linspace(0., 1., 2000)
+    n = dndz(z)
+
+    n_samples = 50
+    tr_wl = ccl.WeakLensingTracer(COSMO, dndz=(z, n), n_samples=n_samples)
+    w, chi = tr_wl.get_kernel(chi=None)
+
+    assert w[0].shape[-1] == n_samples
+    assert chi[0].shape[-1] == n_samples
+
+
+def test_tracer_n_sample_warn():
+    z = np.linspace(0., 1., 50)
+    n = dndz(z)
+
+    with pytest.warns(CCLWarning):
+        _ = ccl.WeakLensingTracer(COSMO, dndz=(z, n))

pyccl/tracers.py

@@ -1,10 +1,14 @@
+import warnings
+
+import numpy as np
+
 from . import ccllib as lib
 from .core import check
 from .background import comoving_radial_distance, growth_rate, \
     growth_factor, scale_factor_of_chi, h_over_h0
+from .errors import CCLWarning
 from .pyutils import _check_array_params, NoneArr, _vectorize_fn6, \
     _get_spline1d_arrays
-import numpy as np
 
 
 def _Sig_MG(cosmo, a, k=None):
@@ -69,7 +73,7 @@ def get_density_kernel(cosmo, dndz):
     return chi, wchi
 
 
-def get_lensing_kernel(cosmo, dndz, mag_bias=None):
+def get_lensing_kernel(cosmo, dndz, mag_bias=None, n_chi=None):
     """This convenience function returns the radial kernel for
     weak-lensing-like. Given an unnormalized redshift distribution
     and an optional magnification bias function, it returns
@@ -96,17 +100,28 @@ def get_lensing_kernel(cosmo, dndz, mag_bias=None):
     z_s, s = _check_array_params(mag_bias, 'mag_bias')
     _check_background_spline_compatibility(cosmo, dndz[0])
 
-    # Calculate number of samples in chi
-    nchi = lib.get_nchi_lensing_kernel_wrapper(z_n)
+    if n_chi is None:
+        # Calculate number of samples in chi
+        n_chi = lib.get_nchi_lensing_kernel_wrapper(z_n)
+
+    if n_chi > len(z_n):
+        warnings.warn(
+            f"The number of samples in the n(z) ({len(z_n)}) is smaller than "
+            f"the number of samples in the lensing kernel ({n_chi}). Consider "
+            f"disabling spline integration for the lensing kernel by setting "
+            f"cosmo.cosmo.gsl_params.LENSING_KERNEL_SPLINE_INTEGRATION "
+            f"= False",
+            category=CCLWarning)
+
     # Compute array of chis
     status = 0
     chi, status = lib.get_chis_lensing_kernel_wrapper(cosmo.cosmo, z_n[-1],
-                                                      nchi, status)
+                                                      n_chi, status)
     # Compute kernel
     wchi, status = lib.get_lensing_kernel_wrapper(cosmo.cosmo,
                                                   z_n, n, z_n[-1],
                                                   int(has_magbias), z_s, s,
-                                                  chi, nchi, status)
+                                                  chi, n_chi, status)
     check(status, cosmo=cosmo)
     return chi, wchi
 
@@ -188,33 +203,47 @@ def get_kernel(self, chi):
         in this `Tracer`.
 
         Args:
-            chi (float or array_like): values of the comoving
-                radial distance in increasing order and in Mpc.
+            chi (float or array_like, optional): values of the comoving
+                radial distance in increasing order and in Mpc. If None,
+                returns the kernel at the internal spline knots.
 
         Returns:
-            array_like: list of radial kernels for each tracer. \
-                The shape will be `(n_tracer, chi.size)`, where \
-                `n_tracer` is the number of tracers. The last \
-                dimension will be squeezed if the input is a \
+            array_like: list of radial kernels for each tracer.
+                The shape will be `(n_tracer, chi.size)`, where
+                `n_tracer` is the number of tracers. The last
+                dimension will be squeezed if the input is a
                 scalar.
+                If no chi was provided, returns two arrays, the radial kernels
+                and the comoving radial distances corresponding to the internal
+                values used for interpolation.
         """
         if not hasattr(self, '_trc'):
             return []
 
-        chi_use = np.atleast_1d(chi)
+        if chi is None:
+            chis = []
+        else:
+            chi_use = np.atleast_1d(chi)
         kernels = []
         for t in self._trc:
-            status = 0
-            w, status = lib.cl_tracer_get_kernel(t, chi_use,
-                                                 chi_use.size,
-                                                 status)
-            check(status)
+            if chi is None:
+                chi_use, w = _get_spline1d_arrays(t.kernel.spline)
+                chis.append(chi_use)
+            else:
+                status = 0
+                w, status = lib.cl_tracer_get_kernel(t, chi_use,
+                                                     chi_use.size,
+                                                     status)
+                check(status)
             kernels.append(w)
-        kernels = np.array(kernels)
-        if np.ndim(chi) == 0:
-            if kernels.shape != (0,):
-                kernels = np.squeeze(kernels, axis=-1)
-        return kernels
+        if chi is None:
+            return kernels, chis
+        else:
+            kernels = np.array(kernels)
+            if np.ndim(chi) == 0:
+                if kernels.shape != (0,):
+                    kernels = np.squeeze(kernels, axis=-1)
+            return kernels
 
     def get_f_ell(self, ell):
         """Get the ell-dependent prefactors for all tracers
@@ -548,8 +577,13 @@ class NumberCountsTracer(Tracer):
             giving the magnification bias as a function of redshift. If
             `None`, the tracer is assumed to not have magnification bias
             terms. Defaults to None.
+        n_samples (int, optional): number of samples over which the
+            magnification lensing kernel is desired. These will be equi-spaced
+            in radial distance. The kernel is quite smooth, so usually O(100)
+            samples is enough.
     """
-    def __init__(self, cosmo, has_rsd, dndz, bias, mag_bias=None):
+    def __init__(self, cosmo, has_rsd, dndz, bias, mag_bias=None,
+                 n_samples=256):
         self._trc = []
 
         # we need the distance functions at the C layer
@@ -583,7 +617,8 @@ def __init__(self, cosmo, has_rsd, dndz, bias, mag_bias=None):
                             transfer_a=t_a, der_bessel=2)
         if mag_bias is not None:  # Has magnification bias
             # Kernel
-            chi, w = get_lensing_kernel(cosmo, dndz, mag_bias=mag_bias)
+            chi, w = get_lensing_kernel(cosmo, dndz, mag_bias=mag_bias,
+                                        n_chi=n_samples)
             # Multiply by -2 for magnification
             kernel_m = (chi, -2 * w)
             if (cosmo['sigma_0'] == 0):
@@ -616,9 +651,13 @@ class WeakLensingTracer(Tracer):
             normalization. Set to False to use the raw input amplitude,
             which will usually be 1 for use with PT IA modeling.
             Defaults to True.
+        n_samples (int, optional): number of samples over which the lensing
+            kernel is desired. These will be equi-spaced in radial distance.
+            The kernel is quite smooth, so usually O(100) samples
+            is enough.
     """
     def __init__(self, cosmo, dndz, has_shear=True, ia_bias=None,
-                 use_A_ia=True):
+                 use_A_ia=True, n_samples=256):
         self._trc = []
 
         # we need the distance functions at the C layer
@@ -630,7 +669,7 @@ def __init__(self, cosmo, dndz, has_shear=True, ia_bias=None,
                               fill_value=0)
 
         if has_shear:
-            kernel_l = get_lensing_kernel(cosmo, dndz)
+            kernel_l = get_lensing_kernel(cosmo, dndz, n_chi=n_samples)
             if (cosmo['sigma_0'] == 0):
                 # GR case
                 self.add_tracer(cosmo, kernel=kernel_l,
@@ -668,7 +707,7 @@ class CMBLensingTracer(Tracer):
     Args:
         cosmo (:class:`~pyccl.core.Cosmology`): Cosmology object.
         z_source (float): Redshift of source plane for CMB lensing.
-        nsamples (int, optional): number of samples over which the kernel
+        n_samples (int, optional): number of samples over which the kernel
             is desired. These will be equi-spaced in radial distance.
             The kernel is quite smooth, so usually O(100) samples
             is enough.