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

classy_szfast/classy_szfast.py

@@ -173,9 +173,16 @@ class Class_szfast(object):
             self.cp_pkl_fftlog_alphas_imag_nn = cp_pkl_fftlog_alphas_imag_nn
 
         self.cosmo_model = 'ede-v2'
+
         self.use_Amod = 0
         self.Amod = 0 
-
+        
+        self.use_pk_z_bins = 0
+        self.pk_z_bins_z1 = 0
+        self.pk_z_bins_z2 = 0
+        self.pk_z_bins_A0 = 0
+        self.pk_z_bins_A1 = 0
+        self.pk_z_bins_A2 = 0
         
 
         if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
@@ -223,6 +230,15 @@ class Class_szfast(object):
                 self.use_Amod = v
                 self.Amod  = params_settings['Amod']
 
+            if k == 'use_pk_z_bins':
+                self.use_pk_z_bins = v
+                self.pk_z_bins_z1 = params_settings['pk_z_bins_z1']
+                self.pk_z_bins_z2 = params_settings['pk_z_bins_z2']
+                self.pk_z_bins_A0 = params_settings['pk_z_bins_A0']
+                self.pk_z_bins_A1 = params_settings['pk_z_bins_A1']
+                self.pk_z_bins_A2 = params_settings['pk_z_bins_A2']
+
+
 
 
         if cosmo_model_dict[params_settings['cosmo_model']] == 'ede-v2':
@@ -535,6 +551,19 @@ class Class_szfast(object):
                 pk_ae  = pkl_p + self.Amod*(pknl_p-pkl_p)
                 predicted_pk_spectrum_z.append(pk_ae)
 
+        elif self.use_pk_z_bins:
+            # print('>>> using pk_z_bins')
+            for zp in z_arr:
+                params_dict_pp = params_dict.copy()
+                params_dict_pp['z_pk_save_nonclass'] = [zp]
+                pkl_p = self.cp_pkl_nn[self.cosmo_model].predictions_np(params_dict_pp)[0]
+                if zp < self.pk_z_bins_z1:
+                    pklp = self.pk_z_bins_A0*pkl_p
+                elif zp < self.pk_z_bins_z2:
+                    pklp = self.pk_z_bins_A1*pkl_p
+                else:
+                    pklp = self.pk_z_bins_A2*pkl_p
+                predicted_pk_spectrum_z.append(pklp)
         else:
 
             for zp in z_arr:

classy_szfast/cosmopower_jax.py

@@ -7,6 +7,7 @@ from .suppress_warnings import suppress_warnings
 from .emulators_meta_data import *
 
 from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ
+from jax.errors import TracerArrayConversionError
 
 
 cp_tt_nn_jax = {}
@@ -28,6 +29,32 @@ class CosmoPowerJAX_custom(CPJ):
         if 'ten_to_predictions' in kwargs.keys():
             self.ten_to_predictions = kwargs['ten_to_predictions']
 
+    def _dict_to_ordered_arr_np(self,
+                               input_dict,
+                               ):
+        """
+        Sort input parameters. Takend verbatim from CP 
+        (https://github.com/alessiospuriomancini/cosmopower/blob/main/cosmopower/cosmopower_NN.py#LL291C1-L308C73)
+
+        Parameters:
+            input_dict (dict [numpy.ndarray]):
+                input dict of (arrays of) parameters to be sorted
+
+        Returns:
+            numpy.ndarray:
+                parameters sorted according to desired order
+        """
+        if self.parameters is not None:
+            try:
+                return np.stack([input_dict[k] for k in self.parameters], axis=1)
+            except TracerArrayConversionError:
+                converted_dict = {k: jnp.array(v) if isinstance(v, list) else v for k, v in input_dict.items()}
+                return jnp.stack([converted_dict[k] for k in self.parameters], axis=1)
+
+        else:
+            return np.stack([input_dict[k] for k in input_dict], axis=1)
+
+
     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.

{classy_szfast-0.0.25.post12.dist-info → classy_szfast-0.0.25.post14.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: classy_szfast
-Version: 0.0.25.post12
+Version: 0.0.25.post14
 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.post12.dist-info → classy_szfast-0.0.25.post14.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=18p1wwBLFKTSXh2iCHoSSVD5NNwsFMyrOOrpbnsBcyk,42915
+classy_szfast/classy_szfast.py,sha256=T2I-exMW1ZYjruf_sKFHchvfZ1inXp0SV5-JJ4GkoLs,44147
 classy_szfast/config.py,sha256=v6DGcBHmfn5JtuO48dKyXCh-Dmn0uwOF_izvVOJFnqw,279
 classy_szfast/cosmopower.py,sha256=ooYK2BDOZSo3XtGHfPtjXHxr5UW-yVngLPkb5gpvTx8,2351
-classy_szfast/cosmopower_jax.py,sha256=NqU_Sw5x0qw16DNHrN4ipCZP8HVdQvUiCq_Y04h8zoo,5468
+classy_szfast/cosmopower_jax.py,sha256=wMKTu3E5OaOR7NX6aPkeUv4ow13xDOcE_Nkkiy-S12c,6543
 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.post12.dist-info/METADATA,sha256=xvy4-nx-2tsgs6ol1CPjWoMyOsBli_PSBR0YSqC088U,572
-classy_szfast-0.0.25.post12.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-classy_szfast-0.0.25.post12.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.25.post12.dist-info/RECORD,,
+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.post12.dist-info → classy_szfast-0.0.25.post14.dist-info}/WHEEL

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