classy-szfast 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

classy_szfast/__init__.py

@@ -6,5 +6,9 @@ from .pks_and_sigmas import *
 
 from pathlib import Path
 import sys
-from .custom_profiles import *
-from .custom_bias import *
+
+
+from .custom_profiles.custom_profiles import *
+
+
+from .custom_bias.custom_bias import *

classy_szfast/classy_sz.py

@@ -10,7 +10,7 @@ import time
 class classy_sz(classy):
     use_class_sz_fast_mode = 0 # this is passed in the yaml file
     use_class_sz_no_cosmo_mode = 0 # this is passed in the yaml file
-    lensing_lkl = "SOLikeT"
+    lensing_lkl = None
     # skip_background_and_thermo = True
     # ell_factor = False # True for pyactlite and bplike, False for clik
 

classy_szfast/classy_szfast.py

@@ -12,6 +12,44 @@ import pickle
 
 H_units_conv_factor = {"1/Mpc": 1, "km/s/Mpc": Const.c_km_s}
 
+
+import logging
+
+def configure_logging(level=logging.INFO):
+    logging.basicConfig(
+        format='%(levelname)s - %(message)s',
+        level=level
+    )
+
+def set_verbosity(verbosity):
+    levels = {
+        'none': logging.CRITICAL,
+        'minimal': logging.INFO,
+        'extensive': logging.DEBUG
+    }
+    # print(f'Setting verbosity to {verbosity}')
+    level = levels.get(verbosity, logging.INFO)
+    configure_logging(level)
+
+
+
+def update_params_with_defaults(params_values, self):
+    """
+    Update params_values with default values if they don't already exist.
+
+    Args:
+    params_values (dict): Dictionary containing parameter values.
+    self (object): The object containing emulator_dict and cosmo_model attributes.
+    """
+    # Retrieve default values
+    default_values = self.emulator_dict[self.cosmo_model]['default']
+    
+    # Update params_values with default values if key does not exist
+    for key, value in default_values.items():
+        if key not in params_values:
+            params_values[key] = value
+
+
 class Class_szfast(object):
     def __init__(self,
                 #lowring=False,  some options if needed
@@ -20,6 +58,15 @@ class Class_szfast(object):
         # self.xy = xy
         # self.lowring = lowring
 
+
+        set_verbosity(kwargs["classy_sz_verbose"])
+        self.logger = logging.getLogger(__name__)
+        
+
+        
+
+
+
         # cosmopower emulators
         self.cp_path_to_cosmopower_organization = path_to_cosmopower_organization + '/'
         self.cp_tt_nn = cp_tt_nn
@@ -33,22 +80,41 @@ class Class_szfast(object):
         self.cp_h_nn = cp_h_nn
         self.cp_s8_nn = cp_s8_nn
 
+        self.emulator_dict = emulator_dict
+
         if dofftlog_alphas == True:
             self.cp_pkl_fftlog_alphas_nus = cp_pkl_fftlog_alphas_nus
             self.cp_pkl_fftlog_alphas_real_nn  = cp_pkl_fftlog_alphas_real_nn
             self.cp_pkl_fftlog_alphas_imag_nn = cp_pkl_fftlog_alphas_imag_nn
 
-        self.cosmo_model = 'lcdm'
+        self.cosmo_model = 'ede-v2'
         self.use_Amod = 0
         self.Amod = 0 
 
         self.cp_lmax = cp_l_max_scalars
         self.cp_ls = np.arange(2,self.cp_lmax+1)
 
-        self.cp_kmax = 50.
-        self.cp_kmin = 1e-4
-        self.cp_nk = 5000
-        self.cp_ndspl_k = 10
+
+        
+        cosmo_model_dict = {0: 'lcdm',
+                            1: 'mnu',
+                            2: 'neff',
+                            3: 'wcdm',
+                            4: 'ede',
+                            5: 'mnu-3states',
+                            6: 'ede-v2'
+                            }
+        
+
+        if cosmo_model_dict[kwargs['cosmo_model']] == 'ede-v2':
+
+            self.cp_ndspl_k = 1
+            self.cp_nk = 1000
+        
+        else:
+        
+            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)
@@ -59,19 +125,43 @@ class Class_szfast(object):
         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_zmax = 5. # current max z of our pk emulators (sept 23)
+        
+
+
+        if (cosmo_model_dict[kwargs['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_kmin = 1e-4
-        self.cszfast_pk_grid_kmax = 50.
+
         self.cszfast_pk_grid_k = np.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_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
+        
         for k,v in kwargs.items():
-            # if k == 'ndim_masses':
-            #     self.cszfast_pk_grid_nk = v
-            #     self.cszfast_pk_grid_k = np.geomspace(self.cszfast_pk_grid_kmin,self.cszfast_pk_grid_kmax,self.cszfast_pk_grid_nk)
+
             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)
@@ -80,23 +170,27 @@ class Class_szfast(object):
                 self.cszfast_pk_grid_pkl_flat = np.zeros(self.cszfast_pk_grid_nz*self.cszfast_pk_grid_nk)
 
             if k == 'cosmo_model':
-                # print('updating cosmo model')
-                cosmo_model_dict = {0: 'lcdm',
-                                    1: 'mnu',
-                                    2: 'neff',
-                                    3: 'wcdm',
-                                    4: 'ede',
-                                    5: 'mnu-3states'}
+
                 self.cosmo_model = cosmo_model_dict[v]
 
             if k == 'use_Amod':
+
                 self.use_Amod = v
                 self.Amod  = kwargs['Amod']
-                # print('self.cosmo_model',self.cosmo_model)
-        # print('self.cszfast_pk_grid_nk',self.cszfast_pk_grid_nk)
-        # print('self.cszfast_pk_grid_nz',self.cszfast_pk_grid_nz)
 
-        # exit(0)
+
+
+        if cosmo_model_dict[kwargs['cosmo_model']] == 'ede-v2':
+        
+            self.pk_power_fac = self.cszfast_pk_grid_k**-3
+        
+        else:
+
+            ls = np.arange(2,self.cp_nk+2)[::self.cp_ndspl_k] # jan 10 ndspl
+            dls = ls*(ls+1.)/2./np.pi
+            self.pk_power_fac= (dls)**-1
+
+
         self.cp_z_interp = np.linspace(0.,20.,5000)
 
         self.csz_base = None
@@ -132,18 +226,21 @@ class Class_szfast(object):
                                                        self.cszfast_gas_pressure_xgrid_xmax,
                                                        self.cszfast_gas_pressure_xgrid_nx)
 
+        self.params_for_emulators = {}
 
     def find_As(self,params_cp):
         # params_cp = self.params_cp
         t0 = time.time()
 
         sigma_8_asked = params_cp["sigma8"]
+        update_params_with_defaults(params_cp, self)
         # print(params_cp)
         def to_root(ln10_10_As_goal):
             params_cp["ln10^{10}A_s"] = ln10_10_As_goal[0]
             params_dict = {}
             for k,v in params_cp.items():
                 params_dict[k]=[v]
+            
             return self.cp_der_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0][1]-sigma_8_asked
 
         lnA_s = optimize.root(to_root,
@@ -221,11 +318,13 @@ class Class_szfast(object):
                       want_ee=True,
                       want_pp=1,
                       **params_values_dict):
+        
         params_values = params_values_dict.copy()
-        # params_values['ln10^{10}A_s'] = params_values.pop("logA")
-        # print('in cmb:',params_values)
+        update_params_with_defaults(params_values, self)
+
+
         params_dict = {}
-        # print('cosmo_model',self.cosmo_model)
+
         for k,v in params_values.items():
             params_dict[k]=[v]
 
@@ -233,16 +332,45 @@ class Class_szfast(object):
             if isinstance(params_dict['m_ncdm'][0],str): 
                 params_dict['m_ncdm'] =  [float(params_dict['m_ncdm'][0].split(',')[0])]
 
-        # print('params_dict',params_dict)
 
-        if want_tt:
-            self.cp_predicted_tt_spectrum = self.cp_tt_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+
+        # if want_tt: for now always want_tt = True
+        self.cp_predicted_tt_spectrum = self.cp_tt_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+
+        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])
+
+        # 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_pp = 1./(lcp*(lcp+1.))**2.        
+
+        self.cp_predicted_tt_spectrum *= factor_ttteee
+        
+
         if want_te:
             self.cp_predicted_te_spectrum = self.cp_te_nn[self.cosmo_model].predictions_np(params_dict)[0]
+            self.cp_predicted_te_spectrum *= factor_ttteee
         if want_ee:
             self.cp_predicted_ee_spectrum = self.cp_ee_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+            self.cp_predicted_ee_spectrum *= factor_ttteee
         if want_pp:
             self.cp_predicted_pp_spectrum = self.cp_pp_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+            self.cp_predicted_pp_spectrum *= factor_pp
+
+        
+
+        
         # print('>>> clssy_szfast.py cmb computed')
 
     def load_cmb_cls_from_file(self,**params_values_dict):
@@ -268,17 +396,16 @@ class Class_szfast(object):
     def calculate_pkl(self,
                       # cosmo_model = self.cosmo_model,
                       **params_values_dict):
-        nz = self.cszfast_pk_grid_nz # number of z-points in redshift data [21oct22] --> set to 80
-        zmax = self.cszfast_pk_grid_zmax # max redshift of redshift data [21oct22] --> set to 4 because boltzmannbase.py wants to extrapolate
-        z_arr = np.linspace(0.,zmax,nz) # z-array of redshift data [21oct22] oct 26 22: nz = 1000, zmax = 20
+ 
+        z_arr = self.cszfast_pk_grid_z
 
-        nk = self.cp_nk
-        ndspl = self.cp_ndspl_k
-        k_arr = np.geomspace(self.cp_kmin,self.cp_kmax,nk)[::ndspl]  # oct 26 22 : (1e-4,50.,5000), jan 10: ndspl
+
+        k_arr = self.cszfast_pk_grid_k 
 
 
         params_values = params_values_dict.copy()
-        # print('in pkl:',params_values)
+        update_params_with_defaults(params_values, self)
+
 
         params_dict = {}
         for k,v in zip(params_values.keys(),params_values.values()):
@@ -293,55 +420,37 @@ class Class_szfast(object):
         predicted_pk_spectrum_z = []
 
         if self.use_Amod:
-            # print(">>> using Amod :",self.Amod)
+
             for zp in z_arr:
+
                 params_dict_pp = params_dict.copy()
                 params_dict_pp['z_pk_save_nonclass'] = [zp]
                 pkl_p = self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0]
                 pknl_p = self.cp_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0]
                 pk_ae  = pkl_p + self.Amod*(pknl_p-pkl_p)
                 predicted_pk_spectrum_z.append(pk_ae)
+
         else:
+
             for zp in z_arr:
+            
                 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])
 
         predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
 
-        # weird scaling to get rid off
-        # scaling factor for the pk emulator:
-        ls = np.arange(2,self.cp_nk+2)[::ndspl] # jan 10 ndspl
-        dls = ls*(ls+1.)/2./np.pi
+
         pk = 10.**predicted_pk_spectrum
-        pk_re =  ((dls)**-1*pk)
+
+        pk_re = pk*self.pk_power_fac
         pk_re = np.transpose(pk_re)
 
-        # print(z_arr,np.log(k_arr),np.log(pk_re))
-        # print(np.log(pk_re).min())
-        # print(np.log(pk_re).max())
-        # print(np.shape(np.log(pk_re)))
-        # print('coordinates')
-        # # z_coords, k_coords = np.meshgrid(z_arr, np.log(k_arr), indexing='ij')
-        # # z_coords, k_coords = z_coords.flatten(), k_coords.flatten()
-        # z_coords, k_coords = z_arr,np.log(k_arr)
-        # print(len(z_coords),len(k_coords))
-        # #
-        # # print(np.shape(list(zip(z_arr,np.log(k_arr)))))
-        # pk_values = np.log(pk_re).T
-        # # pk_values = pk_values.ravel()
-        # print(np.shape(pk_values),pk_values)
-        #
-        # self.pkl_linearnd_interp = LinearNDInterpolator(np.asarray(z_coords),np.asarray(k_coords), pk_values)
-        # exit(0)
-        # # self.pkl_cloughtocher_interp = CloughTocher2DInterpolator(list(zip(z_arr,np.log(k_arr))), np.log(pk_re).T)
-        # exit(0)
-        # self.pkl_cloughtocher_interp = CloughTocher2DInterpolator(list(zip(z_arr,np.log(k_arr))), pk_value)
-        # is this line making things slower?
         self.pkl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
-        # self.pkl_interp = None
+
         self.cszfast_pk_grid_pk = pk_re
         self.cszfast_pk_grid_pkl_flat = pk_re.flatten()
+
         return pk_re, k_arr, z_arr
 
 
@@ -396,6 +505,7 @@ class Class_szfast(object):
                          # cosmo_model = self.cosmo_model,
                          **params_values_dict):
         params_values = params_values_dict.copy()
+        update_params_with_defaults(params_values, self)
         # print('in pkl:',params_values)
 
         params_dict = {}
@@ -408,6 +518,7 @@ class Class_szfast(object):
 
         self.cp_predicted_der = self.cp_der_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
         self.sigma8 = self.cp_predicted_der[1]
+        self.Neff = self.cp_predicted_der[4]
         return 0
 
 
@@ -441,17 +552,16 @@ class Class_szfast(object):
     def calculate_pknl(self,
                       # cosmo_model = self.cosmo_model,
                       **params_values_dict):
-        nz = self.cszfast_pk_grid_nz # number of z-points in redshift data [21oct22] --> set to 80
-        zmax = self.cszfast_pk_grid_zmax # max redshift of redshift data [21oct22] --> set to 4 because boltzmannbase.py wants to extrapolate
-        z_arr = np.linspace(0.,zmax,nz) # z-array of redshift data [21oct22] oct 26 22: nz = 1000, zmax = 20
+        
+        z_arr = self.cszfast_pk_grid_z
+
 
-        nk = self.cp_nk
-        ndspl = self.cp_ndspl_k
-        k_arr = np.geomspace(self.cp_kmin,self.cp_kmax,nk)[::ndspl]  # oct 26 22 : (1e-4,50.,5000), jan 10: ndspl
+        k_arr = self.cszfast_pk_grid_k 
 
 
         params_values = params_values_dict.copy()
-        # print('in pknl:',params_values)
+        update_params_with_defaults(params_values, self)
+
 
         params_dict = {}
         for k,v in zip(params_values.keys(),params_values.values()):
@@ -471,26 +581,117 @@ class Class_szfast(object):
 
         predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
 
-        # weird scaling to get rid off
-        # scaling factor for the pk emulator:
-        ls = np.arange(2,self.cp_nk+2)[::ndspl] # jan 10 ndspl
-        dls = ls*(ls+1.)/2./np.pi
+
         pk = 10.**predicted_pk_spectrum
-        pk_re =  ((dls)**-1*pk)
+
+        pk_re = pk*self.pk_power_fac
         pk_re = np.transpose(pk_re)
 
-        # print(z_arr,np.log(k_arr),np.log(pk_re))
-        # exit(0)
-        # self.pknl_linearnd_interp = LinearNDInterpolator(list(zip(z_arr,np.log(k_arr))), pk_re.T)
-        # self.pknl_cloughtocher_interp = CloughTocher2DInterpolator(list(zip(z_arr,np.log(k_arr))), pk_re.T)
 
-        # is this line making things slower?
         self.pknl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
-        # self.pknl_interp = None
+
+
         self.cszfast_pk_grid_pknl = pk_re
         self.cszfast_pk_grid_pknl_flat = pk_re.flatten()
+
         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)
+
+
+        params_dict = {}
+        for k,v in zip(params_values.keys(),params_values.values()):
+            params_dict[k]=[v]
+
+        if 'm_ncdm' in params_dict.keys():
+            if isinstance(params_dict['m_ncdm'][0],str): 
+                params_dict['m_ncdm'] =  [float(params_dict['m_ncdm'][0].split(',')[0])]
+
+
+        predicted_pk_spectrum_z = []
+
+        z_asked = z_asked
+        params_dict_pp = params_dict.copy()
+        update_params_with_defaults(params_dict_pp, self)
+
+        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])
+
+        predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
+
+
+        pk = 10.**predicted_pk_spectrum
+        pk_re = pk*self.pk_power_fac
+        pk_re = np.transpose(pk_re)
+
+        
+        return pk_re, k_arr
+    
+
+    def calculate_pknl_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)
+
+
+        params_dict = {}
+        for k,v in zip(params_values.keys(),params_values.values()):
+            params_dict[k]=[v]
+
+        if 'm_ncdm' in params_dict.keys():
+            if isinstance(params_dict['m_ncdm'][0],str): 
+                params_dict['m_ncdm'] =  [float(params_dict['m_ncdm'][0].split(',')[0])]
+
+
+        predicted_pk_spectrum_z = []
+
+        z_asked = z_asked
+        params_dict_pp = params_dict.copy()
+        update_params_with_defaults(params_dict_pp, self)
+
+        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)
+
+
+        pk = 10.**predicted_pk_spectrum
+        pk_re = pk*self.pk_power_fac
+        pk_re = np.transpose(pk_re)
 
+        
+        return pk_re, k_arr
+    
 
     def calculate_hubble(self,
                                  # cosmo_model = self.cosmo_model,

classy_szfast/config.py

@@ -1 +1,6 @@
-path_to_cosmopower_organization = '/Users/boris/Work/CLASS-SZ/SO-SZ/class_sz/cosmopower-organization'
+import os
+
+def get_cosmopower_path():
+    return os.getenv('PATH_TO_COSMOPOWER_ORGANIZATION')
+
+path_to_cosmopower_organization = get_cosmopower_path()

classy_szfast/cosmopower.py

@@ -30,7 +30,8 @@ cosmo_model_list = [
     'neff',
     'wcdm',
     'ede',
-    'mnu-3states'
+    'mnu-3states',
+    'ede-v2'
 ]
 
 emulator_dict = {}
@@ -40,6 +41,7 @@ emulator_dict['neff'] = {}
 emulator_dict['wcdm'] = {}
 emulator_dict['ede'] = {}
 emulator_dict['mnu-3states'] = {}
+emulator_dict['ede-v2'] = {}
 
 
 
@@ -56,6 +58,9 @@ emulator_dict['lcdm']['DER'] = 'DER_v1'
 emulator_dict['lcdm']['DAZ'] = 'DAZ_v1'
 emulator_dict['lcdm']['HZ'] = 'HZ_v1'
 emulator_dict['lcdm']['S8Z'] = 'S8Z_v1'
+emulator_dict['lcdm']['default'] = {}
+emulator_dict['lcdm']['default']['tau_reio'] = 0.054
+
 
 emulator_dict['mnu']['TT'] = 'TT_mnu_v1'
 emulator_dict['mnu']['TE'] = 'TE_mnu_v1'
@@ -67,6 +72,8 @@ emulator_dict['mnu']['DER'] = 'DER_mnu_v1'
 emulator_dict['mnu']['DAZ'] = 'DAZ_mnu_v1'
 emulator_dict['mnu']['HZ'] = 'HZ_mnu_v1'
 emulator_dict['mnu']['S8Z'] = 'S8Z_mnu_v1'
+emulator_dict['mnu']['default'] = {}
+emulator_dict['mnu']['default']['tau_reio'] = 0.054
 
 
 emulator_dict['neff']['TT'] = 'TT_neff_v1'
@@ -79,6 +86,8 @@ emulator_dict['neff']['DER'] = 'DER_neff_v1'
 emulator_dict['neff']['DAZ'] = 'DAZ_neff_v1'
 emulator_dict['neff']['HZ'] = 'HZ_neff_v1'
 emulator_dict['neff']['S8Z'] = 'S8Z_neff_v1'
+emulator_dict['neff']['default'] = {}
+emulator_dict['neff']['default']['tau_reio'] = 0.054
 
 
 
@@ -92,6 +101,8 @@ emulator_dict['wcdm']['DER'] = 'DER_w_v1'
 emulator_dict['wcdm']['DAZ'] = 'DAZ_w_v1'
 emulator_dict['wcdm']['HZ'] = 'HZ_w_v1'
 emulator_dict['wcdm']['S8Z'] = 'S8Z_w_v1'
+emulator_dict['wcdm']['default'] = {}
+emulator_dict['wcdm']['default']['tau_reio'] = 0.054
 
 
 emulator_dict['ede']['TT'] = 'TT_v1'
@@ -100,13 +111,21 @@ emulator_dict['ede']['EE'] = 'EE_v1'
 emulator_dict['ede']['PP'] = 'PP_v1'
 emulator_dict['ede']['PKNL'] = 'PKNL_v1'
 emulator_dict['ede']['PKL'] = 'PKL_v1'
-# emulator_dict['ede']['PKLFFTLOG_ALPHAS_REAL'] = 'PKLFFTLOGALPHAS_creal_v1'
-# emulator_dict['ede']['PKLFFTLOG_ALPHAS_IMAG'] = 'PKLFFTLOGALPHAS_cimag_v1'
 emulator_dict['ede']['DER'] = 'DER_v1'
 emulator_dict['ede']['DAZ'] = 'DAZ_v1'
 emulator_dict['ede']['HZ'] = 'HZ_v1'
 emulator_dict['ede']['S8Z'] = 'S8Z_v1'
-
+emulator_dict['ede']['default'] = {}
+emulator_dict['ede']['default']['fEDE'] = 0.001
+emulator_dict['ede']['default']['tau_reio'] = 0.054
+emulator_dict['ede']['default']['log10z_c'] = 3.562 # e.g. from https://github.com/mwt5345/class_ede/blob/master/class/notebooks-ede/2-CMB-Comparison.ipynb
+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
+emulator_dict['ede']['default']['N_ncdm'] = 3
+emulator_dict['ede']['default']['m_ncdm'] = 0.02
+emulator_dict['ede']['default'] = {}
+emulator_dict['ede']['default']['tau_reio'] = 0.054
 
 
 emulator_dict['mnu-3states']['TT'] = 'TT_v1'
@@ -119,6 +138,30 @@ emulator_dict['mnu-3states']['DER'] = 'DER_v1'
 emulator_dict['mnu-3states']['DAZ'] = 'DAZ_v1'
 emulator_dict['mnu-3states']['HZ'] = 'HZ_v1'
 emulator_dict['mnu-3states']['S8Z'] = 'S8Z_v1'
+emulator_dict['mnu-3states']['default'] = {}
+emulator_dict['mnu-3states']['default']['tau_reio'] = 0.054
+
+
+emulator_dict['ede-v2']['TT'] = 'TT_v2'
+emulator_dict['ede-v2']['TE'] = 'TE_v2'
+emulator_dict['ede-v2']['EE'] = 'EE_v2'
+emulator_dict['ede-v2']['PP'] = 'PP_v2'
+emulator_dict['ede-v2']['PKNL'] = 'PKNL_v2'
+emulator_dict['ede-v2']['PKL'] = 'PKL_v2'
+emulator_dict['ede-v2']['DER'] = 'DER_v2'
+emulator_dict['ede-v2']['DAZ'] = 'DAZ_v2'
+emulator_dict['ede-v2']['HZ'] = 'HZ_v2'
+emulator_dict['ede-v2']['S8Z'] = 'S8Z_v2'
+
+emulator_dict['ede-v2']['default'] = {}
+emulator_dict['ede-v2']['default']['fEDE'] = 0.001
+emulator_dict['ede-v2']['default']['tau_reio'] = 0.054
+emulator_dict['ede-v2']['default']['log10z_c'] = 3.562 # e.g. from https://github.com/mwt5345/class_ede/blob/master/class/notebooks-ede/2-CMB-Comparison.ipynb
+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
+emulator_dict['ede-v2']['default']['N_ncdm'] = 3
+emulator_dict['ede-v2']['default']['m_ncdm'] = 0.02
 
 
 cp_tt_nn = {}
@@ -150,13 +193,26 @@ with suppress_warnings():
     tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 
 
+import re
+
+def split_emulator_string(input_string):
+    match = re.match(r"(.+)-v(\d+)", input_string)
+    if match:
+        folder = match.group(1)
+        version = match.group(2)
+        return folder, version
+    else:
+        folder = input_string
+        version = '1'
+        return folder, version
 
 
 
 
 for mp in cosmo_model_list:
-    
-    path_to_emulators = path_to_cosmopower_organization + '/' + mp +'/'
+    folder, version = split_emulator_string(mp)
+    # print(folder, version)
+    path_to_emulators = path_to_cosmopower_organization + '/' + folder +'/'
     
     cp_tt_nn[mp] = cosmopower_NN(restore=True,
                              restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TT'])

{classy_szfast-0.0.7.dist-info → classy_szfast-0.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.7
+Version: 0.0.9
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Home-page: https://github.com/CLASS-SZ/classy_szfast
 Download-URL: https://github.com/CLASS-SZ/classy_szfast

{classy_szfast-0.0.7.dist-info → classy_szfast-0.0.9.dist-info}/RECORD

@@ -1,8 +1,8 @@
-classy_szfast/__init__.py,sha256=Mp1HvqeZ9_XmjVkqljBe8qG8vxpDzQt-CzQDlZFmOIU,247
-classy_szfast/classy_sz.py,sha256=AUCJsWFOTUH1zadbfBcCS8coHLKNEvSlvAUV5DQ3xQ0,25834
-classy_szfast/classy_szfast.py,sha256=gpiGm3WLStyhiT-b-s8wUvlxCcdqjMZkVqpVhszkJ0M,32433
-classy_szfast/config.py,sha256=JwRAmdVL44ftwuJyrtuP6tZKmQC_yMJ-WdPustaFIWY,102
-classy_szfast/cosmopower.py,sha256=O9Jorm7st-z7CKLXgAIrIZmVZx5rS2xjiUtkEIDrkpo,7010
+classy_szfast/__init__.py,sha256=E2thrL0Z9oXFfdzwcsu-xbOytudLFTlRlPqVFGlPPPg,279
+classy_szfast/classy_sz.py,sha256=mp6tbYzeoi27ciShV2N8xGCR9vIrhm1ISgRHjDGEh-E,25829
+classy_szfast/classy_szfast.py,sha256=WogzPos9dGlacklyAOjb58CS4Qbdip8BZyuj0RcGLtA,35452
+classy_szfast/config.py,sha256=OJXwK9gTufJgmOxXUax3XJ6QJ1yF0pJARJBCG-odZQU,151
+classy_szfast/cosmopower.py,sha256=poSQ6zRHJBPedZ-nO3BWjsxXLyd8FSh9aTCx6B-VoqQ,9553
 classy_szfast/cosmosis_classy_szfast_interface.py,sha256=zAnxvFtn73a5yS7jgs59zpWFEYKCIQyraYPs5hQ4Le8,11483
 classy_szfast/pks_and_sigmas.py,sha256=drtuujE1HhlrYY1hY92DyY5lXlYS1uE15MSuVI4uo6k,6625
 classy_szfast/suppress_warnings.py,sha256=6wIBml2Sj9DyRGZlZWhuA9hqvpxqrNyYjuz6BPK_a6E,202
@@ -11,7 +11,7 @@ classy_szfast/custom_bias/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 classy_szfast/custom_bias/custom_bias.py,sha256=53RbMAawCurKuI44nqk6fTTkQbbuJDcZpK0YhNQS3j8,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.7.dist-info/METADATA,sha256=-IehxlVfQsXMIBkODyolN9fmYV4bkk_ruY5N3CFAdSc,713
-classy_szfast-0.0.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-classy_szfast-0.0.7.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.7.dist-info/RECORD,,
+classy_szfast-0.0.9.dist-info/METADATA,sha256=-hPJoxQ6DfJok-Uz8W_RS-77DQ32doSDnDl5_z_-vV4,713
+classy_szfast-0.0.9.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+classy_szfast-0.0.9.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.9.dist-info/RECORD,,