Read non-linear power spectrum from input (#757)

* First step to input distance

* Added function to read input chi of z

* Added input E(a) and computing a(chi) from input chi(a)

* Input arrays depend on scale factor. Removed flake8 warnings. Creating new reversed arrays instead of changing old ones.

* Syntax error fixes

* Added unit test

* Allow input growth factor and growth rate. Unit test created.

* Added unit test for ValueError raises. Dynamic allocation of reversed arrays in C script.

* Check all input arrays are parsed, with unit test. Clean TODOs and finish documentation.

* Flake8 successfully run.

* Creating new function set_background_from_arrays to avoid extra keyword arguments on the Cosmology call. The background_on_input is now set to False initially and converted to True once this new function is called successfully.

* Building input for linear power spectrum

* Towards input linear power spectrum

* Changing set_background_from_arrays to _set_background_from_arrays and tests accordingly.

* Changed background_on_input to _background_on_input

* Adjust range of input a,k arrays. Write unit test.

* Remove white line

* Linear power spectrum from input stored from existing 2d spline. Added transfer function option for when power spectrum is to be input by the user.

* Fixing travis error.

* Fixing travis error.

* Fixing travis

* Parsing input non-linear power spectrum

* Merge from LSSTDESC-CCL

* Docstring and small changes

* Unit test for raises due to input linear power spectrum part of core.py.

* Small fix.

* Unit tests for nonlin power from input.

* Trimming and moving linpower raises test to test_power.

* Making pknl_from_input not part of the public API.

* Trimming and updating error message.

* rename

Co-authored-by: David Alonso <dam.phys@gmail.com>

include/ccl_config.h

@@ -36,7 +36,8 @@ typedef enum matter_power_spectrum_t
     ccl_linear           = 0,
     ccl_halofit          = 1,
     ccl_halo_model       = 3,
-    ccl_emu              = 4
+    ccl_emu              = 4,
+    ccl_pknl_from_input  = 5
 } matter_power_spectrum_t;
 
 /**

pyccl/core.py

@@ -186,9 +186,12 @@ def __init__(
         # This will change to True once the "_set_background_from_arrays"
         # is called.
         self._background_on_input = False
-        # This will change to True once the "set_linear_power_from_arrays"
+        # This will change to True once the "_set_linear_power_from_arrays"
         # is called.
         self._linear_power_on_input = False
+        # This will change to True once the "_set_nonlin_power_from_arrays"
+        # is called.
+        self._nonlinear_power_on_input = False
 
     def _build_cosmo(self):
         """Assemble all of the input data into a valid ccl_cosmology object."""
@@ -825,7 +828,7 @@ def _get_halo_model_nonlin_power(self):
         hmc = hal.HMCalculator(self, hmf, hbf, mdef)
         return hal.halomod_Pk2D(self, hmc, prf, normprof1=True)
 
-    def compute_nonlin_power(self):
+    def _compute_nonlin_power_internal(self):
         """Compute the non-linear power spectrum."""
         if self.has_nonlin_power:
             return
@@ -876,6 +879,33 @@ def compute_nonlin_power(self):
         status = lib.cosmology_compute_nonlin_power(self.cosmo, psp, status)
         check(status, self)
 
+    def compute_nonlin_power(self):
+        """Call the appropriate function to compute the linear power
+        spectrum, either read from input or calculated internally,"""
+        if self._nonlinear_power_on_input:
+            self._compute_nonlin_power_from_arrays()
+        else:
+            self._compute_nonlin_power_internal()
+
+    def _compute_nonlin_power_from_arrays(self):
+        if not self._nonlinear_power_on_input:
+            raise ValueError("Cannot compute non-linear power spectrum from"
+                             "input without input arrays initialized.")
+        pk_lin = Pk2D(pkfunc=None,
+                      a_arr=self.a_array,
+                      lk_arr=np.log(self.k_array),
+                      pk_arr=self.pk_array,
+                      is_logp=False,
+                      extrap_order_lok=1,
+                      extrap_order_hik=2,
+                      cosmo=None)
+
+        psp = pk_lin.psp
+
+        status = 0
+        status = lib.cosmology_compute_nonlin_power(self.cosmo, psp, status)
+        check(status, self)
+
     def compute_sigma(self):
         """Compute the sigma(M) and mass function splines."""
         if self.has_sigma:
@@ -1015,15 +1045,56 @@ def _set_linear_power_from_arrays(self, a_array=None, k_array=None,
         else:
             if ((a_array is None) or (k_array is None)
                     or (pk_array is None)):
-                raise ValueError("One or more input array for a, k,"
-                                 " or Pk is not parsed.")
+                raise ValueError("One or more input arrays for a, k,"
+                                 " or Pk are not parsed.")
             self.cosmo.config.transfer_function_method = lib.pklin_from_input
             self._config_init_kwargs['transfer_function'] = 'pklin_from_input'
             self._linear_power_on_input = True
             self.a_array = a_array
             self.k_array = k_array
             self.pk_array = pk_array
 
+    def _set_nonlin_power_from_arrays(self, a_array=None, k_array=None,
+                                      pk_array=None):
+        """
+        This function initializes the arrays used for parsing
+        a non-linear power spectrum from input. Call this function
+        to have the power spectrum be read from input and not
+        computed by CCL.
+
+        a_array (array): an array holding values of the scale factor
+        k_array (array): an array holding values of the wavenumber
+           in units of Mpc^-1).
+        pk_array (array): a 2D array containing the values of the power
+           spectrum at the values of the scale factor and the wavenumber
+           held by `a_array` and `k_array`. The shape of this array must be
+           `[na,nk]`, where `na` is the size of `a_array` and `nk` is the
+           size of `k_array`. This array can be provided in a flattened
+           form as long as the total size matches `nk*na`.
+           Note that, if you pass your own Pk array, you
+           are responsible of making sure that it is sufficiently well
+           sampled (i.e. the resolution of `a_array` and `k_array` is high
+           enough to sample the main features in the power spectrum).
+           For reference, CCL will use bicubic interpolation to evaluate
+           the power spectrum at any intermediate point in k and a.
+        """
+        if self.has_nonlin_power:
+            raise ValueError("Non-linear power spectrum has been initialized"
+                             "and cannot be reset.")
+        else:
+            if ((a_array is None) or (k_array is None)
+                    or (pk_array is None)):
+                raise ValueError("One or more input arrays for a, k,"
+                                 " or Pk are not parsed.")
+            self.cosmo.config.matter_power_spectrum_method \
+                = lib.pknl_from_input
+            self._config_init_kwargs['matter_power_spectrum']\
+                = 'pknl_from_input'
+            self._nonlinear_power_on_input = True
+            self.a_array = a_array
+            self.k_array = k_array
+            self.pk_array = pk_array
+
 
 class CosmologyVanillaLCDM(Cosmology):
     """A cosmology with typical flat Lambda-CDM parameters (`Omega_c=0.25`,

pyccl/tests/test_power.py

@@ -216,3 +216,47 @@ def test_input_linpower_raises():
                                 A_s=2e-9)
     with pytest.raises(ValueError):
         cosmo_input._compute_linear_power_from_arrays()
+
+
+def test_input_nonlin_power_spectrum():
+    cosmo = ccl.Cosmology(Omega_c=0.27, Omega_b=0.05, h=0.7, n_s=0.965,
+
+                          A_s=2e-9)
+    a_arr = np.linspace(0.1, 1.0, 50)
+    k_arr = np.logspace(np.log10(2e-4), np.log10(1), 1000)
+    pk_arr = np.empty(shape=(len(a_arr), len(k_arr)))
+    for i, a in enumerate(a_arr):
+        pk_arr[i] = ccl.power.nonlin_matter_power(cosmo, k_arr, a)
+
+    chi_from_ccl = ccl.background.comoving_radial_distance(cosmo, a_arr)
+    hoh0_from_ccl = ccl.background.h_over_h0(cosmo, a_arr)
+    growth_from_ccl = ccl.background.growth_factor(cosmo, a_arr)
+    fgrowth_from_ccl = ccl.background.growth_rate(cosmo, a_arr)
+
+    cosmo_input = ccl.Cosmology(Omega_c=0.27, Omega_b=0.05, h=0.7, n_s=0.965,
+                                A_s=2e-9,)
+    cosmo_input._set_background_from_arrays(a_array=a_arr,
+                                            chi_array=chi_from_ccl,
+                                            hoh0_array=hoh0_from_ccl,
+                                            growth_array=growth_from_ccl,
+                                            fgrowth_array=fgrowth_from_ccl)
+    cosmo_input._set_nonlin_power_from_arrays(a_arr, k_arr, pk_arr)
+
+    pk_CCL_input = ccl.power.nonlin_matter_power(cosmo_input, k_arr, 0.5)
+    pk_CCL = ccl.power.nonlin_matter_power(cosmo, k_arr, 0.5)
+
+    assert np.allclose(pk_CCL_input, pk_CCL, atol=0., rtol=1e-5)
+
+
+def test_input_nonlin_raises():
+    cosmo_input = ccl.Cosmology(Omega_c=0.27, Omega_b=0.05, h=0.7, n_s=0.965,
+                                A_s=2e-9)
+    with pytest.raises(ValueError):
+        cosmo_input._set_nonlin_power_from_arrays()
+    with pytest.raises(ValueError):
+        cosmo_input.compute_nonlin_power()
+        cosmo_input._set_nonlin_power_from_arrays()
+    cosmo_input = ccl.Cosmology(Omega_c=0.27, Omega_b=0.05, h=0.7, n_s=0.965,
+                                A_s=2e-9)
+    with pytest.raises(ValueError):
+        cosmo_input._compute_nonlin_power_from_arrays()

src/ccl_power.c

@@ -538,6 +538,10 @@ void ccl_cosmology_compute_nonlin_power(ccl_cosmology* cosmo, ccl_f2d_t *psp_o,
         ccl_cosmology_compute_power_emu(cosmo, status);}
         break;
 
+      case ccl_pknl_from_input: {
+        ccl_cosmology_compute_nonlin_power_from_f2d(cosmo, psp_o, status);}
+        break;
+
     default: {
       *status = CCL_ERROR_INCONSISTENT;
       ccl_cosmology_set_status_message(