pywavelet 0.0.1b0__py3-none-any.whl → 0.1.0__py3-none-any.whl

pywavelet/__init__.py

@@ -2,4 +2,4 @@
 WDM Wavelet transform
 """
 
-__version__ = "0.0.1"
+__version__ = "0.0.2"

pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.1b0'
-__version_tuple__ = version_tuple = (0, 0, 1)
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)

pywavelet/logger.py

@@ -1,15 +1,14 @@
 import sys
 import warnings
+from rich.logging import RichHandler
+import logging
 
-from loguru import logger
-
-logger.add(
-    sys.stderr,
-    format="|<blue>pywavelet</>|{time:DD/MM HH:mm:ss}|{level}| {message} ",
-    colorize=True,
-    level="INFO",
+FORMAT = "%(message)s"
+logging.basicConfig(
+    level="INFO", format=FORMAT, datefmt="[%X]", handlers=[RichHandler(rich_tracebacks=True)]
 )
 
+logger = logging.getLogger("pywavelet")
 
 warnings.filterwarnings("ignore", category=RuntimeWarning)
 warnings.filterwarnings("ignore", category=UserWarning)

pywavelet/transforms/__init__.py

@@ -1,10 +1,10 @@
-from .from_wavelets import (
-    from_wavelet_to_freq,
-    from_wavelet_to_freq_to_time,
-    from_wavelet_to_time,
-)
-from .to_wavelets import (
-    from_freq_to_wavelet,
-    from_time_to_freq_to_wavelet,
-    from_time_to_wavelet,
-)
+from .forward import from_freq_to_wavelet, from_time_to_wavelet, compute_bins
+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",
+    "compute_bins",
+]

pywavelet/transforms/forward/__init__.py

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

pywavelet/transforms/forward/from_freq.py

@@ -0,0 +1,80 @@
+"""helper functions for transform_freq"""
+import numpy as np
+from numba import njit
+from numpy import fft
+
+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
+    DX = np.zeros(Nt, dtype=np.complex128)
+    freq_strain = data.copy()  # Convert
+    __core(Nf, Nt, DX, freq_strain, phif, wave)
+    return wave
+
+# @njit()
+def __core(Nf:int, Nt:int, DX:np.ndarray, freq_strain:np.ndarray, phif:np.ndarray, wave:np.ndarray):
+    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
+        DX_trans = np.fft.ifft(DX, Nt)
+        __fill_wave_2(f_bin, DX_trans, wave, Nt, Nf)
+
+
+
+
+@njit()
+def __fill_wave_1(
+    f_bin: int,
+    Nt: int,
+    Nf: int,
+    DX: np.ndarray,
+    data: np.ndarray,
+    phif: np.ndarray,
+) -> None:
+    """helper for assigning DX in the main loop"""
+    i_base = Nt // 2
+    jj_base = f_bin * Nt // 2
+
+    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
+    else:
+        DX[Nt // 2] = phif[0] * data[f_bin * Nt // 2]
+
+    for jj in range(jj_base + 1 - Nt // 2, jj_base + Nt // 2):
+        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:
+            DX[i] = 0.0
+        elif j == 0:
+            continue
+        else:
+            DX[i] = phif[j] * data[jj]
+
+@njit()
+def __fill_wave_2(
+    f_bin: int, DX_trans: np.ndarray, wave: np.ndarray, Nt: int, Nf: int
+) -> None:
+    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
+        for n in range(0, Nt, 2):
+            wave[n, 0] = DX_trans[n].real * np.sqrt(2)
+    elif f_bin == Nf:
+        for n in range(0, Nt, 2):
+            wave[n + 1, 0] = DX_trans[n].real * np.sqrt(2)
+    else:
+        for n in range(0, Nt):
+            if f_bin % 2:
+                if (n + f_bin) % 2:
+                    wave[n, f_bin] = -DX_trans[n].imag
+                else:
+                    wave[n, f_bin] =  DX_trans[n].real
+            else:
+                if (n + f_bin) % 2:
+                    wave[n, f_bin] = DX_trans[n].imag
+                else:
+                    wave[n, f_bin] = DX_trans[n].real

pywavelet/transforms/forward/from_time.py

@@ -0,0 +1,66 @@
+"""helper functions for transform_time.py"""
+import numpy as np
+from numba import njit
+from numpy import fft
+
+
+def transform_wavelet_time_helper(
+    data: np.ndarray, Nf: int, Nt: int, phi: np.ndarray, mult: int
+) -> np.ndarray:
+    """helper function to do the wavelet transform in the time domain"""
+    # the time domain freqseries stream
+    ND = Nf * Nt
+    K = mult * 2 * Nf
+    assert len(data) == ND, f"len(data)={len(data)} != Nf*Nt={ND}"
+
+    # windowed freqseries packets
+    wdata = np.zeros(K)
+    wave = np.zeros((Nt, Nf))  # wavelet wavepacket transform of the signal
+    data_pad = np.concatenate((data, data[:K]))
+    __core(Nf, Nt, K, ND, wdata, data_pad, phi, wave, mult)
+    return wave
+
+def __core(Nf: int, Nt: int, K: int, ND: int, wdata: np.ndarray, data_pad: np.ndarray, phi: np.ndarray, wave: np.ndarray, mult: int) -> None:
+    for time_bin_i in range(0, Nt):
+        __fill_wave_1(time_bin_i, K, ND, Nf, wdata, data_pad, phi)
+        wdata_trans = np.fft.rfft(wdata, K)
+        __fill_wave_2(time_bin_i, wave, wdata_trans, Nf, mult)
+
+
+@njit()
+def __fill_wave_1(
+    t_bin: int,
+    K: int,
+    ND: int,
+    Nf: int,
+    wdata: np.ndarray,
+    data_pad: np.ndarray,
+    phi: np.ndarray,
+) -> None:
+    """Assign wdata to be FFT'd in a loop with K extra values on the right to loop."""
+    # wrapping the freqseries is needed to make the sum in Eq 13 in Cornish paper from [-K/2, K/2]
+    jj = (t_bin * Nf - K // 2) % ND  # Periodically wrap the freqseries
+    for j in range(K):
+        # Eq 13 from Cornish paper
+        wdata[j] = data_pad[jj] * phi[j]  # Apply the window
+        jj = (jj + 1) % ND  # Periodically wrap the freqseries
+
+
+@njit()
+def __fill_wave_2(
+    t_bin: int, wave: np.ndarray, wdata_trans: np.ndarray, Nf: int, mult: int
+) -> None:
+    # wdata_trans = np.sum(wdata) * np.exp(1j * np.pi * np.arange(0, 1+K//2) / K)
+
+    # pack fft'd wdata into wave array
+    if t_bin % 2 == 0 and t_bin < wave.shape[0] - 1:  # if EVEN t_bin
+        # m=0 value at even Nt and
+        wave[t_bin, 0] = wdata_trans[0].real / np.sqrt(2)
+        wave[t_bin + 1, 0] = wdata_trans[Nf * mult].real / np.sqrt(2)
+
+    # Cnm in eq 13
+    for j in range(1, Nf):
+        if (t_bin + j) % 2:
+            wave[t_bin, j] = -wdata_trans[j * mult].imag
+        else:
+            wave[t_bin, j] = wdata_trans[j * mult].real

pywavelet/transforms/forward/main.py

@@ -0,0 +1,128 @@
+from typing import Union
+
+import numpy as np
+
+from ...logger import logger
+from ..phi_computer import phi_vec, phitilde_vec_norm
+from ..types import FrequencySeries, TimeSeries, Wavelet
+from .from_freq import transform_wavelet_freq_helper
+from .from_time import transform_wavelet_time_helper
+from .wavelet_bins import _get_bins, _preprocess_bins
+
+
+__all__ = ["from_time_to_wavelet", "from_freq_to_wavelet"]
+
+def from_time_to_wavelet(
+    timeseries: TimeSeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+    mult: int = 32,
+) -> Wavelet:
+    """
+    Transform time-domain data to wavelet-domain data.
+
+    This function performs a forward wavelet transform, converting a
+    time-domain signal into a wavelet-domain representation.
+
+    Parameters
+    ----------
+    timeseries : TimeSeries
+        Input time-domain data, represented as a `TimeSeries` object.
+    Nf : int, optional
+        Number of frequency bins for the wavelet transform. Default is None.
+    Nt : int, optional
+        Number of time bins for the wavelet transform. Default is None.
+    nx : float, optional
+        Number of standard deviations for the `phi_vec`, controlling the
+        width of the wavelets. Default is 4.0.
+    mult : int, optional
+        Number of time bins to use for the wavelet transform. Ensure `mult` is
+        not larger than half the number of time bins (`Nt`). Default is 32.
+
+    Returns
+    -------
+    Wavelet
+        A `Wavelet` object representing the transformed wavelet-domain data.
+
+    Warnings
+    --------
+    There can be significant leakage if `mult` is too small. The transform is
+    only approximately exact if `mult = Nt / 2`.
+
+    Notes
+    -----
+    The function warns when the `mult` value is too large, potentially leading
+    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
+
+    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)  # Ensure mult is not larger than ND/2
+    phi = phi_vec(Nf, dt=dt, 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
+    )
+
+
+def from_freq_to_wavelet(
+    freqseries: FrequencySeries,
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+    nx: float = 4.0,
+) -> Wavelet:
+    """
+    Transform frequency-domain data to wavelet-domain data.
+
+    This function performs a forward wavelet transform, converting a
+    frequency-domain signal into a wavelet-domain representation.
+
+    Parameters
+    ----------
+    freqseries : FrequencySeries
+        Input frequency-domain data, represented as a `FrequencySeries` object.
+    Nf : int, optional
+        Number of frequency bins for the wavelet transform. Default is None.
+    Nt : int, optional
+        Number of time bins for the wavelet transform. Default is None.
+    nx : float, optional
+        Number of standard deviations for the `phi_vec`, controlling the
+        width of the wavelets. Default is 4.0.
+
+    Returns
+    -------
+    Wavelet
+        A `Wavelet` object representing the transformed wavelet-domain data.
+
+    Notes
+    -----
+    The function normalizes the wavelet-domain data to ensure consistency
+    during the transformation process.
+    """
+    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)
+    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/forward/wavelet_bins.py

@@ -0,0 +1,58 @@
+from typing import Tuple, Union
+
+import numpy as np
+
+from ..types import FrequencySeries, TimeSeries
+
+
+def _preprocess_bins(
+    data: Union[TimeSeries, FrequencySeries], Nf=None, Nt=None
+) -> Tuple[int, int]:
+    """preprocess the bins"""
+
+
+    if isinstance(data, TimeSeries):
+        N = len(data)
+    elif isinstance(data, FrequencySeries):
+        # len(d) =  N // 2 + 1
+        N = 2 * (len(data) - 1)
+
+    if Nt is not None and Nf is None:
+        assert 1 <= Nt <= N, f"Nt={Nt} must be between 1 and N={N}"
+        Nf = N // Nt
+
+    elif Nf is not None and Nt is None:
+        assert 1 <= Nf <= N, f"Nf={Nf} must be between 1 and N={N}"
+        Nt = N // Nf
+
+    _N = Nf * Nt
+    return Nf, Nt
+
+
+def _get_bins(
+    data: Union[TimeSeries, FrequencySeries],
+    Nf: Union[int, None] = None,
+    Nt: Union[int, None] = None,
+) -> Tuple[np.ndarray, np.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
+
+
+def compute_bins(Nf:int, Nt:int, T:float) -> Tuple[np.ndarray, np.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
+    return  t_bins, f_bins

pywavelet/transforms/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/inverse/main.py

@@ -0,0 +1,96 @@
+import numpy as np
+
+from ...transforms.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
+
+__all__ = ["from_wavelet_to_time", "from_wavelet_to_freq"]
+
+INV_ROOT2 = 1.0 / np.sqrt(2)
+
+def from_wavelet_to_time(
+    wave_in: Wavelet,
+    dt: float,
+    nx: float = 4.0,
+    mult: int = 32,
+) -> TimeSeries:
+    """
+    Perform an inverse wavelet transform to the time domain.
+
+    This function converts a wavelet-domain signal to a time-domain signal using
+    the inverse wavelet transform algorithm.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        Input wavelet, represented by a `Wavelet` object.
+    dt : float
+        Time step of the wavelet data.
+    nx : float, optional
+        Scaling parameter for the phi vector used in the transformation. Default is 4.0.
+    mult : int, optional
+        Multiplier parameter for the phi vector. Ensures that the `mult` value
+        is not larger than half the number of time bins (`wave_in.Nt`). Default is 32.
+
+    Returns
+    -------
+    TimeSeries
+        A `TimeSeries` object containing the signal transformed into the time domain.
+
+    Notes
+    -----
+    The transformation involves normalizing the output by the square root of 2
+    to ensure the proper backwards transformation.
+    """
+
+    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
+    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
+    ts = np.arange(0, wave_in.Nf * wave_in.Nt) * dt
+    return TimeSeries(data=h_t, time=ts)
+
+
+def from_wavelet_to_freq(
+    wave_in: Wavelet,
+    dt: float,
+    nx:float=4.0
+) -> FrequencySeries:
+    """
+    Perform an inverse wavelet transform to the frequency domain.
+
+    This function converts a wavelet-domain signal into a frequency-domain
+    signal using the inverse wavelet transform algorithm.
+
+    Parameters
+    ----------
+    wave_in : Wavelet
+        Input wavelet, represented by a `Wavelet` object.
+    dt : float
+        Time step of the wavelet data.
+    nx : float, optional
+        Scaling parameter for the phi vector used in the transformation. Default is 4.0.
+
+    Returns
+    -------
+    FrequencySeries
+        A `FrequencySeries` object containing the signal transformed into the frequency domain.
+
+    Notes
+    -----
+    The transformation involves normalizing the output by the square root of 2
+    to ensure the proper backwards transformation.
+    """
+
+    phif = phitilde_vec_norm(wave_in.Nf, wave_in.Nt, dt=dt, d=nx)
+    freq_data = inverse_wavelet_freq_helper_fast(
+        wave_in.data, phif, wave_in.Nf, wave_in.Nt
+    )
+
+    freq_data *= INV_ROOT2
+
+    freqs = np.fft.rfftfreq(wave_in.ND, d=dt)
+    return FrequencySeries(data=freq_data, freq=freqs)

pywavelet/transforms/{from_wavelets/inverse_wavelet_freq_funcs.py → inverse/to_freq.py}

@@ -1,29 +1,61 @@
 """functions for computing the inverse wavelet transforms"""
 import numpy as np
 from numba import njit
+from numpy import fft
 
-from ... import fft_funcs as fft
 
-
-# @njit()
-def inverse_wavelet_freq_helper_fast(wave_in, phif, Nf, Nt):
+def inverse_wavelet_freq_helper_fast(
+    wave_in: np.ndarray, phif: np.ndarray, Nf: int, Nt: int
+) -> np.ndarray:
     """jit compatible loop for inverse_wavelet_freq"""
+    wave_in = wave_in.T
     ND = Nf * Nt
 
     prefactor2s = np.zeros(Nt, np.complex128)
-    res = np.zeros(ND // 2 + 1, dtype=np.complex128)
+    res = np.zeros(ND//2 +1, dtype=np.complex128)
+    __core(Nf, Nt, prefactor2s, wave_in, phif, res)
 
-    for m in range(0, Nf + 1):
-        pack_wave_inverse(m, Nt, Nf, prefactor2s, wave_in)
-        # with numba.objmode(fft_prefactor2s="complex128[:]"):
-        fft_prefactor2s = fft.fft(prefactor2s)
-        unpack_wave_inverse(m, Nt, Nf, phif, fft_prefactor2s, res)
 
     return res
 
+def __core(Nf: int, Nt: int, prefactor2s: np.ndarray, wave_in: np.ndarray, phif: np.ndarray, res: np.ndarray) -> None:
+    for m in range(0, Nf + 1):
+        __pack_wave_inverse(m, Nt, Nf, prefactor2s, wave_in)
+        fft_prefactor2s = np.fft.fft(prefactor2s)
+        __unpack_wave_inverse(m, Nt, Nf, phif, fft_prefactor2s, res)
+
 
 @njit()
-def unpack_wave_inverse(m, Nt, Nf, phif, fft_prefactor2s, res):
+def __pack_wave_inverse(
+    m: int, Nt: int, Nf: int, prefactor2s: np.ndarray, wave_in: np.ndarray
+) -> None:
+    """helper for fast frequency domain inverse transform to prepare for fourier transform"""
+    if m == 0:
+        for n in range(0, Nt):
+            prefactor2s[n] = 2 ** (-1 / 2) * wave_in[(2 * n) % Nt, 0]
+    elif m == Nf:
+        for n in range(0, Nt):
+            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
+            if (n + m) % 2:
+                mult2 = -1j
+            else:
+                mult2 = 1
+
+            prefactor2s[n] = mult2 * val
+
+
+@njit()
+def __unpack_wave_inverse(
+    m: int,
+    Nt: int,
+    Nf: int,
+    phif: np.ndarray,
+    fft_prefactor2s: np.ndarray,
+    res: np.ndarray,
+) -> None:
     """helper for unpacking results of frequency domain inverse transform"""
 
     if m == 0 or m == Nf:
@@ -52,25 +84,4 @@ def unpack_wave_inverse(m, Nt, Nf, phif, fft_prefactor2s, res):
                 ind31 = Nt - 1
             if ind32 == Nt:
                 ind32 = 0
-
         res[Nt // 2 * m] = fft_prefactor2s[(Nt // 2 * m) % Nt] * phif[0]
-
-
-@njit()
-def pack_wave_inverse(m, Nt, Nf, prefactor2s, wave_in):
-    """helper for fast frequency domain inverse transform to prepare for fourier transform"""
-    if m == 0:
-        for n in range(0, Nt):
-            prefactor2s[n] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt, 0]
-    elif m == Nf:
-        for n in range(0, Nt):
-            prefactor2s[n] = 1 / np.sqrt(2) * wave_in[(2 * n) % Nt + 1, 0]
-    else:
-        for n in range(0, Nt):
-            val = wave_in[n, m]
-            if (n + m) % 2:
-                mult2 = -1j
-            else:
-                mult2 = 1
-
-            prefactor2s[n] = mult2 * val