classy-szfast 0.0.25.post20__tar.gz → 0.0.25.post22__tar.gz

{classy_szfast-0.0.25.post20 → classy_szfast-0.0.25.post22}/PKG-INFO

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

@@ -13,7 +13,7 @@ from mcfit import TophatVar
 from scipy.interpolate import CubicSpline
 import pickle
 import jax 
-jax.config.update("jax_enable_x64", False)
+jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 import jax.scipy as jscipy
 
@@ -131,7 +131,16 @@ class Class_szfast(object):
             self.pow = jnp.power
 
             self.sigma_B = 2. * self.pow(self.pi,5) * self.pow(Const._k_B_,4) / 15. / self.pow(Const._h_P_,3) / self.pow(Const._c_,2)
-
+            # print('sigma_B',self.sigma_B)
+            # print('pi',self.pi)
+            # print('k_B',Const._k_B_)
+            # print('h_P',Const._h_P_)
+            # print('c',Const._c_)
+            # print('pow(pi,5)',self.pow(self.pi,5)) 
+            # print('pow(k_B,4)',self.pow(Const._k_B_,4)) ## this doesnt work if float64 is disabled
+            # print('15.',15.)
+            # print('pow(h_P,3)',self.pow(Const._h_P_,3)) ## this doesnt work if float64 is disabled
+            # print('pow(c,2)',self.pow(Const._c_,2))
             self.linspace = jnp.linspace
             self.geomspace = jnp.geomspace
             self.arange = jnp.arange
@@ -706,6 +715,8 @@ class Class_szfast(object):
             self.cp_predicted_der = self.cp_der_nn[self.cosmo_model].ten_to_predictions_np(params_dict)[0]
         self.sigma8 = self.cp_predicted_der[1]
         self.Neff = self.cp_predicted_der[4]
+        self.ra_star = self.cp_predicted_der[12]
+
         return 0
 
 
@@ -830,7 +841,7 @@ class Class_szfast(object):
         pk_re = self.transpose(pk_re)
 
         
-        return pk_re, k_arr
+        return pk_re.flatten(), k_arr
     
 
     def calculate_pknl_at_z(self,
@@ -878,7 +889,7 @@ class Class_szfast(object):
         pk_re = self.transpose(pk_re)
 
         
-        return pk_re, k_arr
+        return pk_re.flatten(), k_arr
     
 
     def calculate_hubble(self,

{classy_szfast-0.0.25.post20 → classy_szfast-0.0.25.post22}/classy_szfast/cosmopower_jax.py

@@ -1,7 +1,7 @@
 from .config import path_to_class_sz_data
 import numpy as np
 import jax
-jax.config.update("jax_enable_x64", False)
+jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 from .restore_nn import Restore_NN
 from .restore_nn import Restore_PCAplusNN

{classy_szfast-0.0.25.post20 → classy_szfast-0.0.25.post22}/classy_szfast/utils.py

@@ -14,7 +14,7 @@ from numpy import linalg as LA
 import mcfit
 from mcfit import P2xi
 import jax
-jax.config.update("jax_enable_x64", False)
+jax.config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 # import cosmopower
 # import classy_sz as csz

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

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

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