classy-szfast 0.0.25.post14__py3-none-any.whl → 0.0.25.post15__py3-none-any.whl

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
@@ -320,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:
@@ -692,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/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.post14.dist-info → classy_szfast-0.0.25.post15.dist-info}/METADATA

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

@@ -1,9 +1,9 @@
 classy_szfast/__init__.py,sha256=E2thrL0Z9oXFfdzwcsu-xbOytudLFTlRlPqVFGlPPPg,279
 classy_szfast/classy_sz.py,sha256=QmbwrSXInQLMvCDqsr7KPmtaU0KOiOt1Rb-cTKuulZw,22240
-classy_szfast/classy_szfast.py,sha256=T2I-exMW1ZYjruf_sKFHchvfZ1inXp0SV5-JJ4GkoLs,44147
+classy_szfast/classy_szfast.py,sha256=PFb4Mqy3ACTRwEDAgN3m_-2batQhIZ1z2prT8zSGu5U,44556
 classy_szfast/config.py,sha256=v6DGcBHmfn5JtuO48dKyXCh-Dmn0uwOF_izvVOJFnqw,279
 classy_szfast/cosmopower.py,sha256=ooYK2BDOZSo3XtGHfPtjXHxr5UW-yVngLPkb5gpvTx8,2351
-classy_szfast/cosmopower_jax.py,sha256=wMKTu3E5OaOR7NX6aPkeUv4ow13xDOcE_Nkkiy-S12c,6543
+classy_szfast/cosmopower_jax.py,sha256=Xn8-fUvlE_GoJ-AQVB_wCJ1k9D0pzQMgdlzSE3OZbwI,6440
 classy_szfast/cosmosis_classy_szfast_interface.py,sha256=zAnxvFtn73a5yS7jgs59zpWFEYKCIQyraYPs5hQ4Le8,11483
 classy_szfast/emulators_meta_data.py,sha256=mXG5LQuJw9QBNE_kxXW8Kx0AUCWpbV6uRO9BaBbIfHo,10732
 classy_szfast/pks_and_sigmas.py,sha256=drtuujE1HhlrYY1hY92DyY5lXlYS1uE15MSuVI4uo6k,6625
@@ -14,7 +14,7 @@ classy_szfast/custom_bias/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 classy_szfast/custom_bias/custom_bias.py,sha256=aR2t5RTIwv7P0m2bsEU0Eq6BTkj4pG10AebH6QpG4qM,486
 classy_szfast/custom_profiles/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 classy_szfast/custom_profiles/custom_profiles.py,sha256=4LZwb2XoqwCyWNmW2s24Z7AJdmgVdaRG7yYaBYe-d9Q,1188
-classy_szfast-0.0.25.post14.dist-info/METADATA,sha256=XgfBmaQKt3iW1k61OmIBt2JgThMVnX4D0ycfsRCRUDk,572
-classy_szfast-0.0.25.post14.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-classy_szfast-0.0.25.post14.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.25.post14.dist-info/RECORD,,
+classy_szfast-0.0.25.post15.dist-info/METADATA,sha256=OCfStl2y1zmr2u08imrd_xfrR5xPg7MvRW8VwpV7Iro,572
+classy_szfast-0.0.25.post15.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+classy_szfast-0.0.25.post15.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.25.post15.dist-info/RECORD,,