classy-szfast 0.0.25.post13__tar.gz → 0.0.25.post15__tar.gz

{classy_szfast-0.0.25.post13 → classy_szfast-0.0.25.post15}/PKG-INFO

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

{classy_szfast-0.0.25.post13 → classy_szfast-0.0.25.post15}/classy_szfast/classy_szfast.py

@@ -88,7 +88,7 @@ class Class_szfast(object):
         self.cp_ee_nn = cp_ee_nn
         self.cp_pp_nn = cp_pp_nn
         self.cp_pknl_nn  = cp_pknl_nn
-        self.cp_der_nn = cp_der_nn
+
        
         self.cp_lmax = cp_l_max_scalars
 
@@ -120,6 +120,7 @@ class Class_szfast(object):
             self.cp_h_nn = cp_h_nn_jax
             self.cp_da_nn = cp_da_nn_jax
             self.cp_pkl_nn = cp_pkl_nn_jax
+            self.cp_der_nn = cp_der_nn_jax
 
             self.pi = jnp.pi
             self.transpose = jnp.transpose
@@ -140,7 +141,7 @@ class Class_szfast(object):
             self.cp_h_nn = cp_h_nn
             self.cp_da_nn = cp_da_nn
             self.cp_pkl_nn = cp_pkl_nn
-
+            self.cp_der_nn = cp_der_nn
             self.pi = np.pi
             self.transpose = np.transpose
             self.pow = np.power
@@ -173,9 +174,16 @@ class Class_szfast(object):
             self.cp_pkl_fftlog_alphas_imag_nn = cp_pkl_fftlog_alphas_imag_nn
 
         self.cosmo_model = 'ede-v2'
+
         self.use_Amod = 0
         self.Amod = 0 
-
+        
+        self.use_pk_z_bins = 0
+        self.pk_z_bins_z1 = 0
+        self.pk_z_bins_z2 = 0
+        self.pk_z_bins_A0 = 0
+        self.pk_z_bins_A1 = 0
+        self.pk_z_bins_A2 = 0
         
 
         if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
@@ -223,6 +231,15 @@ class Class_szfast(object):
                 self.use_Amod = v
                 self.Amod  = params_settings['Amod']
 
+            if k == 'use_pk_z_bins':
+                self.use_pk_z_bins = v
+                self.pk_z_bins_z1 = params_settings['pk_z_bins_z1']
+                self.pk_z_bins_z2 = params_settings['pk_z_bins_z2']
+                self.pk_z_bins_A0 = params_settings['pk_z_bins_A0']
+                self.pk_z_bins_A1 = params_settings['pk_z_bins_A1']
+                self.pk_z_bins_A2 = params_settings['pk_z_bins_A2']
+
+
 
 
         if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
@@ -304,6 +321,11 @@ class Class_szfast(object):
         rz = dA*(1.+z)*params_values['h']
         dVdzdOmega = 2.99792458e8/1.0e5*rz*rz/Ez
         return dVdzdOmega
+    
+    def get_sigma8_and_der(self,params_values_dict=None):
+        params_values = self.get_all_relevant_params(params_values_dict)
+        self.calculate_sigma8_and_der(**params_values)
+        return self.cp_predicted_der
 
     def get_r_delta_of_m_delta_at_z(self,delta,m_delta,z,params_values_dict=None):
         if params_values_dict:
@@ -535,6 +557,19 @@ class Class_szfast(object):
                 pk_ae  = pkl_p + self.Amod*(pknl_p-pkl_p)
                 predicted_pk_spectrum_z.append(pk_ae)
 
+        elif self.use_pk_z_bins:
+            # print('>>> using pk_z_bins')
+            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]
+                if zp < self.pk_z_bins_z1:
+                    pklp = self.pk_z_bins_A0*pkl_p
+                elif zp < self.pk_z_bins_z2:
+                    pklp = self.pk_z_bins_A1*pkl_p
+                else:
+                    pklp = self.pk_z_bins_A2*pkl_p
+                predicted_pk_spectrum_z.append(pklp)
         else:
 
             for zp in z_arr:
@@ -663,7 +698,10 @@ 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_der = self.cp_der_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+        if self.jax_mode:
+            self.cp_predicted_der = self.cp_der_nn[self.cosmo_model].predict(params_dict)
+        else:
+            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

{classy_szfast-0.0.25.post13 → classy_szfast-0.0.25.post15}/classy_szfast/cosmopower_jax.py

@@ -135,10 +135,7 @@ for mp in cosmo_model_list:
     cp_pkl_nn_jax[mp] = CosmoPowerJAX_custom(probe='custom_log',filepath=path_to_emulators +'PK/' + emulator_dict[mp]['PKL'] + '.npz')
     cp_pkl_nn_jax[mp].ten_to_predictions = False
 
-    cp_der_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'derived-parameters/' + emulator_dict[mp]['DER'])
-    
-    # cp_da_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['DAZ'])
-    
+    cp_der_nn_jax[mp] = CosmoPowerJAX_custom(probe='custom_log',filepath=path_to_emulators + 'derived-parameters/' + emulator_dict[mp]['DER'] + '.npz')  
     
     cp_da_nn_jax[mp] = CosmoPowerJAX_custom(probe='custom_log',filepath=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['DAZ'] + '.npz')
     if mp != 'ede-v2':

{classy_szfast-0.0.25.post13 → classy_szfast-0.0.25.post15}/classy_szfast.egg-info/PKG-INFO

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

{classy_szfast-0.0.25.post13 → classy_szfast-0.0.25.post15}/pyproject.toml

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