classy-szfast 0.0.25__tar.gz → 0.0.25.post2__tar.gz

{classy_szfast-0.0.25 → classy_szfast-0.0.25.post2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25
+Version: 0.0.25.post2
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
 License: MIT
@@ -13,3 +13,4 @@ Requires-Dist: tensorflow
 Requires-Dist: mcfit
 Requires-Dist: get_cosmopower_emus
 Requires-Dist: class_sz_data
+Requires-Dist: cosmopower-jax

{classy_szfast-0.0.25 → classy_szfast-0.0.25.post2}/classy_szfast/classy_szfast.py

@@ -1,7 +1,9 @@
 from .utils import *
 from .config import *
 import numpy as np
-from .cosmopower import *
+from .emulators_meta_data import emulator_dict, dofftlog_alphas, cp_l_max_scalars
+from .cosmopower import cp_tt_nn, cp_te_nn, cp_ee_nn, cp_ee_nn, cp_pp_nn, cp_pknl_nn, cp_pkl_nn, cp_der_nn, cp_da_nn, cp_h_nn, cp_s8_nn, cp_pkl_fftlog_alphas_real_nn, cp_pkl_fftlog_alphas_imag_nn, cp_pkl_fftlog_alphas_nus
+from .cosmopower_jax import cp_tt_nn_jax, cp_te_nn_jax, cp_ee_nn_jax, cp_ee_nn_jax, cp_pp_nn_jax, cp_pknl_nn_jax, cp_pkl_nn_jax, cp_der_nn_jax, cp_da_nn_jax, cp_h_nn_jax, cp_s8_nn_jax
 from .pks_and_sigmas import *
 import scipy
 import time
@@ -9,7 +11,8 @@ from multiprocessing import Process
 from mcfit import TophatVar
 from scipy.interpolate import CubicSpline
 import pickle
-
+import jax.numpy as jnp
+import jax.scipy as jscipy
 
 H_units_conv_factor = {"1/Mpc": 1, "km/s/Mpc": Const.c_km_s}
 
@@ -69,6 +72,11 @@ class Class_szfast(object):
         except:
             pass
         self.logger = logging.getLogger(__name__)
+
+
+        self.jax_mode = params_settings["jax"]
+
+        # print(f"JAX mode: {self.jax_mode}")
         
 
         
@@ -85,7 +93,12 @@ class Class_szfast(object):
         self.cp_pkl_nn = cp_pkl_nn
         self.cp_der_nn = cp_der_nn
         self.cp_da_nn = cp_da_nn
-        self.cp_h_nn = cp_h_nn
+        
+        if self.jax_mode:
+            self.cp_h_nn = cp_h_nn_jax
+        else:
+            self.cp_h_nn = cp_h_nn
+
         self.cp_s8_nn = cp_s8_nn
 
         self.emulator_dict = emulator_dict
@@ -203,6 +216,7 @@ class Class_szfast(object):
 
 
         self.cp_z_interp = np.linspace(0.,20.,5000)
+        self.cp_z_interp_jax = jnp.linspace(0.,20.,5000)
 
         self.csz_base = None
 
@@ -409,6 +423,11 @@ class Class_szfast(object):
 
         k_arr = self.cszfast_pk_grid_k 
 
+        # print(">>> z_arr:",z_arr)
+        # print(">>> k_arr:",k_arr)
+        # import sys
+        
+
 
         params_values = params_values_dict.copy()
         update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
@@ -445,6 +464,11 @@ class Class_szfast(object):
                 params_dict_pp['z_pk_save_nonclass'] = [zp]
                 predicted_pk_spectrum_z.append(self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0])
 
+                # if abs(zp-0.5) < 0.01:
+                #   print(">>> predicted_pk_spectrum_z:",predicted_pk_spectrum_z[-1])
+                #   import pprint
+                #   pprint.pprint(params_dict_pp)
+       
         predicted_pk_spectrum = np.asarray(predicted_pk_spectrum_z)
 
 
@@ -453,6 +477,10 @@ class Class_szfast(object):
         pk_re = pk*self.pk_power_fac
         pk_re = np.transpose(pk_re)
 
+        # print(">>> pk_re:",pk_re)
+        # import sys
+        # sys.exit(0)
+
         self.pkl_interp = PowerSpectrumInterpolator(z_arr,k_arr,np.log(pk_re).T,logP=True)
 
         self.cszfast_pk_grid_pk = pk_re
@@ -708,17 +736,46 @@ 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]
- 
-        self.hz_interp = scipy.interpolate.interp1d(
-                                            self.cp_z_interp,
-                                            self.cp_predicted_hubble,
-                                            kind='linear',
-                                            axis=-1,
-                                            copy=True,
-                                            bounds_error=None,
-                                            fill_value=np.nan,
-                                            assume_sorted=False)
+
+        if self.jax_mode:
+            # print("JAX MODE in hubble")
+            # self.cp_predicted_hubble = self.cp_h_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+            # print(params_dict)
+            # print("params_dict type:", type(params_dict))
+            self.cp_predicted_hubble = self.cp_h_nn[self.cosmo_model].predict(params_dict)
+            # print("self.cp_predicted_hubble type:", type(self.cp_predicted_hubble))
+            # print("self.cp_predicted_hubble",self.cp_predicted_hubble)
+
+            # self.hz_interp = jscipy.interpolate.interp1d(
+            #                                     self.cp_z_interp_jax,
+            #                                     self.cp_predicted_hubble,
+            #                                     kind='linear',
+            #                                     axis=-1,
+            #                                     copy=True,
+            #                                     bounds_error=None,
+            #                                     fill_value=np.nan,
+            #                                     assume_sorted=False)
+            
+            # Assuming `cp_z_interp` and `cp_predicted_hubble` are JAX arrays
+            def hz_interp(x):
+                return jnp.interp(x, self.cp_z_interp_jax, self.cp_predicted_hubble, left=jnp.nan, right=jnp.nan)
+
+            self.hz_interp = hz_interp
+            # exit()
+        else:
+            self.cp_predicted_hubble = self.cp_h_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+            # print("self.cp_predicted_hubble",self.cp_predicted_hubble)
+            
+
+            self.hz_interp = scipy.interpolate.interp1d(
+                                                self.cp_z_interp,
+                                                self.cp_predicted_hubble,
+                                                kind='linear',
+                                                axis=-1,
+                                                copy=True,
+                                                bounds_error=None,
+                                                fill_value=np.nan,
+                                                assume_sorted=False)
 
     def calculate_chi(self,
                       **params_values_dict):
@@ -819,7 +876,10 @@ class Class_szfast(object):
 
 
     def get_hubble(self, z,units="1/Mpc"):
-        return np.array(self.hz_interp(z)*H_units_conv_factor[units])
+        if self.jax_mode:
+            return jnp.array(self.hz_interp(z)*H_units_conv_factor[units])
+        else:
+            return np.array(self.hz_interp(z)*H_units_conv_factor[units])
 
     def get_chi(self, z):
         return np.array(self.chi_interp(z))

classy_szfast-0.0.25.post2/classy_szfast/cosmopower.py

@@ -0,0 +1,59 @@
+from .config import path_to_class_sz_data
+import numpy as np
+from .restore_nn import Restore_NN
+from .restore_nn import Restore_PCAplusNN
+from .suppress_warnings import suppress_warnings
+from .emulators_meta_data import *
+
+
+cp_tt_nn = {}
+cp_te_nn = {}
+cp_ee_nn = {}
+cp_pp_nn = {}
+cp_pknl_nn = {}
+cp_pkl_nn = {}
+cp_pkl_fftlog_alphas_real_nn = {}
+cp_pkl_fftlog_alphas_imag_nn = {}
+cp_pkl_fftlog_alphas_nus = {}
+cp_der_nn = {}
+cp_da_nn = {}
+cp_h_nn = {}
+cp_s8_nn = {}
+
+
+for mp in cosmo_model_list:
+    folder, version = split_emulator_string(mp)
+    # print(folder, version)
+    path_to_emulators = path_to_class_sz_data + '/' + folder +'/'
+    
+    cp_tt_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TT'])
+    
+    cp_te_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TE'])
+    
+    with suppress_warnings():
+        cp_ee_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['EE'])
+    
+    cp_pp_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PP/' + emulator_dict[mp]['PP'])
+    
+    cp_pknl_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKNL'])
+    
+    cp_pkl_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKL'])
+
+    if (mp == 'lcdm') and (dofftlog_alphas == True):
+        cp_pkl_fftlog_alphas_real_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKLFFTLOG_ALPHAS_REAL']
+                                 )
+        cp_pkl_fftlog_alphas_imag_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKLFFTLOG_ALPHAS_IMAG']
+                                 )
+        cp_pkl_fftlog_alphas_nus[mp] = np.load(path_to_emulators + 'PK/PKL_FFTLog_alphas_nu_v1.npz')
+    
+    cp_der_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'derived-parameters/' + emulator_dict[mp]['DER'])
+    
+    cp_da_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['DAZ'])
+    
+    cp_h_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['HZ'])
+    
+    cp_s8_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['S8Z'])
+    
+
+
+

classy_szfast-0.0.25.post2/classy_szfast/cosmopower_jax.py

@@ -0,0 +1,53 @@
+from .config import path_to_class_sz_data
+import numpy as np
+from .restore_nn import Restore_NN
+from .restore_nn import Restore_PCAplusNN
+from .suppress_warnings import suppress_warnings
+from .emulators_meta_data import *
+
+from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ
+
+
+cp_tt_nn_jax = {}
+cp_te_nn_jax = {}
+cp_ee_nn_jax = {}
+cp_pp_nn_jax = {}
+cp_pknl_nn_jax = {}
+cp_pkl_nn_jax = {}
+cp_der_nn_jax = {}
+cp_da_nn_jax = {}
+cp_h_nn_jax = {}
+cp_s8_nn_jax = {}
+
+
+for mp in cosmo_model_list:
+    folder, version = split_emulator_string(mp)
+    # print(folder, version)
+    path_to_emulators = path_to_class_sz_data + '/' + folder +'/'
+    
+    cp_tt_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TT'])
+    
+    cp_te_nn_jax[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TE'])
+    
+    with suppress_warnings():
+        cp_ee_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['EE'])
+    
+    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_pkl_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKL'])
+    
+    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'])
+    
+    # print(path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['HZ'])
+    emulator_custom = CPJ(probe='custom_log',filepath=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['HZ'] + '.npz')
+    # print(emulator_custom.parameters)
+    # exit()
+
+    cp_h_nn_jax[mp] = CPJ(probe='custom_log',filepath=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['HZ'] + '.npz')
+
+    cp_s8_nn_jax[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['S8Z'])
+    

classy_szfast-0.0.25/classy_szfast/cosmopower.py → classy_szfast-0.0.25.post2/classy_szfast/emulators_meta_data.py

@@ -1,9 +1,20 @@
-from .utils import *
-from .config import *
-
-from .restore_nn import Restore_NN
-from .restore_nn import Restore_PCAplusNN
 from .suppress_warnings import suppress_warnings
+import warnings
+from contextlib import contextmanager
+import logging
+
+# Suppress absl warnings
+import absl.logging
+absl.logging.set_verbosity('error')
+# Suppress TensorFlow warnings
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+with suppress_warnings():
+    import tensorflow as tf
+    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+
+import re
 
 dofftlog_alphas = False
 
@@ -215,36 +226,7 @@ emulator_dict['ede-v2']['default']['N_ncdm'] = 3
 emulator_dict['ede-v2']['default']['m_ncdm'] = 0.02
 
 
-cp_tt_nn = {}
-cp_te_nn = {}
-cp_ee_nn = {}
-cp_pp_nn = {}
-cp_pknl_nn = {}
-cp_pkl_nn = {}
-cp_pkl_fftlog_alphas_real_nn = {}
-cp_pkl_fftlog_alphas_imag_nn = {}
-cp_pkl_fftlog_alphas_nus = {}
-cp_der_nn = {}
-cp_da_nn = {}
-cp_h_nn = {}
-cp_s8_nn = {}
 
-import warnings
-from contextlib import contextmanager
-import logging
-
-# Suppress absl warnings
-import absl.logging
-absl.logging.set_verbosity('error')
-# Suppress TensorFlow warnings
-import os
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-with suppress_warnings():
-    import tensorflow as tf
-    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-
-
-import re
 
 def split_emulator_string(input_string):
     match = re.match(r"(.+)-v(\d+)", input_string)
@@ -258,41 +240,3 @@ def split_emulator_string(input_string):
         return folder, version
 
 
-
-
-for mp in cosmo_model_list:
-    folder, version = split_emulator_string(mp)
-    # print(folder, version)
-    path_to_emulators = path_to_class_sz_data + '/' + folder +'/'
-    
-    cp_tt_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TT'])
-    
-    cp_te_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['TE'])
-    
-    with suppress_warnings():
-        cp_ee_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'TTTEEE/' + emulator_dict[mp]['EE'])
-    
-    cp_pp_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PP/' + emulator_dict[mp]['PP'])
-    
-    cp_pknl_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKNL'])
-    
-    cp_pkl_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKL'])
-
-    if (mp == 'lcdm') and (dofftlog_alphas == True):
-        cp_pkl_fftlog_alphas_real_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKLFFTLOG_ALPHAS_REAL']
-                                 )
-        cp_pkl_fftlog_alphas_imag_nn[mp] = Restore_PCAplusNN(restore_filename=path_to_emulators + 'PK/' + emulator_dict[mp]['PKLFFTLOG_ALPHAS_IMAG']
-                                 )
-        cp_pkl_fftlog_alphas_nus[mp] = np.load(path_to_emulators + 'PK/PKL_FFTLog_alphas_nu_v1.npz')
-    
-    cp_der_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'derived-parameters/' + emulator_dict[mp]['DER'])
-    
-    cp_da_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['DAZ'])
-    
-    cp_h_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['HZ'])
-    
-    cp_s8_nn[mp] = Restore_NN(restore_filename=path_to_emulators + 'growth-and-distances/' + emulator_dict[mp]['S8Z'])
-    
-
-
-

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25
+Version: 0.0.25.post2
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
 License: MIT
@@ -13,3 +13,4 @@ Requires-Dist: tensorflow
 Requires-Dist: mcfit
 Requires-Dist: get_cosmopower_emus
 Requires-Dist: class_sz_data
+Requires-Dist: cosmopower-jax

{classy_szfast-0.0.25 → classy_szfast-0.0.25.post2}/classy_szfast.egg-info/SOURCES.txt

@@ -5,7 +5,9 @@ classy_szfast/classy_sz.py
 classy_szfast/classy_szfast.py
 classy_szfast/config.py
 classy_szfast/cosmopower.py
+classy_szfast/cosmopower_jax.py
 classy_szfast/cosmosis_classy_szfast_interface.py
+classy_szfast/emulators_meta_data.py
 classy_szfast/pks_and_sigmas.py
 classy_szfast/restore_nn.py
 classy_szfast/suppress_warnings.py

{classy_szfast-0.0.25 → classy_szfast-0.0.25.post2}/classy_szfast.egg-info/requires.txt

@@ -4,3 +4,4 @@ tensorflow
 mcfit
 get_cosmopower_emus
 class_sz_data
+cosmopower-jax

{classy_szfast-0.0.25 → classy_szfast-0.0.25.post2}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
-version = "0.0.25"
+version = "0.0.25.post2"
 license = { text = "MIT" }
 name = "classy_szfast"
 maintainers = [{name = "Boris Bolliet",email="bb667@cam.ac.uk"}]
@@ -15,7 +15,8 @@ dependencies = [
     "tensorflow",
     "mcfit",
     "get_cosmopower_emus",
-    "class_sz_data"
+    "class_sz_data",
+    "cosmopower-jax"
 ]
 
 [project.urls]