pywavelet 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

pywavelet/__init__.py

@@ -6,11 +6,10 @@ import importlib
 import os
 
 from . import backend as _backend
+from ._version import __version__
 
-__version__ = "0.0.2"
 
-
-def set_backend(backend: str):
+def set_backend(backend: str, precision: str = "float32") -> None:
     """Set the backend for the wavelet transform.
 
     Parameters
@@ -18,10 +17,23 @@ def set_backend(backend: str):
     backend : str
         Backend to use. Options are "numpy", "jax", "cupy".
     """
-    from . import types
-    from . import transforms
+    from . import transforms, types
+
     os.environ["PYWAVELET_BACKEND"] = backend
+    os.environ["PYWAVELET_PRECISION"] = precision
 
     importlib.reload(_backend)
     importlib.reload(types)
     importlib.reload(transforms)
+    if backend == "cupy":
+        importlib.reload(transforms.cupy)
+        importlib.reload(transforms.cupy.forward)
+        importlib.reload(transforms.cupy.inverse)
+    elif backend == "jax":
+        importlib.reload(transforms.jax)
+        importlib.reload(transforms.jax.forward)
+        importlib.reload(transforms.jax.inverse)
+    else:
+        importlib.reload(transforms.numpy)
+        importlib.reload(transforms.numpy.forward)
+        importlib.reload(transforms.numpy.inverse)

pywavelet/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.6'
-__version_tuple__ = version_tuple = (0, 2, 6)
+__version__ = version = '0.2.8'
+__version_tuple__ = version_tuple = (0, 2, 8)

pywavelet/backend.py

@@ -1,5 +1,10 @@
 import importlib
 import os
+from typing import Tuple
+
+import numpy as np
+from rich.console import Console
+from rich.table import Table, Text
 
 from .logger import logger
 
@@ -7,47 +12,168 @@ JAX = "jax"
 CUPY = "cupy"
 NUMPY = "numpy"
 
+VALID_PRECISIONS = ["float32", "float64"]
+
+
+def cuda_is_available() -> bool:
+    """Check if CUDA is available."""
+    # Check if CuPy is available and CUDA is accessible
+    cupy_available = importlib.util.find_spec("cupy") is not None
+    _cuda_available = False
+    if cupy_available:
+        import cupy
+
+        try:
+            cupy.cuda.runtime.getDeviceCount()  # Check if any CUDA device is available
+            _cuda_available = True
+        except cupy.cuda.runtime.CUDARuntimeError:
+            _cuda_available = False
+    else:
+        _cuda_available = False
+    return _cuda_available
+
+
+def jax_is_available() -> bool:
+    """Check if JAX is available."""
+    return importlib.util.find_spec(JAX) is not None
+
+
+def get_available_backends_table():
+    """Print the available backends as a rich table."""
+
+    jax_avail = jax_is_available()
+    cupy_avail = cuda_is_available()
+    table = Table("Backend", "Available", title="Available backends")
+    true_check = "[green]✓[/green]"
+    false_check = "[red]✗[/red]"
+    table.add_row(JAX, true_check if jax_avail else false_check)
+    table.add_row(CUPY, true_check if cupy_avail else false_check)
+    table.add_row(NUMPY, true_check)
+    console = Console(width=150)
+    with console.capture() as capture:
+        console.print(table)
+    return Text.from_ansi(capture.get())
+
+
+def log_backend(level="info"):
+    """Print the current backend and precision."""
+    backend = os.getenv("PYWAVELET_BACKEND", NUMPY).lower()
+    precision = os.getenv("PYWAVELET_PRECISION", "float32").lower()
+    str = f"Current backend: {backend}[{precision}]"
+    if level == "info":
+        logger.info(str)
+    elif level == "debug":
+        logger.debug(str)
+
 
 def get_backend_from_env():
     """Select and return the appropriate backend module."""
     backend = os.getenv("PYWAVELET_BACKEND", NUMPY).lower()
+    precision = os.getenv("PYWAVELET_PRECISION", "float32").lower()
+
+    if backend == JAX and jax_is_available():
+
+        import jax.numpy as xp
+        from jax.numpy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from jax.scipy.special import betainc
 
+    elif backend == CUPY and cuda_is_available():
+
+        import cupy as xp
+        from cupy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from cupyx.scipy.special import betainc
+
+    elif backend == NUMPY:
+        import numpy as xp
+        from numpy.fft import fft, ifft, irfft, rfft, rfftfreq
+        from scipy.special import betainc
+
+    else:
+        logger.error(f"Backend {backend}[{precision}] is not available. ")
+        print(get_available_backends_table())
+        logger.warning(f"Setting backend to NumPy. ")
+        os.environ["PYWAVELET_BACKEND"] = NUMPY
+        return get_backend_from_env()
+
+    log_backend("debug")
+    return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
+
+
+def get_precision_from_env() -> str:
+    """Get the precision from the environment variable."""
+    precision = os.getenv("PYWAVELET_PRECISION", "float32").lower()
+    if precision not in VALID_PRECISIONS:
+        logger.error(
+            f"Precision {precision} is not supported, defaulting to float32."
+        )
+        precision = "float32"
+    return precision
+
+
+def set_precision(precision: str) -> None:
+    """Set the precision for the backend."""
+    precision = precision.lower()
+    if precision not in VALID_PRECISIONS:
+        logger.error(f"Precision {precision} is not supported.")
+        return
+    else:
+        os.environ["PYWAVELET_PRECISION"] = precision
+        logger.info(f"Setting precision to {precision}.")
+        return
+
+
+def get_dtype_from_env() -> Tuple[np.dtype, np.dtype]:
+    """Get the data type from the environment variable."""
+    precision = get_precision_from_env()
+    backend = get_backend_from_env()[-1]
     if backend == JAX:
-        if importlib.util.find_spec(JAX):
-            import jax.numpy as xp
-            from jax.numpy.fft import fft, ifft, irfft, rfft, rfftfreq
-            from jax.scipy.special import betainc
-
-            logger.info("Using JAX backend")
-            return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
-        else:
-            logger.warning(
-                "JAX backend requested but not installed. Falling back to NumPy."
-            )
+
+        if precision == "float32":
+            import jax
+
+            jax.config.update("jax_enable_x64", False)
+
+            import jax.numpy as jnp
+
+            float_dtype = jnp.float32
+            complex_dtype = jnp.complex64
+        elif precision == "float64":
+            import jax
+
+            jax.config.update("jax_enable_x64", True)
+
+            import jax.numpy as jnp
+
+            float_dtype = jnp.float64
+            complex_dtype = jnp.complex128
 
     elif backend == CUPY:
-        if importlib.util.find_spec(CUPY):
-            import cupy as xp
-            from cupy.fft import fft, ifft, irfft, rfft, rfftfreq
-            from cupyx.scipy.special import betainc
-
-            logger.info("Using CuPy backend")
-            return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
-        else:
-            logger.warning(
-                "CuPy backend requested but not installed. Falling back to NumPy."
-            )
-
-    # Default to NumPy
-    import numpy as xp
-    from numpy.fft import fft, ifft, irfft, rfft, rfftfreq
-    from scipy.special import betainc
-
-    logger.info("Using NumPy+Numba backend")
-    return xp, fft, ifft, irfft, rfft, rfftfreq, betainc, backend
+        import cupy as cp
+
+        if precision == "float32":
+            float_dtype = cp.float32
+            complex_dtype = cp.complex64
+        elif precision == "float64":
+            float_dtype = cp.float64
+            complex_dtype = cp.complex128
 
+    else:
+        if precision == "float32":
+            float_dtype = np.float32
+            complex_dtype = np.complex64
+        elif precision == "float64":
+            float_dtype = np.float64
+            complex_dtype = np.complex128
+
+    return float_dtype, complex_dtype
+
+
+cuda_available = cuda_is_available()
 
 # Get the chosen backend
 xp, fft, ifft, irfft, rfft, rfftfreq, betainc, current_backend = (
     get_backend_from_env()
 )
+
+# Get the chosen precision
+float_dtype, complex_dtype = get_dtype_from_env()

pywavelet/logger.py

@@ -5,14 +5,14 @@ import warnings
 from rich.logging import RichHandler
 
 FORMAT = "%(message)s"
-logging.basicConfig(
-    level="INFO",
-    format=FORMAT,
-    datefmt="[%X]",
-    handlers=[RichHandler(rich_tracebacks=True)],
-)
 
 logger = logging.getLogger("pywavelet")
+if not logger.hasHandlers():
+    handler = RichHandler(rich_tracebacks=True)
+    formatter = logging.Formatter(fmt=FORMAT, datefmt="[%X]")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
 
 warnings.filterwarnings("ignore", category=RuntimeWarning)
 warnings.filterwarnings("ignore", category=UserWarning)

pywavelet/transforms/__init__.py

@@ -1,4 +1,8 @@
 from ..backend import current_backend
+from ..logger import logger
+
+logger.debug(f"Using {current_backend} backend")
+
 
 if current_backend == "jax":
     from .jax import (

pywavelet/transforms/cupy/forward/from_freq.py

@@ -1,14 +1,18 @@
+import logging
+
 import cupy as cp
 from cupyx.scipy.fft import ifft
 
-
-import logging
-
-logger = logging.getLogger('pywavelet')
+logger = logging.getLogger("pywavelet")
 
 
 def transform_wavelet_freq_helper(
-    data: cp.ndarray, Nf: int, Nt: int, phif: cp.ndarray
+    data: cp.ndarray,
+    Nf: int,
+    Nt: int,
+    phif: cp.ndarray,
+    float_dtype=cp.float64,
+    complex_dtype=cp.complex128,
 ) -> cp.ndarray:
     """
     Transforms input data from the frequency domain to the wavelet domain using a
@@ -24,69 +28,62 @@ def transform_wavelet_freq_helper(
     - wave (cp.ndarray): 2D array of wavelet-transformed data with shape (Nf, Nt).
     """
 
-    logger.debug(f"[CUPY TRANSFORM] Input types [data:{type(data)}, phif:{type(phif)}]")
-
-    # Initialize the wavelet output array with zeros (time-rows, frequency-columns)
-    wave = cp.zeros((Nt, Nf))
-    f_bins = cp.arange(Nf)  # Frequency bin indices
-
-    # Compute base indices for time (i_base) and frequency (jj_base)
-    i_base = Nt // 2
-    jj_base = f_bins * Nt // 2
-
-    # Set initial values for the center of the transformation
-    initial_values = cp.where(
-        (f_bins == 0)
-        | (f_bins == Nf),  # Edge cases: DC (f=0) and Nyquist (f=Nf)
-        phif[0] * data[f_bins * Nt // 2] / 2.0,  # Adjust for symmetry
-        phif[0] * data[f_bins * Nt // 2],
+    logger.debug(
+        f"[CUPY TRANSFORM] Input types [data:{type(data)}, phif:{type(phif)}, precision:{data.dtype}]"
     )
 
-    # Initialize a 2D array to store intermediate FFT input values
-    DX = cp.zeros((Nf, Nt), dtype=cp.complex64)
-    DX[:, Nt // 2] = (
-        initial_values  # Set initial values at the center of the transformation (2 sided FFT)
+    half = Nt // 2
+    f_bins = cp.arange(Nf + 1)  # [0, 1, ..., Nf]
+
+    # 1) Build (Nf+1, Nt) DX
+    # — center:
+    center = phif[0] * data[f_bins * half]
+    center = cp.where((f_bins == 0) | (f_bins == Nf), center * 0.5, center)
+    DX = cp.zeros((Nf + 1, Nt), complex_dtype)
+    DX[:, half] = center
+
+    # — off-center (j = ±1...±(half-1))
+    offs = cp.arange(1 - half, half)  # length Nt-1
+    jj = f_bins[:, None] * half + offs[None, :]  # (Nf+1, Nt-1)
+    ii = half + offs  # (Nt-1,)
+    mask = ((f_bins[:, None] == Nf) & (offs[None, :] > 0)) | (
+        (f_bins[:, None] == 0) & (offs[None, :] < 0)
+    )
+    vals = phif[cp.abs(offs)] * data[jj]
+    vals = cp.where(mask, 0.0, vals)
+    DX[:, ii] = vals
+
+    # 2) IFFT along time
+    DXt = ifft(DX, axis=1)
+
+    # 3) Unpack into wave (Nt, Nf)
+    wave = cp.zeros((Nt, Nf), float_dtype)
+    sqrt2 = cp.sqrt(2.0)
+
+    # 3a) DC into col 0, even rows
+    evens = cp.arange(0, Nt, 2)
+    wave[evens, 0] = cp.real(DXt[0, evens]) * sqrt2
+
+    # 3b) Nyquist into col 0, odd rows
+    odds = cp.arange(1, Nt, 2)
+    wave[odds, 0] = cp.real(DXt[Nf, evens]) * sqrt2
+
+    # 3c) intermediate bins 1...Nf-1
+    mids = cp.arange(1, Nf)  # [1...Nf-1]
+    Dt_mid = DXt[mids, :]  # (Nf-1, Nt)
+    real_m = cp.real(Dt_mid).T  # (Nt, Nf-1)
+    imag_m = cp.imag(Dt_mid).T  # (Nt, Nf-1)
+
+    odd_f = (mids % 2) == 1  # shape (Nf-1,)
+    n_idx = cp.arange(Nt)[:, None]  # (Nt,1)
+    odd_nf = ((n_idx + mids[None, :]) % 2) == 1
+
+    # Select real vs imag and sign exactly as in __fill_wave_2
+    mid_vals = cp.where(
+        odd_nf,
+        cp.where(odd_f, -imag_m, imag_m),
+        cp.where(odd_f, real_m, real_m),
     )
+    wave[:, 1:Nf] = mid_vals
 
-    # Compute time indices for all offsets around the midpoint
-    j_range = cp.arange(
-        1 - Nt // 2, Nt // 2
-    )  # Time offsets (centered around zero)
-    j = cp.abs(j_range)  # Absolute offset indices
-    i = i_base + j_range  # Time indices in output array
-
-    # Compute conditions for edge cases
-    cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)  # Nyquist
-    cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)  # DC
-    cond3 = j[None, :] == 0  # Center of the transformation (no offset)
-
-    # Compute frequency indices for the input data
-    jj = jj_base[:, None] + j_range[None, :]  # Frequency offsets
-    val = cp.where(
-        cond1 | cond2, 0.0, phif[j] * data[jj]
-    )  # Wavelet filter application
-    DX[:, i] = cp.where(cond3, DX[:, i], val)  # Update DX with computed values
-
-    # Perform the inverse FFT along the time dimension
-    DX_trans = ifft(DX, axis=1)
-
-    # Fill the wavelet output array based on the inverse FFT results
-    n_range = cp.arange(Nt)  # Time indices
-    cond1 = (
-        n_range[:, None] + f_bins[None, :]
-    ) % 2 == 1  # Odd/even alternation
-    cond2 = cp.expand_dims(f_bins % 2 == 1, axis=-1)  # Odd frequency bins
-
-    # Assign real and imaginary parts based on conditions
-    real_part = cp.where(cond2, -cp.imag(DX_trans), cp.real(DX_trans))
-    imag_part = cp.where(cond2, cp.real(DX_trans), cp.imag(DX_trans))
-    wave = cp.where(cond1, imag_part.T, real_part.T)
-
-    # Special cases for frequency bins 0 (DC) and Nf (Nyquist)
-    wave[::2, 0] = cp.real(DX_trans[0, ::2] * cp.sqrt(2))  # DC component
-    wave[1::2, -1] = cp.real(
-        DX_trans[-1, ::2] * cp.sqrt(2)
-    )  # Nyquist component
-
-    # Return the wavelet-transformed array (transposed for freq-major layout)
-    return wave.T  # (Nt, Nf) -> (Nf, Nt)
+    return wave.T

pywavelet/transforms/cupy/forward/main.py

@@ -2,6 +2,7 @@ from typing import Union
 
 import cupy as cp
 
+from .... import backend
 from ....logger import logger
 from ....types import FrequencySeries, TimeSeries, Wavelet
 from ....types.wavelet_bins import _get_bins, _preprocess_bins
@@ -16,7 +17,6 @@ def from_time_to_wavelet(
     Nt: Union[int, None] = None,
     nx: float = 4.0,
     mult: int = 32,
-    **kwargs,
 ) -> Wavelet:
     """Transforms time-domain data to wavelet-domain data.
 
@@ -45,7 +45,6 @@ def from_time_to_wavelet(
 
     """
     Nf, Nt = _preprocess_bins(timeseries, Nf, Nt)
-    dt = timeseries.dt
     t_bins, f_bins = _get_bins(timeseries, Nf, Nt)
 
     ND = Nf * Nt
@@ -73,7 +72,6 @@ def from_freq_to_wavelet(
     Nf: Union[int, None] = None,
     Nt: Union[int, None] = None,
     nx: float = 4.0,
-    **kwargs,
 ) -> Wavelet:
     """Transforms frequency-domain data to wavelet-domain data.
 
@@ -98,9 +96,15 @@ def from_freq_to_wavelet(
     """
     Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
     t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
-    phif = cp.array(phitilde_vec_norm(Nf, Nt, d=nx))
+    phif = cp.array(phitilde_vec_norm(Nf, Nt, d=nx), dtype=backend.float_dtype)
+    data = cp.array(freqseries.data, dtype=backend.complex_dtype)
     wave = transform_wavelet_freq_helper(
-        cp.array(freqseries.data), Nf=Nf, Nt=Nt, phif=phif
+        data,
+        Nf=Nf,
+        Nt=Nt,
+        phif=phif,
+        float_dtype=backend.float_dtype,
+        complex_dtype=backend.complex_dtype,
     )
-
-    return Wavelet((2 / Nf) * wave * cp.sqrt(2), time=t_bins, freq=f_bins)
+    factor = (2 / Nf) * cp.sqrt(2)
+    return Wavelet(factor * wave, time=t_bins, freq=f_bins)

pywavelet/transforms/cupy/inverse/to_freq.py

@@ -6,57 +6,61 @@ def inverse_wavelet_freq_helper(
     wave_in: cp.ndarray, phif: cp.ndarray, Nf: int, Nt: int
 ) -> cp.ndarray:
     """CuPy vectorized function for inverse_wavelet_freq"""
-    wave_in = wave_in.T
-    ND = Nf * Nt
+    wave = wave_in.T
+    ND2 = (Nf * Nt) // 2
+    half = Nt // 2
 
+    # === STEP 1: build prefactor2s[m, n] ===
     m_range = cp.arange(Nf + 1)
-    prefactor2s = cp.zeros((Nf + 1, Nt), dtype=cp.complex128)
-
-    n_range = cp.arange(Nt)
-
-    # m == 0 case
-    prefactor2s[0] = 2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt, 0]
-
-    # m == Nf case
-    prefactor2s[Nf] = 2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt + 1, 0]
-
-    # Other m cases
-    m_mid = m_range[1:Nf]
-    n_grid, m_grid = cp.meshgrid(n_range, m_mid)
-    val = wave_in[n_grid, m_grid]
-    mult2 = cp.where((n_grid + m_grid) % 2, -1j, 1)
-    prefactor2s[1:Nf] = mult2 * val
-
-    # Vectorized FFT
-    fft_prefactor2s = fft(prefactor2s, axis=1)
-
-    # Vectorized __unpack_wave_inverse
-    res = cp.zeros(ND, dtype=cp.complex128)
-
-    # m == 0 or m == Nf cases
-    i_ind_range = cp.arange(Nt // 2)
-    i_0 = cp.abs(i_ind_range)
-    i_Nf = cp.abs(Nf * Nt // 2 - i_ind_range)
-    ind3_0 = (2 * i_0) % Nt
-    ind3_Nf = (2 * i_Nf) % Nt
-
-    res[i_0] += fft_prefactor2s[0, ind3_0] * phif[i_ind_range]
-    res[i_Nf] += fft_prefactor2s[Nf, ind3_Nf] * phif[i_ind_range]
-
-    # Special case for m == Nf
-    res[Nf * Nt // 2] += fft_prefactor2s[Nf, 0] * phif[Nt // 2]
-
-    # Other m cases
-    m_mid = m_range[1:Nf]
-    i_ind_range = cp.arange(Nt // 2 + 1)
-    m_grid, i_ind_grid = cp.meshgrid(m_mid, i_ind_range)
-
-    i1 = Nt // 2 * m_grid - i_ind_grid
-    i2 = Nt // 2 * m_grid + i_ind_grid
-    ind31 = (Nt // 2 * m_grid - i_ind_grid) % Nt
-    ind32 = (Nt // 2 * m_grid + i_ind_grid) % Nt
-
-    res[i1] += fft_prefactor2s[m_grid, ind31] * phif[i_ind_grid]
-    res[i2] += fft_prefactor2s[m_grid, ind32] * phif[i_ind_grid]
+    pref2 = cp.zeros((Nf + 1, Nt), dtype=cp.complex128)
+    n = cp.arange(Nt)
+
+    # m=0 and m=Nf, with 1/√2
+    pref2[0, :] = wave[(2 * n) % Nt, 0] * (2 ** (-0.5))
+    pref2[Nf, :] = wave[((2 * n) % Nt) + 1, 0] * (2 ** (-0.5))
+
+    # middle m=1...Nf-1
+    m_mid = cp.arange(1, Nf)
+    # build meshgrids (m_mid rows, n cols)
+    mm, nn = cp.meshgrid(m_mid, n, indexing="ij")
+    vals = wave[nn, mm]
+    signs = cp.where(((nn + mm) % 2) == 1, -1j, 1)
+    pref2[1:Nf, :] = signs * vals
+
+    # === STEP 2: FFT along time axis ===
+    F = fft(pref2, axis=1)  # shape (Nf+1, Nt)
+
+    # === STEP 3: unpack back into half-spectrum res[0...ND2] ===
+    res = cp.zeros(ND2 + 1, dtype=cp.complex128)
+    idx = cp.arange(half)
+
+    # 3a) contribution from m=0
+    res[idx] += F[0, (2 * idx) % Nt] * phif[idx]
+
+    # 3b) contribution from m=Nf
+    iNf = cp.abs(Nf * half - idx)
+    res[iNf] += F[Nf, (2 * idx) % Nt] * phif[idx]
+
+    # special Nyquist‐folding term
+    special = cp.abs(Nf * half - half)
+    res[special] += F[Nf, 0] * phif[half]
+
+    # 3c) middle m cases
+    m_mid = cp.arange(1, Nf)
+    i_mid = cp.arange(half)  # 0...half-1
+    mm, ii = cp.meshgrid(m_mid, i_mid, indexing="ij")
+    i1 = (half * mm - ii) % (ND2 + 1)
+    i2 = (half * mm + ii) % (ND2 + 1)
+    ind1 = (half * mm - ii) % Nt
+    ind2 = (half * mm + ii) % Nt
+
+    # accumulate
+    res[i1] += F[mm, ind1] * phif[ii]
+    res[i2] += F[mm, ind2] * phif[ii]
+
+    # override the "center" points (j=0) exactly
+    centers = half * m_mid
+    fft_idx = centers % Nt
+    res[centers] = F[m_mid, fft_idx] * phif[0]
 
     return res

pywavelet/transforms/jax/__init__.py

@@ -4,6 +4,30 @@ from .inverse import from_wavelet_to_freq, from_wavelet_to_time
 
 logger.warning("JAX SUBPACKAGE NOT FULLY TESTED")
 
+
+def _log_jax_info():
+    """Log JAX backend and precision information.
+
+    backend : str
+        JAX backend. ["cpu", "gpu", "tpu"]
+    precision : str
+        JAX precision. ["32bit", "64bit"]
+    """
+    import jax
+
+    _backend = jax.default_backend()
+    _precision = "64bit" if jax.config.jax_enable_x64 else "32bit"
+
+    logger.info(f"Jax running on {_backend} [{_precision} precision].")
+    if _precision == "32bit":
+        logger.warning(
+            "Jax is not running in 64bit precision. "
+            "To change, use jax.config.update('jax_enable_x64', True)."
+        )
+
+
+_log_jax_info()
+
 __all__ = [
     "from_wavelet_to_time",
     "from_wavelet_to_freq",