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

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)

pywavelet/transforms/jax/inverse/to_freq.py

@@ -1,70 +1,88 @@
+from functools import partial
+
 import jax
 import jax.numpy as jnp
 from jax import jit
 from jax.numpy.fft import fft
 
-from functools import partial
 
-
-@partial(jit, static_argnames=('Nf', 'Nt'))
+@partial(jit, static_argnames=("Nf", "Nt"))
 def inverse_wavelet_freq_helper(
-        wave_in: jnp.ndarray, phif: jnp.ndarray, Nf: int, Nt: int
+    wave_in: jnp.ndarray, phif: jnp.ndarray, Nf: int, Nt: int
 ) -> jnp.ndarray:
-    """JAX vectorized function for inverse_wavelet_freq"""
+    """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
-    prefactor2s = prefactor2s.at[0].set(2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt, 0])
+    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])
+    prefactor2s = prefactor2s.at[Nf].set(
+        2 ** (-1 / 2) * wave_in[(2 * n_range) % Nt + 1, 0]
+    )
 
-    # Other m cases
+    # Other m cases: use meshgrid for vectorization.
     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.
     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)
 
-    # Vectorized FFT
+    # Apply FFT along axis 1 for all m.
     fft_prefactor2s = fft(prefactor2s, axis=1)
 
-    # Vectorized __unpack_wave_inverse
-    ## TODO: Check with Giorgio
-    # ND or ND // 2 + 1?
-    # https://github.com/pywavelet/pywavelet/blob/63151a47cde9edc14f1e7e0bf17f554e78ad257c/src/pywavelet/transforms/from_wavelets/inverse_wavelet_freq_funcs.py
-    res = jnp.zeros(ND, dtype=jnp.complex128)
+    # Allocate the result array with corrected shape.
+    res = jnp.zeros(ND // 2 + 1, dtype=jnp.complex128)
 
-    # m == 0 or m == Nf cases
+    # Unpacking for m == 0 and m == Nf cases:
     i_ind_range = jnp.arange(Nt // 2)
-    i_0 = jnp.abs(i_ind_range)
-    i_Nf = jnp.abs(Nf * Nt // 2 - i_ind_range)
+    i_0 = jnp.abs(i_ind_range)  # for m == 0: i = i_ind_range
+    i_Nf = jnp.abs(Nf * (Nt // 2) - 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])
     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])
 
-    # Special case for m == Nf
-    res = res.at[Nf * Nt // 2].add(fft_prefactor2s[Nf, 0] * phif[Nt // 2])
-
-    # Other m cases
+    # Unpacking for m in (1, ..., Nf-1)
     m_mid = m_range[1:Nf]
-    i_ind_range = jnp.arange(Nt // 2 + 1)
-    m_grid, i_ind_grid = jnp.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 = res.at[i1].add(fft_prefactor2s[m_grid, ind31] * phif[i_ind_grid])
-    res = res.at[i2].add(fft_prefactor2s[m_grid, ind32] * phif[i_ind_grid])
-
-    return res
+    # 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]
+    )
+    res = res.at[i2].add(
+        fft_prefactor2s[m_grid_mid, ind32] * phif[i_ind_grid_mid]
+    )
+
+    return res

pywavelet/transforms/numpy/__init__.py

@@ -1,10 +1,9 @@
 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/numpy/forward/from_freq.py

@@ -1,32 +1,69 @@
 """helper functions for transform_freq"""
 
+import logging
+
 import numpy as np
 from numba import njit
 
+# import rocket_fft.special as rfft  # JIT‐able FFT routines
+
+
+logger = logging.getLogger("pywavelet")
+
 
 def transform_wavelet_freq_helper(
     data: np.ndarray, Nf: int, Nt: int, phif: np.ndarray
 ) -> np.ndarray:
-    """helper to do the wavelet transform using the fast wavelet domain transform"""
-    wave = np.zeros((Nt, Nf))  # wavelet wavepacket transform of the signal
+    """
+    Forward wavelet transform helper using the fast wavelet domain transform,
+    with a JIT-able FFT (rocket-fft) so that the whole transform is jittable.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        Input frequency-domain data (1D array).
+    Nf : int
+        Number of frequency bins.
+    Nt : int
+        Number of time bins.
+    phif : np.ndarray
+        Fourier-domain phase factors (complex-valued array of length Nt//2 + 1).
+
+    Returns
+    -------
+    wave : np.ndarray
+        The wavelet transform output of shape (Nt, Nf). Note that contributions from
+        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)}]"
+    )
+    wave = np.zeros((Nt, Nf), dtype=np.float64)
     DX = np.zeros(Nt, dtype=np.complex128)
-    freq_strain = data.copy()  # Convert
+    # Create a copy of the input data (if needed).
+    freq_strain = data.copy()
     __core(Nf, Nt, DX, freq_strain, phif, wave)
     return wave
 
 
-# @njit()
+@njit()
 def __core(
     Nf: int,
     Nt: int,
     DX: np.ndarray,
-    freq_strain: np.ndarray,
+    data: np.ndarray,
     phif: np.ndarray,
     wave: np.ndarray,
 ):
+    """
+    Process each frequency bin (f_bin) to compute the temporary array DX,
+    perform the inverse FFT using rocket-fft, and then unpack the result into wave.
+
+    This function is fully jittable.
+    """
     for f_bin in range(0, Nf + 1):
-        __fill_wave_1(f_bin, Nt, Nf, DX, freq_strain, phif)
-        # Numba doesn't support np.ifft so we cant jit this
+        __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)
         __fill_wave_2(f_bin, DX_trans, wave, Nt, Nf)
 
@@ -40,24 +77,32 @@ def __fill_wave_1(
     data: np.ndarray,
     phif: np.ndarray,
 ) -> None:
-    """helper for assigning DX in the main loop"""
+    """
+    Fill the temporary complex array DX for the given frequency bin (f_bin)
+    based on the input data and the phase factors phif.
+
+    The computation is performed over a window of indices defined by the current f_bin.
+    """
     i_base = Nt // 2
-    jj_base = f_bin * Nt // 2
+    jj_base = f_bin * (Nt // 2)
 
+    # Special center assignment:
     if f_bin == 0 or f_bin == Nf:
-        # NOTE this term appears to be needed to recover correct constant (at least for m=0), but was previously missing
-        DX[Nt // 2] = phif[0] * data[f_bin * Nt // 2] / 2.0
+        DX[i_base] = phif[0] * data[jj_base] / 2.0
     else:
-        DX[Nt // 2] = phif[0] * data[f_bin * Nt // 2]
+        DX[i_base] = phif[0] * data[jj_base]
 
-    for jj in range(jj_base + 1 - Nt // 2, jj_base + Nt // 2):
+    # Determine the window of indices.
+    start = jj_base + 1 - (Nt // 2)
+    end = jj_base + (Nt // 2)
+    for jj in range(start, end):
         j = np.abs(jj - jj_base)
         i = i_base - jj_base + jj
-        if f_bin == Nf and jj > jj_base:
-            DX[i] = 0.0
-        elif f_bin == 0 and jj < jj_base:
+        # For the highest frequency (f_bin==Nf) or the lowest (f_bin==0), zero out the out-of-range values.
+        if (f_bin == Nf and jj > jj_base) or (f_bin == 0 and jj < jj_base):
             DX[i] = 0.0
         elif j == 0:
+            # Center already assigned.
             continue
         else:
             DX[i] = phif[j] * data[jj]
@@ -67,21 +112,34 @@ def __fill_wave_1(
 def __fill_wave_2(
     f_bin: int, DX_trans: np.ndarray, wave: np.ndarray, Nt: int, Nf: int
 ) -> None:
+    """
+    Unpack the inverse FFT output (DX_trans) into the output wave array.
+
+    For f_bin==0 and f_bin==Nf, the results are stored in column 0 of wave,
+    using even- or odd-indexed rows respectively. For intermediate f_bin values,
+    the values are stored in column f_bin with a sign and component (real or imag)
+    determined by parity.
+    """
+    sqrt2 = np.sqrt(2.0)
     if f_bin == 0:
-        # half of lowest and highest frequency bin pixels are redundant, so store them in even and odd components of f_bin=0 respectively
+        # f_bin==0: assign even-indexed rows of column 0.
         for n in range(0, Nt, 2):
-            wave[n, 0] = DX_trans[n].real * np.sqrt(2)
+            wave[n, 0] = DX_trans[n].real * sqrt2
     elif f_bin == Nf:
+        # f_bin==Nf: assign odd-indexed rows of column 0.
         for n in range(0, Nt, 2):
-            wave[n + 1, 0] = DX_trans[n].real * np.sqrt(2)
+            wave[n + 1, 0] = DX_trans[n].real * sqrt2
     else:
+        # For intermediate f_bin, assign values to column f_bin.
         for n in range(0, Nt):
             if f_bin % 2:
+                # For odd f_bin: use -imag when (n+f_bin) is odd; otherwise use real.
                 if (n + f_bin) % 2:
                     wave[n, f_bin] = -DX_trans[n].imag
                 else:
                     wave[n, f_bin] = DX_trans[n].real
             else:
+                # For even f_bin: use imag when (n+f_bin) is odd; otherwise use real.
                 if (n + f_bin) % 2:
                     wave[n, f_bin] = DX_trans[n].imag
                 else:

pywavelet/transforms/numpy/forward/main.py

@@ -56,7 +56,6 @@ def from_time_to_wavelet(
     to inaccurate results.
     """
     Nf, Nt = _preprocess_bins(timeseries, Nf, Nt)
-    dt = timeseries.dt
     t_bins, f_bins = _get_bins(timeseries, Nf, Nt)
 
     ND = Nf * Nt
@@ -113,7 +112,7 @@ 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)
+    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/numpy/inverse/main.py

@@ -1,14 +1,12 @@
 import numpy as np
 
-from ...phi_computer import phi_vec, phitilde_vec_norm
 from ....types import FrequencySeries, TimeSeries, Wavelet
+from ...phi_computer import phi_vec, phitilde_vec_norm
 from .to_freq import inverse_wavelet_freq_helper_fast
 from .to_time import inverse_wavelet_time_helper_fast
 
 __all__ = ["from_wavelet_to_time", "from_wavelet_to_freq"]
 
-INV_ROOT2 = 1.0 / np.sqrt(2)
-
 
 def from_wavelet_to_time(
     wave_in: Wavelet,
@@ -50,7 +48,7 @@ def from_wavelet_to_time(
     h_t = inverse_wavelet_time_helper_fast(
         wave_in.data.T, phi, wave_in.Nf, wave_in.Nt, mult
     )
-    h_t *= INV_ROOT2  # Normalize to get proper backward transformation
+    h_t *= 1.0 / np.sqrt(2)  # Normalize to get proper backward transformation
     ts = np.arange(0, wave_in.Nf * wave_in.Nt) * dt
     return TimeSeries(data=h_t, time=ts)
 
@@ -84,12 +82,12 @@ def from_wavelet_to_freq(
     to ensure the proper backwards transformation.
     """
 
-    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt,  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
     )
 
-    freq_data *= INV_ROOT2
+    freq_data *= 1.0 / np.sqrt(2)
 
     freqs = np.fft.rfftfreq(wave_in.ND, d=dt)
     return FrequencySeries(data=freq_data, freq=freqs)

pywavelet/transforms/numpy/inverse/to_freq.py

@@ -46,7 +46,7 @@ def __pack_wave_inverse(
             prefactor2s[n] = 2 ** (-1 / 2) * wave_in[(2 * n) % Nt + 1, 0]
     else:
         for n in range(0, Nt):
-            val = wave_in[n, m]  # bug is here
+            val = wave_in[n, m]
             if (n + m) % 2:
                 mult2 = -1j
             else:
@@ -93,3 +93,66 @@ def __unpack_wave_inverse(
             if ind32 == Nt:
                 ind32 = 0
         res[Nt // 2 * m] = fft_prefactor2s[(Nt // 2 * m) % Nt] * phif[0]
+
+
+#
+# # @njit
+# def inverse_wavelet_freq_helper_fast_version2(
+#         wave_in: np.ndarray, phif: np.ndarray, Nf: int, Nt: int
+# ) -> np.ndarray:
+#     wave_in = wave_in.T
+#     ND = Nf * Nt
+#     prefactor2s = np.zeros((Nf + 1, Nt), dtype=np.complex128)
+#     n_range = np.arange(Nt)
+#
+#     # m == 0 case
+#     indices = (2 * n_range) % Nt
+#     prefactor2s[0] = (2 ** (-0.5)) * wave_in[indices, 0]
+#
+#     # m == Nf case
+#     indices = ((2 * n_range) % Nt) + 1
+#     prefactor2s[Nf] = (2 ** (-0.5)) * wave_in[indices, 0]
+#
+#     # For m = 1, ..., Nf-1
+#     m_mid = np.arange(1, Nf)
+#     m_grid, n_grid = np.meshgrid(m_mid, n_range, indexing='ij')
+#     val = wave_in[n_grid, m_grid]
+#     mult2 = np.where(((n_grid + m_grid) % 2) != 0, -1j, 1)
+#     prefactor2s[1:Nf] = mult2 * val
+#
+#     fft_prefactor2s = np.fft.fft(prefactor2s, axis=1)
+#
+#     res = np.zeros(ND // 2 + 1, dtype=np.complex128)
+#
+#     # Unpacking for m == 0 and m == Nf
+#     for m in [0, Nf]:
+#         i_ind_range = np.arange(Nt // 2 + 1 if m == Nf else Nt // 2)
+#         i = np.abs(m * Nt // 2 - i_ind_range)
+#         ind3 = (2 * i) % Nt
+#         res[i] += fft_prefactor2s[m, ind3] * phif[i_ind_range]
+#
+#     # Unpacking for m = 1,..., Nf-1
+#     for m in range(1, Nf):
+#         ind31 = (Nt // 2 * m) % Nt
+#         ind32 = ind31
+#         for i_ind in range(Nt // 2):
+#             i1 = Nt // 2 * m - i_ind
+#             i2 = Nt // 2 * m + i_ind
+#             res[i1] += fft_prefactor2s[m, ind31] * phif[i_ind]
+#             res[i2] += fft_prefactor2s[m, ind32] * phif[i_ind]
+#             ind31 = (ind31 - 1) % Nt
+#             ind32 = (ind32 + 1) % Nt
+#         res[Nt // 2 * m] += fft_prefactor2s[m, (Nt // 2 * m) % Nt] * phif[0]
+#
+#     return res
+#
+# #
+# #
+# # if __name__ == '__main__':
+# #     phif = np.array(np.random.rand(64))
+# #     wave_in = np.array(np.random.rand(64, 64))
+# #     Nf = 64
+# #     Nt = 64
+# #     res = inverse_wavelet_freq_helper_fast(wave_in, phif, Nf, Nt)
+# #     res2 = inverse_wavelet_freq_helper_fast_version2(wave_in, phif, Nf, Nt)
+# #     assert np.allclose(res, res2), "Results do not match!"

pywavelet/transforms/phi_computer.py

@@ -1,10 +1,62 @@
-from ..backend import xp, PI, betainc, ifft
+import numpy as np
+from jaxtyping import Array, Float
 
+from ..backend import betainc, ifft, xp
 
+__all__ = ["phitilde_vec_norm", "phi_vec", "omega"]
 
-def phitilde_vec(
-    omega: xp.ndarray, Nf: int, d: float = 4.0
-) -> xp.ndarray:
+
+def omega(Nf: int, Nt: int) -> Float[Array, " 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)
+
+
+def phitilde_vec_norm(Nf: int, Nt: int, d: float) -> Float[Array, " 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)
+
+
+def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> Float[Array, " 2*q*Nf"]:
+    """get time domain phi as fourier transform of _phitilde_vec
+    q: number of Nf bins over which the window extends?
+
+    """
+    insDOM = 1.0 / np.sqrt(np.pi / Nf)
+    K = q * 2 * 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
+    DX[0] = insDOM
+
+    DX = DX.copy()
+    # postive frequencies
+    DX[1 : half_K + 1] = _phitilde_vec(dom * np.arange(1, half_K + 1), Nf, d)
+    # negative frequencies
+    DX[half_K + 1 :] = _phitilde_vec(
+        -dom * np.arange(half_K - 1, 0, -1), Nf, d
+    )
+    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])
+
+    nrm = np.sqrt(2.0) / np.sqrt(K / dom)  # *xp.linalg.norm(phi)
+
+    phi *= nrm
+    return xp.array(phi)
+
+
+def _phitilde_vec(
+    omega: Float[Array, " Nt//2+1"], Nf: int, d: float = 4.0
+) -> Float[Array, " Nt//2+1"]:
     """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)
@@ -33,25 +85,25 @@ def phitilde_vec(
 
     """
     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 / xp.sqrt(dOmega)
+    dOmega = 2 * np.pi * dF  # Near Eq 10 # 2 pi times DF
+    inverse_sqrt_dOmega = 1.0 / np.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 = 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 = np.zeros(omega.size)
+    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[xp.abs(omega) < A] = inverse_sqrt_dOmega
+    phi[np.abs(omega) < A] = inverse_sqrt_dOmega
     return phi
 
 
-def __nu_d(
-    omega: xp.ndarray, A: float, B: float, d: float = 4.0
-) -> xp.ndarray:
+def _nu_d(
+    omega: Float[Array, " Nt//2+1"], A: float, B: float, d: float = 4.0
+) -> Float[Array, " Nt//2+1"]:
     """Compute the normalized incomplete beta function.
 
     Parameters
@@ -74,76 +126,5 @@ def __nu_d(
     https://docs.scipy.org/doc/scipy-1.7.1/reference/reference/generated/scipy.special.betainc.html
 
     """
-    x = (xp.abs(omega) - A) / B
-    return betainc(d, d, x) / betainc(d, d, 1)
-
-
-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 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)
-
-    # Calculate the unnormalized phitilde (u_phit)
-    u_phit = phitilde_vec(omegas, Nf,  d)
-
-    # Normalize the phitilde
-    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
-    # frequency domain. First note that
-
-    # old normalising factor -- This factor is absolutely ridiculous. Why!?
-    # Matt_normalising_factor = np.sqrt(
-    #     (2 * np.sum(u_phit[1:] ** 2) + u_phit[0] ** 2) * 2 * PI / ND
-    # )
-    # Matt_normalising_factor /= PI**(3/2)/PI
-
-    # The expression above is equal to np.pi**(-1/2) after working through the maths.
-    # I have pulled (2/Nf) from __init__.py (from freq to wavelet) into the normalsiing
-    # factor here. I thnk it's cleaner to have ONE normalising constant. Avoids confusion
-    # and it is much easier to track.
-
-    # TODO: understand the following:
-    # (2 * np.sum(u_phit[1:] ** 2) + u_phit[0] ** 2) = 0.5 * Nt / dOmega
-    # Matt_normalising_factor is equal to 1/sqrt(pi)... why is this computed?
-    # in such a stupid way?
-
-    return u_phit / (normalising_factor)
-
-
-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 / xp.sqrt(PI / Nf)
-    K = q * 2 * Nf
-    half_K = q * Nf  # xp.int64(K/2)
-
-    dom = 2 * PI / K  # max frequency is K/2*dom = pi/dt = OM
-
-    DX = xp.zeros(K, dtype=xp.complex128)
-
-    # zero frequency
-    DX[0] = insDOM
-
-    DX = DX.copy()
-    # postive frequencies
-    DX[1 : half_K + 1] = phitilde_vec(
-        dom * xp.arange(1, half_K + 1), Nf, d
-    )
-    # negative frequencies
-    DX[half_K + 1 :] = phitilde_vec(
-        -dom * xp.arange(half_K - 1, 0, -1), Nf,  d
-    )
-    DX = K * ifft(DX, 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 = xp.sqrt(K / dom)  # *xp.linalg.norm(phi)
-
-    fac = xp.sqrt(2.0) / nrm
-    phi *= fac
-    return phi
+    x = (np.abs(omega) - A) / B
+    return betainc(d, d, x)

pywavelet/types/common.py

@@ -1,6 +1,7 @@
-from typing import Tuple, Union, Callable
-from ..logger import logger
+from typing import Callable, Tuple, Union
+
 from ..backend import xp
+from ..logger import logger
 
 
 def _len_check(d):
@@ -8,7 +9,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:Callable):
+def is_documented_by(original: Callable):
     def wrapper(target):
         target.__doc__ = original.__doc__
         return target

pywavelet/types/frequencyseries.py

@@ -2,7 +2,7 @@ from typing import Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 
-from ..backend import xp, irfft
+from ..backend import irfft, xp
 from .common import fmt_pow2, fmt_time, is_documented_by
 from .plotting import plot_freqseries, plot_periodogram