classy-szfast 0.0.25.post3__tar.gz → 0.0.25.post5__tar.gz

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post3
+Version: 0.0.25.post5
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
 License: MIT

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/classy_szfast/classy_szfast.py

@@ -1,7 +1,8 @@
 from .utils import *
+from .utils import Const
 from .config import *
 import numpy as np
-from .emulators_meta_data import emulator_dict, dofftlog_alphas, cp_l_max_scalars
+from .emulators_meta_data import emulator_dict, dofftlog_alphas, cp_l_max_scalars, cosmo_model_list
 from .cosmopower import cp_tt_nn, cp_te_nn, cp_ee_nn, cp_ee_nn, cp_pp_nn, cp_pknl_nn, cp_pkl_nn, cp_der_nn, cp_da_nn, cp_h_nn, cp_s8_nn, cp_pkl_fftlog_alphas_real_nn, cp_pkl_fftlog_alphas_imag_nn, cp_pkl_fftlog_alphas_nus
 from .cosmopower_jax import cp_tt_nn_jax, cp_te_nn_jax, cp_ee_nn_jax, cp_ee_nn_jax, cp_pp_nn_jax, cp_pknl_nn_jax, cp_pkl_nn_jax, cp_der_nn_jax, cp_da_nn_jax, cp_h_nn_jax, cp_s8_nn_jax
 from .pks_and_sigmas import *
@@ -76,13 +77,10 @@ class Class_szfast(object):
 
         self.jax_mode = params_settings["jax"]
 
+ 
         # print(f"JAX mode: {self.jax_mode}")
         
 
-        
-
-
-
         # cosmopower emulators
         # self.cp_path_to_cosmopower_organization = path_to_cosmopower_organization + '/'
         self.cp_tt_nn = cp_tt_nn
@@ -90,16 +88,80 @@ class Class_szfast(object):
         self.cp_ee_nn = cp_ee_nn
         self.cp_pp_nn = cp_pp_nn
         self.cp_pknl_nn  = cp_pknl_nn
-        self.cp_pkl_nn = cp_pkl_nn
         self.cp_der_nn = cp_der_nn
        
+        self.cp_lmax = cp_l_max_scalars
+
+        cosmo_model_dict = {i: model for i, model in enumerate(cosmo_model_list)}
+
+        if (cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2'):
+        
+            self.cszfast_pk_grid_zmax = 20.
+            self.cszfast_pk_grid_kmin = 5e-4
+            self.cszfast_pk_grid_kmax = 10. 
+            self.cp_kmax = self.cszfast_pk_grid_kmax
+            self.cp_kmin = self.cszfast_pk_grid_kmin
+            # self.logger.info(f">>> using kmin = {self.cp_kmin}")
+            # self.logger.info(f">>> using kmax = {self.cp_kmax}")
+            # self.logger.info(f">>> using zmax = {self.cszfast_pk_grid_zmax}")
+        
+        else:
+        
+            self.cszfast_pk_grid_zmax = 5. #  max z of our pk emulators (sept 23)
+            self.cszfast_pk_grid_kmin = 1e-4
+            self.cszfast_pk_grid_kmax = 50.
+            self.cp_kmax = self.cszfast_pk_grid_kmax
+            self.cp_kmin = self.cszfast_pk_grid_kmin
+            # self.logger.info(f">>> using kmin = {self.cp_kmin}")
+            # self.logger.info(f">>> using kmax = {self.cp_kmax}")
+            # self.logger.info(f">>> using zmax = {self.cszfast_pk_grid_zmax}")
         
         if self.jax_mode:
             self.cp_h_nn = cp_h_nn_jax
             self.cp_da_nn = cp_da_nn_jax
+            self.cp_pkl_nn = cp_pkl_nn_jax
+
+            self.pi = jnp.pi
+            self.transpose = jnp.transpose
+            self.asarray = jnp.asarray
+            self.log = jnp.log
+            self.pow = jnp.power
+
+            self.sigma_B = 2. * self.pow(self.pi,5) * self.pow(Const._k_B_,4) / 15. / self.pow(Const._h_P_,3) / self.pow(Const._c_,2)
+
+            self.linspace = jnp.linspace
+            self.geomspace = jnp.geomspace
+            self.arange = jnp.arange
+            self.zeros = jnp.zeros
+            self.gradient = jnp.gradient
+
+
         else:
             self.cp_h_nn = cp_h_nn
             self.cp_da_nn = cp_da_nn
+            self.cp_pkl_nn = cp_pkl_nn
+
+            self.pi = np.pi
+            self.transpose = np.transpose
+            self.pow = np.power
+
+            self.sigma_B = 2. * self.pow(self.pi,5) * self.pow(Const._k_B_,4) / 15. / self.pow(Const._h_P_,3) / self.pow(Const._c_,2)
+
+
+            self.linspace = np.linspace
+            self.geomspace = np.geomspace
+            self.arange = np.arange
+            self.zeros = np.zeros
+            self.asarray = np.asarray
+            self.log = np.log
+            self.gradient = np.gradient
+
+
+        self.cp_ls = self.arange(2,self.cp_lmax+1)
+        self.cp_predicted_tt_spectrum =self.zeros(self.cp_lmax)
+        self.cp_predicted_te_spectrum =self.zeros(self.cp_lmax)
+        self.cp_predicted_ee_spectrum =self.zeros(self.cp_lmax)
+        self.cp_predicted_pp_spectrum =self.zeros(self.cp_lmax)
 
         self.cp_s8_nn = cp_s8_nn
 
@@ -114,20 +176,6 @@ class Class_szfast(object):
         self.use_Amod = 0
         self.Amod = 0 
 
-        self.cp_lmax = cp_l_max_scalars
-        self.cp_ls = np.arange(2,self.cp_lmax+1)
-
-
-
-        
-        cosmo_model_dict = {0: 'lcdm',
-                            1: 'mnu',
-                            2: 'neff',
-                            3: 'wcdm',
-                            4: 'ede',
-                            5: 'mnu-3states',
-                            6: 'ede-v2'
-                            }
         
 
         if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
@@ -140,60 +188,31 @@ class Class_szfast(object):
             self.cp_ndspl_k = 10
             self.cp_nk = 5000
 
-        self.cp_predicted_tt_spectrum =np.zeros(self.cp_lmax)
-        self.cp_predicted_te_spectrum =np.zeros(self.cp_lmax)
-        self.cp_predicted_ee_spectrum =np.zeros(self.cp_lmax)
-        self.cp_predicted_pp_spectrum =np.zeros(self.cp_lmax)
-
 
         self.cszfast_ldim = 20000 # used for the cls arrays
-
         self.cszfast_pk_grid_nz = 100 # has to be same as narraySZ, i.e., ndim_redshifts; it is setup hereafter if ndim_redshifts is passed
         
 
+        self.cszfast_pk_grid_z = self.linspace(0.,self.cszfast_pk_grid_zmax,self.cszfast_pk_grid_nz)
+        self.cszfast_pk_grid_ln1pz = self.log(1.+self.cszfast_pk_grid_z)
 
-        if (cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2'):
-        
-            self.cszfast_pk_grid_zmax = 20.
-            self.cszfast_pk_grid_kmin = 5e-4
-            self.cszfast_pk_grid_kmax = 10. 
-            self.cp_kmax = self.cszfast_pk_grid_kmax
-            self.cp_kmin = self.cszfast_pk_grid_kmin
-            # self.logger.info(f">>> using kmin = {self.cp_kmin}")
-            # self.logger.info(f">>> using kmax = {self.cp_kmax}")
-            # self.logger.info(f">>> using zmax = {self.cszfast_pk_grid_zmax}")
-        
-        else:
-        
-            self.cszfast_pk_grid_zmax = 5. #  max z of our pk emulators (sept 23)
-            self.cszfast_pk_grid_kmin = 1e-4
-            self.cszfast_pk_grid_kmax = 50.
-            self.cp_kmax = self.cszfast_pk_grid_kmax
-            self.cp_kmin = self.cszfast_pk_grid_kmin
-            # self.logger.info(f">>> using kmin = {self.cp_kmin}")
-            # self.logger.info(f">>> using kmax = {self.cp_kmax}")
-            # self.logger.info(f">>> using zmax = {self.cszfast_pk_grid_zmax}")
 
-        self.cszfast_pk_grid_z = np.linspace(0.,self.cszfast_pk_grid_zmax,self.cszfast_pk_grid_nz)
-        self.cszfast_pk_grid_ln1pz = np.log(1.+self.cszfast_pk_grid_z)
-
-
-        self.cszfast_pk_grid_k = np.geomspace(self.cp_kmin,self.cp_kmax,self.cp_nk)[::self.cp_ndspl_k]
+        self.cszfast_pk_grid_k = self.geomspace(self.cp_kmin,self.cp_kmax,self.cp_nk)[::self.cp_ndspl_k]
         
-        self.cszfast_pk_grid_lnk = np.log(self.cszfast_pk_grid_k)
+        self.cszfast_pk_grid_lnk = self.log(self.cszfast_pk_grid_k)
         
-        self.cszfast_pk_grid_nk = len(np.geomspace(self.cp_kmin,self.cp_kmax,self.cp_nk)[::self.cp_ndspl_k]) # has to be same as ndimSZ, and the same as dimension of cosmopower pk emulators
+        self.cszfast_pk_grid_nk = len(self.geomspace(self.cp_kmin,self.cp_kmax,self.cp_nk)[::self.cp_ndspl_k]) # has to be same as ndimSZ, and the same as dimension of cosmopower pk emulators
         
         for k,v in params_settings.items():
 
             if k == 'ndim_redshifts':
                 
                 self.cszfast_pk_grid_nz = v
-                self.cszfast_pk_grid_z = np.linspace(0.,self.cszfast_pk_grid_zmax,self.cszfast_pk_grid_nz)
-                self.cszfast_pk_grid_ln1pz = np.log(1.+self.cszfast_pk_grid_z)
+                self.cszfast_pk_grid_z = self.linspace(0.,self.cszfast_pk_grid_zmax,self.cszfast_pk_grid_nz)
+                self.cszfast_pk_grid_ln1pz = self.log(1.+self.cszfast_pk_grid_z)
 
-                self.cszfast_pk_grid_pknl_flat = np.zeros(self.cszfast_pk_grid_nz*self.cszfast_pk_grid_nk)
-                self.cszfast_pk_grid_pkl_flat = np.zeros(self.cszfast_pk_grid_nz*self.cszfast_pk_grid_nk)
+                self.cszfast_pk_grid_pknl_flat = self.zeros(self.cszfast_pk_grid_nz*self.cszfast_pk_grid_nk)
+                self.cszfast_pk_grid_pkl_flat = self.zeros(self.cszfast_pk_grid_nz*self.cszfast_pk_grid_nk)
 
             if k == 'cosmo_model':
 
@@ -212,13 +231,12 @@ class Class_szfast(object):
         
         else:
 
-            ls = np.arange(2,self.cp_nk+2)[::self.cp_ndspl_k] # jan 10 ndspl
-            dls = ls*(ls+1.)/2./np.pi
+            ls = self.arange(2,self.cp_nk+2)[::self.cp_ndspl_k] # jan 10 ndspl
+            dls = ls*(ls+1.)/2./self.pi
             self.pk_power_fac= (dls)**-1
 
 
-        self.cp_z_interp = np.linspace(0.,20.,5000)
-        self.cp_z_interp_jax = jnp.linspace(0.,20.,5000)
+        self.cp_z_interp = self.linspace(0.,20.,5000)
 
         self.csz_base = None
 
@@ -226,7 +244,7 @@ class Class_szfast(object):
         self.cszfast_zgrid_zmin = 0.
         self.cszfast_zgrid_zmax = 4.
         self.cszfast_zgrid_nz = 250
-        self.cszfast_zgrid = np.linspace(self.cszfast_zgrid_zmin,
+        self.cszfast_zgrid = self.linspace(self.cszfast_zgrid_zmin,
                                          self.cszfast_zgrid_zmax,
                                          self.cszfast_zgrid_nz)
 
@@ -234,19 +252,41 @@ class Class_szfast(object):
         self.cszfast_mgrid_mmin = 1e10
         self.cszfast_mgrid_mmax = 1e15
         self.cszfast_mgrid_nm = 50
-        self.cszfast_mgrid = np.geomspace(self.cszfast_mgrid_mmin,
+        self.cszfast_mgrid = self.geomspace(self.cszfast_mgrid_mmin,
                                           self.cszfast_mgrid_mmax,
                                           self.cszfast_mgrid_nm)
 
         self.cszfast_gas_pressure_xgrid_xmin = 1e-2
         self.cszfast_gas_pressure_xgrid_xmax = 1e2
         self.cszfast_gas_pressure_xgrid_nx = 100
-        self.cszfast_gas_pressure_xgrid = np.geomspace(self.cszfast_gas_pressure_xgrid_xmin,
+        self.cszfast_gas_pressure_xgrid = self.geomspace(self.cszfast_gas_pressure_xgrid_xmin,
                                                        self.cszfast_gas_pressure_xgrid_xmax,
                                                        self.cszfast_gas_pressure_xgrid_nx)
         
         self.params_for_emulators = {}
 
+
+    def get_all_relevant_params(self,params_values_dict=None):
+        if params_values_dict:
+            params_values = params_values_dict.copy()
+        else:
+            params_values = self.params_for_emulators
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+        params_values['h'] = params_values['H0']/100.
+        params_values['Omega_b'] = params_values['omega_b']/params_values['h']**2.
+        params_values['Omega_cdm'] = params_values['omega_cdm']/params_values['h']**2.
+        params_values['Omega0_g'] = (4.*self.sigma_B/Const._c_*pow(params_values['T_cmb'],4.)) / (3.*Const._c_*Const._c_*1.e10*params_values['h']*params_values['h']/Const._Mpc_over_m_/Const._Mpc_over_m_/8./self.pi/Const._G_)
+        params_values['Omega0_ur'] = params_values['N_ur']*7./8.*self.pow(4./11.,4./3.)*params_values['Omega0_g']
+        params_values['Omega0_ncdm'] = params_values['deg_ncdm']*params_values['m_ncdm']/(93.14*params_values['h']*params_values['h']) ## valid only in standard cases, default T_ncdm etc
+        params_values['Omega_Lambda'] = 1. - params_values['Omega0_g'] - params_values['Omega_b'] - params_values['Omega_cdm'] - params_values['Omega0_ncdm'] - params_values['Omega0_ur']
+        params_values['Omega0_m'] = params_values['Omega_cdm'] + params_values['Omega_b'] + params_values['Omega0_ncdm']
+        params_values['Omega0_r'] = params_values['Omega0_ur']+params_values['Omega0_g']
+        params_values['Omega0_m_nonu'] = params_values['Omega0_m'] - params_values['Omega0_ncdm']
+        params_values['Omega0_cb'] = params_values['Omega0_m_nonu'] 
+        return params_values
+
+
+
     def find_As(self,params_cp):
 
         sigma_8_asked = params_cp["sigma8"]
@@ -312,7 +352,7 @@ class Class_szfast(object):
         creal = predicted_testing_alphas_creal
         cimag = predicted_testing_alphas_cimag
         Nmax = len(self.cszfast_pk_grid_k)
-        cnew = np.zeros(Nmax+1,dtype=complex)
+        cnew = self.zeros(Nmax+1,dtype=complex)
         for i in range(Nmax+1):
             if i<int(Nmax/2):
                 cnew[i] = complex(creal[i],cimag[i])
@@ -362,20 +402,20 @@ class Class_szfast(object):
 
         nl = len(self.cp_predicted_tt_spectrum)
         cls = {}
-        cls['ell'] = np.arange(20000)
-        cls['tt'] = np.zeros(20000)
-        cls['te'] = np.zeros(20000)
-        cls['ee'] = np.zeros(20000)
-        cls['pp'] = np.zeros(20000)
-        cls['bb'] = np.zeros(20000)
-        lcp = np.asarray(cls['ell'][2:nl+2])
+        cls['ell'] = self.arange(20000)
+        cls['tt'] = self.zeros(20000)
+        cls['te'] = self.zeros(20000)
+        cls['ee'] = self.zeros(20000)
+        cls['pp'] = self.zeros(20000)
+        cls['bb'] = self.zeros(20000)
+        lcp = self.asarray(cls['ell'][2:nl+2])
 
         # print('cosmo_model:',self.cosmo_model,nl)
         if self.cosmo_model == 'ede-v2':
             factor_ttteee = 1./lcp**2 
             factor_pp = 1./lcp**3
         else:
-            factor_ttteee = 1./(lcp*(lcp+1.)/2./np.pi)
+            factor_ttteee = 1./(lcp*(lcp+1.)/2./self.pi)
             factor_pp = 1./(lcp*(lcp+1.))**2.        
 
         self.cp_predicted_tt_spectrum *= factor_ttteee
@@ -402,7 +442,7 @@ class Class_szfast(object):
             cmb_cls_loaded = pickle.load(handle)
         nl_cls_file = len(cmb_cls_loaded['ell'])
         cls_ls = cmb_cls_loaded['ell']
-        dlfac = cls_ls*(cls_ls+1.)/2./np.pi
+        dlfac = cls_ls*(cls_ls+1.)/2./self.pi
         cls_tt = cmb_cls_loaded['tt']*dlfac
         cls_te = cmb_cls_loaded['te']*dlfac
         cls_ee = cmb_cls_loaded['ee']*dlfac
@@ -421,8 +461,6 @@ class Class_szfast(object):
                       **params_values_dict):
  
         z_arr = self.cszfast_pk_grid_z
-
-
         k_arr = self.cszfast_pk_grid_k 
 
         # print(">>> z_arr:",z_arr)
@@ -464,26 +502,32 @@ class Class_szfast(object):
             
                 params_dict_pp = params_dict.copy()
                 params_dict_pp['z_pk_save_nonclass'] = [zp]
-                predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
+                if self.jax_mode:
+                    predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predict(params_dict_pp))
+                else:
+                    predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
                 # if abs(zp-0.5) < 0.01:
                 #   print(">>> predicted_pk_spectrum_z:",predicted_pk_spectrum_z[-1])
                 #   import pprint
                 #   pprint.pprint(params_dict_pp)
        
-        predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
+        predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 
 
         pk = 10.**predicted_pk_spectrum
 
         pk_re = pk*self.pk_power_fac
-        pk_re = np.transpose(pk_re)
+        pk_re = self.transpose(pk_re)
 
         # print(">>> pk_re:",pk_re)
         # import sys
         # sys.exit(0)
 
-        self.pkl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
+        if self.jax_mode:
+            self.pkl_interp = None
+        else:
+            self.pkl_interp = PowerSpectrumInterpolator(z_arr,k_arr,self.log(pk_re).T,logP=True)
 
         self.cszfast_pk_grid_pk = pk_re
         self.cszfast_pk_grid_pkl_flat = pk_re.flatten()
@@ -491,9 +535,7 @@ class Class_szfast(object):
         return pk_re, k_arr, z_arr
 
 
-    def calculate_sigma(self,
-                        
-                        **params_values_dict):
+    def calculate_sigma(self,**params_values_dict):
 
         params_values = params_values_dict.copy()
 
@@ -509,7 +551,7 @@ class Class_szfast(object):
 
             R, var[:,iz] = TophatVar(k, lowring=True)(P[:,iz], extrap=True)
 
-            dvar[:,iz] = np.gradient(var[:,iz], R)
+            dvar[:,iz] = self.gradient(var[:,iz], R)
 
         # print(k)
         # print(R)
@@ -517,11 +559,11 @@ class Class_szfast(object):
         # exit(0)
  
 
-        self.cszfast_pk_grid_lnr = np.log(R)
+        self.cszfast_pk_grid_lnr = self.log(R)
         self.cszfast_pk_grid_sigma2 = var
  
         self.cszfast_pk_grid_sigma2_flat = var.flatten()
-        self.cszfast_pk_grid_lnsigma2_flat = 0.5*np.log(var.flatten())
+        self.cszfast_pk_grid_lnsigma2_flat = 0.5*self.log(var.flatten())
  
         self.cszfast_pk_grid_dsigma2 = dvar
         self.cszfast_pk_grid_dsigma2_flat = dvar.flatten()
@@ -570,7 +612,7 @@ class Class_szfast(object):
         s8z  = self.cp_s8_nn[self.cosmo_model].predictions_np(params_dict)
         # print(self.s8z)
         self.s8z_interp = scipy.interpolate.interp1d(
-                                                    np.linspace(0.,20.,5000),
+                                                    self.linspace(0.,20.,5000),
                                                     s8z[0],
                                                     kind='linear',
                                                     axis=-1,
@@ -609,16 +651,16 @@ class Class_szfast(object):
             params_dict_pp['z_pk_save_nonclass'] = [zp]
             predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
-        predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
+        predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 
 
         pk = 10.**predicted_pk_spectrum
 
         pk_re = pk*self.pk_power_fac
-        pk_re = np.transpose(pk_re)
+        pk_re = self.transpose(pk_re)
 
 
-        self.pknl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
+        self.pknl_interp = PowerSpectrumInterpolator(z_arr,k_arr,self.log(pk_re).T,logP=True)
 
 
         self.cszfast_pk_grid_pknl = pk_re
@@ -627,20 +669,18 @@ class Class_szfast(object):
         return pk_re, k_arr, z_arr
     
 
+
+
     def calculate_pkl_at_z(self,
                            z_asked,
                           params_values_dict=None):
 
-        z_arr = self.cszfast_pk_grid_z
 
         k_arr = self.cszfast_pk_grid_k 
 
         if params_values_dict:
-
             params_values = params_values_dict.copy()
-
         else:
-
             params_values = self.params_for_emulators
 
         update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
@@ -657,19 +697,21 @@ class Class_szfast(object):
 
         predicted_pk_spectrum_z = []
 
-        z_asked = z_asked
         params_dict_pp = params_dict.copy()
         update_params_with_defaults(params_dict_pp, self.emulator_dict[self.cosmo_model]['default'])
 
         params_dict_pp['z_pk_save_nonclass'] = [z_asked]
-        predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
+        if self.jax_mode:
+            predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predict(params_dict_pp))
+        else:
+            predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
-        predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
+        predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 
 
         pk = 10.**predicted_pk_spectrum
         pk_re = pk*self.pk_power_fac
-        pk_re = np.transpose(pk_re)
+        pk_re = self.transpose(pk_re)
 
         
         return pk_re, k_arr
@@ -712,12 +754,12 @@ class Class_szfast(object):
         params_dict_pp['z_pk_save_nonclass'] = [z_asked]
         predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
-        predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
+        predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 
 
         pk = 10.**predicted_pk_spectrum
         pk_re = pk*self.pk_power_fac
-        pk_re = np.transpose(pk_re)
+        pk_re = self.transpose(pk_re)
 
         
         return pk_re, k_arr
@@ -748,19 +790,10 @@ class Class_szfast(object):
             # print("self.cp_predicted_hubble type:", type(self.cp_predicted_hubble))
             # print("self.cp_predicted_hubble",self.cp_predicted_hubble)
 
-            # self.hz_interp = jscipy.interpolate.interp1d(
-            #                                     self.cp_z_interp_jax,
-            #                                     self.cp_predicted_hubble,
-            #                                     kind='linear',
-            #                                     axis=-1,
-            #                                     copy=True,
-            #                                     bounds_error=None,
-            #                                     fill_value=np.nan,
-            #                                     assume_sorted=False)
             
             # Assuming `cp_z_interp` and `cp_predicted_hubble` are JAX arrays
             def hz_interp(x):
-                return jnp.interp(x, self.cp_z_interp_jax, self.cp_predicted_hubble, left=jnp.nan, right=jnp.nan)
+                return jnp.interp(x, self.cp_z_interp, self.cp_predicted_hubble, left=jnp.nan, right=jnp.nan)
 
             self.hz_interp = hz_interp
             # exit()
@@ -803,10 +836,10 @@ class Class_szfast(object):
             if self.cosmo_model == 'ede-v2':
                 # print('ede-v2 case')
                 self.cp_predicted_da = jnp.insert(self.cp_predicted_da, 0, 0)
-            self.cp_predicted_da *= (1.+self.cp_z_interp_jax)
+            self.cp_predicted_da *= (1.+self.cp_z_interp)
     
             def chi_interp(x):
-                return jnp.interp(x, self.cp_z_interp_jax, self.cp_predicted_da, left=jnp.nan, right=jnp.nan)
+                return jnp.interp(x, self.cp_z_interp, self.cp_predicted_da, left=jnp.nan, right=jnp.nan)
 
             self.chi_interp = chi_interp
 
@@ -835,12 +868,12 @@ class Class_szfast(object):
     def get_cmb_cls(self,ell_factor=True,Tcmb_uk = Tcmb_uk):
 
         cls = {}
-        cls['ell'] = np.arange(self.cszfast_ldim)
-        cls['tt'] = np.zeros(self.cszfast_ldim)
-        cls['te'] = np.zeros(self.cszfast_ldim)
-        cls['ee'] = np.zeros(self.cszfast_ldim)
-        cls['pp'] = np.zeros(self.cszfast_ldim)
-        cls['bb'] = np.zeros(self.cszfast_ldim)
+        cls['ell'] = self.arange(self.cszfast_ldim)
+        cls['tt'] = self.zeros(self.cszfast_ldim)
+        cls['te'] = self.zeros(self.cszfast_ldim)
+        cls['ee'] = self.zeros(self.cszfast_ldim)
+        cls['pp'] = self.zeros(self.cszfast_ldim)
+        cls['bb'] = self.zeros(self.cszfast_ldim)
         cls['tt'][2:self.cp_lmax+1] = (Tcmb_uk)**2.*self.cp_predicted_tt_spectrum.copy()
         cls['te'][2:self.cp_lmax+1] = (Tcmb_uk)**2.*self.cp_predicted_te_spectrum.copy()
         cls['ee'][2:self.cp_lmax+1] = (Tcmb_uk)**2.*self.cp_predicted_ee_spectrum.copy()
@@ -848,21 +881,21 @@ class Class_szfast(object):
 
 
         if ell_factor==False:
-            fac_l = np.zeros(self.cszfast_ldim)
-            fac_l[2:self.cp_lmax+1] = 1./(cls['ell'][2:self.cp_lmax+1]*(cls['ell'][2:self.cp_lmax+1]+1.)/2./np.pi)
+            fac_l = self.zeros(self.cszfast_ldim)
+            fac_l[2:self.cp_lmax+1] = 1./(cls['ell'][2:self.cp_lmax+1]*(cls['ell'][2:self.cp_lmax+1]+1.)/2./self.pi)
             cls['tt'][2:self.cp_lmax+1] *= fac_l[2:self.cp_lmax+1]
             cls['te'][2:self.cp_lmax+1] *= fac_l[2:self.cp_lmax+1]
             cls['ee'][2:self.cp_lmax+1] *= fac_l[2:self.cp_lmax+1]
-            # cls['bb'] = np.zeros(self.cszfast_ldim)
+            # cls['bb'] = self.zeros(self.cszfast_ldim)
         return cls
 
 
     def get_pknl_at_k_and_z(self,k_asked,z_asked):
     # def get_pkl_at_k_and_z(self,k_asked,z_asked,method = 'cloughtocher'):
         # if method == 'linear':
-        #     pk = self.pkl_linearnd_interp(z_asked,np.log(k_asked))
+        #     pk = self.pkl_linearnd_interp(z_asked,self.log(k_asked))
         # elif method == 'cloughtocher':
-        #     pk = self.pkl_cloughtocher_interp(z_asked,np.log(k_asked))
+        #     pk = self.pkl_cloughtocher_interp(z_asked,self.log(k_asked))
         # return np.exp(pk)
         return self.pknl_interp.P(z_asked,k_asked)
 
@@ -870,18 +903,18 @@ class Class_szfast(object):
     # def get_pkl_at_k_and_z(self,k_asked,z_asked,method = 'cloughtocher'):
     def get_pkl_at_k_and_z(self,k_asked,z_asked):
         # if method == 'linear':
-        #     pk = self.pknl_linearnd_interp(z_asked,np.log(k_asked))
+        #     pk = self.pknl_linearnd_interp(z_asked,self.log(k_asked))
         # elif method == 'cloughtocher':
-        #     pk = self.pknl_cloughtocher_interp(z_asked,np.log(k_asked))
+        #     pk = self.pknl_cloughtocher_interp(z_asked,self.log(k_asked))
         # return np.exp(pk)
         return self.pkl_interp.P(z_asked,k_asked)
 
     # function used to overwrite the classy function in fast mode.
     def get_sigma_at_r_and_z(self,r_asked,z_asked):
         # if method == 'linear':
-        #     pk = self.pknl_linearnd_interp(z_asked,np.log(k_asked))
+        #     pk = self.pknl_linearnd_interp(z_asked,self.log(k_asked))
         # elif method == 'cloughtocher':
-        #     pk = self.pknl_cloughtocher_interp(z_asked,np.log(k_asked))
+        #     pk = self.pknl_cloughtocher_interp(z_asked,self.log(k_asked))
         # return np.exp(pk)
         k = self.cszfast_pk_grid_k
         P_at_z = self.get_pkl_at_k_and_z(k,z_asked)
@@ -899,7 +932,10 @@ class Class_szfast(object):
             return np.array(self.hz_interp(z)*H_units_conv_factor[units])
 
     def get_chi(self, z):
-        return np.array(self.chi_interp(z))
+        if self.jax_mode:
+            return jnp.array(self.chi_interp(z))
+        else:
+            return np.array(self.chi_interp(z))
 
     def get_gas_pressure_profile_x(self,z,m,x):
         return 0#np.vectorize(self.csz_base.get_pressure_P_over_P_delta_at_x_M_z_b12_200c)(x,m,z)
@@ -913,7 +949,7 @@ class Class_szfast(object):
 
 
     def tabulate_gas_pressure_profile_k(self):
-        z_asked,m_asked,x_asked = 0.2,3e14,np.geomspace(1e-3,1e2,500)
+        z_asked,m_asked,x_asked = 0.2,3e14,self.geomspace(1e-3,1e2,500)
         start = time.time()
         px = self.get_gas_pressure_profile_x(z_asked,m_asked,x_asked)
         end = time.time()

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/classy_szfast/cosmopower_jax.py

@@ -105,8 +105,9 @@ for mp in cosmo_model_list:
     
     cp_pknl_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKNL'])
     
-    cp_pkl_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKL'])
-    
+    cp_pkl_nn_jax[mp] = CosmoPowerJAX_custom(probe='custom_log',filepath=path_to_emulators +'PK/' + emulator_dict[mp]['PKL'] + '.npz')
+    cp_pkl_nn_jax[mp].ten_to_predictions = False
+
     cp_der_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'derived-parameters/' + emulator_dict[mp]['DER'])
     
     # cp_da_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['DAZ'])

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/classy_szfast/emulators_meta_data.py

@@ -33,7 +33,7 @@ cosmopower_derived_params_names = ['100*theta_s',
                                    'ra_star',
                                    'rs_drag']
 
-cp_l_max_scalars = 11000 # max multipole of train ing data
+cp_l_max_scalars = 11000 # max multipole of training data
 
 cosmo_model_list = [
     'lcdm',
@@ -87,6 +87,8 @@ emulator_dict['lcdm']['default']['n_s'] = 0.9665
 emulator_dict['lcdm']['default']['N_ur'] = 2.0328
 emulator_dict['lcdm']['default']['N_ncdm'] = 1
 emulator_dict['lcdm']['default']['m_ncdm'] = 0.06
+emulator_dict['lcdm']['default']['deg_ncdm'] = 1
+emulator_dict['lcdm']['default']['T_cmb'] = 2.7255
 
 emulator_dict['mnu']['TT'] = 'TT_mnu_v1'
 emulator_dict['mnu']['TE'] = 'TE_mnu_v1'
@@ -107,7 +109,9 @@ emulator_dict['mnu']['default']['omega_cdm'] = 0.11933
 emulator_dict['mnu']['default']['n_s'] = 0.9665
 emulator_dict['mnu']['default']['N_ur'] = 2.0328
 emulator_dict['mnu']['default']['N_ncdm'] = 1
+emulator_dict['mnu']['default']['deg_ncdm'] = 1
 emulator_dict['mnu']['default']['m_ncdm'] = 0.06
+emulator_dict['mnu']['default']['T_cmb'] = 2.7255
 
 emulator_dict['neff']['TT'] = 'TT_neff_v1'
 emulator_dict['neff']['TE'] = 'TE_neff_v1'
@@ -128,8 +132,9 @@ emulator_dict['neff']['default']['omega_cdm'] = 0.11933
 emulator_dict['neff']['default']['n_s'] = 0.9665
 emulator_dict['neff']['default']['N_ur'] = 2.0328 # this is the default value in class v2 to get Neff = 3.046
 emulator_dict['neff']['default']['N_ncdm'] = 1
+emulator_dict['neff']['default']['deg_ncdm'] = 1
 emulator_dict['neff']['default']['m_ncdm'] = 0.06
-
+emulator_dict['neff']['default']['T_cmb'] = 2.7255
 
 emulator_dict['wcdm']['TT'] = 'TT_w_v1'
 emulator_dict['wcdm']['TE'] = 'TE_w_v1'
@@ -151,6 +156,9 @@ emulator_dict['wcdm']['default']['n_s'] = 0.9665
 emulator_dict['wcdm']['default']['N_ur'] = 2.0328 # this is the default value in class v2 to get Neff = 3.046
 emulator_dict['wcdm']['default']['N_ncdm'] = 1
 emulator_dict['wcdm']['default']['m_ncdm'] = 0.06
+emulator_dict['wcdm']['default']['deg_ncdm'] = 1
+emulator_dict['wcdm']['default']['T_cmb'] = 2.7255
+
 
 emulator_dict['ede']['TT'] = 'TT_v1'
 emulator_dict['ede']['TE'] = 'TE_v1'
@@ -174,9 +182,10 @@ emulator_dict['ede']['default']['log10z_c'] = 3.562 # e.g. from https://github.c
 emulator_dict['ede']['default']['thetai_scf'] = 2.83 # e.g. from https://github.com/mwt5345/class_ede/blob/master/class/notebooks-ede/2-CMB-Comparison.ipynb
 emulator_dict['ede']['default']['r'] = 0.
 emulator_dict['ede']['default']['N_ur'] = 0.00641 # this is the default value in class v2 to get Neff = 3.046
-emulator_dict['ede']['default']['N_ncdm'] = 3
+emulator_dict['ede']['default']['N_ncdm'] = 1 ### use equivalence with deg_ncdm = 3
+emulator_dict['ede']['default']['deg_ncdm'] = 3 ### use equivalence with deg_ncdm = 3 for faster computation
 emulator_dict['ede']['default']['m_ncdm'] = 0.02
-
+emulator_dict['ede']['default']['T_cmb'] = 2.7255
 
 emulator_dict['mnu-3states']['TT'] = 'TT_v1'
 emulator_dict['mnu-3states']['TE'] = 'TE_v1'
@@ -196,8 +205,10 @@ emulator_dict['mnu-3states']['default']['omega_b'] = 0.02242
 emulator_dict['mnu-3states']['default']['omega_cdm'] = 0.11933
 emulator_dict['mnu-3states']['default']['n_s'] = 0.9665
 emulator_dict['mnu-3states']['default']['N_ur'] = 0.00641 # this is the default value in class v2 to get Neff = 3.046
-emulator_dict['mnu-3states']['default']['N_ncdm'] = 3
+emulator_dict['mnu-3states']['default']['N_ncdm'] = 1 ### use equivalence with deg_ncdm = 3
+emulator_dict['mnu-3states']['default']['deg_ncdm'] = 3 ### use equivalence with deg_ncdm = 3 for faster computation
 emulator_dict['mnu-3states']['default']['m_ncdm'] = 0.02
+emulator_dict['mnu-3states']['default']['T_cmb'] = 2.7255
 
 emulator_dict['ede-v2']['TT'] = 'TT_v2'
 emulator_dict['ede-v2']['TE'] = 'TE_v2'
@@ -222,9 +233,10 @@ emulator_dict['ede-v2']['default']['log10z_c'] = 3.562 # e.g. from https://githu
 emulator_dict['ede-v2']['default']['thetai_scf'] = 2.83 # e.g. from https://github.com/mwt5345/class_ede/blob/master/class/notebooks-ede/2-CMB-Comparison.ipynb
 emulator_dict['ede-v2']['default']['r'] = 0.
 emulator_dict['ede-v2']['default']['N_ur'] = 0.00441 # this is the default value in class v3 to get Neff = 3.044
-emulator_dict['ede-v2']['default']['N_ncdm'] = 3
+emulator_dict['ede-v2']['default']['N_ncdm'] = 1 ### use equivalence with deg_ncdm = 3
+emulator_dict['ede-v2']['default']['deg_ncdm'] = 3 ### use equivalence with deg_ncdm = 3 for faster computation
 emulator_dict['ede-v2']['default']['m_ncdm'] = 0.02
-
+emulator_dict['ede-v2']['default']['T_cmb'] = 2.7255
 
 
 

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/classy_szfast/utils.py

@@ -60,3 +60,13 @@ class Const:
     c_km_s = 299792.458  # speed of light
     h_J_s = 6.626070040e-34  # Planck's constant
     kB_J_K = 1.38064852e-23  # Boltzmann constant
+
+    _c_ = 2.99792458e8      # c in m/s 
+    _Mpc_over_m_ = 3.085677581282e22  # conversion factor from meters to megaparsecs 
+    _Gyr_over_Mpc_ = 3.06601394e2 # conversion factor from megaparsecs to gigayears
+    _G_ = 6.67428e-11             # Newton constant in m^3/Kg/s^2 
+    _eV_ = 1.602176487e-19        # 1 eV expressed in J 
+
+    # parameters entering in Stefan-Boltzmann constant sigma_B 
+    _k_B_ = 1.3806504e-23
+    _h_P_ = 6.62606896e-34

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/classy_szfast.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post3
+Version: 0.0.25.post5
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
 License: MIT

{classy_szfast-0.0.25.post3 → classy_szfast-0.0.25.post5}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
-version = "0.0.25.post3"
+version = "0.0.25.post5"
 license = { text = "MIT" }
 name = "classy_szfast"
 maintainers = [{name = "Boris Bolliet",email="bb667@cam.ac.uk"}]