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

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 = None
-    # 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

@@ -10,6 +10,7 @@ from mcfit import TophatVar
 from scipy.interpolate import CubicSpline
 import pickle
 
+
 H_units_conv_factor = {"1/Mpc": 1, "km/s/Mpc": Const.c_km_s}
 
 
@@ -33,16 +34,15 @@ def set_verbosity(verbosity):
 
 
 
-def update_params_with_defaults(params_values, self):
+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.
-    self (object): The object containing emulator_dict and cosmo_model attributes.
+    default_values (dict): Dictionary containing default parameter values.
     """
-    # 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():
@@ -52,6 +52,7 @@ def update_params_with_defaults(params_values, self):
 
 class Class_szfast(object):
     def __init__(self,
+                 params_settings = {},
                 #lowring=False,  some options if needed
                  **kwargs):
         # some parameters
@@ -59,7 +60,11 @@ class Class_szfast(object):
         # self.lowring = lowring
 
 
-        set_verbosity(kwargs["classy_sz_verbose"])
+        self.A_s_fast = 0 
+        self.logA_fast = 0
+
+
+        set_verbosity(params_settings["classy_sz_verbose"])
         self.logger = logging.getLogger(__name__)
         
 
@@ -95,6 +100,7 @@ class Class_szfast(object):
         self.cp_ls = np.arange(2,self.cp_lmax+1)
 
 
+
         
         cosmo_model_dict = {0: 'lcdm',
                             1: 'mnu',
@@ -106,7 +112,7 @@ class Class_szfast(object):
                             }
         
 
-        if cosmo_model_dict[kwargs['cosmo_model']] == 'ede-v2':
+        if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
 
             self.cp_ndspl_k = 1
             self.cp_nk = 1000
@@ -128,16 +134,16 @@ class Class_szfast(object):
         
 
 
-        if (cosmo_model_dict[kwargs['cosmo_model']] == 'ede-v2'):
+        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}")
+            # 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:
         
@@ -146,9 +152,9 @@ class Class_szfast(object):
             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.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)
@@ -160,12 +166,14 @@ class Class_szfast(object):
         
         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():
+        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)
 
@@ -176,11 +184,11 @@ class Class_szfast(object):
             if k == 'use_Amod':
 
                 self.use_Amod = v
-                self.Amod  = kwargs['Amod']
+                self.Amod  = params_settings['Amod']
 
 
 
-        if cosmo_model_dict[kwargs['cosmo_model']] == 'ede-v2':
+        if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
         
             self.pk_power_fac = self.cszfast_pk_grid_k**-3
         
@@ -194,14 +202,6 @@ class Class_szfast(object):
         self.cp_z_interp = np.linspace(0.,20.,5000)
 
         self.csz_base = None
-        # self.csz_base.compute()
-
-
-        # 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
 
 
         self.cszfast_zgrid_zmin = 0.
@@ -225,16 +225,15 @@ 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"]
-        update_params_with_defaults(params_cp, self)
-        # 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 = {}
@@ -247,14 +246,18 @@ class Class_szfast(object):
                               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):
@@ -262,7 +265,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,
@@ -319,9 +322,10 @@ class Class_szfast(object):
                       want_pp=1,
                       **params_values_dict):
         
+
         params_values = params_values_dict.copy()
-        update_params_with_defaults(params_values, self)
 
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
         params_dict = {}
 
@@ -404,7 +408,7 @@ class Class_szfast(object):
 
 
         params_values = params_values_dict.copy()
-        update_params_with_defaults(params_values, self)
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
 
         params_dict = {}
@@ -455,57 +459,50 @@ class Class_szfast(object):
 
 
     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)
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
+        
         # print('in pkl:',params_values)
 
         params_dict = {}
@@ -560,7 +557,7 @@ class Class_szfast(object):
 
 
         params_values = params_values_dict.copy()
-        update_params_with_defaults(params_values, self)
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
 
         params_dict = {}
@@ -613,7 +610,7 @@ class Class_szfast(object):
 
             params_values = self.params_for_emulators
 
-        update_params_with_defaults(params_values, self)
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
 
         params_dict = {}
@@ -629,7 +626,7 @@ class Class_szfast(object):
 
         z_asked = z_asked
         params_dict_pp = params_dict.copy()
-        update_params_with_defaults(params_dict_pp, self)
+        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])
@@ -661,7 +658,7 @@ class Class_szfast(object):
 
             params_values = self.params_for_emulators
 
-        update_params_with_defaults(params_values, self)
+        update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
 
         params_dict = {}
@@ -677,7 +674,7 @@ class Class_szfast(object):
 
         z_asked = z_asked
         params_dict_pp = params_dict.copy()
-        update_params_with_defaults(params_dict_pp, self)
+        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])
@@ -694,10 +691,12 @@ class Class_szfast(object):
     
 
     def calculate_hubble(self,
-                                 # cosmo_model = self.cosmo_model,
-                                 **params_values_dict):
+                         **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]
@@ -706,11 +705,8 @@ class Class_szfast(object):
             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,
@@ -722,36 +718,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),
@@ -763,6 +756,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)
@@ -774,13 +768,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/config.py

@@ -1,6 +1,10 @@
 import os
 
+import get_cosmopower_emus
+
+
 def get_cosmopower_path():
+    get_cosmopower_emus.set()
     return os.getenv('PATH_TO_COSMOPOWER_ORGANIZATION')
 
 path_to_cosmopower_organization = get_cosmopower_path()

classy_szfast/cosmopower.py

@@ -60,7 +60,9 @@ 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'
@@ -74,7 +76,9 @@ 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'
@@ -88,7 +92,9 @@ 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'
@@ -103,7 +109,9 @@ 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'
@@ -124,8 +132,6 @@ 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'
@@ -140,7 +146,9 @@ 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'
@@ -156,6 +164,10 @@ 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.

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.9.dist-info → classy_szfast-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.9
+Version: 0.0.11
 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
@@ -16,4 +16,5 @@ Requires-Dist: tensorflow ==2.13.0
 Requires-Dist: tensorflow-probability ==0.21.0
 Requires-Dist: cosmopower
 Requires-Dist: mcfit
+Requires-Dist: get-cosmopower-emus
 

classy_szfast-0.0.11.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=TBmm98jDakvLRbrDiI6m_youZXOONyhXAV2MyU41gFc,33729
+classy_szfast/config.py,sha256=4CvejtLcFOQR30bJ8tlEeBHhu3Rr7LakeLO6dbFgPSU,210
+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.11.dist-info/METADATA,sha256=q7a06RaiRRY8j-mjhJGEVeQUoZx-l-JJ1IpDnjL_UGw,749
+classy_szfast-0.0.11.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+classy_szfast-0.0.11.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.11.dist-info/RECORD,,

{classy_szfast-0.0.9.dist-info → classy_szfast-0.0.11.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.9.dist-info/RECORD

@@ -1,17 +0,0 @@
-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
-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.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,,