classy-szfast 0.0.25.post4__py3-none-any.whl → 0.0.25.post5__py3-none-any.whl

classy_szfast/classy_szfast.py

@@ -1,4 +1,5 @@
 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, cosmo_model_list
@@ -123,11 +124,16 @@ class Class_szfast(object):
             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:
@@ -137,13 +143,18 @@ class Class_szfast(object):
 
             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)
@@ -183,12 +194,12 @@ class Class_szfast(object):
         
 
         self.cszfast_pk_grid_z = self.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_ln1pz = self.log(1.+self.cszfast_pk_grid_z)
 
 
         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(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
         
@@ -198,7 +209,7 @@ class Class_szfast(object):
                 
                 self.cszfast_pk_grid_nz = v
                 self.cszfast_pk_grid_z = self.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_ln1pz = self.log(1.+self.cszfast_pk_grid_z)
 
                 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)
@@ -254,6 +265,28 @@ class Class_szfast(object):
         
         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"]
@@ -494,7 +527,7 @@ class Class_szfast(object):
         if self.jax_mode:
             self.pkl_interp = None
         else:
-            self.pkl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
+            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()
@@ -502,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()
 
@@ -520,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)
@@ -528,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()
@@ -629,7 +660,7 @@ class Class_szfast(object):
         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
@@ -638,6 +669,8 @@ class Class_szfast(object):
         return pk_re, k_arr, z_arr
     
 
+
+
     def calculate_pkl_at_z(self,
                            z_asked,
                           params_values_dict=None):
@@ -860,9 +893,9 @@ class Class_szfast(object):
     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)

classy_szfast/emulators_meta_data.py

@@ -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/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.post4.dist-info → classy_szfast-0.0.25.post5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post4
+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.post4.dist-info → classy_szfast-0.0.25.post5.dist-info}/RECORD

@@ -1,20 +1,20 @@
 classy_szfast/__init__.py,sha256=E2thrL0Z9oXFfdzwcsu-xbOytudLFTlRlPqVFGlPPPg,279
 classy_szfast/classy_sz.py,sha256=QmbwrSXInQLMvCDqsr7KPmtaU0KOiOt1Rb-cTKuulZw,22240
-classy_szfast/classy_szfast.py,sha256=H3DPD4_ZvVJ_FCVbEzc4rW8DYQHPZCCaGQZtQ1f-7oQ,37646
+classy_szfast/classy_szfast.py,sha256=Fajoc14gdED3LpY0dMFhh1A_NzRqFVCCIof5pElCA4Y,39791
 classy_szfast/config.py,sha256=cd7Z62-qnX_4FJWfUNqcyJVh-AdBiXrF8DcQGpyAUZM,274
 classy_szfast/cosmopower.py,sha256=ooYK2BDOZSo3XtGHfPtjXHxr5UW-yVngLPkb5gpvTx8,2351
 classy_szfast/cosmopower_jax.py,sha256=NvwY1a9YoizOV0zKp0SEtiHQkhWRO-seGFzO3FIzLl0,5436
 classy_szfast/cosmosis_classy_szfast_interface.py,sha256=zAnxvFtn73a5yS7jgs59zpWFEYKCIQyraYPs5hQ4Le8,11483
-classy_szfast/emulators_meta_data.py,sha256=faA0iQqzfC5lOj1hu7wSEZUFfhBGLJqX9-xovMfTbr0,9723
+classy_szfast/emulators_meta_data.py,sha256=mXG5LQuJw9QBNE_kxXW8Kx0AUCWpbV6uRO9BaBbIfHo,10732
 classy_szfast/pks_and_sigmas.py,sha256=drtuujE1HhlrYY1hY92DyY5lXlYS1uE15MSuVI4uo6k,6625
 classy_szfast/restore_nn.py,sha256=DqA9thhTRiGBDVb9zjhqcbF2W4V0AU0vrjJFhnLboU4,21075
 classy_szfast/suppress_warnings.py,sha256=6wIBml2Sj9DyRGZlZWhuA9hqvpxqrNyYjuz6BPK_a6E,202
-classy_szfast/utils.py,sha256=VdaRsJK2ttHI9zkyxVhergxHPC6t99usrlycblyqcP8,1464
+classy_szfast/utils.py,sha256=Et90Yo3EOjSaOTorqPmhLE3kVMF7qc2ZJHOD5NRY5vU,1916
 classy_szfast/custom_bias/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 classy_szfast/custom_bias/custom_bias.py,sha256=aR2t5RTIwv7P0m2bsEU0Eq6BTkj4pG10AebH6QpG4qM,486
 classy_szfast/custom_profiles/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 classy_szfast/custom_profiles/custom_profiles.py,sha256=4LZwb2XoqwCyWNmW2s24Z7AJdmgVdaRG7yYaBYe-d9Q,1188
-classy_szfast-0.0.25.post4.dist-info/METADATA,sha256=oNj9IJ8zELtcPCEOOi4wQzx9SfxVaMmdIl6QjyAFdkY,548
-classy_szfast-0.0.25.post4.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-classy_szfast-0.0.25.post4.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.25.post4.dist-info/RECORD,,
+classy_szfast-0.0.25.post5.dist-info/METADATA,sha256=nhGFJzuwxdDWzS3UQpZMNk736flVpLyjjRB68r-wxbY,548
+classy_szfast-0.0.25.post5.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+classy_szfast-0.0.25.post5.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.25.post5.dist-info/RECORD,,