classy-szfast 0.0.8__py3-none-any.whl → 0.0.10__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

@@ -8,59 +8,69 @@ import numpy as np
 import time
 
 class classy_sz(classy):
-    use_class_sz_fast_mode = 0 # this is passed in the yaml file
+
+    use_class_sz_fast_mode = 1 # this is passed in the yaml file
     use_class_sz_no_cosmo_mode = 0 # this is passed in the yaml file
-    lensing_lkl = "SOLikeT"
-    # skip_background_and_thermo = True
-    # ell_factor = False # True for pyactlite and bplike, False for clik
+    lensing_lkl = 'ACT'
+
+    baos = None 
 
     def initialize(self):
         """Importing CLASS from the correct path, if given, and if not, globally."""
+
         self.classy_module = self.is_installed()
         if not self.classy_module:
             raise NotInstalledError(
                 self.log, "Could not find CLASS_SZ. Check error message above.")
+        
         from classy_sz import Class, CosmoSevereError, CosmoComputationError
+        
         global CosmoComputationError, CosmoSevereError
+
         self.classy = Class()
         super(classy,self).initialize()
+
+
         # Add general CLASS stuff
         self.extra_args["output"] = self.extra_args.get("output", "")
+
         if "sBBN file" in self.extra_args:
             self.extra_args["sBBN file"] = (
                 self.extra_args["sBBN file"].format(classy=self.path))
+            
         # Derived parameters that may not have been requested, but will be necessary later
         self.derived_extra = []
         self.log.info("Initialized!")
 
+
+        ## rename some parameters to avoid conflices
+        classy_sz_renames = {
+
+            'omega_m':'Omega_m',
+            'Omegam':'Omega_m',
+            'Omega_m':'Omega_m'
+        }
+        self.renames.update(classy_sz_renames)
+
+
         if self.use_class_sz_no_cosmo_mode == 1:
+
             self.log.info("Initializing cosmology part!")
+
             initial_parameters = self.extra_args.copy()
-            # print("initial_parameters:",initial_parameters)
+
 
             self.classy.set(initial_parameters)
-            self.classy.compute_class_szfast()
+            self.classy.compute_class_szfast(likelihood_mode=True)
             self.log.info("cosmology part initialized!")
 
 
-        # print(self.lensing_lkl)
-        # exit(0)
-
-        # # class_sz default params for lkl
-        # self.extra_args["output"] = 'tSZ_1h'
-        # self.extra_args["multipoles_sz"] = 'P15'
-        # self.extra_args['nlSZ'] = 18
 
 
-    # # here modify if you want to bypass stuff in the class computation
-    # def calculate(self, state, want_derived=True, **params_values_dict):
-    #     print("Bypassing class_sz")
-
+    def must_provide(self, **requirements):
 
 
 
-    def must_provide(self, **requirements):
-
         if "Cl_sz" in requirements:
             # make sure cobaya still runs as it does for standard classy
             requirements.pop("Cl_sz")
@@ -246,25 +256,25 @@ class classy_sz(classy):
                         args=[])
         super().must_provide(**requirements)
 
+
     # get the required new observable
     def get_Cl(self, ell_factor=False, units="FIRASmuK2"):
+        
         if self.use_class_sz_fast_mode:
+
             return self.get_Clfast(ell_factor=ell_factor)
+        
         else:
+        
             return self._get_Cl(ell_factor=ell_factor, units=units, lensed=True)
 
     def get_Clfast(self,ell_factor = False):
-        # print('ell_factor:',self.ell_factor)
-        # exit(0)
+
         cls = {}
         cls = deepcopy(self._current_state["Cl"])
-        # ell_factor = self.ell_factor
-        # print(cls)
-        # exit(0)
-        # print('in get clfast:',ell_factor)
-        # print(cls)
+
         lcp = np.asarray(cls['ell'])
-        # print(self.lensing_lkl)
+
         if ell_factor==True:
             cls['tt'] *= (2.7255e6)**2.*(lcp*(lcp+1.))/2./np.pi
             cls['te'] *= (2.7255e6)**2.*(lcp*(lcp+1.))/2./np.pi
@@ -274,17 +284,18 @@ class classy_sz(classy):
             cls['tt'] *= (2.7255e6)**2.
             cls['te'] *= (2.7255e6)**2.
             cls['ee'] *= (2.7255e6)**2.
-        # print(cls['tt'][1230])
-        # print(cls['te'][1230])
-        # print(cls['ee'][1230])
-        # exit(0)
+
         if self.lensing_lkl ==  "SOLikeT":
             cls['pp'] *= (lcp*(lcp+1.))**2./4.
+
         elif self.lensing_lkl == "ACT":
-            cls['pp'] *= 1.#(lcp*(lcp+1.))**2./4.
+            cls['pp'] *= 1.
+             
         else: # here for the planck lensing lkl, using lfactor option gives:
             cls['pp'] *= (lcp*(lcp+1.))**2.*1./2./np.pi
+
         return cls
+    
     # get the required new observable
     def get_Cl_sz(self):
         cls = {}
@@ -362,8 +373,9 @@ class classy_sz(classy):
 
     # get the required new observable
     def get_sz_unbinned_cluster_counts(self):
+
         cls = deepcopy(self._current_state["sz_unbinned_cluster_counts"])
-        # print(cls)
+
         return cls['loglike'],cls['ntot'],cls['rates']
 
 
@@ -380,37 +392,53 @@ class classy_sz(classy):
 
     # IMPORTANT: this method is imported from cobaya and modified to accomodate the emulators
     def calculate(self, state, want_derived=True, **params_values_dict):
-        # Set parameters
+
+        
         params_values = params_values_dict.copy()
-        # print('\n\n')
-        # print('>>> class_sz.py: class/class_sz using params:',params_values)
+        # if baos are requested we need to update the relevant flags
+        if self.baos:
+            params_values.update({'skip_chi':0,'skip_hubble':0})
+
         if 'N_ncdm' in self.extra_args.keys():
+
             if self.extra_args['N_ncdm'] ==  3:
+            
                 str_mncdm = str(params_values['m_ncdm'])
                 params_values['m_ncdm'] = str_mncdm+','+str_mncdm+','+str_mncdm
-        # print('>>> class_sz.py: class/class_sz using params:',params_values)
-        # exit(0)
+
         try:
+            
             params_values['ln10^{10}A_s'] = params_values.pop("logA")
+            
             self.set(params_values)
+        
         except KeyError:
+        
             self.set(params_values)
+
         # Compute!
         try:
             if self.use_class_sz_fast_mode == 1:
-                # start = time.perf_counter()
+
+
                 if self.use_class_sz_no_cosmo_mode == 1:
-                    # print(params_values)
+
+                    start_time = time.time()
                     self.classy.compute_class_sz(params_values)
+                    end_time = time.time()
+                    # self.log.info("Execution time of class_sz: {:.5f} seconds".format(end_time - start_time))
+
                 else:
-                    self.classy.compute_class_szfast()
-                # end = time.perf_counter()
-                # print('classy_szfast took:',end-start)
-            # self.classy.compute_class_szfast()
-            # elif self.use_class_sz_no_cosmo_mode == 1:
-            #     self.classy.compute_class_sz(params_values)
+
+                    start_time = time.time()
+                    self.classy.compute_class_szfast(likelihood_mode=True)
+                    end_time = time.time()
+                    # self.log.info("Execution time of class_szfast: {:.5f} seconds".format(end_time - start_time))
+                    # print('pars in classy',self.classy.pars)
             else:
+
                 self.classy.compute()
+
         # "Valid" failure of CLASS: parameters too extreme -> log and report
         except self.classy_module.CosmoComputationError as e:
             if self.stop_at_error:
@@ -492,81 +520,15 @@ class classy_sz(classy):
             if collector.post:
                 state[product] = collector.post(*state[product])
         # Prepare derived parameters
+
         d, d_extra = self._get_derived_all(derived_requested=want_derived)
+
         if want_derived:
+
             state["derived"] = {p: d.get(p) for p in self.output_params}
-            # Prepare necessary extra derived parameters
+
+
         state["derived_extra"] = deepcopy(d_extra)
-        # exit(0)
-
-
-    # # get the required new observable
-    # def get_Cl(self,ell_factor=True,units="FIRASmuK2"):
-    #     ell_factor=self.ell_factor
-    #     # if self.tsz.use_class_sz_fast_mode == 1:
-    #     cls = {}
-    #     cls['ell'] = np.arange(20000)
-    #     # print(cls['ell'])
-    #     cls['tt'] = np.zeros(20000)
-    #     cls['te'] = np.zeros(20000)
-    #     cls['ee'] = np.zeros(20000)
-    #     cls['pp'] = np.zeros(20000)
-    #     # if self.tt_spectra is not None:
-    #     nl = len(self.classy.class_szfast.cp_predicted_tt_spectrum)
-    #     # print('nl:',nl)
-    #     cls['tt'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_tt_spectrum
-    #     if ell_factor==False:
-    #         lcp = np.asarray(cls['ell'][2:nl+2])
-    #         cls['tt'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
-    #
-    # # if self.te_spectra is not None:
-    #     cls['te'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_te_spectrum
-    #     if ell_factor==False:
-    #         lcp = np.asarray(cls['ell'][2:nl+2])
-    #         cls['te'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
-    # # if self.ee_spectra is not None:
-    #     cls['ee'][2:nl+2] = (2.7255e6)**2.*self.classy.class_szfast.cp_predicted_ee_spectrum
-    #     if ell_factor==False:
-    #         lcp = np.asarray(cls['ell'][2:nl+2])
-    #         cls['ee'][2:nl+2] *= 1./(lcp*(lcp+1.)/2./np.pi)
-    #     # if self.pp_spectra is not None:
-    #         # nl = len(self.pp_spectra[0])
-    #     if self.lensing_lkl ==  "SOLikeT":
-    #         cls['pp'][2:nl+2] = self.classy.class_szfast.cp_predicted_pp_spectrum/4. ## this is clkk... works for so lensinglite lkl
-    #     else:
-    #         # here for the planck lensing lkl, using lfactor option gives:
-    #         lcp = np.asarray(cls['ell'][2:nl+2])
-    #         cls['pp'][2:nl+2] = self.classy.class_szfast.cp_predicted_pp_spectrum/(lcp*(lcp+1.))**2.
-    #         cls['pp'][2:nl+2] *= (lcp*(lcp+1.))**2./2./np.pi
-    #     return cls
-
-    # def check_ranges(self, z, k):
-    #     return 1
-
-    # IMPORTANT: copied from cobaya and changed.
-    def get_param(self, p):
-        translated = self.translate_param(p)
-        for pool in ["params", "derived", "derived_extra"]:
-            value = (self.current_state[pool] or {}).get(translated, None)
-            if p == 'omegam':
-                print('getting omegam in get_param')
-                # print(translated)
-                # print(self.classy.Omega_m())
-                # exit(0)
-                return self.classy.Omega_m()
-            if value is not None:
-                return value
-
-        raise LoggedError(self.log, "Parameter not known: '%s'", p)
-### ORIGINAL FUNCTION:
-#    def get_param(self, p):
-#        translated = self.translate_param(p)
-#        for pool in ["params", "derived", "derived_extra"]:
-#            value = (self.current_state[pool] or {}).get(translated, None)
-#            if value is not None:
-#                return value
-#
-#        raise LoggedError(self.log, "Parameter not known: '%s'", p)
 
 
 

classy_szfast/classy_szfast.py

@@ -12,14 +12,65 @@ 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, default_values):
+    """
+    Update params_values with default values if they don't already exist.
+
+    Args:
+    params_values (dict): Dictionary containing parameter values.
+    default_values (dict): Dictionary containing default parameter values.
+    """
+
+    
+    # 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,
+                 params_settings = {},
                 #lowring=False,  some options if needed
                  **kwargs):
         # some parameters
         # self.xy = xy
         # self.lowring = lowring
 
+
+        self.A_s_fast = 0 
+        self.logA_fast = 0
+
+
+        set_verbosity(params_settings["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 +84,42 @@ 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[params_settings['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,55 +130,77 @@ 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[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_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)
+        
+        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_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)
 
             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)
+                self.Amod  = params_settings['Amod']
 
-        # exit(0)
-        self.cp_z_interp = np.linspace(0.,20.,5000)
 
-        self.csz_base = None
-        # self.csz_base.compute()
 
+        if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
+        
+            self.pk_power_fac = self.cszfast_pk_grid_k**-3
+        
+        else:
 
-        # z_arr = np.linspace(0.,zmax,nz) # z-array of redshift data [21oct22] oct 26 22: nz = 1000, zmax = 20
-        #
-        # 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
+            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
 
 
         self.cszfast_zgrid_zmin = 0.
@@ -131,33 +224,39 @@ class Class_szfast(object):
         self.cszfast_gas_pressure_xgrid = np.geomspace(self.cszfast_gas_pressure_xgrid_xmin,
                                                        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"]
-        # print(params_cp)
+
+        update_params_with_defaults(params_cp, self.emulator_dict[self.cosmo_model]['default'])
+
         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,
                               x0=3.046,
                               #tol = 1e-10,
                               method="hybr")
-        params_cp['ln10^{10}A_s'] = lnA_s.x[0]# .x[0]
+        
+        params_cp['ln10^{10}A_s'] = lnA_s.x[0]
+
         params_cp.pop('sigma8')
-        # print("T total in find As",time.time()-t0)#self.t_total)
-        # print(params_cp)
+
         return 1
 
 
     def get_H0_from_thetas(self,params_values):
+
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+
         # print(params_values)
         theta_s_asked = params_values['100*theta_s']
         def fzero(H0_goal):
@@ -165,7 +264,7 @@ class Class_szfast(object):
           params_dict = {}
           for k,v in params_values.items():
               params_dict[k]=[v]
-          # print(params_dict)
+
           predicted_der_params = self.cp_der_nn[self.cosmo_model].ten_to_predictions_np(params_dict)
           return predicted_der_params[0][0]-theta_s_asked
         sol = optimize.root(fzero,
@@ -221,11 +320,14 @@ 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.emulator_dict[self.cosmo_model]['default'])
+
         params_dict = {}
-        # print('cosmo_model',self.cosmo_model)
+
         for k,v in params_values.items():
             params_dict[k]=[v]
 
@@ -233,16 +335,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 +399,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.emulator_dict[self.cosmo_model]['default'])
+
 
         params_dict = {}
         for k,v in zip(params_values.keys(),params_values.values()):
@@ -293,109 +423,85 @@ 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
 
 
     def calculate_sigma(self,
-                        # cosmo_model = self.cosmo_model,
-                        # z_asked = None,
-                        # r_asked = None,
+                        
                         **params_values_dict):
+
         params_values = params_values_dict.copy()
+
         k = self.cszfast_pk_grid_k
-        # self.cszfast_pk_grid_z
-        # print(self.cszfast_pk_grid_pk,np.shape(self.cszfast_pk_grid_pk))
+
         P = self.cszfast_pk_grid_pk
+
         var = P.copy()
+
         dvar = P.copy()
+
         for iz,zp in enumerate(self.cszfast_pk_grid_z):
+
             R, var[:,iz] = TophatVar(k, lowring=True)(P[:,iz], extrap=True)
-            # dvar[:,iz] = np.gradient(var[:,iz], np.log(R))
-            # if params_values_dict['sigma_derivative'] == 1: ## need more points here !!
-            #     rds,dvar[:,iz] =  TophatVar(k,lowring=True,deriv=1)(P[:,iz]*k,extrap=True)
-            # else:
-            #     dvar[:,iz] = np.gradient(var[:,iz], R)
 
             dvar[:,iz] = np.gradient(var[:,iz], R)
-            # from inigo: R_vec,self.dvar = TophatVar(k,lowring=True,deriv=1)(P[:,iz]*k,extrap=True)
-            # from inigo: self.dsigma_vec = self.dvar/(2.*self.sigma_vec)
-            # print('R:',R,np.shape(R))
-            # varR = CubicSpline(R, var[:,iz])
-            # print(zp,np.sqrt(varR(8)))
-        # print(params_values)
-        # print('in sigma:',params_values)
-        # h = params_values['H0']/100.
-        # var = var.T
-        # dvar = dvar.T
+
+        # print(k)
+        # print(R)
+        # print(k*R)
+        # exit(0)
+ 
 
         self.cszfast_pk_grid_lnr = np.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 = self.cszfast_pk_grid_lnsigma2_flat.T
+ 
         self.cszfast_pk_grid_dsigma2 = dvar
         self.cszfast_pk_grid_dsigma2_flat = dvar.flatten()
-        # if z_asked != None and r_asked != None:
-        #     print(z_asked[0],r_asked[0])
-        #     return z_asked, r_asked
-        # else:
-        #     return 0
+ 
         return 0
 
 
     def calculate_sigma8_and_der(self,
                          # cosmo_model = self.cosmo_model,
                          **params_values_dict):
+        
         params_values = params_values_dict.copy()
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+        
         # print('in pkl:',params_values)
 
         params_dict = {}
@@ -408,6 +514,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 +548,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.emulator_dict[self.cosmo_model]['default'])
+
 
         params_dict = {}
         for k,v in zip(params_values.keys(),params_values.values()):
@@ -471,31 +577,40 @@ 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.emulator_dict[self.cosmo_model]['default'])
 
-    def calculate_hubble(self,
-                                 # cosmo_model = self.cosmo_model,
-                                 **params_values_dict):
-        params_values = params_values_dict.copy()
 
         params_dict = {}
         for k,v in zip(params_values.keys(),params_values.values()):
@@ -506,10 +621,91 @@ class Class_szfast(object):
                 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.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])
+
+        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.emulator_dict[self.cosmo_model]['default'])
+
+
+        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.emulator_dict[self.cosmo_model]['default'])
+
+        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,
+                         **params_values_dict):
+        
+        params_values = params_values_dict.copy()
+
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+
+        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])]
 
         self.cp_predicted_hubble = self.cp_h_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
-        # print(self.cp_predicted_hubble)
-        # z_interp =
+ 
         self.hz_interp = scipy.interpolate.interp1d(
                                             self.cp_z_interp,
                                             self.cp_predicted_hubble,
@@ -521,36 +717,33 @@ class Class_szfast(object):
                                             assume_sorted=False)
 
     def calculate_chi(self,
-                      # cosmo_model = self.cosmo_model,
                       **params_values_dict):
-        # def test_integrand_func(x, alpha, beta, i, j, k, l):
-        #     return x * alpha * beta + i * j * k
-        #
-        #
-        # grid = np.random.random((5000, 2))
-        #
-        # res, err = pyquad.quad_grid(test_integrand_func, 0, 1, grid, (1.0, 1.0, 1.0, 1.0))
-        #
-        # print(res,err)
-        # def integrand_chi(z,alpha, beta, i, j, k, l):
-        #     z = z-1.
-        #     return 1./self.get_Hubble(z)
-        # zmax = 1.
-        # grid = np.random.random((10000000, 2))
-        # chiz,err = pyquad.quad_grid(integrand_chi, 1., 1.+zmax, grid, (1.0, 1.0, 1.0, 1.0))
-        # print(chiz)
+
         params_values = params_values_dict.copy()
 
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+
         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])]
 
+        # deal with different scaling of DA in different model from emulator training
+        if self.cosmo_model == 'ede-v2':
+
+            self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+            self.cp_predicted_da = np.insert(self.cp_predicted_da, 0, 0)
+        
+        else:
+        
+            self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].predictions_np(params_dict)[0]
+        
 
-        self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].predictions_np(params_dict)[0]
         self.chi_interp = scipy.interpolate.interp1d(
                                                     self.cp_z_interp,
                                                     self.cp_predicted_da*(1.+self.cp_z_interp),
@@ -562,6 +755,7 @@ class Class_szfast(object):
                                                     assume_sorted=False)
 
     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)
@@ -573,13 +767,6 @@ class Class_szfast(object):
         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()
         cls['pp'][2:self.cp_lmax+1] = self.cp_predicted_pp_spectrum.copy()/4. ## this is clkk... works for so lensinglite lkl
-        # cls['bb'][2:self.cp_lmax+1] = self.cp_predicted_pp_spectrum.copy()/4. ## this is clkk... works for so lensinglite lkl
-        # print('doing gets')
-        # For planck likelihood:
-        # lcp = np.asarray(cls['ell'][2:nl+2])
-        # cls['pp'][2:nl+2] = self.pp_spectra[0].copy()/(lcp*(lcp+1.))**2.
-        # cls['pp'][2:nl+2] *= (lcp*(lcp+1.))**2./2./np.pi
-
 
 
         if ell_factor==False:

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,11 @@ 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['lcdm']['default']['N_ur'] = 2.0328
+emulator_dict['lcdm']['default']['N_ncdm'] = 1
+emulator_dict['lcdm']['default']['m_ncdm'] = 0.06
 
 emulator_dict['mnu']['TT'] = 'TT_mnu_v1'
 emulator_dict['mnu']['TE'] = 'TE_mnu_v1'
@@ -67,7 +74,11 @@ 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['mnu']['default']['N_ur'] = 2.0328
+emulator_dict['mnu']['default']['N_ncdm'] = 1
+emulator_dict['mnu']['default']['m_ncdm'] = 0.06
 
 emulator_dict['neff']['TT'] = 'TT_neff_v1'
 emulator_dict['neff']['TE'] = 'TE_neff_v1'
@@ -79,7 +90,11 @@ 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
+emulator_dict['neff']['default']['N_ur'] = 2.0328
+emulator_dict['neff']['default']['N_ncdm'] = 1
+emulator_dict['neff']['default']['m_ncdm'] = 0.06
 
 
 emulator_dict['wcdm']['TT'] = 'TT_w_v1'
@@ -92,7 +107,11 @@ 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['wcdm']['default']['N_ur'] = 2.0328
+emulator_dict['wcdm']['default']['N_ncdm'] = 1
+emulator_dict['wcdm']['default']['m_ncdm'] = 0.06
 
 emulator_dict['ede']['TT'] = 'TT_v1'
 emulator_dict['ede']['TE'] = 'TE_v1'
@@ -100,13 +119,19 @@ 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['mnu-3states']['TT'] = 'TT_v1'
@@ -119,6 +144,36 @@ 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['mnu-3states']['default']['N_ur'] = 0.00641
+emulator_dict['mnu-3states']['default']['N_ncdm'] = 3
+emulator_dict['mnu-3states']['default']['m_ncdm'] = 0.02
+
+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']['H0'] = 67.66
+emulator_dict['ede-v2']['default']['ln10^{10}A_s'] = 3.047
+emulator_dict['ede-v2']['default']['omega_b'] = 0.02242
+emulator_dict['ede-v2']['default']['omega_cdm'] = 0.11933
+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 +205,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/custom_bias/custom_bias.py

@@ -17,4 +17,4 @@ def custom1_b(z,classy_sz,*params):
     # else:
     #     w = 3./2.*(H0/h)**2/Rho_crit_0*(chi/(1.+z))**-1.*(chi_star-chi)/chi_star # 
 
-    return 2.# + 2*z
+    return 4.# + 2*z

{classy_szfast-0.0.8.dist-info → classy_szfast-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.8
+Version: 0.0.10
 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.10.dist-info/RECORD

@@ -0,0 +1,17 @@
+classy_szfast/__init__.py,sha256=E2thrL0Z9oXFfdzwcsu-xbOytudLFTlRlPqVFGlPPPg,279
+classy_szfast/classy_sz.py,sha256=QmbwrSXInQLMvCDqsr7KPmtaU0KOiOt1Rb-cTKuulZw,22240
+classy_szfast/classy_szfast.py,sha256=YiT4xHKmgYvEo5LphwLHooPc5H4kr2FuDSJG--H7cdo,33728
+classy_szfast/config.py,sha256=OJXwK9gTufJgmOxXUax3XJ6QJ1yF0pJARJBCG-odZQU,151
+classy_szfast/cosmopower.py,sha256=u9iLdt51Qss5KwOQ637lHDBWsUz0H2u-53VUMyP5foY,10435
+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
+classy_szfast/utils.py,sha256=ZZxujm1yBM0KIeVVLOuoNqUVkXIZt817QDi7U_Fz_IM,1462
+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.10.dist-info/METADATA,sha256=9xenjOGGj3pgWbw9Ygm8Wkq9IFUBkHcBNHUByCRherY,714
+classy_szfast-0.0.10.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+classy_szfast-0.0.10.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.10.dist-info/RECORD,,

{classy_szfast-0.0.8.dist-info → classy_szfast-0.0.10.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (71.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

classy_szfast-0.0.8.dist-info/RECORD

@@ -1,17 +0,0 @@
-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=OJXwK9gTufJgmOxXUax3XJ6QJ1yF0pJARJBCG-odZQU,151
-classy_szfast/cosmopower.py,sha256=O9Jorm7st-z7CKLXgAIrIZmVZx5rS2xjiUtkEIDrkpo,7010
-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
-classy_szfast/utils.py,sha256=ZZxujm1yBM0KIeVVLOuoNqUVkXIZt817QDi7U_Fz_IM,1462
-classy_szfast/custom_bias/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-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.8.dist-info/METADATA,sha256=aTR5RDlPBTAxaplz-G8x-Erlw8NVPFH6lNusCCxjNO4,713
-classy_szfast-0.0.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-classy_szfast-0.0.8.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.8.dist-info/RECORD,,