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

pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.2'
-__version_tuple__ = version_tuple = (0, 2, 2)
+__version__ = version = '0.2.4'
+__version_tuple__ = version_tuple = (0, 2, 4)

pywavelet/backend.py

@@ -1,5 +1,7 @@
 import os
 
+from .logger import logger
+
 try:
     import jax
 
@@ -13,14 +15,17 @@ 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.numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
     from jax.scipy.special import betainc  # type: ignore
 
+    logger.info("Using JAX backend")
 
 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
+    from numpy.fft import fft, ifft, irfft, rfft, rfftfreq  # type: ignore
+    from scipy.special import betainc  # type: ignore
+
+    logger.info("Using NumPy+numba backend")
 
 
-PI = xp.pi
+PI = xp.pi

pywavelet/transforms/__init__.py

@@ -1,10 +1,21 @@
-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
+from ..backend import use_jax
+
+if use_jax:
+    from .jax import (
+        from_freq_to_wavelet,
+        from_time_to_wavelet,
+        from_wavelet_to_freq,
+        from_wavelet_to_time,
+    )
+else:
+    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, phitilde_vec_norm
 
 __all__ = [
     "from_wavelet_to_time",

pywavelet/transforms/jax/__init__.py

@@ -0,0 +1,12 @@
+from ...logger import logger
+from .forward import from_freq_to_wavelet, from_time_to_wavelet
+from .inverse import from_wavelet_to_freq, from_wavelet_to_time
+
+logger.warning("JAX SUBPACKAGE NOT FULLY TESTED")
+
+__all__ = [
+    "from_wavelet_to_time",
+    "from_wavelet_to_freq",
+    "from_time_to_wavelet",
+    "from_freq_to_wavelet",
+]

pywavelet/transforms/jax/forward/__init__.py

@@ -1,6 +1,3 @@
 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

@@ -1,56 +1,97 @@
-import jax
-import jax.numpy as jnp
 from functools import partial
+
+import jax.numpy as jnp
 from jax import jit
 from jax.numpy.fft import ifft
 
-@partial(jit, static_argnames=('Nf', 'Nt'))
+
+@partial(jit, static_argnames=("Nf", "Nt"))
 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
 ) -> jnp.ndarray:
-    # Initially all wrk being done in time-rws, freq-cols
+    """
+    Transforms input data from the frequency domain to the wavelet domain using a
+    pre-computed wavelet filter (`phif`) and performs an efficient inverse FFT.
+
+    Parameters:
+    - data (jnp.ndarray): 1D array representing the input data in the frequency domain.
+    - Nf (int): Number of frequency bins.
+    - Nt (int): Number of time bins. (Note: Nt * Nf == len(data))
+    - phif (jnp.ndarray): Pre-computed wavelet filter for frequency components.
+
+    Returns:
+    - wave (jnp.ndarray): 2D array of wavelet-transformed data with shape (Nf, Nt).
+    """
+
+    # Initialize the wavelet output array with zeros (time-rows, frequency-columns)
     wave = jnp.zeros((Nt, Nf))
-    f_bins = jnp.arange(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),
-        phif[0] * data[f_bins * Nt // 2] / 2.0,
-        phif[0] * data[f_bins * Nt // 2]
+        (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],
     )
 
-    DX = jnp.zeros((Nf, Nt), dtype=jnp.complex64)
-    DX = DX.at[:, Nt // 2].set(initial_values)
+    # 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)
 
-    j_range = jnp.arange(1 - Nt // 2, Nt // 2)
-    j = jnp.abs(j_range)
-    i = i_base + j_range
+    # 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
 
-    cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)
-    cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)
-    cond3 = j[None, :] == 0
+    # 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)
 
-    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))
+    # 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)
 
-    # Vectorized ifft
+    # Perform the inverse FFT along the time dimension
     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)
+    # 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 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)))
+    # 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 wave.T
+    # Return the wavelet-transformed array (transposed for freq-major layout)
+    return wave.T  # (Nt, Nf) -> (Nf, Nt)

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

@@ -1,10 +1,9 @@
 import jax.numpy as jnp
 from jax.numpy.fft import rfftfreq
 
-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
-# from .inverse_wavelet_time_funcs import inverse_wavelet_time_helper
 
 
 def from_wavelet_to_time(
@@ -65,5 +64,5 @@ def from_wavelet_to_freq(
         -1 / 2
     )  # Normalise to get the proper backwards transformation
 
-    freqs = rfftfreq(wave_in.ND*2, d=dt)[1:]
-    return FrequencySeries(data=freq_data, freq=freqs)
+    freqs = rfftfreq(wave_in.ND * 2, d=dt)[1:]
+    return FrequencySeries(data=freq_data, freq=freqs)

pywavelet/types/wavelet.py

@@ -1,6 +1,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 .plotting import plot_wavelet_grid, plot_wavelet_trend
@@ -343,14 +344,18 @@ class Wavelet:
 
     def __mul__(self, other):
         """Element-wise multiplication of two Wavelet objects."""
-        if isinstance(other, Wavelet):
-            return Wavelet(
-                data=self.data * other.data, time=self.time, freq=self.freq
-            )
+        if isinstance(other, WaveletMask):
+            data = self.data.copy()
+            data[~other.mask] = np.nan
+            return Wavelet(data=data, time=self.time, freq=self.freq)
         elif isinstance(other, float):
             return Wavelet(
                 data=self.data * other, time=self.time, freq=self.freq
             )
+        elif isinstance(other, WaveletMask):
+            return Wavelet(
+                data=self.data * other.data, time=self.time, freq=self.freq
+            )
 
     def __truediv__(self, other):
         """Element-wise division of two Wavelet objects."""
@@ -445,11 +450,40 @@ class WaveletMask(Wavelet):
         return rpr
 
     @classmethod
-    def from_frange(
-        cls, time_grid: xp.ndarray, freq_grid: xp.ndarray, frange: List[float]
+    def from_restrictions(
+        cls,
+        time_grid: xp.ndarray,
+        freq_grid: xp.ndarray,
+        frange: List[float],
+        tgaps: List[Tuple[float, float]] = [],
     ):
+        """
+        Create a WaveletMask object from restrictions on time and frequency.
+
+        Parameters
+        ----------
+        time_grid : xp.ndarray
+            Array of time points.
+        freq_grid : xp.ndarray
+            Array of corresponding frequency points.
+        frange : List[float]
+            Frequency range to include.
+        tgaps : List[Tuple[float, float]]
+            List of time gaps to exclude.
+
+        Returns
+        -------
+        WaveletMask
+            A WaveletMask object with the specified restrictions.
+        """
         self = cls.zeros_from_grid(time_grid, freq_grid)
         self.data[
             (freq_grid >= frange[0]) & (freq_grid <= frange[1]), :
         ] = True
+
+        for tgap in tgaps:
+            self.data[
+                :, (time_grid >= tgap[0]) & (time_grid <= tgap[1])
+            ] = False
+        self.data = self.data.astype(bool)
         return self

{pywavelet-0.2.2.dist-info → pywavelet-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.2
+Version: 0.2.4
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/

{pywavelet-0.2.2.dist-info → pywavelet-0.2.4.dist-info}/RECORD

@@ -1,17 +1,17 @@
 pywavelet/__init__.py,sha256=zcK3Qj4wTrGZF1rU3aT6yA9LvliAOD4DVOY7gNfHhCI,53
-pywavelet/_version.py,sha256=RrHB9KG1O3GPm--rbTedqmZbdDrbgeRLXBmT4OBUqqI,411
-pywavelet/backend.py,sha256=k4pDi6f4cwNY6HsUIx1xfuga9f2wLnFr_FIb7Fs1Mds,553
+pywavelet/_version.py,sha256=4gL0W4-u58XR5lRLpeoIPrGhcewTk0-527de6uTNmkg,411
+pywavelet/backend.py,sha256=SmpgIBHvTO1rtIAQQN_zpVB8i6R-x23FNKJG6_JlrNs,666
 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/__init__.py,sha256=EYX8glRWojYbrjtbgrjS4vigYTRi7FOtIV3D1UwI5fY,604
 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/__init__.py,sha256=D_f-JgFAzOIJ-EuQZhTMziD4MT6lVWS3XV9s51Cu7Kg,335
+pywavelet/transforms/jax/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
+pywavelet/transforms/jax/forward/from_freq.py,sha256=tKEdqPyEvX8ZKVQf16wGxN3d6gkcjm_RtAHQuWHUzy4,3764
 pywavelet/transforms/jax/forward/from_time.py,sha256=xNeoZq54B6Gi3TdTTYLr_euaFeJcwpms-lSyCG53AdI,1726
 pywavelet/transforms/jax/forward/main.py,sha256=mm0R4m0pXcnzZB0jCckAc4ynG8STH5mldCmHyyU_PGo,3091
-pywavelet/transforms/jax/inverse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pywavelet/transforms/jax/inverse/main.py,sha256=ZK8NyfMI6oFYMKcALatETWnCystH0LWjEYInvnMMmh0,1714
+pywavelet/transforms/jax/inverse/__init__.py,sha256=J4KIzPzbHNg_8fV_c1MpPq3slSqHQV0j3VFrjfd1Nog,121
+pywavelet/transforms/jax/inverse/main.py,sha256=-HVOOBsYo3GJvGNCsQLbNPnt9s14JvbB2bGAd9LOr3A,1647
 pywavelet/transforms/jax/inverse/to_freq.py,sha256=ASNARcDBJQr4EizAP_77e5ai36iPwP6hzfvwGbZQ6BM,2295
 pywavelet/transforms/numpy/__init__.py,sha256=qFLpGpW3VJSbDp2JpD0Gx7PdwDjH-wrW_aO84ASkIgA,255
 pywavelet/transforms/numpy/forward/__init__.py,sha256=E_A8plyfTSKDRXlAAvdiRMTe9f3Y6MbK3pXMHFg8mr0,121
@@ -27,9 +27,9 @@ pywavelet/types/common.py,sha256=aIcYq-0KOLHnPQjrVbVmw_TQ3Xm5a7xA30rSgwt3rk4,127
 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.py,sha256=uHJzTS2ZXTRr7I7NHWv3qNjknSBhQUpcED3jM6ti7UM,13587
 pywavelet/types/wavelet_bins.py,sha256=GoQGKeZlPc-KbYY7LoxAhB-HI4diHpPcTABBXRfUTLA,1459
-pywavelet-0.2.2.dist-info/METADATA,sha256=JqDqYCarCAWF5DTmbKzY1u1zUlvlB7OexbrrcLuCRRM,2241
-pywavelet-0.2.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-pywavelet-0.2.2.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
-pywavelet-0.2.2.dist-info/RECORD,,
+pywavelet-0.2.4.dist-info/METADATA,sha256=Thhhz8I2XTKr0mVuf09UpcvjeEGKUnVUX0jxENu6gEQ,2241
+pywavelet-0.2.4.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+pywavelet-0.2.4.dist-info/top_level.txt,sha256=g0Ezt0Rg0X-nrd-a0pAXKVRkuWNsF2M9Ynsjb9b2UYQ,10
+pywavelet-0.2.4.dist-info/RECORD,,