classy-szfast 0.0.25.post25__tar.gz → 0.0.25.post27__tar.gz

{classy_szfast-0.0.25.post25 → classy_szfast-0.0.25.post27}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: classy_szfast
-Version: 0.0.25.post25
+Version: 0.0.25.post27
 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.post25 → classy_szfast-0.0.25.post27}/classy_szfast/classy_szfast.py

@@ -13,7 +13,7 @@ from mcfit import TophatVar
 from scipy.interpolate import CubicSpline
 import pickle
 import jax 
-jax.config.update("jax_enable_x64", False)
+jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 import jax.scipy as jscipy
 
@@ -122,6 +122,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_pknl_nn = cp_pknl_nn_jax
             self.cp_der_nn = cp_der_nn_jax
 
             self.pi = jnp.pi
@@ -130,12 +131,13 @@ class Class_szfast(object):
             self.log = jnp.log
             self.pow = jnp.power
 
-            self.sigma_B = 2. * self.pow(self.pi,5) * self.pow(Const._k_B_,4) / 15. / self.pow(Const._h_P_,3) / self.pow(Const._c_,2)
+            # self.sigma_B = 2. * self.pow(self.pi,5) * self.pow(Const._k_B_,4) / 15. / self.pow(Const._h_P_,3) / self.pow(Const._c_,2)
+            self.sigma_B = 5.6704004737209545e-08
             # print('sigma_B',self.sigma_B)
             # print('pi',self.pi)
             # print('k_B',Const._k_B_)
             # print('h_P',Const._h_P_)
-            # print('c',Const._c_)
+            # # print('c',Const._c_)
             # print('pow(pi,5)',self.pow(self.pi,5)) 
             # print('pow(k_B,4)',self.pow(Const._k_B_,4)) ## this doesnt work if float64 is disabled
             # print('15.',15.)
@@ -152,6 +154,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_pknl_nn = cp_pknl_nn
             self.cp_der_nn = cp_der_nn
             self.pi = np.pi
             self.transpose = np.transpose
@@ -778,7 +781,10 @@ class Class_szfast(object):
         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_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
+            if self.jax_mode:
+                predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predict(params_dict_pp))
+            else:
+                predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
         predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 
@@ -789,7 +795,10 @@ class Class_szfast(object):
         pk_re = self.transpose(pk_re)
 
 
-        self.pknl_interp = PowerSpectrumInterpolator(z_arr,k_arr,self.log(pk_re).T,logP=True)
+        if self.jax_mode:
+            self.pknl_interp = None
+        else:
+            self.pknl_interp = PowerSpectrumInterpolator(z_arr,k_arr,self.log(pk_re).T,logP=True)
 
 
         self.cszfast_pk_grid_pknl = pk_re
@@ -876,12 +885,14 @@ class Class_szfast(object):
 
         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])
+        if self.jax_mode:
+            predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predict(params_dict_pp))
+        else:
+            predicted_pk_spectrum_z.append(self.cp_pknl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
         predicted_pk_spectrum = self.asarray(predicted_pk_spectrum_z)
 

{classy_szfast-0.0.25.post25 → classy_szfast-0.0.25.post27}/classy_szfast/cosmopower_jax.py

@@ -132,8 +132,12 @@ for mp in cosmo_model_list:
     
     cp_pp_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PP/' + emulator_dict[mp]['PP'])
     
-    cp_pknl_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKNL'])
-    
+    # cp_pknl_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKNL'])
+
+    cp_pknl_nn_jax[mp] = CosmoPowerJAX_custom(probe='custom_log',filepath=path_to_emulators +'PK/' + emulator_dict[mp]['PKNL'] + '.npz')
+    cp_pknl_nn_jax[mp].ten_to_predictions = False
+
+
     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
 

{classy_szfast-0.0.25.post25 → classy_szfast-0.0.25.post27}/classy_szfast.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: classy_szfast
-Version: 0.0.25.post25
+Version: 0.0.25.post27
 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.post25 → classy_szfast-0.0.25.post27}/pyproject.toml

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