Change to use jax.Array and jax.typing

eelregit · Oct 22, 2023 · 9e231c9 · 9e231c9
1 parent e830c2f
commit 9e231c9
Show file tree

Hide file tree

Showing 13 changed files with 96 additions and 103 deletions.
diff --git a/pmwd/boltzmann.py b/pmwd/boltzmann.py
@@ -34,14 +34,14 @@ def transfer_fit(k, cosmo, conf):
 
     Parameters
     ----------
-    k : array_like
+    k : ArrayLike
         Wavenumbers in [1/L].
     cosmo : Cosmology
     conf : Configuration
 
     Returns
     -------
-    T : jax.numpy.ndarray of (k * 1.).dtype
+    T : jax.Array of (k * 1.).dtype
         Matter transfer function.
 
     .. _Transfer Function:
@@ -129,14 +129,14 @@ def transfer(k, cosmo, conf):
 
     Parameters
     ----------
-    k : array_like
+    k : ArrayLike
         Wavenumbers in [1/L].
     cosmo : Cosmology
     conf : Configuration
 
     Returns
     -------
-    T : jax.numpy.ndarray of (k * 1.).dtype
+    T : jax.Array of (k * 1.).dtype
         Matter transfer function.
 
     Raises
@@ -238,7 +238,7 @@ def growth(a, cosmo, conf, order=1, deriv=0):
 
     Parameters
     ----------
-    a : array_like
+    a : ArrayLike
         Scale factors.
     cosmo : Cosmology
     conf : Configuration
@@ -249,7 +249,7 @@ def growth(a, cosmo, conf, order=1, deriv=0):
 
     Returns
     -------
-    D : jax.numpy.ndarray of (a * 1.).dtype
+    D : jax.Array of (a * 1.).dtype
         Growth functions or derivatives.
 
     Raises
@@ -297,16 +297,16 @@ def varlin(R, a, cosmo, conf):
 
     Parameters
     ----------
-    R : array_like
+    R : ArrayLike
         Scales in [L].
-    a : array_like or None
+    a : ArrayLike or None
         Scale factors. If None, output is not scaled by growth.
     cosmo : Cosmology
     conf : Configuration
 
     Returns
     -------
-    sigma2 : jax.numpy.ndarray of (k * a * 1.).dtype
+    sigma2 : jax.Array of (k * a * 1.).dtype
         Linear matter overdensity variance.
 
     Raises
@@ -401,16 +401,16 @@ def linear_power(k, a, cosmo, conf):
 
     Parameters
     ----------
-    k : array_like
+    k : ArrayLike
         Wavenumbers in [1/L].
-    a : array_like or None
+    a : ArrayLike or None
         Scale factors. If None, output is not scaled by growth.
     cosmo : Cosmology
     conf : Configuration
 
     Returns
     -------
-    Plin : jax.numpy.ndarray of (k * a * 1.).dtype
+    Plin : jax.Array of (k * a * 1.).dtype
         Linear matter power spectrum in [L^3].
 
     Raises

diff --git a/pmwd/configuration.py b/pmwd/configuration.py
@@ -2,9 +2,9 @@
 import math
 from typing import ClassVar, Optional, Tuple, Union
 
-from numpy.typing import DTypeLike
 import jax
 from jax import ensure_compile_time_eval
+from jax.typing import DTypeLike
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 from mcfit import TophatVar

diff --git a/pmwd/cosmology.py b/pmwd/cosmology.py
@@ -1,20 +1,17 @@
 from dataclasses import field
 from functools import partial
 from operator import add, sub
-from typing import ClassVar, Optional, Union
+from typing import ClassVar, Optional
 
-import numpy as np
-from jax import value_and_grad
+from jax import Array, value_and_grad
+from jax.typing import ArrayLike
 import jax.numpy as jnp
 from jax.tree_util import tree_map
 
 from pmwd.tree_util import pytree_dataclass
 from pmwd.configuration import Configuration
 
 
-FloatParam = Union[float, jnp.ndarray]
-
-
 @partial(pytree_dataclass, aux_fields="conf", frozen=True)
 class Cosmology:
     """Cosmological and configuration parameters, "immutable" as a frozen dataclass.
@@ -34,45 +31,45 @@ class Cosmology:
     ----------
     conf : Configuration
         Configuration parameters.
-    A_s_1e9 : float or jax.numpy.ndarray
+    A_s_1e9 : float ArrayLike
         Primordial scalar power spectrum amplitude, multiplied by 1e9.
-    n_s : float or jax.numpy.ndarray
+    n_s : float ArrayLike
         Primordial scalar power spectrum spectral index.
-    Omega_m : float or jax.numpy.ndarray
+    Omega_m : float ArrayLike
         Total matter density parameter today.
-    Omega_b : float or jax.numpy.ndarray
+    Omega_b : float ArrayLike
         Baryonic matter density parameter today.
-    Omega_k_ : None, float, or jax.numpy.ndarray, optional
+    Omega_k_ : None or float ArrayLike, optional
         Spatial curvature density parameter today. Default is None.
-    w_0_ : None, float, or jax.numpy.ndarray, optional
+    w_0_ : None or float ArrayLike, optional
         Dark energy equation of state constant parameter. Default is None.
-    w_a_ : None, float, or jax.numpy.ndarray, optional
+    w_a_ : None or float ArrayLike, optional
         Dark energy equation of state linear parameter. Default is None.
-    h : float or jax.numpy.ndarray
+    h : float ArrayLike
         Hubble constant in unit of 100 [km/s/Mpc].
 
     """
 
     conf: Configuration = field(repr=False)
 
-    A_s_1e9: FloatParam
-    n_s: FloatParam
-    Omega_m: FloatParam
-    Omega_b: FloatParam
-    h: FloatParam
+    A_s_1e9: ArrayLike
+    n_s: ArrayLike
+    Omega_m: ArrayLike
+    Omega_b: ArrayLike
+    h: ArrayLike
 
-    Omega_k_: Optional[FloatParam] = None
+    Omega_k_: Optional[ArrayLike] = None
     Omega_k_fixed: ClassVar[float] = 0
-    w_0_: Optional[FloatParam] = None
+    w_0_: Optional[ArrayLike] = None
     w_0_fixed: ClassVar[float] = -1
-    w_a_: Optional[FloatParam] = None
+    w_a_: Optional[ArrayLike] = None
     w_a_fixed: ClassVar[float] = 0
 
-    transfer: Optional[jnp.ndarray] = field(default=None, compare=False)
+    transfer: Optional[Array] = field(default=None, compare=False)
 
-    growth: Optional[jnp.ndarray] = field(default=None, compare=False)
+    growth: Optional[Array] = field(default=None, compare=False)
 
-    varlin: Optional[jnp.ndarray] = field(default=None, compare=False)
+    varlin: Optional[Array] = field(default=None, compare=False)
 
     def __post_init__(self):
         if self._is_transforming():
@@ -191,13 +188,13 @@ def E2(a, cosmo):
 
     Parameters
     ----------
-    a : array_like
+    a : ArrayLike
         Scale factors.
     cosmo : Cosmology
 
     Returns
     -------
-    E2 : jax.numpy.ndarray of cosmo.conf.cosmo_dtype
+    E2 : jax.Array of cosmo.conf.cosmo_dtype
         Squared Hubble parameter time scaling factors.
 
     Notes
@@ -229,13 +226,13 @@ def H_deriv(a, cosmo):
 
     Parameters
     ----------
-    a : array_like
+    a : ArrayLike
         Scale factors.
     cosmo : Cosmology
 
     Returns
     -------
-    dlnH_dlna : jax.numpy.ndarray of cosmo.conf.cosmo_dtype
+    dlnH_dlna : jax.Array of cosmo.conf.cosmo_dtype
         Hubble parameter derivatives.
 
     """
@@ -250,13 +247,13 @@ def Omega_m_a(a, cosmo):
 
     Parameters
     ----------
-    a : array_like
+    a : ArrayLike
         Scale factors.
     cosmo : Cosmology
 
     Returns
     -------
-    Omega : jax.numpy.ndarray of cosmo.conf.cosmo_dtype
+    Omega : jax.Array of cosmo.conf.cosmo_dtype
         Matter density parameters.
 
     Notes

diff --git a/pmwd/gather.py b/pmwd/gather.py
@@ -12,18 +12,18 @@ def gather(ptcl, conf, mesh, val=0, offset=0, cell_size=None):
     ----------
     ptcl : Particles
     conf : Configuration
-    mesh : array_like
+    mesh : ArrayLike
         Input mesh.
-    val : array_like, optional
+    val : ArrayLike, optional
         Input values, can be 0D.
-    offset : array_like, optional
+    offset : ArrayLike, optional
         Offset of mesh to particle grid. If 0D, the value is used in each dimension.
     cell_size : float, optional
         Mesh cell size in [L]. Default is ``conf.cell_size``.
 
     Returns
     -------
-    val : jax.numpy.ndarray
+    val : jax.Array
         Output values.
 
     """

diff --git a/pmwd/lpt.py b/pmwd/lpt.py
@@ -81,7 +81,7 @@ def levi_civita(indices):
 
     Parameters
     ----------
-    indices : array_like
+    indices : ArrayLike
 
     Returns
     -------
@@ -140,7 +140,7 @@ def lpt(modes, cosmo, conf):
 
     Parameters
     ----------
-    modes : jax.numpy.ndarray
+    modes : jax.Array
         Linear matter overdensity Fourier modes in [L^3].
     cosmo : Cosmology
     conf : Configuration

diff --git a/pmwd/modes.py b/pmwd/modes.py
@@ -27,7 +27,7 @@ def white_noise(seed, conf, real=False, unit_abs=False, negate=False):
 
     Returns
     -------
-    modes : jax.numpy.ndarray of conf.float_dtype
+    modes : jax.Array of conf.float_dtype
         White noise modes.
 
     """
@@ -75,7 +75,7 @@ def linear_modes(modes, cosmo, conf, a=None):
 
     Parameters
     ----------
-    modes : jax.numpy.ndarray
+    modes : jax.Array
         Fourier or real modes with white noise prior.
     cosmo : Cosmology
     conf : Configuration
@@ -84,7 +84,7 @@ def linear_modes(modes, cosmo, conf, a=None):
 
     Returns
     -------
-    modes : jax.numpy.ndarray of conf.float_dtype
+    modes : jax.Array of conf.float_dtype
         Linear matter overdensity Fourier modes in [L^3].
 
     Notes