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

pywavelet/transforms/jax/forward/from_freq.py

@@ -1,17 +1,28 @@
 from functools import partial
 
+import jax
 import jax.numpy as jnp
 from jax import jit
 from jax.numpy.fft import ifft
 
+X64_PRECISION = jax.config.jax_enable_x64
+
+CMPLX_DTYPE = jnp.complex128 if X64_PRECISION else jnp.complex64
+
+
 import logging
 
-logger = logging.getLogger('pywavelet')
+logger = logging.getLogger("pywavelet")
 
 
-@partial(jit, static_argnames=("Nf", "Nt"))
+@partial(jit, static_argnames=("Nf", "Nt", "float_dtype", "complex_dtype"))
 def transform_wavelet_freq_helper(
-    data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray
+    data: jnp.ndarray,
+    Nf: int,
+    Nt: int,
+    phif: jnp.ndarray,
+    float_dtype=jnp.float64,
+    complex_dtype=jnp.complex128,
 ) -> jnp.ndarray:
     """
     Transforms input data from the frequency domain to the wavelet domain using a
@@ -26,78 +37,60 @@ def transform_wavelet_freq_helper(
     Returns:
     - 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
-
-    # 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 = jnp.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"[JAX TRANSFORM] Input types [data:{type(data)},{data.dtype}, phif:{type(phif)},{phif.dtype}]"
+    )
+    half = Nt // 2
+    f_bins = jnp.arange(Nf + 1)  # [0,1,...,Nf]
+
+    # --- 1) build the full (Nf+1, Nt) DX array ---
+    # center (j = 0):
+    center = phif[0] * data[f_bins * half]
+    center = jnp.where((f_bins == 0) | (f_bins == Nf), center / 2.0, center)
+    DX = jnp.zeros((Nf + 1, Nt), complex_dtype)
+    DX = DX.at[:, half].set(center)
+
+    # off-center (j = +/-1...+/-(half−1))
+    offs = jnp.arange(1 - half, half)  # length Nt−1
+    jj = f_bins[:, None] * half + offs[None, :]  # shape (Nf+1, Nt−1)
+    ii = half + offs  # shape (Nt−1,)
+    mask = ((f_bins[:, None] == Nf) & (offs[None, :] > 0)) | (
+        (f_bins[:, None] == 0) & (offs[None, :] < 0)
+    )
+    vals = phif[jnp.abs(offs)] * data[jj]
+    vals = jnp.where(mask, 0.0, vals)
+    DX = DX.at[:, ii].set(vals)
+
+    # --- 2) ifft along time axis ---
+    DXt = jnp.fft.ifft(DX, n=Nt, axis=1)
+
+    # --- 3) unpack into wave (Nt, Nf) ---
+    wave = jnp.zeros((Nt, Nf), float_dtype)
+    sqrt2 = jnp.sqrt(2.0)
+
+    # 3a) DC into col 0, even rows
+    evens = jnp.arange(0, Nt, 2)
+    wave = wave.at[evens, 0].set(jnp.real(DXt[0, evens]) * sqrt2)
+
+    # 3b) Nyquist into col 0, odd rows
+    odds = jnp.arange(1, Nt, 2)
+    wave = wave.at[odds, 0].set(jnp.real(DXt[Nf, evens]) * sqrt2)
+
+    # 3c) intermediate bins 1...Nf−1
+    mids = jnp.arange(1, Nf)  # [1...Nf-1]
+    Dt_mid = DXt[mids, :]  # shape (Nf-1, Nt)
+    real_m = jnp.real(Dt_mid).T  # (Nt, Nf-1)
+    imag_m = jnp.imag(Dt_mid).T  # (Nt, Nf-1)
+
+    odd_f = (mids % 2) == 1  # shape (Nf-1,)
+    n_idx = jnp.arange(Nt)[:, None]  # (Nt,1)
+    odd_n_f = ((n_idx + mids[None, :]) % 2) == 1  # (Nt, Nf-1)
+
+    mid_vals = jnp.where(
+        odd_n_f,
+        jnp.where(odd_f, -imag_m, imag_m),
+        jnp.where(odd_f, real_m, real_m),
     )
+    wave = wave.at[:, 1:Nf].set(mid_vals)
 
-    # Initialize a 2D array to store intermediate FFT input values
-    DX = jnp.zeros(
-        (Nf, Nt), dtype=jnp.complex64
-    )  # TODO: Check dtype -- is complex64 sufficient?
-    DX = DX.at[:, Nt // 2].set(
-        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 = jnp.arange(
-        1 - Nt // 2, Nt // 2
-    )  # Time offsets (centered around zero)
-    j = jnp.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 = jnp.where(
-        cond1 | cond2, 0.0, phif[j] * data[jj]
-    )  # Wavelet filter application
-    DX = DX.at[:, i].set(
-        jnp.where(cond3, DX[:, i], val)
-    )  # Update DX with computed values
-    # At this point, DX contains the data FFT'd with the wavelet filter
-    # (each row is a frequency bin, each column is a time bin)
-
-    # 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 = jnp.arange(Nt)  # Time indices
-    cond1 = (
-        n_range[:, None] + f_bins[None, :]
-    ) % 2 == 1  # Odd/even alternation
-    cond2 = jnp.expand_dims(f_bins % 2 == 1, axis=-1)  # Odd frequency bins
-
-    # Assign real and imaginary parts based on conditions
-    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 frequency bins 0 (DC) and Nf (Nyquist)
-    wave = wave.at[::2, 0].set(
-        jnp.real(DX_trans[0, ::2] * jnp.sqrt(2))
-    )  # DC component
-    wave = wave.at[1::2, -1].set(
-        jnp.real(DX_trans[-1, ::2] * jnp.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/jax/forward/main.py

@@ -2,6 +2,7 @@ from typing import Union
 
 import jax.numpy as jnp
 
+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.
 
@@ -100,7 +98,12 @@ def from_freq_to_wavelet(
     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
+        freqseries.data,
+        Nf=Nf,
+        Nt=Nt,
+        phif=phif,
+        float_dtype=backend.float_dtype,
+        complex_dtype=backend.complex_dtype,
     )
-
-    return Wavelet((2 / Nf) * wave * jnp.sqrt(2), time=t_bins, freq=f_bins)
+    factor = (2 / Nf) * jnp.sqrt(2)
+    return Wavelet(factor * wave, time=t_bins, freq=f_bins)

pywavelet/transforms/jax/inverse/to_freq.py

@@ -11,73 +11,56 @@ def inverse_wavelet_freq_helper(
     wave_in: jnp.ndarray, phif: jnp.ndarray, Nf: int, Nt: int
 ) -> jnp.ndarray:
     """JAX vectorized function for inverse_wavelet_freq with corrected shapes and ranges."""
-    # Transpose to match the NumPy version.
     wave_in = wave_in.T
     ND = Nf * Nt
 
-    # Allocate prefactor2s for each m (shape: (Nf+1, Nt)).
     m_range = jnp.arange(Nf + 1)
     prefactor2s = jnp.zeros((Nf + 1, Nt), dtype=jnp.complex128)
     n_range = jnp.arange(Nt)
 
-    # m == 0 case
+    # Handle m=0 and m=Nf cases
     prefactor2s = prefactor2s.at[0].set(
         2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt, 0]
     )
-
-    # m == Nf case
     prefactor2s = prefactor2s.at[Nf].set(
         2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt + 1, 0]
     )
 
-    # Other m cases: use meshgrid for vectorization.
+    # Handle middle m cases
     m_mid = m_range[1:Nf]
-    # Create grids: n_grid (columns) and m_grid (rows)
-    n_grid, m_grid = jnp.meshgrid(n_range, m_mid)
-    # Index the transposed wave_in using (n, m) as in the NumPy version.
+    n_grid, m_grid = jnp.meshgrid(n_range, m_mid, indexing="ij")
     val = wave_in[n_grid, m_grid]
-    # Apply the alternating multiplier based on (n+m) parity.
     mult2 = jnp.where((n_grid + m_grid) % 2, -1j, 1)
-    prefactor2s = prefactor2s.at[1:Nf].set(mult2 * val)
+    prefactor2s = prefactor2s.at[1:Nf].set((mult2 * val).T)
 
-    # Apply FFT along axis 1 for all m.
     fft_prefactor2s = fft(prefactor2s, axis=1)
 
-    # Allocate the result array with corrected shape.
     res = jnp.zeros(ND // 2 + 1, dtype=jnp.complex128)
 
-    # Unpacking for m == 0 and m == Nf cases:
+    # Unpack for m=0 and m=Nf
     i_ind_range = jnp.arange(Nt // 2)
-    i_0 = jnp.abs(i_ind_range)  # for m == 0: i = i_ind_range
-    i_Nf = jnp.abs(Nf * (Nt // 2) - i_ind_range)
+    i_0 = i_ind_range
     ind3_0 = (2 * i_0) % Nt
-    ind3_Nf = (2 * i_Nf) % Nt
-
     res = res.at[i_0].add(fft_prefactor2s[0, ind3_0] * phif[i_ind_range])
+
+    i_Nf = jnp.abs(Nf * (Nt // 2) - i_ind_range)
+    ind3_Nf = (2 * i_Nf) % Nt
     res = res.at[i_Nf].add(fft_prefactor2s[Nf, ind3_Nf] * phif[i_ind_range])
-    # Special case for m == Nf (ensure the Nyquist frequency is updated correctly)
+
     special_index = jnp.abs(Nf * (Nt // 2) - (Nt // 2))
     res = res.at[special_index].add(fft_prefactor2s[Nf, 0] * phif[Nt // 2])
 
-    # Unpacking for m in (1, ..., Nf-1)
+    # Unpack for middle m values
     m_mid = m_range[1:Nf]
-    # Use range [0, Nt//2) to match the loop in NumPy version.
     i_ind_range_mid = jnp.arange(Nt // 2)
-    # Create meshgrid for vectorized computation.
     m_grid_mid, i_ind_grid_mid = jnp.meshgrid(
         m_mid, i_ind_range_mid, indexing="ij"
     )
-
-    # Compute indices i1 and i2 following the NumPy logic.
     i1 = (Nt // 2) * m_grid_mid - i_ind_grid_mid
     i2 = (Nt // 2) * m_grid_mid + i_ind_grid_mid
-    # Compute the wrapped indices for FFT results.
     ind31 = i1 % Nt
     ind32 = i2 % Nt
 
-    # Update result array using vectorized adds.
-    # Note: You might need to adjust this further if your target res shape is non-trivial,
-    # because here we assume that i1 and i2 indices fall within the allocated result shape.
     res = res.at[i1].add(
         fft_prefactor2s[m_grid_mid, ind31] * phif[i_ind_grid_mid]
     )
@@ -85,4 +68,10 @@ def inverse_wavelet_freq_helper(
         fft_prefactor2s[m_grid_mid, ind32] * phif[i_ind_grid_mid]
     )
 
+    # Correct the center points for middle m's
+    center_indices = (Nt // 2) * m_mid
+    fft_indices = center_indices % Nt
+    values = fft_prefactor2s[m_mid, fft_indices] * phif[0]
+    res = res.at[center_indices].set(values)
+
     return res

pywavelet/transforms/numpy/forward/from_freq.py

@@ -12,7 +12,12 @@ logger = logging.getLogger("pywavelet")
 
 
 def transform_wavelet_freq_helper(
-    data: np.ndarray, Nf: int, Nt: int, phif: np.ndarray
+    data: np.ndarray,
+    Nf: int,
+    Nt: int,
+    phif: np.ndarray,
+    float_dtype: np.dtype = np.float64,
+    complex_dtype: np.dtype = np.complex128,
 ) -> np.ndarray:
     """
     Forward wavelet transform helper using the fast wavelet domain transform,
@@ -36,14 +41,14 @@ def transform_wavelet_freq_helper(
         f_bin==0 and f_bin==Nf are both stored in column 0.
     """
     logger.debug(
-        f"[NUMPY TRANSFORM] Input types [data:{type(data)}, phif:{type(phif)}]"
+        f"[NUMPY TRANSFORM] Input types [data:{type(data)},{data.dtype}, phif:{type(phif)},{phif.dtype}]"
     )
-    wave = np.zeros((Nt, Nf), dtype=np.float64)
-    DX = np.zeros(Nt, dtype=np.complex128)
+    wave = np.zeros((Nt, Nf), dtype=float_dtype)
+    DX = np.zeros(Nt, dtype=complex_dtype)
     # Create a copy of the input data (if needed).
     freq_strain = data.copy()
     __core(Nf, Nt, DX, freq_strain, phif, wave)
-    return wave
+    return wave.T
 
 
 @njit()
@@ -64,7 +69,10 @@ def __core(
     for f_bin in range(0, Nf + 1):
         __fill_wave_1(f_bin, Nt, Nf, DX, data, phif)
         # Use rocket-fft's ifft (which is JIT-able) instead of np.fft.ifft.
-        DX_trans = np.fft.ifft(DX, Nt)
+        DX_trans = np.fft.ifft(
+            DX,
+            Nt,
+        )
         __fill_wave_2(f_bin, DX_trans, wave, Nt, Nf)
 
 

pywavelet/transforms/numpy/forward/main.py

@@ -2,6 +2,7 @@ from typing import Union
 
 import numpy as np
 
+from .... import backend
 from ....logger import logger
 from ....types import FrequencySeries, TimeSeries, Wavelet
 from ....types.wavelet_bins import _get_bins, _preprocess_bins
@@ -112,7 +113,15 @@ def from_freq_to_wavelet(
     """
     Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
     t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
-    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)
+    phif = np.array(phitilde_vec_norm(Nf, Nt, d=nx), dtype=backend.float_dtype)
+    data = np.array(freqseries.data, dtype=backend.complex_dtype)
+    wave = transform_wavelet_freq_helper(
+        data,
+        Nf,
+        Nt,
+        phif,
+        float_dtype=backend.float_dtype,
+        complex_dtype=backend.complex_dtype,
+    )
+    factor = (backend.float_dtype)((2 / Nf) * np.sqrt(2))
+    return Wavelet(factor * wave, time=t_bins, freq=f_bins)

pywavelet/transforms/phi_computer.py

@@ -1,25 +1,43 @@
-import numpy as np
-from jaxtyping import Array, Float
+"""
+This module contains functions to compute the Fourier transform of the
+wavelet function and its normalization. The wavelet function is defined
+in the frequency domain and is used to transform time-domain data into
+the wavelet domain.
+
+Everything in this module is retured as a npfloat64 array.
+"""
 
-from ..backend import betainc, ifft, xp
+from functools import lru_cache
+
+import numpy as np
+from jaxtyping import Float64
+from numpy.fft import ifft
+from scipy.special import betainc
 
 __all__ = ["phitilde_vec_norm", "phi_vec", "omega"]
 
 
-def omega(Nf: int, Nt: int) -> Float[Array, " Nt//2+1"]:
+@lru_cache(maxsize=None)
+def omega(Nf: int, Nt: int) -> Float64[np.ndarray, "{Nt}//2+1"]:
     """Get the angular frequencies of the time domain signal."""
     df = 2 * np.pi / (Nf * Nt)
-    return df * np.arange(0, Nt // 2 + 1)
+    return df * np.arange(0, Nt // 2 + 1, dtype=np.float64)
 
 
-def phitilde_vec_norm(Nf: int, Nt: int, d: float) -> Float[Array, " Nt//2+1"]:
+@lru_cache(maxsize=None)
+def phitilde_vec_norm(
+    Nf: int, Nt: int, d: float
+) -> Float64[np.ndarray, "{Nt}//2+1"]:
     """Normalize phitilde for inverse frequency domain transform."""
     omegas = omega(Nf, Nt)
     _phi_t = _phitilde_vec(omegas, Nf, d) * np.sqrt(np.pi)
-    return xp.array(_phi_t)
+    return np.array(_phi_t)
 
 
-def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> Float[Array, " 2*q*Nf"]:
+@lru_cache(maxsize=None)
+def phi_vec(
+    Nf: int, d: float = 4.0, q: int = 16
+) -> Float64[np.ndarray, "2*{q}*{Nf}"]:
     """get time domain phi as fourier transform of _phitilde_vec
     q: number of Nf bins over which the window extends?
 
@@ -29,7 +47,6 @@ def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> Float[Array, " 2*q*Nf"]:
     half_K = q * Nf  # xp.int64(K/2)
 
     dom = 2 * np.pi / K  # max frequency is K/2*dom = pi/dt = OM
-
     DX = np.zeros(K, dtype=np.complex128)
 
     # zero frequency
@@ -51,12 +68,12 @@ def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> Float[Array, " 2*q*Nf"]:
     nrm = np.sqrt(2.0) / np.sqrt(K / dom)  # *xp.linalg.norm(phi)
 
     phi *= nrm
-    return xp.array(phi)
+    return np.array(phi)
 
 
 def _phitilde_vec(
-    omega: Float[Array, " Nt//2+1"], Nf: int, d: float = 4.0
-) -> Float[Array, " Nt//2+1"]:
+    omega: Float64[np.ndarray, "dim"], Nf: int, d: float = 4.0
+) -> Float64[np.ndarray, "dim"]:
     """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)
@@ -67,8 +84,6 @@ def _phitilde_vec(
 
     Where nu_d = normalized incomplete beta function
 
-
-
     Parameters
     ----------
     ω : xp.ndarray
@@ -93,22 +108,22 @@ def _phitilde_vec(
     if B <= 0:
         raise ValueError("B must be greater than 0")
 
-    phi = np.zeros(omega.size)
+    phi = np.zeros(omega.size, dtype=np.float64)
     mask = (A <= np.abs(omega)) & (np.abs(omega) < A + B)  # Minor changes
     vd = (np.pi / 2.0) * _nu_d(omega[mask], A, B, d=d)  # different from paper
-    phi[mask] = inverse_sqrt_dOmega * xp.cos(vd)
+    phi[mask] = inverse_sqrt_dOmega * np.cos(vd)
     phi[np.abs(omega) < A] = inverse_sqrt_dOmega
     return phi
 
 
 def _nu_d(
-    omega: Float[Array, " Nt//2+1"], A: float, B: float, d: float = 4.0
-) -> Float[Array, " Nt//2+1"]:
+    omega: Float64[np.ndarray, "dim"], A: float, B: float, d: float = 4.0
+) -> Float64[np.ndarray, "dim"]:
     """Compute the normalized incomplete beta function.
 
     Parameters
     ----------
-    ω : xp.ndarray
+    omega : np.ndarray
         Array of angular frequencies
     A : float
         Lower bound for the beta function
@@ -119,7 +134,7 @@ def _nu_d(
 
     Returns
     -------
-    xp.ndarray
+    np.ndarray
         Array of ν_d values
 
     scipy.special.betainc

pywavelet/types/common.py

@@ -1,6 +1,6 @@
 from typing import Callable, Tuple, Union
 
-from ..backend import xp
+from ..backend import complex_dtype, float_dtype, xp
 from ..logger import logger
 
 

pywavelet/types/plotting.py

@@ -303,7 +303,7 @@ def plot_periodogram(
     flow = np.min(np.abs(freq))
     ax.set_xlabel("Frequency [Hz]")
     ax.set_ylabel("Periodigram")
-    # ax.set_xlim(left=flow, right=nyquist_frequency / 2)
+    ax.set_xlim(left=flow, right=nyquist_frequency / 2)
     return ax.figure, ax
 
 

pywavelet/types/timeseries.py

@@ -1,6 +1,7 @@
 from typing import Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
+from astropy.utils.metadata.utils import dtype
 from scipy.signal import butter, sosfiltfilt
 from scipy.signal.windows import tukey
 

pywavelet/types/wavelet.py

@@ -3,7 +3,7 @@ from typing import List, Tuple
 import matplotlib.pyplot as plt
 import numpy as np
 
-from .common import fmt_timerange, is_documented_by, xp
+from .common import float_dtype, fmt_timerange, is_documented_by, xp
 from .plotting import plot_wavelet_grid, plot_wavelet_trend
 from .wavelet_bins import compute_bins
 
@@ -70,7 +70,9 @@ class Wavelet:
             A Wavelet object with zero-filled data array.
         """
         Nf, Nt = len(freq), len(time)
-        return cls(data=xp.zeros((Nf, Nt)), time=time, freq=freq)
+        return cls(
+            data=xp.zeros((Nf, Nt), dtype=float_dtype), time=time, freq=freq
+        )
 
     @classmethod
     def zeros(cls, Nf: int, Nt: int, T: float) -> "Wavelet":

pywavelet/types/wavelet_bins.py

@@ -1,6 +1,6 @@
 from typing import Tuple, Union
 
-from ..backend import xp
+from ..backend import complex_dtype, float_dtype, xp
 from .frequencyseries import FrequencySeries
 from .timeseries import TimeSeries
 
@@ -35,13 +35,7 @@ def _get_bins(
 ) -> Tuple[xp.ndarray, xp.ndarray]:
     T = data.duration
     t_bins, f_bins = compute_bins(Nf, Nt, T)
-
-    # N = len(data)
-    # fs = N / T
-    # assert delta_f == fmax / Nf, f"delta_f={delta_f} != fmax/Nf={fmax/Nf}"
-
     t_bins += data.t0
-
     return t_bins, f_bins
 
 
@@ -51,6 +45,6 @@ def compute_bins(Nf: int, Nt: int, T: float) -> Tuple[xp.ndarray, xp.ndarray]:
     """
     delta_T = T / Nt
     delta_F = 1 / (2 * delta_T)
-    t_bins = xp.arange(0, Nt) * delta_T
-    f_bins = xp.arange(0, Nf) * delta_F
+    t_bins = xp.arange(0, Nt, dtype=float_dtype) * delta_T
+    f_bins = xp.arange(0, Nf, dtype=float_dtype) * delta_F
     return t_bins, f_bins

pywavelet/utils/__init__.py

@@ -0,0 +1,6 @@
+from .analysis import (
+    compute_likelihood,
+    compute_snr,
+    evolutionary_psd_from_stationary_psd,
+    noise_weighted_inner_product,
+)

pywavelet/{utils.py → utils/analysis.py}

@@ -3,7 +3,7 @@ from typing import Union
 import numpy as np
 from scipy.interpolate import interp1d
 
-from .types import FrequencySeries, TimeSeries, Wavelet, WaveletMask
+from ..types import FrequencySeries, TimeSeries, Wavelet, WaveletMask
 
 DATA_TYPE = Union[TimeSeries, FrequencySeries, Wavelet]