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

pywavelet/transforms/{from_wavelets/inverse_wavelet_time_funcs.py → inverse/to_time.py}

@@ -1,39 +1,48 @@
 """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
 
-
-def inverse_wavelet_time_helper_fast(wave_in, phi, Nf, Nt, mult):
+def inverse_wavelet_time_helper_fast(
+    wave_in: np.ndarray, phi: np.ndarray, Nf: int, Nt: int, mult: int
+) -> np.ndarray:
     """helper loop for fast inverse wavelet transform"""
     ND = Nf * Nt
     K = mult * 2 * Nf
-    # res = np.zeros(ND)
 
     # extend this array, we can use wrapping boundary conditions at end
     res = np.zeros(ND + K + Nf)
 
     afins = np.zeros(2 * Nf, dtype=np.complex128)
 
+    __core(Nf, Nt, K, ND, wave_in, phi, res, afins)
+
+    return res[:ND]
+
+
+def __core(Nf: int, Nt: int, K: int, ND: int, wave_in: np.ndarray, phi: np.ndarray, res: np.ndarray, afins:np.ndarray) -> None:
     for n in range(0, Nt):
         if n % 2 == 0:
             pack_wave_time_helper_compact(n, Nf, Nt, wave_in, afins)
-            ffts_fin = fft.fft(afins)
+            ffts_fin = np.fft.fft(afins)
             unpack_time_wave_helper_compact(n, Nf, Nt, K, phi, ffts_fin, res)
 
     # wrap boundary conditions
-    res[: min(K + Nf, ND)] += res[ND : min(ND + K + Nf, 2 * ND)]
+    res[: min(K + Nf, ND)] += res[ND: min(ND + K + Nf, 2 * ND)]
     if K + Nf > ND:
-        res[: K + Nf - ND] += res[2 * ND : ND + K * Nf]
-
-    res = res[:ND]
-
-    return res
-
-
-@njit()
-def unpack_time_wave_helper(n, Nf, Nt, K, phis, fft_fin_real, res):
+        res[: K + Nf - ND] += res[2 * ND: ND + K * Nf]
+
+
+def unpack_time_wave_helper(
+    n: int,
+    Nf: int,
+    Nt: int,
+    K: int,
+    phis: np.ndarray,
+    fft_fin_real: np.ndarray,
+    res: np.ndarray,
+) -> None:
     """helper for time domain wavelet transform to unpack wavelet domain coefficients"""
     ND = Nf * Nt
 
@@ -52,10 +61,19 @@ def unpack_time_wave_helper(n, Nf, Nt, K, phis, fft_fin_real, res):
             k = 0
 
 
-@njit()
-def unpack_time_wave_helper_compact(n, Nf, Nt, K, phis, fft_fin, res):
+def unpack_time_wave_helper_compact(
+    n: int,
+    Nf: int,
+    Nt: int,
+    K: int,
+    phis: np.ndarray,
+    fft_fin: np.ndarray,
+    res: np.ndarray,
+):
     """helper for time domain wavelet transform to unpack wavelet domain coefficients
     in compact representation where cosine and sine parts are real and imaginary parts
+
+    IN-PLACE EDITS res
     """
     ND = Nf * Nt
     fft_fin_real = np.zeros(4 * Nf)
@@ -79,8 +97,13 @@ def unpack_time_wave_helper_compact(n, Nf, Nt, K, phis, fft_fin, res):
 
 
 @njit()
-def pack_wave_time_helper(n, Nf, Nt, wave_in, afins):
-    """helper for time domain transform to pack wavelet domain coefficients"""
+def pack_wave_time_helper(
+    n: int, Nf: int, Nt: int, wave_in: np.ndarray, afins: np.ndarray
+) -> None:
+    """helper for time domain transform to pack wavelet domain coefficients
+
+    IN-PLACE EDITS afins
+    """
     if n % 2 == 0:
         # assign highest and lowest bin correctly
         afins[0] = np.sqrt(2) * wave_in[n, 0]
@@ -108,9 +131,14 @@ def pack_wave_time_helper(n, Nf, Nt, wave_in, afins):
 
 
 @njit()
-def pack_wave_time_helper_compact(n, Nf, Nt, wave_in, afins):
-    """helper for time domain transform to pack wavelet domain coefficients
+def pack_wave_time_helper_compact(
+    n: int, Nf: int, Nt: int, wave_in: np.ndarray, afins: np.ndarray
+) -> None:
+    """
+    Helper for time domain transform to pack wavelet domain coefficients
     in packed representation with odd and even coefficients in real and imaginary pars
+
+    IN-PLACE EDITS afins
     """
     afins[0] = np.sqrt(2) * wave_in[n, 0]
     if n + 1 < Nt:

pywavelet/transforms/phi_computer.py

@@ -0,0 +1,152 @@
+import numpy as np
+from numpy import fft
+from scipy.special import betainc
+
+PI = np.pi
+
+
+def phitilde_vec(
+    omega: np.ndarray, Nf: int, dt: float, d: float = 4.0
+) -> np.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)
+
+    phi(omega_i) =
+        1/sqrt(2π∆F) if |omega_i| < A
+        1/sqrt(2π∆F) cos(nu_d π/2 * |omega|-A / B) if A < |omega_i| < A + B
+
+    Where nu_d = normalized incomplete beta function
+
+
+
+    Parameters
+    ----------
+    ω : np.ndarray
+        Array of angular frequencies
+    Nf : int
+        Number of frequency bins
+    d : float, optional
+        Number of standard deviations for the gaussian wavelet, by default 4.
+
+    Returns
+    -------
+    np.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)
+
+    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
+    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
+    return phi
+
+
+def __nu_d(
+    omega: np.ndarray, A: float, B: float, d: float = 4.0
+) -> np.ndarray:
+    """Compute the normalized incomplete beta function.
+
+    Parameters
+    ----------
+    ω : np.ndarray
+        Array of angular frequencies
+    A : float
+        Lower bound for the beta function
+    B : float
+        Upper bound for the beta function
+    d : float, optional
+        Number of standard deviations for the gaussian wavelet, by default 4.
+
+    Returns
+    -------
+    np.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
+    return betainc(d, d, x) / betainc(d, d, 1)
+
+
+def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.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)
+
+    # Calculate the unnormalized phitilde (u_phit)
+    u_phit = phitilde_vec(omegas, Nf, dt, d)
+
+    # Normalize the phitilde
+    normalising_factor = np.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, dt, d: float = 4.0, q: int = 16) -> np.ndarray:
+    """get time domain phi as fourier transform of phitilde_vec"""
+    insDOM = 1.0 / np.sqrt(PI / Nf)
+    K = q * 2 * Nf
+    half_K = q * Nf  # np.int64(K/2)
+
+    dom = 2 * 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, dt, d
+    )
+    # negative frequencies
+    DX[half_K + 1 :] = phitilde_vec(
+        -dom * np.arange(half_K - 1, 0, -1), Nf, dt, d
+    )
+    DX = K * fft.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(K / dom)  # *np.linalg.norm(phi)
+
+    fac = np.sqrt(2.0) / nrm
+    phi *= fac
+    return phi

pywavelet/transforms/types/__init__.py

@@ -0,0 +1,4 @@
+from .frequencyseries import FrequencySeries
+from .timeseries import TimeSeries
+from .wavelet import Wavelet
+from .wavelet_mask import WaveletMask

pywavelet/transforms/types/common.py

@@ -0,0 +1,53 @@
+from typing import Literal, Tuple
+
+import numpy as xp
+from numpy.fft import irfft, fft, rfft, rfftfreq
+
+from ...logger import logger
+
+
+def _len_check(d):
+    if not xp.log2(len(d)).is_integer():
+        logger.warning(f"Data length {len(d)} is suggested to be a power of 2")
+
+
+def is_documented_by(original):
+    def wrapper(target):
+        target.__doc__ = original.__doc__
+        return target
+
+    return wrapper
+
+
+def fmt_time(seconds: float, units=False) -> Tuple[str, str]:
+    """Returns formatted time and units [ms, s, min, hr, day]"""
+    t, u = "", ""
+    if seconds < 1e-3:
+        t, u = f"{seconds * 1e6:.2f}", "µs"
+    elif seconds < 1:
+        t, u = f"{seconds * 1e3:.2f}", "ms"
+    elif seconds < 60:
+        t, u = f"{seconds:.2f}", "s"
+    elif seconds < 60 * 60:
+        t, u = f"{seconds / 60:.2f}", "min"
+    elif seconds < 60 * 60 * 24:
+        t, u = f"{seconds / 3600:.2f}", "hr"
+    else:
+        t, u = f"{seconds / 86400:.2f}", "day"
+
+    if units:
+        return t, u
+    return t
+
+
+def fmt_timerange(trange):
+    t0 = fmt_time(trange[0])
+    tend, units = fmt_time(trange[1], units = True)
+    return f"[{t0}, {tend}] {units}"
+
+
+def fmt_pow2(n:float)->str:
+    pow2 = xp.log2(n)
+    if pow2.is_integer():
+        return f"2^{int(pow2)}"
+    return f"{n:,}"

pywavelet/transforms/types/frequencyseries.py

@@ -0,0 +1,237 @@
+import matplotlib.pyplot as plt
+from typing import Tuple, Union, Optional
+
+from .common import is_documented_by, xp, irfft, fmt_time, fmt_pow2
+from .plotting import plot_freqseries, plot_periodogram
+
+__all__ = ["FrequencySeries"]
+
+class FrequencySeries:
+    """
+    A class to represent a one-sided frequency series, with various methods for
+    plotting and converting the series to a time-domain representation.
+
+    Attributes
+    ----------
+    data : xp.ndarray
+        Frequency domain data.
+    freq : xp.ndarray
+        Corresponding frequencies (must be non-negative).
+    """
+
+    def __init__(self, data: xp.ndarray, freq: xp.ndarray, t0: float = 0):
+        """
+        Initialize the FrequencySeries with data and frequencies.
+
+        Parameters
+        ----------
+        data : xp.ndarray
+            Frequency domain data.
+        freq : xp.ndarray
+            Array of frequencies. Must be non-negative.
+        t0 : float, optional
+            Initial time of the time domain signal (default is 0).
+            (This is not used in this class, but is included for compatibility with TimeSeries.)
+
+        Raises
+        ------
+        ValueError
+            If any frequency is negative or if `data` and `freq` do not have the same length.
+        """
+        if xp.any(freq < 0):
+            raise ValueError("FrequencySeries must be one-sided (only non-negative frequencies)")
+        if len(data) != len(freq):
+            raise ValueError(f"data and freq must have the same length ({len(data)} != {len(freq)})")
+        self.data = data
+        self.freq = freq
+        self.t0 = t0
+
+    @is_documented_by(plot_freqseries)
+    def plot(self, ax=None, **kwargs) -> Tuple[plt.Figure, plt.Axes]:
+        return plot_freqseries(
+            self.data, self.freq, self.nyquist_frequency, ax=ax, **kwargs
+        )
+
+    @is_documented_by(plot_periodogram)
+    def plot_periodogram(self, ax=None, **kwargs) -> Tuple[plt.Figure, plt.Axes]:
+        return plot_periodogram(
+            self.data, self.freq, self.fs, ax=ax, **kwargs
+        )
+
+    def __len__(self):
+        """Return the length of the frequency series."""
+        return len(self.data)
+
+    def __getitem__(self, item):
+        """Return the data point at the specified index."""
+        return self.data[item]
+
+    @property
+    def df(self) -> float:
+        """Return the frequency resolution (Δf)."""
+        return float(self.freq[1] - self.freq[0])
+
+    @property
+    def dt(self) -> float:
+        """Return the time resolution (Δt)."""
+        return 1 / self.fs
+
+    @property
+    def sample_rate(self) -> float:
+        """Return the sample rate (fs)."""
+        return 2 * float(self.freq[-1])
+
+    @property
+    def fs(self) -> float:
+        """Return the sample rate (fs)."""
+        return self.sample_rate
+
+    @property
+    def nyquist_frequency(self) -> float:
+        """Return the Nyquist frequency (fs/2)."""
+        return self.sample_rate / 2
+
+    @property
+    def duration(self) -> float:
+        """Return the duration of the time domain signal."""
+        return (2 * (len(self) - 1)) / self.fs
+
+    @property
+    def minimum_frequency(self) -> float:
+        """Return the minimum frequency in the frequency series."""
+        return float(xp.abs(self.freq).min())
+
+    @property
+    def maximum_frequency(self) -> float:
+        """Return the maximum frequency in the frequency series."""
+        return float(xp.abs(self.freq).max())
+
+    @property
+    def range(self) -> Tuple[float, float]:
+        """Return the frequency range (minimum and maximum frequencies)."""
+        return (self.minimum_frequency, self.maximum_frequency)
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        """Return the shape of the data array."""
+        return self.data.shape
+
+    @property
+    def ND(self) -> int:
+        """Return the number of data points in the time domain signal."""
+        return 2 * (len(self) - 1)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the FrequencySeries."""
+        dur = fmt_time(self.duration)
+        n = fmt_pow2(len(self))
+        return f"FrequencySeries(n={n}, frange=[{self.range[0]:.2f}, {self.range[1]:.2f}] Hz, T={dur}, fs={self.fs:.2f} Hz)"
+
+    def noise_weighted_inner_product(self, other: "FrequencySeries", psd:"FrequencySeries") -> float:
+        """
+        Compute the noise-weighted inner product of two FrequencySeries.
+
+        Parameters
+        ----------
+        other : FrequencySeries
+            The other FrequencySeries.
+        psd : FrequencySeries
+            The power spectral density (PSD) of the noise.
+
+        Returns
+        -------
+        float
+            The noise-weighted inner product of the two FrequencySeries.
+        """
+        integrand = xp.real(xp.conj(self.data) * other.data / psd.data)
+        return (4 * self.dt/self.ND) * xp.nansum(integrand)
+
+    def matched_filter_snr(self, other: "FrequencySeries", psd: "FrequencySeries") -> float:
+        """
+        Compute the signal-to-noise ratio (SNR) of a matched filter.
+
+        Parameters
+        ----------
+        other : FrequencySeries
+            The other FrequencySeries.
+        psd : FrequencySeries
+            The power spectral density (PSD) of the noise.
+
+        Returns
+        -------
+        float
+            The SNR of the matched filter.
+        """
+        return xp.sqrt(self.noise_weighted_inner_product(other, psd))
+
+    def optimal_snr(self, psd: "FrequencySeries") -> float:
+        """
+        Compute the optimal signal-to-noise ratio (SNR) of a FrequencySeries.
+
+        Parameters
+        ----------
+        psd : FrequencySeries
+            The power spectral density (PSD) of the noise.
+
+        Returns
+        -------
+        float
+            The optimal SNR of the FrequencySeries.
+        """
+        return xp.sqrt(self.noise_weighted_inner_product(self, psd))
+
+    def to_timeseries(self) -> "TimeSeries":
+        """
+        Convert the frequency series to a time series using inverse Fourier transform.
+
+        Returns
+        -------
+        TimeSeries
+            The corresponding time domain signal.
+        """
+        # Perform the inverse FFT
+        time_data = irfft(self.data, n=2 * (len(self) - 1))
+
+        # Calculate the time array
+        dt = 1 / (2 * self.nyquist_frequency)
+        time = xp.arange(len(time_data)) * dt
+        time += self.t0
+
+        # Create and return a TimeSeries object
+        from .timeseries import TimeSeries
+        return TimeSeries(time_data, time)
+
+
+    def to_wavelet(
+            self,
+            Nf: Union[int, None] = None,
+            Nt: Union[int, None] = None,
+            nx: Optional[float] = 4.0,
+        )->"Wavelet":
+        """
+        Convert the frequency series to a wavelet using inverse Fourier transform.
+
+        Returns
+        -------
+        Wavelet
+            The corresponding wavelet.
+        """
+        from ..forward import from_freq_to_wavelet
+        return from_freq_to_wavelet(self, Nf=Nf, Nt=Nt, nx=nx)
+
+
+    def __eq__(self, other):
+        """Check if two FrequencySeries objects are equal."""
+        data_same = xp.allclose(self.data, other.data)
+        freq_same = xp.allclose(self.freq, other.freq)
+        return data_same and freq_same
+
+    def __copy__(self):
+        return FrequencySeries(
+            xp.copy(self.data),
+            xp.copy(self.freq),
+            t0=self.t0
+        )
+
+    def copy(self):
+        return self.__copy__()