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

classy_szfast/classy_szfast.py

@@ -1,5 +1,5 @@
 from .utils import *
-from .utils import Const
+from .utils import Const, jax_gradient
 from .config import *
 import numpy as np
 from .emulators_meta_data import emulator_dict, dofftlog_alphas, cp_l_max_scalars, cosmo_model_list
@@ -133,7 +133,7 @@ class Class_szfast(object):
             self.geomspace = jnp.geomspace
             self.arange = jnp.arange
             self.zeros = jnp.zeros
-            self.gradient = jnp.gradient
+            self.gradient = jax_gradient
 
 
         else:
@@ -591,8 +591,8 @@ class Class_szfast(object):
 
             R, var[:,iz] = TophatVar(k, lowring=True)(P[:,iz], extrap=True)
 
-            # dvar[:,iz] = self.gradient(var[:,iz], R) ## old form
-            _, dvar[:,iz] = TophatVar(k, lowring=True, deriv=1)(P[:,iz]*k, extrap=True) # new form 
+            dvar[:,iz] = self.gradient(var[:,iz], R) ## old form with jax gradient from inigo zubeldia if needed. 
+            # _, dvar[:,iz] = TophatVar(k, lowring=True, deriv=1)(P[:,iz]*k, extrap=True) # new form doesnt seem to work accurately
 
         # dvar = dvar/(2.*np.sqrt(var))
         # print(dvar_grads/dvar)

classy_szfast/utils.py

@@ -13,6 +13,8 @@ import math
 from numpy import linalg as LA
 import mcfit
 from mcfit import P2xi
+import jax
+import jax.numpy as jnp
 # import cosmopower
 # import classy_sz as csz
 
@@ -71,3 +73,150 @@ class Const:
     _k_B_ = 1.3806504e-23
     _h_P_ = 6.62606896e-34
     _M_sun_ =  1.98855e30 # solar mass in kg
+
+
+
+def jax_gradient(f, *varargs, axis=None, edge_order=1):
+    f = jnp.asarray(f)
+    N = f.ndim  # number of dimensions
+    if axis is None:
+        axes = tuple(range(N))
+    else:
+        axes = jax.numpy._normalize_axis_index(axis, N)
+    len_axes = len(axes)
+    n = len(varargs)
+    if n == 0:
+        # no spacing argument - use 1 in all axes
+        dx = [1.0] * len_axes
+    elif n == 1 and jnp.ndim(varargs[0]) == 0:
+        # single scalar for all axes
+        dx = varargs * len_axes
+    elif n == len_axes:
+        # scalar or 1d array for each axis
+        dx = list(varargs)
+        for i, distances in enumerate(dx):
+            distances = jnp.asarray(distances)
+            if distances.ndim == 0:
+                continue
+            elif distances.ndim != 1:
+                raise ValueError("distances must be either scalars or 1D")
+            if len(distances) != f.shape[axes[i]]:
+                raise ValueError("when 1D, distances must match "
+                                 "the length of the corresponding dimension")
+            if jnp.issubdtype(distances.dtype, jnp.integer):
+                # Convert jax integer types to float64 to avoid modular
+                # arithmetic in np.diff(distances).
+                distances = distances.astype(jnp.float64)
+            diffx = jnp.diff(distances)
+            # if distances are constant reduce to the scalar case
+            # since it brings a consistent speedup
+            if (diffx == diffx[0]).all():
+                diffx = diffx[0]
+            dx[i] = diffx
+    else:
+        raise TypeError("invalid number of arguments")
+    if edge_order > 2:
+        raise ValueError("'edge_order' greater than 2 not supported")
+    outvals = []
+    slice1 = [slice(None)] * N
+    slice2 = [slice(None)] * N
+    slice3 = [slice(None)] * N
+    slice4 = [slice(None)] * N
+    otype = f.dtype
+    # All other types convert to floating point.
+    if jnp.issubdtype(otype, jnp.integer):
+        f = f.astype(jnp.float64)
+    otype = jnp.float64
+    for axis, ax_dx in zip(axes, dx):
+        if f.shape[axis] < edge_order + 1:
+            raise ValueError(
+                "Shape of array too small to calculate a numerical gradient, "
+                "at least (edge_order + 1) elements are required.")
+        # result allocation
+        out = jnp.empty_like(f, dtype=otype)
+        uniform_spacing = jnp.ndim(ax_dx) == 0
+        # Numerical differentiation: 2nd order interior
+        slice1[axis] = slice(1, -1)
+        slice2[axis] = slice(None, -2)
+        slice3[axis] = slice(1, -1)
+        slice4[axis] = slice(2, None)
+        if uniform_spacing:
+            out = out.at[tuple(slice1)].set(
+                (f[tuple(slice4)] - f[tuple(slice2)]) / (2.0 * ax_dx)
+            )
+        else:
+            dx1 = ax_dx[:-1]
+            dx2 = ax_dx[1:]
+            a = -(dx2) / (dx1 * (dx1 + dx2))
+            b = (dx2 - dx1) / (dx1 * dx2)
+            c = dx1 / (dx2 * (dx1 + dx2))
+            shape = [1] * N
+            shape[axis] = -1
+            a = a.reshape(shape)
+            b = b.reshape(shape)
+            c = c.reshape(shape)
+            out = out.at[tuple(slice1)].set(
+                a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
+            )
+        # Numerical differentiation: 1st order edges
+        if edge_order == 1:
+            slice1[axis] = 0
+            slice2[axis] = 1
+            slice3[axis] = 0
+            dx_0 = ax_dx if uniform_spacing else ax_dx[0]
+            out = out.at[tuple(slice1)].set(
+                (f[tuple(slice2)] - f[tuple(slice3)]) / dx_0
+            )
+            slice1[axis] = -1
+            slice2[axis] = -1
+            slice3[axis] = -2
+            dx_n = ax_dx if uniform_spacing else ax_dx[-1]
+            out = out.at[tuple(slice1)].set(
+                (f[tuple(slice2)] - f[tuple(slice3)]) / dx_n
+            )
+        # Numerical differentiation: 2nd order edges
+        else:
+            slice1[axis] = 0
+            slice2[axis] = 0
+            slice3[axis] = 1
+            slice4[axis] = 2
+            if uniform_spacing:
+                a = -1.5 / ax_dx
+                b = 2. / ax_dx
+                c = -0.5 / ax_dx
+            else:
+                dx1 = ax_dx[0]
+                dx2 = ax_dx[1]
+                a = -(2. * dx1 + dx2) / (dx1 * (dx1 + dx2))
+                b = (dx1 + dx2) / (dx1 * dx2)
+                c = -dx1 / (dx2 * (dx1 + dx2))
+            out = out.at[tuple(slice1)].set(
+                a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
+            )
+            slice1[axis] = -1
+            slice2[axis] = -3
+            slice3[axis] = -2
+            slice4[axis] = -1
+            if uniform_spacing:
+                a = 0.5 / ax_dx
+                b = -2. / ax_dx
+                c = 1.5 / ax_dx
+            else:
+                dx1 = ax_dx[-2]
+                dx2 = ax_dx[-1]
+                a = dx2 / (dx1 * (dx1 + dx2))
+                b = -(dx2 + dx1) / (dx1 * dx2)
+                c = (2. * dx2 + dx1) / (dx2 * (dx1 + dx2))
+            out = out.at[tuple(slice1)].set(
+                a * f[tuple(slice2)] + b * f[tuple(slice3)] + c * f[tuple(slice4)]
+            )
+        outvals.append(out)
+        # reset the slice object in this dimension to ":"
+        slice1[axis] = slice(None)
+        slice2[axis] = slice(None)
+        slice3[axis] = slice(None)
+        slice4[axis] = slice(None)
+    if len_axes == 1:
+        return outvals[0]
+    ret = tuple(outvals)
+    print("return",ret)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: classy_szfast
-Version: 0.0.25.post10
+Version: 0.0.25.post12
 Summary: The accelerator of the class_sz code from https://github.com/CLASS-SZ
 Maintainer-email: Boris Bolliet <bb667@cam.ac.uk>
 License: MIT
@@ -10,7 +10,7 @@ Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.19.0
 Requires-Dist: Cython>=0.29.21
 Requires-Dist: tensorflow
-Requires-Dist: mcfit
-Requires-Dist: get_cosmopower_emus==0.0.15
-Requires-Dist: class_sz_data
+Requires-Dist: mcfit>=0.0.22
+Requires-Dist: get_cosmopower_emus>=0.0.15
+Requires-Dist: class_sz_data>=0.0.13
 Requires-Dist: cosmopower-jax

{classy_szfast-0.0.25.post10.dist-info → classy_szfast-0.0.25.post12.dist-info}/RECORD

@@ -1,6 +1,6 @@
 classy_szfast/__init__.py,sha256=E2thrL0Z9oXFfdzwcsu-xbOytudLFTlRlPqVFGlPPPg,279
 classy_szfast/classy_sz.py,sha256=QmbwrSXInQLMvCDqsr7KPmtaU0KOiOt1Rb-cTKuulZw,22240
-classy_szfast/classy_szfast.py,sha256=ccKynjQks6nrPXa-eLREmCoIz_cl6e6dVJxdcdqavAE,42821
+classy_szfast/classy_szfast.py,sha256=18p1wwBLFKTSXh2iCHoSSVD5NNwsFMyrOOrpbnsBcyk,42915
 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
@@ -9,12 +9,12 @@ classy_szfast/emulators_meta_data.py,sha256=mXG5LQuJw9QBNE_kxXW8Kx0AUCWpbV6uRO9B
 classy_szfast/pks_and_sigmas.py,sha256=drtuujE1HhlrYY1hY92DyY5lXlYS1uE15MSuVI4uo6k,6625
 classy_szfast/restore_nn.py,sha256=DqA9thhTRiGBDVb9zjhqcbF2W4V0AU0vrjJFhnLboU4,21075
 classy_szfast/suppress_warnings.py,sha256=6wIBml2Sj9DyRGZlZWhuA9hqvpxqrNyYjuz6BPK_a6E,202
-classy_szfast/utils.py,sha256=pDZSIgmII3J5bMUYAD8ErvngyeJEEhtla12kCIAWkMg,1962
+classy_szfast/utils.py,sha256=hdA4_ZqCTuO6YODptpROVTns4jYxJqX6rDU9FHmJ1NA,7502
 classy_szfast/custom_bias/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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.post10.dist-info/METADATA,sha256=RycC7LMaSonaPaonweilq5GM_AfGs4sNbcZrkZ9QMNM,556
-classy_szfast-0.0.25.post10.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-classy_szfast-0.0.25.post10.dist-info/top_level.txt,sha256=hRgqpilUck4lx2KkaWI2y9aCDKqF6pFfGHfNaoPFxv0,14
-classy_szfast-0.0.25.post10.dist-info/RECORD,,
+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,,