classy-szfast 0.0.25.post12__tar.gz → 0.0.25.post13__tar.gz

{classy_szfast-0.0.25.post12 → classy_szfast-0.0.25.post13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post12
+Version: 0.0.25.post13
 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 → classy_szfast-0.0.25.post13}/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 → classy_szfast-0.0.25.post13}/classy_szfast.egg-info/PKG-INFO

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

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