pywavelet 0.0.5__tar.gz → 0.1.0__tar.gz

{pywavelet-0.0.5 → pywavelet-0.1.0}/CHANGELOG.rst

@@ -5,6 +5,19 @@ CHANGELOG
 =========
 
 
+.. _changelog-v0.1.0:
+
+v0.1.0 (2025-01-15)
+===================
+
+Features
+--------
+
+* feat: add wavelet mask and more tests (`e009903`_)
+
+.. _e009903: https://github.com/pywavelet/pywavelet/commit/e00990300d9c013438580c2bc47ea93570fd95be
+
+
 .. _changelog-v0.0.5:
 
 v0.0.5 (2024-12-12)
@@ -15,7 +28,13 @@ Bug Fixes
 
 * fix: update changelog generator (`884c87b`_)
 
+Chores
+------
+
+* chore(release): 0.0.5 (`4ed6b03`_)
+
 .. _884c87b: https://github.com/pywavelet/pywavelet/commit/884c87bcd36b5d21eb1a8e10ee9e0edf6f65d744
+.. _4ed6b03: https://github.com/pywavelet/pywavelet/commit/4ed6b03618347cc179195feec57b05e04a004100
 
 
 .. _changelog-v0.0.4:

{pywavelet-0.0.5 → pywavelet-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.0.5
+Version: 0.1.0
 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.0.5 → pywavelet-0.1.0}/src/pywavelet/_version.py

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

{pywavelet-0.0.5 → pywavelet-0.1.0}/src/pywavelet/transforms/types/__init__.py

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

pywavelet-0.1.0/src/pywavelet/transforms/types/wavelet_mask.py

@@ -0,0 +1,34 @@
+from typing import List
+
+from .common import fmt_timerange, is_documented_by, xp
+
+
+class WaveletMask:
+    def __init__(
+        self,
+        mask: xp.ndarray,
+        time: xp.ndarray,
+        freq: xp.ndarray,
+    ):
+        self.mask = mask
+        self.time = time
+        self.freq = freq
+
+    def __repr__(self):
+        return f"WaveletMask({self.mask.shape}, {fmt_timerange(self.time)}, {self.freq})"
+
+    @classmethod
+    def from_grid(cls, time_grid: xp.ndarray, freq_grid: xp.ndarray):
+        nt, nf = len(time_grid), len(freq_grid)
+        mask = xp.zeros((nf, nt), dtype=bool)
+        return cls(mask, time_grid, freq_grid)
+
+    @classmethod
+    def from_frange(
+        cls, time_grid: xp.ndarray, freq_grid: xp.ndarray, frange: List[float]
+    ):
+        self = cls.from_grid(time_grid, freq_grid)
+        self.mask[
+            (freq_grid >= frange[0]) & (freq_grid <= frange[1]), :
+        ] = True
+        return self

{pywavelet-0.0.5 → pywavelet-0.1.0}/src/pywavelet/utils.py

@@ -3,7 +3,7 @@ from typing import Union
 import numpy as np
 from scipy.interpolate import interp1d
 
-from .transforms.types import FrequencySeries, TimeSeries, Wavelet
+from .transforms.types import FrequencySeries, TimeSeries, Wavelet, WaveletMask
 
 DATA_TYPE = Union[TimeSeries, FrequencySeries, Wavelet]
 
@@ -34,10 +34,13 @@ def evolutionary_psd_from_stationary_psd(
     return Wavelet(psd_grid.T, time=t_grid, freq=f_grid)
 
 
-def noise_weighted_inner_product(d: Wavelet, h: Wavelet, PSD: Wavelet) -> float:
+def noise_weighted_inner_product(
+    d: Wavelet, h: Wavelet, PSD: Wavelet
+) -> float:
     return np.nansum((d.data * h.data) / PSD.data)
 
-def compute_snr(d:Wavelet, h: Wavelet, PSD: Wavelet) -> float:
+
+def compute_snr(d: Wavelet, h: Wavelet, PSD: Wavelet) -> float:
     """Compute the SNR of a model h[ti,fi] given freqseries d[ti,fi] and PSD[ti,fi].
 
     SNR(h) = Sum_{ti,fi} [ h_hat[ti,fi] d[ti,fi] / PSD[ti,fi]
@@ -60,5 +63,14 @@ def compute_snr(d:Wavelet, h: Wavelet, PSD: Wavelet) -> float:
     return np.sqrt(noise_weighted_inner_product(d, h, PSD))
 
 
-def compute_likelihood(data:Wavelet, template:Wavelet, psd:Wavelet) -> float:
-    return -0.5 * np.nansum((data.data - template.data) ** 2 / psd.data)
+def compute_likelihood(
+    data: Wavelet, template: Wavelet, psd: Wavelet, mask: WaveletMask = None
+) -> float:
+    d = data.data
+    h = template.data
+    p = psd.data
+    if mask is not None:
+        m = mask.mask
+        d, h, p = d * m, h * m, p * m
+
+    return -0.5 * np.nansum((d - h) ** 2 / p)

{pywavelet-0.0.5 → pywavelet-0.1.0}/src/pywavelet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.0.5
+Version: 0.1.0
 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.0.5 → pywavelet-0.1.0}/src/pywavelet.egg-info/SOURCES.txt

@@ -42,7 +42,10 @@ src/pywavelet/transforms/types/frequencyseries.py
 src/pywavelet/transforms/types/plotting.py
 src/pywavelet/transforms/types/timeseries.py
 src/pywavelet/transforms/types/wavelet.py
+src/pywavelet/transforms/types/wavelet_mask.py
 tests/conftest.py
+tests/test_lnl.py
+tests/test_mask.py
 tests/test_phi.py
 tests/test_psd.py
 tests/test_roundtrip_conversion.py

{pywavelet-0.0.5 → pywavelet-0.1.0}/tests/conftest.py

@@ -1,13 +1,16 @@
 import os
 
-import pytest
 import numpy as np
+import pytest
 from utils import (
+    BRANCH,
     generate_chirp_time_domain_signal,
     generate_sine_time_domain_signal,
-    BRANCH
 )
 
+from pywavelet.transforms import compute_bins
+from pywavelet.transforms.types import Wavelet
+
 # set global env var "NUMBA_DISABLE_JIT=1"
 os.environ["NUMBA_DISABLE_JIT"] = "1"
 
@@ -18,8 +21,8 @@ fs = 1024
 fmax = 50
 frange = [10, fmax]
 dt = 1 / fs
-Nt = 2 ** 6
-Nf = 2 ** 7
+Nt = 2**6
+Nf = 2**7
 mult = 16
 ND = Nt * Nf
 ts = np.arange(0, ND) * dt
@@ -49,4 +52,24 @@ def sine_time():
 
 @pytest.fixture()
 def sine_freq():
-    return generate_sine_time_domain_signal(ts, ND, f_true=10).to_frequencyseries()
+    return generate_sine_time_domain_signal(
+        ts, ND, f_true=10
+    ).to_frequencyseries()
+
+
+def monochromatic_wnm(
+    f0: float = 20,
+    dt: float = 0.0125,
+    A: float = 2,
+    Nt: int = 128,
+    Nf: int = 256,
+):
+    T = Nt * Nf * dt
+    N = Nt * Nf
+    t_bins, f_bins = compute_bins(Nf, Nt, T)
+    wnm = np.zeros((Nt, Nf))
+    m0 = int(f0 * N * dt)
+    f0_bin_idx = int(2 * m0 / Nt)
+    odd_t_indices = np.arange(Nt) % 2 != 0
+    wnm[odd_t_indices, f0_bin_idx] = A * np.sqrt(2 * Nf)
+    return Wavelet(wnm.T, t_bins, f_bins)

pywavelet-0.1.0/tests/test_lnl.py

@@ -0,0 +1,32 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from conftest import monochromatic_wnm
+
+from pywavelet.transforms import (
+    compute_bins,
+    from_freq_to_wavelet,
+    from_time_to_wavelet,
+)
+from pywavelet.transforms.types import (
+    FrequencySeries,
+    TimeSeries,
+    Wavelet,
+    WaveletMask,
+)
+from pywavelet.utils import (
+    compute_likelihood,
+    compute_snr,
+    evolutionary_psd_from_stationary_psd,
+)
+
+
+def test_lnl(plot_dir):
+    f0 = 20
+    d = monochromatic_wnm(f0=f0)
+    h = monochromatic_wnm(f0=f0)
+    h2 = monochromatic_wnm(f0=f0 + 1)
+    psd = Wavelet(np.ones((h.Nf, h.Nt)), h.time, h.freq)
+    lnl = compute_likelihood(d, h, psd)
+    lnl2 = compute_likelihood(d, h2, psd)
+    assert lnl == 0
+    assert lnl > lnl2

pywavelet-0.1.0/tests/test_mask.py

@@ -0,0 +1,51 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from conftest import monochromatic_wnm
+
+from pywavelet.transforms import (
+    compute_bins,
+    from_freq_to_wavelet,
+    from_time_to_wavelet,
+)
+from pywavelet.transforms.types import (
+    FrequencySeries,
+    TimeSeries,
+    Wavelet,
+    WaveletMask,
+)
+from pywavelet.utils import (
+    compute_likelihood,
+    compute_snr,
+    evolutionary_psd_from_stationary_psd,
+)
+
+
+def test_mask(plot_dir):
+    f0 = 20
+    d = monochromatic_wnm(f0=f0)
+    h = monochromatic_wnm(f0=f0)
+    psd = Wavelet(np.ones((h.Nf, h.Nt)), h.time, h.freq)
+    assert compute_likelihood(d, h, psd) == 0
+
+    mask = WaveletMask.from_frange(h.time, h.freq, [f0 - 0.5, f0 + 0.5])
+    assert np.isclose(compute_likelihood(d, h, psd, mask), 0)
+
+    mask1 = WaveletMask.from_frange(h.time, h.freq, [f0 + 0.5, f0 + 1.5])
+    # assert np.isclose(compute_likelihood(d, h, psd, mask1), 0) == False
+
+    # plt the 3 differnet datasets
+    fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)
+    kwgs = dict(
+        cmap="viridis",
+        aspect="auto",
+        origin="lower",
+        extent=[d.time[0], d.time[-1], d.freq[0], d.freq[-1]],
+    )
+    axes[0].imshow(d.data, **kwgs)
+    axes[0].set_title("data")
+    axes[1].imshow(d.data * mask.mask, **kwgs)
+    axes[1].set_title("data*mask[f0-0.5, f0+0.5]")
+    axes[2].imshow(d.data * mask1.mask, **kwgs)
+    axes[2].set_title("data*mask[f0+0.5, f0+1.5]")
+    plt.tight_layout()
+    fig.savefig(f"{plot_dir}/test_mask.png")

{pywavelet-0.0.5 → pywavelet-0.1.0}/tests/test_snr.py

@@ -1,19 +1,14 @@
+import matplotlib.pyplot as plt
 import numpy as np
-
-from pywavelet.transforms import from_time_to_wavelet, from_freq_to_wavelet, compute_bins
-from pywavelet.transforms.types import TimeSeries, Wavelet, FrequencySeries
+from conftest import monochromatic_wnm
+
+from pywavelet.transforms import (
+    compute_bins,
+    from_freq_to_wavelet,
+    from_time_to_wavelet,
+)
+from pywavelet.transforms.types import FrequencySeries, TimeSeries, Wavelet
 from pywavelet.utils import compute_snr, evolutionary_psd_from_stationary_psd
-import matplotlib.pyplot as plt
-
-def monochromatic_wnm(f0:float, dt:float, A:float, T:float, Nt:int, Nf:int):
-    N = Nt * Nf
-    t_bins, f_bins = compute_bins(Nf, Nt, T)
-    wnm = np.zeros((Nt, Nf))
-    m0 = int(f0 * N * dt)
-    f0_bin_idx = int( 2 * m0 / Nt)
-    odd_t_indices = np.arange(Nt) % 2 != 0
-    wnm[odd_t_indices, f0_bin_idx] = A * np.sqrt(2 * Nf)
-    return Wavelet(wnm.T, t_bins, f_bins)
 
 
 def test_toy_model_snr(plot_dir):
@@ -59,7 +54,7 @@ def test_toy_model_snr(plot_dir):
     ########################################
     ND = len(y)
     Nt = ND // Nf
-    assert Nt > 1 , f"Nt={Nt} must be greater than 1 (ND={ND}, Nf={Nf})"
+    assert Nt > 1, f"Nt={Nt} must be greater than 1 (ND={ND}, Nf={Nf})"
     signal_timeseries = TimeSeries(y, t)
 
     # time --> wavelet
@@ -71,13 +66,17 @@ def test_toy_model_snr(plot_dir):
         t_grid=signal_wavelet.time,
         dt=dt,
     )
-    time2wavelet_snr2 = compute_snr(signal_wavelet, signal_wavelet, psd_wavelet_time) ** 2
+    time2wavelet_snr2 = (
+        compute_snr(signal_wavelet, signal_wavelet, psd_wavelet_time) ** 2
+    )
 
     # freq --> wavelet
     signal_freq = signal_timeseries.to_frequencyseries()
-    psd_freq = FrequencySeries(PSD_AMP * np.ones(len(signal_freq)), signal_freq.freq)
+    psd_freq = FrequencySeries(
+        PSD_AMP * np.ones(len(signal_freq)), signal_freq.freq
+    )
     np.testing.assert_almost_equal(
-        signal_freq.optimal_snr(psd_freq)**2, SNR2_f
+        signal_freq.optimal_snr(psd_freq) ** 2, SNR2_f
     )
 
     # assert len(signal_freq) == (ND // 2 ) + 1 , f"Not one sided spectrum {len(signal_freq)}!={(ND // 2 ) + 1}"
@@ -89,12 +88,15 @@ def test_toy_model_snr(plot_dir):
         t_grid=signal_wavelet_f.time,
         dt=dt,
     )
-    freq2wavelet_snr2 = compute_snr(signal_wavelet_f, signal_wavelet_f, psd_wavelet_freq) ** 2
+    freq2wavelet_snr2 = (
+        compute_snr(signal_wavelet_f, signal_wavelet_f, psd_wavelet_freq) ** 2
+    )
 
     # analytical wavelet
-    analytical_wnm = monochromatic_wnm(f0, dt, A, T, Nt, Nf)
-    analytical_wavelet_snr2 = compute_snr(analytical_wnm, analytical_wnm, psd_wavelet_time) ** 2
-
+    analytical_wnm = monochromatic_wnm(f0, dt, A, Nt, Nf)
+    analytical_wavelet_snr2 = (
+        compute_snr(analytical_wnm, analytical_wnm, psd_wavelet_time) ** 2
+    )
 
     assert np.isclose(
         time2wavelet_snr2, SNR2_t, atol=1e-2
@@ -106,7 +108,6 @@ def test_toy_model_snr(plot_dir):
         analytical_wavelet_snr2, SNR2_t, atol=1e-2
     ), f"SNRs dont match {analytical_wavelet_snr2:.2f}!={SNR2_t:.2f} (factor:{SNR2_t/analytical_wavelet_snr2:.2f})"
 
-
     #### PLOTTING
     fig, axes = plt.subplots(1, 3, figsize=(10, 5))
     signal_wavelet.plot(ax=axes[0], absolute=False)
@@ -123,6 +124,3 @@ def test_toy_model_snr(plot_dir):
     # axes[0].set_ylabel("Frequency [Hz]")
     plt.tight_layout()
     fig.savefig(f"{plot_dir}/snr_comparison.png")
-
-
-