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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post2
+Version: 0.0.25.post3
 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.post2 → classy_szfast-0.0.25.post3}/classy_szfast/classy_szfast.py

@@ -92,12 +92,14 @@ class Class_szfast(object):
         self.cp_pknl_nn  = cp_pknl_nn
         self.cp_pkl_nn = cp_pkl_nn
         self.cp_der_nn = cp_der_nn
-        self.cp_da_nn = cp_da_nn
+       
         
         if self.jax_mode:
             self.cp_h_nn = cp_h_nn_jax
+            self.cp_da_nn = cp_da_nn_jax
         else:
             self.cp_h_nn = cp_h_nn
+            self.cp_da_nn = cp_da_nn
 
         self.cp_s8_nn = cp_s8_nn
 
@@ -785,35 +787,50 @@ class Class_szfast(object):
         update_params_with_defaults(params_values, self.emulator_dict[self.cosmo_model]['default'])
 
         params_dict = {}
-
         for k,v in zip(params_values.keys(),params_values.values()):
-        
             params_dict[k]=[v]
 
         if 'm_ncdm' in params_dict.keys():
             if isinstance(params_dict['m_ncdm'][0],str): 
                 params_dict['m_ncdm'] =  [float(params_dict['m_ncdm'][0].split(',')[0])]
 
-        # deal with different scaling of DA in different model from emulator training
-        if self.cosmo_model == 'ede-v2':
+        if self.jax_mode: 
+            # print("JAX MODE in chi")
+            # self.cp_da_nn[self.cosmo_model].log = False
+            self.cp_predicted_da = self.cp_da_nn[self.cosmo_model].predict(params_dict)
+            # print("self.cp_predicted_da",self.cp_predicted_da)
+
+            if self.cosmo_model == 'ede-v2':
+                # print('ede-v2 case')
+                self.cp_predicted_da = jnp.insert(self.cp_predicted_da, 0, 0)
+            self.cp_predicted_da *= (1.+self.cp_z_interp_jax)
+    
+            def chi_interp(x):
+                return jnp.interp(x, self.cp_z_interp_jax, self.cp_predicted_da, left=jnp.nan, right=jnp.nan)
+
+            self.chi_interp = chi_interp
 
-            self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
-            self.cp_predicted_da = np.insert(self.cp_predicted_da, 0, 0)
-        
         else:
-        
-            self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].predictions_np(params_dict)[0]
-        
+            # deal with different scaling of DA in different model from emulator training
+            if self.cosmo_model == 'ede-v2':
 
-        self.chi_interp = scipy.interpolate.interp1d(
-                                                    self.cp_z_interp,
-                                                    self.cp_predicted_da*(1.+self.cp_z_interp),
-                                                    kind='linear',
-                                                    axis=-1,
-                                                    copy=True,
-                                                    bounds_error=None,
-                                                    fill_value=np.nan,
-                                                    assume_sorted=False)
+                self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
+                self.cp_predicted_da = np.insert(self.cp_predicted_da, 0, 0)
+            
+            else:
+            
+                self.cp_predicted_da  = self.cp_da_nn[self.cosmo_model].predictions_np(params_dict)[0]
+            
+
+            self.chi_interp = scipy.interpolate.interp1d(
+                                                        self.cp_z_interp,
+                                                        self.cp_predicted_da*(1.+self.cp_z_interp),
+                                                        kind='linear',
+                                                        axis=-1,
+                                                        copy=True,
+                                                        bounds_error=None,
+                                                        fill_value=np.nan,
+                                                        assume_sorted=False)
 
     def get_cmb_cls(self,ell_factor=True,Tcmb_uk = Tcmb_uk):
 

classy_szfast-0.0.25.post3/classy_szfast/cosmopower_jax.py

@@ -0,0 +1,127 @@
+from .config import path_to_class_sz_data
+import numpy as np
+import jax.numpy as jnp
+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 = {}
+
+
+class CosmoPowerJAX_custom(CPJ):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ten_to_predictions = True
+        if 'ten_to_predictions' in kwargs.keys():
+            self.ten_to_predictions = kwargs['ten_to_predictions']
+
+    def _predict(self, weights, hyper_params, param_train_mean, param_train_std,
+                 feature_train_mean, feature_train_std, input_vec):
+        """ Forward pass through pre-trained network.
+        In its current form, it does not make use of high-level frameworks like
+        FLAX et similia; rather, it simply loops over the network layers.
+        In future work this can be improved, especially if speed is a problem.
+        
+        Parameters
+        ----------
+        weights : array
+            The stored weights of the neural network.
+        hyper_params : array
+            The stored hyperparameters of the activation function for each layer.
+        param_train_mean : array
+            The stored mean of the training cosmological parameters.
+        param_train_std : array
+            The stored standard deviation of the training cosmological parameters.
+        feature_train_mean : array
+            The stored mean of the training features.
+        feature_train_std : array
+            The stored  standard deviation of the training features.
+        input_vec : array of shape (n_samples, n_parameters) or (n_parameters)
+            The cosmological parameters given as input to the network.
+            
+        Returns
+        -------
+        predictions : array
+            The prediction of the trained neural network.
+        """        
+        act = []
+        # Standardise
+        layer_out = [(input_vec - param_train_mean)/param_train_std]
+
+        # Loop over layers
+        for i in range(len(weights[:-1])):
+            w, b = weights[i]
+            alpha, beta = hyper_params[i]
+            act.append(jnp.dot(layer_out[-1], w.T) + b)
+            layer_out.append(self._activation(act[-1], alpha, beta))
+
+        # Final layer prediction (no activations)
+        w, b = weights[-1]
+        if self.probe == 'custom_log' or self.probe == 'custom_pca':
+            # in original CP models, we assumed a full final bias vector...
+            preds = jnp.dot(layer_out[-1], w.T) + b
+        else:   
+            # ... unlike in cpjax, where we used only a single bias vector
+            preds = jnp.dot(layer_out[-1], w.T) + b[-1]
+
+        # Undo the standardisation
+        preds = preds * feature_train_std + feature_train_mean
+        if self.log == True:
+            if self.ten_to_predictions:
+                preds = 10**preds
+        else:
+            preds = (preds@self.pca_matrix)*self.training_std + self.training_mean
+            if self.probe == 'cmb_pp':
+                preds = 10**preds
+        predictions = preds.squeeze()
+        return predictions
+
+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'])
+    
+    
+    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':
+        cp_da_nn_jax[mp].ten_to_predictions = False
+    # print(cp_da_nn_jax[mp].parameters)
+    # 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] = CosmoPowerJAX_custom(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.post2 → classy_szfast-0.0.25.post3}/classy_szfast.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post2
+Version: 0.0.25.post3
 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.post2 → classy_szfast-0.0.25.post3}/pyproject.toml

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

classy_szfast-0.0.25.post2/classy_szfast/cosmopower_jax.py

@@ -1,53 +0,0 @@
-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'])
-