Change to use jax.Array and jax.typing

eelregit · Oct 23, 2023 · 62a9aa7 · 62a9aa7
1 parent e830c2f
commit 62a9aa7
Show file tree

Hide file tree

Showing 13 changed files with 103 additions and 110 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
@@ -33,11 +33,11 @@ class Configuration:
         ratio, to determine the mesh shape from that of the particle grid. The mesh grid
         cannot be smaller than the particle grid (int or float values must not be
         smaller than 1) and the two grids must have the same aspect ratio.
-    cosmo_dtype : dtype_like, optional
+    cosmo_dtype : DTypeLike, optional
         Float dtype for Cosmology and Configuration.
-    pmid_dtype : dtype_like, optional
+    pmid_dtype : DTypeLike, optional
         Signed integer dtype for particle or mesh grid indices.
-    float_dtype : dtype_like, optional
+    float_dtype : DTypeLike, optional
         Float dtype for other particle and mesh quantities.
     k_pivot_Mpc : float, optional
         Primordial scalar power spectrum pivot scale in 1/Mpc.

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