pywavelet 0.2.3__tar.gz → 0.2.5__tar.gz

{pywavelet-0.2.3 → pywavelet-0.2.5}/CHANGELOG.rst

@@ -5,6 +5,48 @@ CHANGELOG
 =========
 
 
+.. _changelog-v0.2.5:
+
+v0.2.5 (2025-01-28)
+===================
+
+Bug Fixes
+---------
+
+* fix: filter first then window (#20)
+
+* fix: filter first then window
+
+* [pre-commit.ci] auto fixes from pre-commit.com hooks
+
+for more information, see https://pre-commit.ci
+
+---------
+
+Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (`3228f7b`_)
+
+.. _3228f7b: https://github.com/pywavelet/pywavelet/commit/3228f7be0d7efb48812920e822a14e795ebac57f
+
+
+.. _changelog-v0.2.4:
+
+v0.2.4 (2025-01-24)
+===================
+
+Chores
+------
+
+* chore(release): 0.2.4 (`8ea5b2f`_)
+
+Unknown
+-------
+
+* Merge branch 'main' of github.com:avivajpeyi/pywavelet into main (`d2c84d9`_)
+
+.. _8ea5b2f: https://github.com/pywavelet/pywavelet/commit/8ea5b2fc0ffa5a1cb274273ec48e164f8bd2064a
+.. _d2c84d9: https://github.com/pywavelet/pywavelet/commit/d2c84d980b1701baf99e40ba6191cbd9336cfa59
+
+
 .. _changelog-v0.2.3:
 
 v0.2.3 (2025-01-24)
@@ -13,12 +55,21 @@ v0.2.3 (2025-01-24)
 Bug Fixes
 ---------
 
+* fix: add masks to filter out gaps (`26fe40a`_)
+
 * fix: add backend check for os.environ (`98c0818`_)
 
 * fix: add test for jax (`1940394`_)
 
+Chores
+------
+
+* chore(release): 0.2.3 (`d067461`_)
+
+.. _26fe40a: https://github.com/pywavelet/pywavelet/commit/26fe40ace80d5f9d598e1efeba2f8ca4a6f1043b
 .. _98c0818: https://github.com/pywavelet/pywavelet/commit/98c0818078190d829a23734f932f1f93c9932167
 .. _1940394: https://github.com/pywavelet/pywavelet/commit/194039437a3a9b3ada303d101b4e2573ab7d0afd
+.. _d067461: https://github.com/pywavelet/pywavelet/commit/d0674615df328774a0d80eb224b5c503fbd8f332
 
 
 .. _changelog-v0.2.2:

{pywavelet-0.2.3 → pywavelet-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.3
+Version: 0.2.5
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/
@@ -76,10 +76,10 @@ First set up a conda environment with the latest version of python.
 Test code
 ---------
 
-Locate directory /tests from root directory. run 
+Locate directory /tests from root directory. run
 
 .. code-block::
 
     $ pytest .
 
-Hopefully everything should run fine. 
+Hopefully everything should run fine.

{pywavelet-0.2.3 → pywavelet-0.2.5}/README.rst

@@ -39,10 +39,10 @@ First set up a conda environment with the latest version of python.
 Test code
 ---------
 
-Locate directory /tests from root directory. run 
+Locate directory /tests from root directory. run
 
 .. code-block::
 
     $ pytest .
 
-Hopefully everything should run fine. 
+Hopefully everything should run fine.

{pywavelet-0.2.3 → pywavelet-0.2.5}/docs/api.rst

@@ -32,4 +32,3 @@ Types
 .. autoclass:: pywavelet.transforms.types.Wavelet
     :members:
     :undoc-members:
-

{pywavelet-0.2.3 → pywavelet-0.2.5}/docs/example.ipynb

@@ -70,6 +70,7 @@
     "    ax.ticklabel_format(axis=\"x\", style=\"sci\", scilimits=(0, 0))\n",
     "    return ax\n",
     "\n",
+    "\n",
     "# Sizes\n",
     "dt = 1 / 512\n",
     "Nt, Nf = 2**6, 2**6\n",
@@ -88,7 +89,7 @@
     "# transform back to time domain\n",
     "h_reconstructed = from_wavelet_to_time(h_wavelet, dt=h_time.dt, mult=mult)\n",
     "\n",
-    "# Plots \n",
+    "# Plots\n",
     "fig, axes = plt.subplots(1, 4, figsize=(18, 4))\n",
     "_ = h_time.plot_spectrogram(ax=axes[0])\n",
     "_ = h_wavelet.plot(ax=axes[1], absolute=True, cmap=\"Reds\")\n",
@@ -146,7 +147,7 @@
     "freq = np.fft.rfftfreq(N, dt)\n",
     "hf = np.zeros_like(freq, dtype=np.complex128)\n",
     "hf[np.argmin(np.abs(freq - f0))] = 1.0\n",
-    "    \n",
+    "\n",
     "\n",
     "h_freq = FrequencySeries(data=hf, freq=freq)\n",
     "h_wavelet = from_freq_to_wavelet(h_freq, Nf=Nf, Nt=Nt)\n",
@@ -154,10 +155,10 @@
     "\n",
     "\n",
     "fig, axes = plt.subplots(1, 2, figsize=(9, 4))\n",
-    "_ = h_freq.plot(ax=axes[0], label='Original')\n",
+    "_ = h_freq.plot(ax=axes[0], label=\"Original\")\n",
     "_ = h_wavelet.plot(ax=axes[1], absolute=True, cmap=\"Reds\")\n",
-    "_ = h_reconstructed.plot(ax=axes[0], ls=\":\", label='Reconstructed')\n",
-    "axes[1].set_ylim(f0-5, f0+5)\n",
+    "_ = h_reconstructed.plot(ax=axes[0], ls=\":\", label=\"Reconstructed\")\n",
+    "axes[1].set_ylim(f0 - 5, f0 + 5)\n",
     "axes[0].legend()\n",
     "fig.savefig(\"roundtrip_freq.png\")\n",
     "plt.close()"

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.3'
-__version_tuple__ = version_tuple = (0, 2, 3)
+__version__ = version = '0.2.5'
+__version_tuple__ = version_tuple = (0, 2, 5)

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/jax/forward/from_time.py

@@ -1,12 +1,14 @@
+from functools import partial
+
 import jax
 import jax.numpy as jnp
 from jax import jit
 from jax.numpy.fft import rfft
-from functools import partial
 
-@partial(jit, static_argnames=('Nf', 'Nt', 'mult'))
+
+@partial(jit, static_argnames=("Nf", "Nt", "mult"))
 def transform_wavelet_time_helper(
-    data: jnp.ndarray, phi: jnp.ndarray, Nf: int, Nt: int,  mult: int
+    data: jnp.ndarray, phi: jnp.ndarray, Nf: int, Nt: int, mult: int
 ) -> jnp.ndarray:
     """Helper function to do the wavelet transform in the time domain using JAX"""
     # Define constants
@@ -35,17 +37,23 @@ def transform_wavelet_time_helper(
     even_indices = jnp.nonzero(even_mask, size=even_mask.shape[0])[0]
 
     # Update wave for m=0 using even time bins
-    wave = wave.at[even_indices, 0].set(jnp.real(wdata_trans[even_indices, 0]) / jnp.sqrt(2))
-    wave = wave.at[even_indices + 1, 0].set(jnp.real(wdata_trans[even_indices, Nf * mult]) / jnp.sqrt(2))
+    wave = wave.at[even_indices, 0].set(
+        jnp.real(wdata_trans[even_indices, 0]) / jnp.sqrt(2)
+    )
+    wave = wave.at[even_indices + 1, 0].set(
+        jnp.real(wdata_trans[even_indices, Nf * mult]) / jnp.sqrt(2)
+    )
 
     # Handle other cases (j > 0) using vectorized operations
     j_range = jnp.arange(1, Nf)
-    odd_condition = ((time_bins[:, None] + j_range[None, :]) % 2 == 1)
+    odd_condition = (time_bins[:, None] + j_range[None, :]) % 2 == 1
 
     wave = wave.at[:, 1:].set(
-        jnp.where(odd_condition,
-                  -jnp.imag(wdata_trans[:, j_range * mult]),
-                  jnp.real(wdata_trans[:, j_range * mult]))
+        jnp.where(
+            odd_condition,
+            -jnp.imag(wdata_trans[:, j_range * mult]),
+            jnp.real(wdata_trans[:, j_range * mult]),
+        )
     )
 
-    return wave.T
+    return wave.T

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/jax/forward/main.py

@@ -62,10 +62,10 @@ def from_time_to_wavelet(
 
     mult = min(mult, Nt // 2)  # make sure K isn't bigger than ND
     phi = jnp.array(phi_vec(Nf, d=nx, q=mult))
-    wave = transform_wavelet_time_helper(timeseries.data, Nf=Nf, Nt=Nt, phi=phi, mult=mult)
-    return Wavelet(
-        wave* jnp.sqrt(2), time=t_bins, freq=f_bins
+    wave = transform_wavelet_time_helper(
+        timeseries.data, Nf=Nf, Nt=Nt, phi=phi, mult=mult
     )
+    return Wavelet(wave * jnp.sqrt(2), time=t_bins, freq=f_bins)
 
 
 def from_freq_to_wavelet(
@@ -98,13 +98,9 @@ def from_freq_to_wavelet(
     """
     Nf, Nt = _preprocess_bins(freqseries, Nf, Nt)
     t_bins, f_bins = _get_bins(freqseries, Nf, Nt)
-    phif = jnp.array(phitilde_vec_norm(Nf, Nt,  d=nx))
-    wave =  transform_wavelet_freq_helper(
+    phif = jnp.array(phitilde_vec_norm(Nf, Nt, d=nx))
+    wave = transform_wavelet_freq_helper(
         freqseries.data, Nf=Nf, Nt=Nt, phif=phif
     )
 
-    return Wavelet(
-        (2 / Nf) * wave * jnp.sqrt(2),
-        time=t_bins,
-        freq=f_bins
-    )
+    return Wavelet((2 / Nf) * wave * jnp.sqrt(2), time=t_bins, freq=f_bins)

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/jax/inverse/to_freq.py

@@ -1,14 +1,14 @@
+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"""
     wave_in = wave_in.T
@@ -20,10 +20,14 @@ def inverse_wavelet_freq_helper(
     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
     m_mid = m_range[1:Nf]
@@ -67,4 +71,4 @@ def inverse_wavelet_freq_helper(
     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
+    return res

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/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-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/numpy/forward/main.py

@@ -113,7 +113,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-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/numpy/inverse/main.py

@@ -1,7 +1,7 @@
 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
 
@@ -84,7 +84,7 @@ 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
     )

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/transforms/phi_computer.py

@@ -1,10 +1,7 @@
-from ..backend import xp, PI, betainc, ifft
+from ..backend import PI, betainc, ifft, xp
 
 
-
-def phitilde_vec(
-    omega: xp.ndarray, Nf: int, d: float = 4.0
-) -> xp.ndarray:
+def phitilde_vec(omega: xp.ndarray, Nf: int, d: float = 4.0) -> xp.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)
@@ -78,7 +75,7 @@ def __nu_d(
     return betainc(d, d, x) / betainc(d, d, 1)
 
 
-def phitilde_vec_norm(Nf: int, Nt: int,  d: float) -> xp.ndarray:
+def phitilde_vec_norm(Nf: int, Nt: int, d: float) -> xp.ndarray:
     """Normalize phitilde for inverse frequency domain transform."""
 
     # Calculate the frequency values
@@ -86,7 +83,7 @@ def phitilde_vec_norm(Nf: int, Nt: int,  d: float) -> xp.ndarray:
     omegas = 2 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)
 
     # Calculate the unnormalized phitilde (u_phit)
-    u_phit = phitilde_vec(omegas, Nf,  d)
+    u_phit = phitilde_vec(omegas, Nf, d)
 
     # Normalize the phitilde
     normalising_factor = PI ** (-1 / 2)  # Ollie's normalising factor
@@ -129,13 +126,9 @@ def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> xp.ndarray:
 
     DX = DX.copy()
     # postive frequencies
-    DX[1 : half_K + 1] = phitilde_vec(
-        dom * xp.arange(1, half_K + 1), Nf, d
-    )
+    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[half_K + 1 :] = phitilde_vec(-dom * xp.arange(half_K - 1, 0, -1), Nf, d)
     DX = K * ifft(DX, K)
 
     phi = xp.zeros(K)

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/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-0.2.3 → pywavelet-0.2.5}/src/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
 

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/types/timeseries.py

@@ -4,15 +4,9 @@ import matplotlib.pyplot as plt
 from scipy.signal import butter, sosfiltfilt
 from scipy.signal.windows import tukey
 
+from ..backend import rfft, rfftfreq, xp
 from ..logger import logger
-from .common import (
-    fmt_pow2,
-    fmt_time,
-    fmt_timerange,
-    is_documented_by,
-)
-from ..backend import xp, rfftfreq, rfft
-
+from .common import fmt_pow2, fmt_time, fmt_timerange, is_documented_by
 from .plotting import plot_spectrogram, plot_timeseries
 
 __all__ = ["TimeSeries"]
@@ -272,9 +266,9 @@ class TimeSeries:
         sos = butter(
             bandpass_order, Wn=fmin, btype="highpass", output="sos", fs=self.fs
         )
-        window = tukey(self.ND, alpha=tukey_window_alpha)
         data = self.data.copy()
-        data = sosfiltfilt(sos, data * window)
+        data = sosfiltfilt(sos, data)
+        data = data * tukey(self.ND, alpha=tukey_window_alpha)
         return TimeSeries(data, self.time)
 
     def __copy__(self):

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/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.3 → pywavelet-0.2.5}/src/pywavelet/types/wavelet_bins.py

@@ -1,7 +1,6 @@
 from typing import Tuple, Union
 
 from ..backend import xp
-
 from .frequencyseries import FrequencySeries
 from .timeseries import TimeSeries
 

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet/utils.py

@@ -71,6 +71,8 @@ def compute_likelihood(
     p = psd.data
     if mask is not None:
         m = mask.mask
+        # convert mask to numbers -- 0 for False, 1 for True
+        m = m.astype(int)
         d, h, p = d * m, h * m, p * m
 
     return -0.5 * np.nansum((d - h) ** 2 / p)

{pywavelet-0.2.3 → pywavelet-0.2.5}/src/pywavelet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.2.3
+Version: 0.2.5
 Summary: WDM wavelet transform your time/freq series!
 Author-email: Pywavelet Team <avi.vajpeyi@gmail.com>
 Project-URL: Homepage, https://pywavelet.github.io/pywavelet/
@@ -76,10 +76,10 @@ First set up a conda environment with the latest version of python.
 Test code
 ---------
 
-Locate directory /tests from root directory. run 
+Locate directory /tests from root directory. run
 
 .. code-block::
 
     $ pytest .
 
-Hopefully everything should run fine. 
+Hopefully everything should run fine.

{pywavelet-0.2.3 → pywavelet-0.2.5}/tests/test_mask.py

@@ -13,25 +13,26 @@ def test_mask(plot_dir):
     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])
+    mask = WaveletMask.from_restrictions(h.time, h.freq, [f0 - 0.5, f0 + 0.5])
+    dmasked = d * mask
     assert np.isclose(compute_likelihood(d, h, psd, mask), 0)
+    assert np.isclose(compute_likelihood(dmasked, h, psd), 0)
+    # number of nans in dmasked.data
+    assert np.isnan(dmasked.data).sum() != 0
 
-    mask1 = WaveletMask.from_frange(h.time, h.freq, [f0 + 0.5, f0 + 1.5])
+    mask1 = WaveletMask.from_restrictions(
+        h.time, h.freq, [f0 + 0.5, f0 + 1.5], tgaps=[[1.7 * 60, 3.4 * 60]]
+    )
+    dmasked1 = d * mask1
     # assert np.isclose(compute_likelihood(d, h, psd_analysis, 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)
+    d.plot(ax=axes[0])
     axes[0].set_title("data")
-    axes[1].imshow(d.data * mask.mask, **kwgs)
+    (d * mask).plot(ax=axes[1])
     axes[1].set_title("data*mask[f0-0.5, f0+0.5]")
-    axes[2].imshow(d.data * mask1.mask, **kwgs)
+    (d * mask1).plot(ax=axes[2])
     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.2.3 → pywavelet-0.2.5}/tests/test_phi.py

@@ -1,11 +1,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from pywavelet.transforms import (
-    phi_vec,
-    phitilde_vec,
-    phitilde_vec_norm,
-)
+from pywavelet.transforms import phi_vec, phitilde_vec, phitilde_vec_norm
 
 
 def test_phi(plot_dir):