classy-szfast 0.0.25.post26__py3-none-any.whl → 0.0.25.post27__py3-none-any.whl

classy_szfast/classy_szfast.py

@@ -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
@@ -153,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
@@ -779,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)
 
@@ -790,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
@@ -877,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/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.post26.dist-info → classy_szfast-0.0.25.post27.dist-info}/METADATA

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

@@ -1,9 +1,9 @@
 classy_szfast/__init__.py,sha256=8XVvg_q0kf97P_glT4o-6Jnss4AG4RgHat4CEslWDKs,305
 classy_szfast/classy_sz.py,sha256=SCzZYCUSR0c-DAGheX-OlgzzOxOc24d5T4b-ezdhubU,25512
-classy_szfast/classy_szfast.py,sha256=MgxgFUvI77SOVvgdHlUtm2YFSEwzOmXqEPf1FxiCTes,45332
+classy_szfast/classy_szfast.py,sha256=x-svNHmiyhsfH3p0MAoCe86fs0EghqQpLkqIJLyWGp8,45776
 classy_szfast/config.py,sha256=v6DGcBHmfn5JtuO48dKyXCh-Dmn0uwOF_izvVOJFnqw,279
 classy_szfast/cosmopower.py,sha256=ooYK2BDOZSo3XtGHfPtjXHxr5UW-yVngLPkb5gpvTx8,2351
-classy_szfast/cosmopower_jax.py,sha256=aw-DcTVgBIOxz43tkt90gPzhXi4LbEx29YC-P2vfTvI,6493
+classy_szfast/cosmopower_jax.py,sha256=CWrfXG0W4FQ2FbzRz7X7ZZoh6FbDT56XU0g2QFunna4,6680
 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.post26.dist-info/METADATA,sha256=hOT3xPzZ6hBHdIiM9Xv8yK7YY0K_H5vPwLcbJuY3Vd0,579
-classy_szfast-0.0.25.post26.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-classy_szfast-0.0.25.post26.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.25.post26.dist-info/RECORD,,
+classy_szfast-0.0.25.post27.dist-info/METADATA,sha256=FUM3Nf3QS2nf9c4CM4FW9vwCAXQl1OBJ6j8K6XrGwIg,579
+classy_szfast-0.0.25.post27.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+classy_szfast-0.0.25.post27.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
+classy_szfast-0.0.25.post27.dist-info/RECORD,,

{classy_szfast-0.0.25.post26.dist-info → classy_szfast-0.0.25.post27.dist-info}/WHEEL

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