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

pywavelet/__init__.py

@@ -2,4 +2,26 @@
 WDM Wavelet transform
 """
 
+import importlib
+import os
+
+from . import backend as _backend
+
 __version__ = "0.0.2"
+
+
+def set_backend(backend: str):
+    """Set the backend for the wavelet transform.
+
+    Parameters
+    ----------
+    backend : str
+        Backend to use. Options are "numpy", "jax", "cupy".
+    """
+    from . import types
+    from . import transforms
+    os.environ["PYWAVELET_BACKEND"] = backend
+
+    importlib.reload(_backend)
+    importlib.reload(types)
+    importlib.reload(transforms)

pywavelet/_version.py

@@ -1,8 +1,13 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
 else:
     VERSION_TUPLE = object
@@ -12,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.5'
-__version_tuple__ = version_tuple = (0, 2, 5)
+__version__ = version = '0.2.6'
+__version_tuple__ = version_tuple = (0, 2, 6)

pywavelet/backend.py

@@ -1,31 +1,53 @@
+import importlib
 import os
 
 from .logger import logger
 
-try:
-    import jax
-
-    jax_available = True
-
-
-except ImportError:
-    jax_available = False
-
-use_jax = jax_available and os.getenv("PYWAVELET_JAX", "0") == "1"
-
-if use_jax:
-    import jax.numpy as xp  # type: ignore
-    from jax.numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
-    from jax.scipy.special import betainc  # type: ignore
-
-    logger.info("Using JAX backend")
-
-else:
-    import numpy as xp  # type: ignore
-    from numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
-    from scipy.special import betainc  # type: ignore
-
-    logger.info("Using NumPy+numba backend")
-
-
-PI = xp.pi
+JAX = "jax"
+CUPY = "cupy"
+NUMPY = "numpy"
+
+
+def get_backend_from_env():
+    """Select and return the appropriate backend module."""
+    backend = os.getenv("PYWAVELET_BACKEND", NUMPY).lower()
+
+    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."
+            )
+
+    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
+
+
+# Get the chosen backend
+xp, fft, ifft, irfft, rfft, rfftfreq, betainc, current_backend = (
+    get_backend_from_env()
+)

pywavelet/transforms/__init__.py

@@ -1,12 +1,19 @@
-from ..backend import use_jax
+from ..backend import current_backend
 
-if use_jax:
+if current_backend == "jax":
     from .jax import (
         from_freq_to_wavelet,
         from_time_to_wavelet,
         from_wavelet_to_freq,
         from_wavelet_to_time,
     )
+elif current_backend == "cupy":
+    from .cupy import (
+        from_freq_to_wavelet,
+        from_time_to_wavelet,
+        from_wavelet_to_freq,
+        from_wavelet_to_time,
+    )
 else:
     from .numpy import (
         from_wavelet_to_time,
@@ -15,7 +22,7 @@ else:
         from_freq_to_wavelet,
     )
 
-from .phi_computer import phi_vec, phitilde_vec, phitilde_vec_norm
+from .phi_computer import omega, phi_vec, phitilde_vec_norm
 
 __all__ = [
     "from_wavelet_to_time",
@@ -24,5 +31,4 @@ __all__ = [
     "from_freq_to_wavelet",
     "phitilde_vec_norm",
     "phi_vec",
-    "phitilde_vec",
 ]

pywavelet/transforms/cupy/__init__.py

@@ -0,0 +1,12 @@
+from ...logger import logger
+from .forward import from_freq_to_wavelet, from_time_to_wavelet
+from .inverse import from_wavelet_to_freq, from_wavelet_to_time
+
+logger.warning("CUPY SUBPACKAGE NOT FULLY TESTED")
+
+__all__ = [
+    "from_wavelet_to_time",
+    "from_wavelet_to_freq",
+    "from_time_to_wavelet",
+    "from_freq_to_wavelet",
+]

pywavelet/transforms/cupy/forward/__init__.py

@@ -0,0 +1,3 @@
+from .main import from_freq_to_wavelet, from_time_to_wavelet
+
+__all__ = ["from_time_to_wavelet", "from_freq_to_wavelet"]

pywavelet/transforms/cupy/forward/from_freq.py

@@ -0,0 +1,92 @@
+import cupy as cp
+from cupyx.scipy.fft import ifft
+
+
+import logging
+
+logger = logging.getLogger('pywavelet')
+
+
+def transform_wavelet_freq_helper(
+    data: cp.ndarray, Nf: int, Nt: int, phif: cp.ndarray
+) -> cp.ndarray:
+    """
+    Transforms input data from the frequency domain to the wavelet domain using a
+    pre-computed wavelet filter (`phif`) and performs an efficient inverse FFT.
+
+    Parameters:
+    - data (cp.ndarray): 1D array representing the input data in the frequency domain.
+    - Nf (int): Number of frequency bins.
+    - Nt (int): Number of time bins. (Note: Nt * Nf == len(data))
+    - phif (cp.ndarray): Pre-computed wavelet filter for frequency components.
+
+    Returns:
+    - 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],
+    )
+
+    # 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)
+    )
+
+    # 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)

pywavelet/transforms/cupy/forward/from_time.py

@@ -0,0 +1,50 @@
+import cupy as cp
+from cupyx.scipy.fft import rfft
+
+
+def transform_wavelet_time_helper(
+    data: cp.ndarray, phi: cp.ndarray, Nf: int, Nt: int, mult: int
+) -> cp.ndarray:
+    """Helper function to do the wavelet transform in the time domain using CuPy"""
+    # Define constants
+    ND = Nf * Nt
+    K = mult * 2 * Nf
+
+    # Pad the data with K extra values
+    data_pad = cp.concatenate((data, data[:K]))
+
+    # Generate time bin indices
+    time_bins = cp.arange(Nt)
+    jj_base = (time_bins[:, None] * Nf - K // 2) % ND
+    jj = (jj_base + cp.arange(K)[None, :]) % ND
+
+    # Apply the window (phi) to the data
+    wdata = data_pad[jj] * phi[None, :]
+
+    # Perform FFT on the windowed data
+    wdata_trans = rfft(wdata, axis=1)
+
+    # Initialize the wavelet transform result
+    wave = cp.zeros((Nt, Nf))
+
+    # Handle m=0 case for even time bins
+    even_mask = (time_bins % 2 == 0) & (time_bins < Nt - 1)
+    even_indices = cp.nonzero(even_mask)[0]
+
+    # Update wave for m=0 using even time bins
+    wave[even_indices, 0] = cp.real(wdata_trans[even_indices, 0]) / cp.sqrt(2)
+    wave[even_indices + 1, 0] = cp.real(
+        wdata_trans[even_indices, Nf * mult]
+    ) / cp.sqrt(2)
+
+    # Handle other cases (j > 0) using vectorized operations
+    j_range = cp.arange(1, Nf)
+    odd_condition = (time_bins[:, None] + j_range[None, :]) % 2 == 1
+
+    wave[:, 1:] = cp.where(
+        odd_condition,
+        -cp.imag(wdata_trans[:, j_range * mult]),
+        cp.real(wdata_trans[:, j_range * mult]),
+    )
+
+    return wave.T

pywavelet/transforms/cupy/forward/main.py

@@ -0,0 +1,106 @@
+from typing import Union
+
+import cupy as cp
+
+from ....logger import logger
+from ....types import FrequencySeries, TimeSeries, Wavelet
+from ....types.wavelet_bins import _get_bins, _preprocess_bins
+from ...phi_computer import phi_vec, phitilde_vec_norm
+from .from_freq import transform_wavelet_freq_helper
+from .from_time import transform_wavelet_time_helper
+
+
+def from_time_to_wavelet(
+    timeseries: TimeSeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+    mult: int = 32,
+    **kwargs,
+) -> Wavelet:
+    """Transforms time-domain data to wavelet-domain data.
+
+    Warning: there can be significant leakage if mult is too small and the
+    transform is only approximately exact if mult=Nt/2
+
+    Parameters
+    ----------
+    timeseries : TimeSeries
+        Time domain data
+    Nf : int
+        Number of frequency bins
+    Nt : int
+        Number of time bins
+    nx : float, optional
+        Number of standard deviations for the phi_vec, by default 4.
+    mult : int, optional
+        Number of time bins to use for the wavelet transform, by default 32
+    **kwargs:
+        Additional keyword arguments passed to the Wavelet.from_data constructor.
+
+    Returns
+    -------
+    Wavelet
+        Wavelet domain data
+
+    """
+    Nf, Nt = _preprocess_bins(timeseries, Nf, Nt)
+    dt = timeseries.dt
+    t_bins, f_bins = _get_bins(timeseries, Nf, Nt)
+
+    ND = Nf * Nt
+
+    if len(timeseries) != ND:
+        logger.warning(
+            f"len(freqseries)={len(timeseries)} != Nf*Nt={ND}. Truncating to freqseries[:{ND}]"
+        )
+        timeseries = timeseries[:ND]
+    if mult > Nt / 2:
+        logger.warning(
+            f"mult={mult} is too large for Nt={Nt}. This may lead to bogus results."
+        )
+
+    mult = min(mult, Nt // 2)  # make sure K isn't bigger than ND
+    phi = cp.array(phi_vec(Nf, d=nx, q=mult))
+    wave = transform_wavelet_time_helper(
+        cp.array(timeseries.data), Nf=Nf, Nt=Nt, phi=phi, mult=mult
+    )
+    return Wavelet(wave * cp.sqrt(2), time=t_bins, freq=f_bins)
+
+
+def from_freq_to_wavelet(
+    freqseries: FrequencySeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+    **kwargs,
+) -> Wavelet:
+    """Transforms frequency-domain data to wavelet-domain data.
+
+    Parameters
+    ----------
+    freqseries : FrequencySeries
+        Frequency domain data
+    Nf : int
+        Number of frequency bins
+    Nt : int
+        Number of time bins
+    nx : float, optional
+        Number of standard deviations for the phi_vec, by default 4.
+    **kwargs:
+        Additional keyword arguments passed to the Wavelet.from_data constructor.
+
+    Returns
+    -------
+    Wavelet
+        Wavelet domain data
+
+    """
+    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))
+    wave = transform_wavelet_freq_helper(
+        cp.array(freqseries.data), Nf=Nf, Nt=Nt, phif=phif
+    )
+
+    return Wavelet((2 / Nf) * wave * cp.sqrt(2), time=t_bins, freq=f_bins)

pywavelet/transforms/cupy/inverse/__init__.py

@@ -0,0 +1,3 @@
+from .main import from_wavelet_to_freq, from_wavelet_to_time
+
+__all__ = ["from_wavelet_to_time", "from_wavelet_to_freq"]

pywavelet/transforms/cupy/inverse/main.py

@@ -0,0 +1,67 @@
+import cupy as cp
+from cupyx.scipy.fft import rfftfreq
+
+from ....types import FrequencySeries, TimeSeries, Wavelet
+from ...phi_computer import phi_vec, phitilde_vec_norm
+from .to_freq import inverse_wavelet_freq_helper
+
+
+def from_wavelet_to_time(
+    wave_in: Wavelet,
+    dt: float,
+    nx: float = 4.0,
+    mult: int = 32,
+) -> TimeSeries:
+    """Inverse wavelet transform to time domain.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        input wavelet
+    dt : float
+        time step
+    nx : float, optional
+        parameter for phi_vec, by default 4.0
+    mult : int, optional
+        parameter for phi_vec, by default 32
+
+    Returns
+    -------
+    TimeSeries
+        Time domain signal
+    """
+    freq = from_wavelet_to_freq(wave_in, dt=dt, nx=nx)
+    return freq.to_timeseries()
+
+
+def from_wavelet_to_freq(
+    wave_in: Wavelet, dt: float, nx=4.0
+) -> FrequencySeries:
+    """Inverse wavelet transform to frequency domain.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        input wavelet
+    dt : float
+        time step
+    nx : float, optional
+        parameter for phitilde_vec_norm, by default 4.0
+
+    Returns
+    -------
+    FrequencySeries
+        Frequency domain signal
+
+    """
+    phif = cp.array(phitilde_vec_norm(wave_in.Nf, wave_in.Nt, d=nx))
+    freq_data = inverse_wavelet_freq_helper(
+        wave_in.data, phif=phif, Nf=wave_in.Nf, Nt=wave_in.Nt
+    )
+
+    freq_data *= 2 ** (
+        -1 / 2
+    )  # Normalise to get the proper backwards transformation
+
+    freqs = rfftfreq(wave_in.ND * 2, d=dt)[1:]
+    return FrequencySeries(data=freq_data, freq=freqs)

pywavelet/transforms/cupy/inverse/to_freq.py

@@ -0,0 +1,62 @@
+import cupy as cp
+from cupyx.scipy.fft import fft
+
+
+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
+
+    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]
+
+    return res

pywavelet/transforms/jax/forward/from_freq.py

@@ -4,6 +4,10 @@ import jax.numpy as jnp
 from jax import jit
 from jax.numpy.fft import ifft
 
+import logging
+
+logger = logging.getLogger('pywavelet')
+
 
 @partial(jit, static_argnames=("Nf", "Nt"))
 def transform_wavelet_freq_helper(
@@ -23,6 +27,8 @@ def transform_wavelet_freq_helper(
     - wave (jnp.ndarray): 2D array of wavelet-transformed data with shape (Nf, Nt).
     """
 
+    logger.debug(f"[JAX TRANSFORM] Input types [data:{type(data)}, phif:{type(phif)}]")
+
     # Initialize the wavelet output array with zeros (time-rows, frequency-columns)
     wave = jnp.zeros((Nt, Nf))
     f_bins = jnp.arange(Nf)  # Frequency bin indices

pywavelet/transforms/jax/inverse/main.py

@@ -10,7 +10,7 @@ def from_wavelet_to_time(
     wave_in: Wavelet,
     dt: float,
     nx: float = 4.0,
-    mult: int = 32,
+    mult: int = None,
 ) -> TimeSeries:
     """Inverse wavelet transform to time domain.
 
@@ -55,14 +55,12 @@ def from_wavelet_to_freq(
         Frequency domain signal
 
     """
-    phif = jnp.array(phitilde_vec_norm(wave_in.Nf, wave_in.Nt, dt=dt, d=nx))
+    phif = jnp.array(phitilde_vec_norm(wave_in.Nf, wave_in.Nt, d=nx))
     freq_data = inverse_wavelet_freq_helper(
         wave_in.data, phif=phif, Nf=wave_in.Nf, Nt=wave_in.Nt
     )
 
-    freq_data *= 2 ** (
-        -1 / 2
-    )  # Normalise to get the proper backwards transformation
+    freq_data *= 1.0 / jnp.sqrt(2)
 
-    freqs = rfftfreq(wave_in.ND * 2, d=dt)[1:]
+    freqs = rfftfreq(wave_in.ND, d=dt)
     return FrequencySeries(data=freq_data, freq=freqs)