pywavelet 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.1'
-__version_tuple__ = version_tuple = (0, 1, 1)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

pywavelet/backend.py

@@ -0,0 +1,26 @@
+import os
+
+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.scipy.fft import fft, ifft, rfft, irfft, rfftfreq  # type: ignore
+    from jax.scipy.special import betainc  # type: ignore
+
+
+else:
+    import numpy as xp  # type: ignore
+    from numpy.fft import fft, ifft, rfft, irfft, rfftfreq  # type: ignore
+    from scipy.special import betainc # type: ignore
+
+
+PI = xp.pi

pywavelet/transforms/__init__.py

@@ -1,9 +1,17 @@
-from .forward import from_freq_to_wavelet, from_time_to_wavelet
-from .inverse import from_wavelet_to_freq, from_wavelet_to_time
+from .numpy import (
+    from_wavelet_to_time,
+    from_wavelet_to_freq,
+    from_time_to_wavelet,
+    from_freq_to_wavelet,
+)
+from .phi_computer import phi_vec, phitilde_vec_norm, phitilde_vec
 
 __all__ = [
     "from_wavelet_to_time",
     "from_wavelet_to_freq",
     "from_time_to_wavelet",
     "from_freq_to_wavelet",
+    "phitilde_vec_norm",
+    "phi_vec",
+    "phitilde_vec",
 ]

pywavelet/transforms/jax/forward/__init__.py

@@ -0,0 +1,6 @@
+from .main import from_freq_to_wavelet, from_time_to_wavelet
+from ....logger import logger
+
+logger.warning("JAX SUBPACKAGE NOT YET TESTED")
+
+__all__ = ["from_time_to_wavelet", "from_freq_to_wavelet"]

pywavelet/transforms/jax/forward/from_freq.py

@@ -0,0 +1,56 @@
+import jax
+import jax.numpy as jnp
+from functools import partial
+from jax import jit
+from jax.numpy.fft import ifft
+
+@partial(jit, static_argnames=('Nf', 'Nt'))
+def transform_wavelet_freq_helper(
+        data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray
+) -> jnp.ndarray:
+    # Initially all wrk being done in time-rws, freq-cols
+    wave = jnp.zeros((Nt, Nf))
+    f_bins = jnp.arange(Nf)
+
+    i_base = Nt // 2
+    jj_base = f_bins * Nt // 2
+
+    initial_values = jnp.where(
+        (f_bins == 0) | (f_bins == Nf),
+        phif[0] * data[f_bins * Nt // 2] / 2.0,
+        phif[0] * data[f_bins * Nt // 2]
+    )
+
+    DX = jnp.zeros((Nf, Nt), dtype=jnp.complex64)
+    DX = DX.at[:, Nt // 2].set(initial_values)
+
+    j_range = jnp.arange(1 - Nt // 2, Nt // 2)
+    j = jnp.abs(j_range)
+    i = i_base + j_range
+
+    cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)
+    cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)
+    cond3 = j[None, :] == 0
+
+    jj = jj_base[:, None] + j_range[None, :]
+    val = jnp.where(cond1 | cond2, 0.0, phif[j] * data[jj])
+    DX = DX.at[:, i].set(jnp.where(cond3, DX[:, i], val))
+
+    # Vectorized ifft
+    DX_trans = ifft(DX, axis=1)
+
+    # Vectorized __fill_wave_2_jax
+    n_range = jnp.arange(Nt)
+    cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1
+    cond2 = jnp.expand_dims(f_bins % 2 == 1, axis=-1) # shape: (Nf, 1)
+
+    real_part = jnp.where(cond2, -jnp.imag(DX_trans), jnp.real(DX_trans))
+    imag_part = jnp.where(cond2, jnp.real(DX_trans), jnp.imag(DX_trans))
+
+    wave = jnp.where(cond1, imag_part.T, real_part.T)
+
+    ## Special cases for f_bin 0 and Nf
+    wave = wave.at[::2, 0].set(jnp.real(DX_trans[0, ::2] * jnp.sqrt(2)))
+    wave = wave.at[1::2, -1].set(jnp.real(DX_trans[-1, ::2] * jnp.sqrt(2)))
+
+    return wave.T

pywavelet/transforms/jax/forward/from_time.py

@@ -0,0 +1,51 @@
+import jax
+import jax.numpy as jnp
+from jax import jit
+from jax.numpy.fft import rfft
+from functools import partial
+
+@partial(jit, static_argnames=('Nf', 'Nt', 'mult'))
+def transform_wavelet_time_helper(
+    data: jnp.ndarray, phi: jnp.ndarray, Nf: int, Nt: int,  mult: int
+) -> jnp.ndarray:
+    """Helper function to do the wavelet transform in the time domain using JAX"""
+    # Define constants
+    ND = Nf * Nt
+    K = mult * 2 * Nf
+
+    # Pad the data with K extra values
+    data_pad = jnp.concatenate((data, data[:K]))
+
+    # Generate time bin indices
+    time_bins = jnp.arange(Nt)
+    jj_base = (time_bins[:, None] * Nf - K // 2) % ND
+    jj = (jj_base + jnp.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 = jnp.zeros((Nt, Nf))
+
+    # Handle m=0 case for even time bins
+    even_mask = (time_bins % 2 == 0) & (time_bins < Nt - 1)
+    even_indices = jnp.nonzero(even_mask, size=even_mask.shape[0])[0]
+
+    # Update wave for m=0 using even time bins
+    wave = wave.at[even_indices, 0].set(jnp.real(wdata_trans[even_indices, 0]) / jnp.sqrt(2))
+    wave = wave.at[even_indices + 1, 0].set(jnp.real(wdata_trans[even_indices, Nf * mult]) / jnp.sqrt(2))
+
+    # Handle other cases (j > 0) using vectorized operations
+    j_range = jnp.arange(1, Nf)
+    odd_condition = ((time_bins[:, None] + j_range[None, :]) % 2 == 1)
+
+    wave = wave.at[:, 1:].set(
+        jnp.where(odd_condition,
+                  -jnp.imag(wdata_trans[:, j_range * mult]),
+                  jnp.real(wdata_trans[:, j_range * mult]))
+    )
+
+    return wave.T

pywavelet/transforms/jax/forward/main.py

@@ -0,0 +1,110 @@
+from typing import Union
+
+import jax.numpy as jnp
+
+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 = jnp.array(phi_vec(Nf, d=nx, q=mult))
+    wave = transform_wavelet_time_helper(timeseries.data, Nf=Nf, Nt=Nt, phi=phi, mult=mult)
+    return Wavelet(
+        wave* jnp.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 = jnp.array(phitilde_vec_norm(Nf, Nt,  d=nx))
+    wave =  transform_wavelet_freq_helper(
+        freqseries.data, Nf=Nf, Nt=Nt, phif=phif
+    )
+
+    return Wavelet(
+        (2 / Nf) * wave * jnp.sqrt(2),
+        time=t_bins,
+        freq=f_bins
+    )

pywavelet/transforms/numpy/__init__.py

@@ -0,0 +1,10 @@
+from .forward import from_freq_to_wavelet, from_time_to_wavelet
+from .inverse import from_wavelet_to_freq, from_wavelet_to_time
+
+
+__all__ = [
+    "from_wavelet_to_time",
+    "from_wavelet_to_freq",
+    "from_time_to_wavelet",
+    "from_freq_to_wavelet",
+]

pywavelet/transforms/{forward → numpy/forward}/from_freq.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 from numba import njit
-from numpy import fft
 
 
 def transform_wavelet_freq_helper(

pywavelet/transforms/{forward → numpy/forward}/main.py

@@ -2,10 +2,10 @@ from typing import Union
 
 import numpy as np
 
-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 ....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
 
@@ -72,7 +72,7 @@ def from_time_to_wavelet(
         )
 
     mult = min(mult, Nt // 2)  # Ensure mult is not larger than ND/2
-    phi = phi_vec(Nf, dt=dt, d=nx, q=mult)
+    phi = phi_vec(Nf, d=nx, q=mult)
     wave = transform_wavelet_time_helper(timeseries.data, Nf, Nt, phi, mult).T
     return Wavelet(wave * np.sqrt(2), time=t_bins, freq=f_bins)
 
@@ -113,8 +113,7 @@ def from_freq_to_wavelet(
     """
     Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
     t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
-    dt = freqseries.dt
-    phif = phitilde_vec_norm(Nf, Nt, dt=dt, d=nx)
+    phif = phitilde_vec_norm(Nf, Nt,  d=nx)
     wave = transform_wavelet_freq_helper(freqseries.data, Nf, Nt, phif)
 
     return Wavelet((2 / Nf) * wave.T * np.sqrt(2), time=t_bins, freq=f_bins)

pywavelet/transforms/{inverse → numpy/inverse}/main.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from ...transforms.phi_computer import phi_vec, phitilde_vec_norm
-from ...types import FrequencySeries, TimeSeries, Wavelet
+from ...phi_computer import phi_vec, phitilde_vec_norm
+from ....types import FrequencySeries, TimeSeries, Wavelet
 from .to_freq import inverse_wavelet_freq_helper_fast
 from .to_time import inverse_wavelet_time_helper_fast
 
@@ -46,7 +46,7 @@ def from_wavelet_to_time(
     """
 
     mult = min(mult, wave_in.Nt // 2)  # Ensure mult is not larger than ND/2
-    phi = phi_vec(wave_in.Nf, d=nx, q=mult, dt=dt) / 2
+    phi = phi_vec(wave_in.Nf, d=nx, q=mult) / 2
     h_t = inverse_wavelet_time_helper_fast(
         wave_in.data.T, phi, wave_in.Nf, wave_in.Nt, mult
     )
@@ -84,7 +84,7 @@ def from_wavelet_to_freq(
     to ensure the proper backwards transformation.
     """
 
-    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt, dt=dt, d=nx)
+    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt,  d=nx)
     freq_data = inverse_wavelet_freq_helper_fast(
         wave_in.data, phif, wave_in.Nf, wave_in.Nt
     )

pywavelet/transforms/phi_computer.py

@@ -1,13 +1,10 @@
-import numpy as np
-from numpy import fft
-from scipy.special import betainc
+from ..backend import xp, PI, betainc, ifft
 
-PI = np.pi
 
 
 def phitilde_vec(
-    omega: np.ndarray, Nf: int, dt: float, d: float = 4.0
-) -> np.ndarray:
+    omega: xp.ndarray, Nf: int, d: float = 4.0
+) -> xp.ndarray:
     """Compute phi_tilde(omega_i) array, nx is filter steepness, defaults to 4.
 
     Eq 11 of https://arxiv.org/pdf/2009.00043.pdf (Cornish et al. 2020)
@@ -22,7 +19,7 @@ def phitilde_vec(
 
     Parameters
     ----------
-    ω : np.ndarray
+    ω : xp.ndarray
         Array of angular frequencies
     Nf : int
         Number of frequency bins
@@ -31,35 +28,35 @@ def phitilde_vec(
 
     Returns
     -------
-    np.ndarray
+    xp.ndarray
         Array of phi_tilde(omega_i) values
 
     """
     dF = 1.0 / (2 * Nf)  # NOTE: missing 1/dt?
     dOmega = 2 * PI * dF  # Near Eq 10 # 2 pi times DF
-    inverse_sqrt_dOmega = 1.0 / np.sqrt(dOmega)
+    inverse_sqrt_dOmega = 1.0 / xp.sqrt(dOmega)
 
     A = dOmega / 4
     B = dOmega - 2 * A  # Cannot have B \leq 0.
     if B <= 0:
         raise ValueError("B must be greater than 0")
 
-    phi = np.zeros(omega.size)
-    mask = (A <= np.abs(omega)) & (np.abs(omega) < A + B)  # Minor changes
+    phi = xp.zeros(omega.size)
+    mask = (A <= xp.abs(omega)) & (xp.abs(omega) < A + B)  # Minor changes
     vd = (PI / 2.0) * __nu_d(omega[mask], A, B, d=d)  # different from paper
-    phi[mask] = inverse_sqrt_dOmega * np.cos(vd)
-    phi[np.abs(omega) < A] = inverse_sqrt_dOmega
+    phi[mask] = inverse_sqrt_dOmega * xp.cos(vd)
+    phi[xp.abs(omega) < A] = inverse_sqrt_dOmega
     return phi
 
 
 def __nu_d(
-    omega: np.ndarray, A: float, B: float, d: float = 4.0
-) -> np.ndarray:
+    omega: xp.ndarray, A: float, B: float, d: float = 4.0
+) -> xp.ndarray:
     """Compute the normalized incomplete beta function.
 
     Parameters
     ----------
-    ω : np.ndarray
+    ω : xp.ndarray
         Array of angular frequencies
     A : float
         Lower bound for the beta function
@@ -70,29 +67,29 @@ def __nu_d(
 
     Returns
     -------
-    np.ndarray
+    xp.ndarray
         Array of ν_d values
 
     scipy.special.betainc
     https://docs.scipy.org/doc/scipy-1.7.1/reference/reference/generated/scipy.special.betainc.html
 
     """
-    x = (np.abs(omega) - A) / B
+    x = (xp.abs(omega) - A) / B
     return betainc(d, d, x) / betainc(d, d, 1)
 
 
-def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.ndarray:
+def phitilde_vec_norm(Nf: int, Nt: int,  d: float) -> xp.ndarray:
     """Normalize phitilde for inverse frequency domain transform."""
 
     # Calculate the frequency values
     ND = Nf * Nt
-    omegas = 2 * np.pi / ND * np.arange(0, Nt // 2 + 1)
+    omegas = 2 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)
 
     # Calculate the unnormalized phitilde (u_phit)
-    u_phit = phitilde_vec(omegas, Nf, dt, d)
+    u_phit = phitilde_vec(omegas, Nf,  d)
 
     # Normalize the phitilde
-    normalising_factor = np.pi ** (-1 / 2)  # Ollie's normalising factor
+    normalising_factor = PI ** (-1 / 2)  # Ollie's normalising factor
 
     # Notes: this is the overall normalising factor that is different from Cornish's paper
     # It is the only way I can force this code to be consistent with our work in the
@@ -117,15 +114,15 @@ def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.ndarray:
     return u_phit / (normalising_factor)
 
 
-def phi_vec(Nf: int, dt, d: float = 4.0, q: int = 16) -> np.ndarray:
+def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> xp.ndarray:
     """get time domain phi as fourier transform of phitilde_vec"""
-    insDOM = 1.0 / np.sqrt(PI / Nf)
+    insDOM = 1.0 / xp.sqrt(PI / Nf)
     K = q * 2 * Nf
-    half_K = q * Nf  # np.int64(K/2)
+    half_K = q * Nf  # xp.int64(K/2)
 
     dom = 2 * PI / K  # max frequency is K/2*dom = pi/dt = OM
 
-    DX = np.zeros(K, dtype=np.complex128)
+    DX = xp.zeros(K, dtype=xp.complex128)
 
     # zero frequency
     DX[0] = insDOM
@@ -133,20 +130,20 @@ def phi_vec(Nf: int, dt, d: float = 4.0, q: int = 16) -> np.ndarray:
     DX = DX.copy()
     # postive frequencies
     DX[1 : half_K + 1] = phitilde_vec(
-        dom * np.arange(1, half_K + 1), Nf, dt, d
+        dom * xp.arange(1, half_K + 1), Nf, d
     )
     # negative frequencies
     DX[half_K + 1 :] = phitilde_vec(
-        -dom * np.arange(half_K - 1, 0, -1), Nf, dt, d
+        -dom * xp.arange(half_K - 1, 0, -1), Nf,  d
     )
-    DX = K * fft.ifft(DX, K)
+    DX = K * ifft(DX, K)
 
-    phi = np.zeros(K)
-    phi[0:half_K] = np.real(DX[half_K:K])
-    phi[half_K:] = np.real(DX[0:half_K])
+    phi = xp.zeros(K)
+    phi[0:half_K] = xp.real(DX[half_K:K])
+    phi[half_K:] = xp.real(DX[0:half_K])
 
-    nrm = np.sqrt(K / dom)  # *np.linalg.norm(phi)
+    nrm = xp.sqrt(K / dom)  # *xp.linalg.norm(phi)
 
-    fac = np.sqrt(2.0) / nrm
+    fac = xp.sqrt(2.0) / nrm
     phi *= fac
     return phi

pywavelet/types/common.py

@@ -1,9 +1,6 @@
-from typing import Tuple, Union
-
-import numpy as xp
-from numpy.fft import fft, irfft, rfft, rfftfreq  # type: ignore
-
+from typing import Tuple, Union, Callable
 from ..logger import logger
+from ..backend import xp
 
 
 def _len_check(d):
@@ -11,7 +8,7 @@ def _len_check(d):
         logger.warning(f"Data length {len(d)} is suggested to be a power of 2")
 
 
-def is_documented_by(original):
+def is_documented_by(original:Callable):
     def wrapper(target):
         target.__doc__ = original.__doc__
         return target

pywavelet/types/frequencyseries.py

@@ -2,7 +2,8 @@ from typing import Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 
-from .common import fmt_pow2, fmt_time, irfft, is_documented_by, xp
+from ..backend import xp, irfft
+from .common import fmt_pow2, fmt_time, is_documented_by
 from .plotting import plot_freqseries, plot_periodogram
 
 __all__ = ["FrequencySeries"]
@@ -226,7 +227,7 @@ class FrequencySeries:
         Wavelet
             The corresponding wavelet.
         """
-        from ..transforms.forward import from_freq_to_wavelet
+        from ..transforms import from_freq_to_wavelet
 
         return from_freq_to_wavelet(self, Nf=Nf, Nt=Nt, nx=nx)
 

pywavelet/types/timeseries.py

@@ -10,10 +10,9 @@ from .common import (
     fmt_time,
     fmt_timerange,
     is_documented_by,
-    rfft,
-    rfftfreq,
-    xp,
 )
+from ..backend import xp, rfftfreq, rfft
+
 from .plotting import plot_spectrogram, plot_timeseries
 
 __all__ = ["TimeSeries"]

pywavelet/types/wavelet.py

@@ -1,7 +1,6 @@
-from typing import List, Optional, Tuple
+from typing import List, Tuple
 
 import matplotlib.pyplot as plt
-import numpy as np
 
 from .common import fmt_timerange, is_documented_by, xp
 from .plotting import plot_wavelet_grid, plot_wavelet_trend
@@ -286,7 +285,7 @@ class Wavelet:
         TimeSeries
             A `TimeSeries` object representing the time-domain signal.
         """
-        from ..transforms.inverse import from_wavelet_to_time
+        from ..transforms import from_wavelet_to_time
 
         return from_wavelet_to_time(self, dt=self.delta_t, nx=nx, mult=mult)
 
@@ -299,7 +298,7 @@ class Wavelet:
         FrequencySeries
             A `FrequencySeries` object representing the frequency-domain signal.
         """
-        from ..transforms.inverse import from_wavelet_to_freq
+        from ..transforms import from_wavelet_to_freq
 
         return from_wavelet_to_freq(self, dt=self.delta_t, nx=nx)
 
@@ -315,8 +314,8 @@ class Wavelet:
 
         frange = ",".join([f"{f:.2e}" for f in (self.freq[0], self.freq[-1])])
         trange = fmt_timerange((self.t0, self.tend))
-        Nfpow2 = int(np.log2(self.shape[0]))
-        Ntpow2 = int(np.log2(self.shape[1]))
+        Nfpow2 = int(xp.log2(self.shape[0]))
+        Ntpow2 = int(xp.log2(self.shape[1]))
         shapef = f"NfxNf=[2^{Nfpow2}, 2^{Ntpow2}]"
         return f"Wavelet({shapef}, [{frange}]Hz, {trange})"
 
@@ -441,14 +440,16 @@ class WaveletMask(Wavelet):
         return self.data
 
     def __repr__(self):
-        return f"WaveletMask({self.mask.shape}, {fmt_timerange(self.time)}, {self.freq})"
+        rpr = super().__repr__()
+        rpr = rpr.replace("Wavelet", "WaveletMask")
+        return rpr
 
     @classmethod
     def from_frange(
         cls, time_grid: xp.ndarray, freq_grid: xp.ndarray, frange: List[float]
     ):
         self = cls.zeros_from_grid(time_grid, freq_grid)
-        self.mask[
+        self.data[
             (freq_grid >= frange[0]) & (freq_grid <= frange[1]), :
         ] = True
         return self

pywavelet/types/wavelet_bins.py

@@ -1,6 +1,6 @@
 from typing import Tuple, Union
 
-import numpy as np
+from ..backend import xp
 
 from .frequencyseries import FrequencySeries
 from .timeseries import TimeSeries
@@ -33,7 +33,7 @@ def _get_bins(
     data: Union[TimeSeries, FrequencySeries],
     Nf: Union[int, None] = None,
     Nt: Union[int, None] = None,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple[xp.ndarray, xp.ndarray]:
     T = data.duration
     t_bins, f_bins = compute_bins(Nf, Nt, T)
 
@@ -46,12 +46,12 @@ def _get_bins(
     return t_bins, f_bins
 
 
-def compute_bins(Nf: int, Nt: int, T: float) -> Tuple[np.ndarray, np.ndarray]:
+def compute_bins(Nf: int, Nt: int, T: float) -> Tuple[xp.ndarray, xp.ndarray]:
     """Get the bins for the wavelet transform
     Eq 4-6 in Wavelets paper
     """
     delta_T = T / Nt
     delta_F = 1 / (2 * delta_T)
-    t_bins = np.arange(0, Nt) * delta_T
-    f_bins = np.arange(0, Nf) * delta_F
+    t_bins = xp.arange(0, Nt) * delta_T
+    f_bins = xp.arange(0, Nf) * delta_F
     return t_bins, f_bins

{pywavelet-0.1.1.dist-info → pywavelet-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.1.1
+Version: 0.2.0
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/
@@ -32,4 +32,6 @@ Requires-Dist: isort; extra == "dev"
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: jupyter-book; extra == "dev"
 Requires-Dist: GitPython; extra == "dev"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
 

pywavelet-0.2.0.dist-info/RECORD

@@ -0,0 +1,32 @@
+pywavelet/__init__.py,sha256=zcK3Qj4wTrGZF1rU3aT6yA9LvliAOD4DVOY7gNfHhCI,53
+pywavelet/_version.py,sha256=H-qsvrxCpdhaQzyddR-yajEqI71hPxLa4KxzpP3uS1g,411
+pywavelet/backend.py,sha256=k4pDi6f4cwNY6HsUIx1xfuga9f2wLnFr_FIb7Fs1Mds,553
+pywavelet/logger.py,sha256=DyKC-pJ_N9GlVeXL00E1D8hUd8GceBg-pnn7g1YPKcM,391
+pywavelet/utils.py,sha256=l47C643nGlV9q4a0G7wtKzuas0Ou4En2e1FTATCgwlw,1907
+pywavelet/transforms/__init__.py,sha256=uc1fKbBGQgEDafJHk6GEVCc0G_EXL5CtFTKCoFsewoM,381
+pywavelet/transforms/phi_computer.py,sha256=ppFSGJwtNnO2flaiok9ms3WXlAxGQikvA7eNfLgriNQ,4461
+pywavelet/transforms/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pywavelet/transforms/jax/forward/__init__.py,sha256=Ki2RJCfkE9Zy59mqT3oEtGK9Ro9kS5kAz9duZFbxyZo,200
+pywavelet/transforms/jax/forward/from_freq.py,sha256=PsUC7RfrN6pRWWkMSXYHk9z5lxCXW3DfF0m-Rd1GOBE,1785
+pywavelet/transforms/jax/forward/from_time.py,sha256=xNeoZq54B6Gi3TdTTYLr_euaFeJcwpms-lSyCG53AdI,1726
+pywavelet/transforms/jax/forward/main.py,sha256=mm0R4m0pXcnzZB0jCckAc4ynG8STH5mldCmHyyU_PGo,3091
+pywavelet/transforms/numpy/__init__.py,sha256=qFLpGpW3VJSbDp2JpD0Gx7PdwDjH-wrW_aO84ASkIgA,255
+pywavelet/transforms/numpy/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
+pywavelet/transforms/numpy/forward/from_freq.py,sha256=JmJyjrNSb64WnpP50VZRt0BICP64iZJP5QAZTZoexkw,2675
+pywavelet/transforms/numpy/forward/from_time.py,sha256=-Y6VEKwDCYBAHAjLdO46vT-6alpM5fXTgTZ_xkYxqA8,2381
+pywavelet/transforms/numpy/forward/main.py,sha256=3y-YCnhpvN7M4N7xy3CVts7n3QQPwDcJ6mkklX1QbFM,3973
+pywavelet/transforms/numpy/inverse/__init__.py,sha256=J4KIzPzbHNg_8fV_c1MpPq3slSqHQV0j3VFrjfd1Nog,121
+pywavelet/transforms/numpy/inverse/main.py,sha256=-11U5tnDizIssHk824rpYrzbJRl6WFpH6K2KKpVpDnU,2989
+pywavelet/transforms/numpy/inverse/to_freq.py,sha256=so_TDbwdS1N8sd1QcpeAEkI10XFDtoFJGohtD4YulZM,2809
+pywavelet/transforms/numpy/inverse/to_time.py,sha256=w5vmImdsb_4YeInZtXh0llsThLTxS0tmYDlNGJ-IUew,5080
+pywavelet/types/__init__.py,sha256=5YptzQvYBnRfC8N5lpOBf9I1lzpJ0pw0QMnvIcwP3YI,122
+pywavelet/types/common.py,sha256=aIcYq-0KOLHnPQjrVbVmw_TQ3Xm5a7xA30rSgwt3rk4,1275
+pywavelet/types/frequencyseries.py,sha256=hrtLaIUaRrqXw8l00yFe2tPJwpksDa_4n1z6R8XSPPQ,7531
+pywavelet/types/plotting.py,sha256=JNDxeP-fB8U09E90J-rVT-h5yCGA_tGRHtctbgINiRo,10625
+pywavelet/types/timeseries.py,sha256=u35bIqFo3QdlQRBEu6maeWA7DePS11LQ6WMiLjZPcWo,9456
+pywavelet/types/wavelet.py,sha256=el48oyAfwtSw2tCQLUb85F9lKr0qMSRJPUmAUU8TS50,12552
+pywavelet/types/wavelet_bins.py,sha256=GoQGKeZlPc-KbYY7LoxAhB-HI4diHpPcTABBXRfUTLA,1459
+pywavelet-0.2.0.dist-info/METADATA,sha256=rg9LNZxrykv39lKIGuX65v7WzuY93_D0oY268mMe5iw,1362
+pywavelet-0.2.0.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+pywavelet-0.2.0.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
+pywavelet-0.2.0.dist-info/RECORD,,

pywavelet-0.1.1.dist-info/RECORD

@@ -1,25 +0,0 @@
-pywavelet/__init__.py,sha256=zcK3Qj4wTrGZF1rU3aT6yA9LvliAOD4DVOY7gNfHhCI,53
-pywavelet/_version.py,sha256=PKIMyjdUACH4-ONvtunQCnYE2UhlMfp9su83e3HXl5E,411
-pywavelet/logger.py,sha256=DyKC-pJ_N9GlVeXL00E1D8hUd8GceBg-pnn7g1YPKcM,391
-pywavelet/utils.py,sha256=l47C643nGlV9q4a0G7wtKzuas0Ou4En2e1FTATCgwlw,1907
-pywavelet/transforms/__init__.py,sha256=1Ibsup9UwMajeZ9NCQ4BN15qZTeJ_EHkgGu8XNFdA18,255
-pywavelet/transforms/phi_computer.py,sha256=vo1PK9Z70kKV-1lfyRoxWdhSYqwIgJK5CRCCJVei3xI,4545
-pywavelet/transforms/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
-pywavelet/transforms/forward/from_freq.py,sha256=wCiyLpzJE3rGxYjQBdXlwkxPIRYhQWjKq0C_8zYlmDk,2697
-pywavelet/transforms/forward/from_time.py,sha256=-Y6VEKwDCYBAHAjLdO46vT-6alpM5fXTgTZ_xkYxqA8,2381
-pywavelet/transforms/forward/main.py,sha256=Gfy0sp-woy_3ihKMzuk2WuZ7dRk-Mm6sp5dVpYrSvj4,4005
-pywavelet/transforms/inverse/__init__.py,sha256=J4KIzPzbHNg_8fV_c1MpPq3slSqHQV0j3VFrjfd1Nog,121
-pywavelet/transforms/inverse/main.py,sha256=Q7wUaRjB1sgqdB7dniWQGbPTWYQNnIsIrYtjsaHJEdE,3012
-pywavelet/transforms/inverse/to_freq.py,sha256=so_TDbwdS1N8sd1QcpeAEkI10XFDtoFJGohtD4YulZM,2809
-pywavelet/transforms/inverse/to_time.py,sha256=w5vmImdsb_4YeInZtXh0llsThLTxS0tmYDlNGJ-IUew,5080
-pywavelet/types/__init__.py,sha256=5YptzQvYBnRfC8N5lpOBf9I1lzpJ0pw0QMnvIcwP3YI,122
-pywavelet/types/common.py,sha256=OSAW6GqLTgqJ-RYEv__XbzsfFd8AFo5w-ctXQ4XAFZo,1317
-pywavelet/types/frequencyseries.py,sha256=UqcE6UQfw5HZm4na2q9k-X-mfqO-BCiTAvGjaYpSrwc,7518
-pywavelet/types/plotting.py,sha256=JNDxeP-fB8U09E90J-rVT-h5yCGA_tGRHtctbgINiRo,10625
-pywavelet/types/timeseries.py,sha256=6DPO0xLi4Dq2srhJLmavFMf4fYIC3wwdbyMU7lMdjTo,9446
-pywavelet/types/wavelet.py,sha256=ptTEnq6nRZiW2x6g_NV_FuoeVNOsbNOu6caSxYDZNgk,12583
-pywavelet/types/wavelet_bins.py,sha256=SC9nhyigWvOfs2TbH8-Ck_iS1Mrz6sG-8EaIFYLIHuk,1453
-pywavelet-0.1.1.dist-info/METADATA,sha256=OgKTqqZQKfhCCcsYpcXnrut-2EAH42c3tCpCGc0PIDg,1307
-pywavelet-0.1.1.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-pywavelet-0.1.1.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
-pywavelet-0.1.1.dist-info/RECORD,,