pywavelet 0.0.5__tar.gz → 0.1.1__tar.gz

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

@@ -5,6 +5,44 @@ CHANGELOG
 =========
 
 
+.. _changelog-v0.1.1:
+
+v0.1.1 (2025-01-16)
+===================
+
+Unknown
+-------
+
+* Merge branch 'main' of github.com:avivajpeyi/pywavelet into main (`69eefa2`_)
+
+.. _69eefa2: https://github.com/pywavelet/pywavelet/commit/69eefa29b7873c30fcb74ad1e051eb20101a277a
+
+
+.. _changelog-v0.1.0:
+
+v0.1.0 (2025-01-15)
+===================
+
+Bug Fixes
+---------
+
+* fix: refactor type outside transforms (`efb8878`_)
+
+Chores
+------
+
+* chore(release): 0.1.0 (`c5a3fde`_)
+
+Features
+--------
+
+* feat: add wavelet mask and more tests (`e009903`_)
+
+.. _efb8878: https://github.com/pywavelet/pywavelet/commit/efb88789f8468ff18f99abaf6168bb8fc0f5947b
+.. _c5a3fde: https://github.com/pywavelet/pywavelet/commit/c5a3fdea455c16478f04049f14bc35dfcf4efb15
+.. _e009903: https://github.com/pywavelet/pywavelet/commit/e00990300d9c013438580c2bc47ea93570fd95be
+
+
 .. _changelog-v0.0.5:
 
 v0.0.5 (2024-12-12)
@@ -15,7 +53,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.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pywavelet
-Version: 0.0.5
+Version: 0.1.1
 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.1}/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.1'
+__version_tuple__ = version_tuple = (0, 1, 1)

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/logger.py

@@ -1,11 +1,15 @@
+import logging
 import sys
 import warnings
+
 from rich.logging import RichHandler
-import logging
 
 FORMAT = "%(message)s"
 logging.basicConfig(
-    level="INFO", format=FORMAT, datefmt="[%X]", handlers=[RichHandler(rich_tracebacks=True)]
+    level="INFO",
+    format=FORMAT,
+    datefmt="[%X]",
+    handlers=[RichHandler(rich_tracebacks=True)],
 )
 
 logger = logging.getLogger("pywavelet")

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

@@ -1,4 +1,4 @@
-from .forward import from_freq_to_wavelet, from_time_to_wavelet, compute_bins
+from .forward import from_freq_to_wavelet, from_time_to_wavelet
 from .inverse import from_wavelet_to_freq, from_wavelet_to_time
 
 __all__ = [
@@ -6,5 +6,4 @@ __all__ = [
     "from_wavelet_to_freq",
     "from_time_to_wavelet",
     "from_freq_to_wavelet",
-    "compute_bins",
 ]

pywavelet-0.1.1/src/pywavelet/transforms/forward/__init__.py

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

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/forward/from_freq.py

@@ -1,8 +1,10 @@
 """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:
@@ -13,8 +15,16 @@ def transform_wavelet_freq_helper(
     __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):
+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
@@ -22,8 +32,6 @@ def __core(Nf:int, Nt:int, DX:np.ndarray, freq_strain:np.ndarray, phif:np.ndarra
         __fill_wave_2(f_bin, DX_trans, wave, Nt, Nf)
 
 
-
-
 @njit()
 def __fill_wave_1(
     f_bin: int,
@@ -55,6 +63,7 @@ def __fill_wave_1(
         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
@@ -72,7 +81,7 @@ def __fill_wave_2(
                 if (n + f_bin) % 2:
                     wave[n, f_bin] = -DX_trans[n].imag
                 else:
-                    wave[n, f_bin] =  DX_trans[n].real
+                    wave[n, f_bin] = DX_trans[n].real
             else:
                 if (n + f_bin) % 2:
                     wave[n, f_bin] = DX_trans[n].imag

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/forward/from_time.py

@@ -1,4 +1,5 @@
 """helper functions for transform_time.py"""
+
 import numpy as np
 from numba import njit
 from numpy import fft
@@ -20,7 +21,18 @@ def transform_wavelet_time_helper(
     __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:
+
+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)

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/forward/main.py

@@ -3,15 +3,15 @@ from typing import Union
 import numpy as np
 
 from ...logger import logger
+from ...types import FrequencySeries, TimeSeries, Wavelet
+from ...types.wavelet_bins import _get_bins, _preprocess_bins
 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,
@@ -74,9 +74,7 @@ def from_time_to_wavelet(
     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
-    )
+    return Wavelet(wave * np.sqrt(2), time=t_bins, freq=f_bins)
 
 
 def from_freq_to_wavelet(
@@ -117,12 +115,6 @@ def from_freq_to_wavelet(
     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
-    )
+    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.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/inverse/main.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from ...transforms.phi_computer import phi_vec, phitilde_vec_norm
-from ..types import FrequencySeries, TimeSeries, Wavelet
+from ...types import FrequencySeries, TimeSeries, Wavelet
 from .to_freq import inverse_wavelet_freq_helper_fast
 from .to_time import inverse_wavelet_time_helper_fast
 
@@ -9,6 +9,7 @@ __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,
@@ -55,9 +56,7 @@ def from_wavelet_to_time(
 
 
 def from_wavelet_to_freq(
-    wave_in: Wavelet,
-    dt: float,
-    nx:float=4.0
+    wave_in: Wavelet, dt: float, nx: float = 4.0
 ) -> FrequencySeries:
     """
     Perform an inverse wavelet transform to the frequency domain.

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/inverse/to_freq.py

@@ -1,4 +1,5 @@
 """functions for computing the inverse wavelet transforms"""
+
 import numpy as np
 from numba import njit
 from numpy import fft
@@ -12,13 +13,20 @@ def inverse_wavelet_freq_helper_fast(
     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)
 
-
     return res
 
-def __core(Nf: int, Nt: int, prefactor2s: np.ndarray, wave_in: np.ndarray, phif: np.ndarray, res: np.ndarray) -> None:
+
+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)

{pywavelet-0.0.5 → pywavelet-0.1.1}/src/pywavelet/transforms/inverse/to_time.py

@@ -1,4 +1,5 @@
 """functions for computing the inverse wavelet transforms"""
+
 import numpy as np
 from numba import njit
 from numpy import fft
@@ -21,7 +22,16 @@ def inverse_wavelet_time_helper_fast(
     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:
+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)
@@ -29,9 +39,9 @@ def __core(Nf: int, Nt: int, K: int, ND: int, wave_in: np.ndarray, phi: np.ndarr
             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[: K + Nf - ND] += res[2 * ND : ND + K * Nf]
 
 
 def unpack_time_wave_helper(

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

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

{pywavelet-0.0.5/src/pywavelet/transforms → pywavelet-0.1.1/src/pywavelet}/types/common.py

@@ -1,9 +1,9 @@
-from typing import Literal, Tuple
+from typing import Tuple, Union
 
 import numpy as xp
-from numpy.fft import irfft, fft, rfft, rfftfreq
+from numpy.fft import fft, irfft, rfft, rfftfreq  # type: ignore
 
-from ...logger import logger
+from ..logger import logger
 
 
 def _len_check(d):
@@ -19,7 +19,7 @@ def is_documented_by(original):
     return wrapper
 
 
-def fmt_time(seconds: float, units=False) -> Tuple[str, str]:
+def fmt_time(seconds: float, units=False) -> Union[str, Tuple[str, str]]:
     """Returns formatted time and units [ms, s, min, hr, day]"""
     t, u = "", ""
     if seconds < 1e-3:
@@ -42,12 +42,12 @@ def fmt_time(seconds: float, units=False) -> Tuple[str, str]:
 
 def fmt_timerange(trange):
     t0 = fmt_time(trange[0])
-    tend, units = fmt_time(trange[1], units = True)
+    tend, units = fmt_time(trange[1], units=True)
     return f"[{t0}, {tend}] {units}"
 
 
-def fmt_pow2(n:float)->str:
+def fmt_pow2(n: float) -> str:
     pow2 = xp.log2(n)
     if pow2.is_integer():
         return f"2^{int(pow2)}"
-    return f"{n:,}"
+    return f"{n:,}"

{pywavelet-0.0.5/src/pywavelet/transforms → pywavelet-0.1.1/src/pywavelet}/types/frequencyseries.py

@@ -1,11 +1,13 @@
+from typing import Optional, Tuple, Union
+
 import matplotlib.pyplot as plt
-from typing import Tuple, Union, Optional
 
-from .common import is_documented_by, xp, irfft, fmt_time, fmt_pow2
+from .common import fmt_pow2, fmt_time, irfft, is_documented_by, xp
 from .plotting import plot_freqseries, plot_periodogram
 
 __all__ = ["FrequencySeries"]
 
+
 class FrequencySeries:
     """
     A class to represent a one-sided frequency series, with various methods for
@@ -39,9 +41,13 @@ class FrequencySeries:
             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)")
+            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)})")
+            raise ValueError(
+                f"data and freq must have the same length ({len(data)} != {len(freq)})"
+            )
         self.data = data
         self.freq = freq
         self.t0 = t0
@@ -53,10 +59,10 @@ class FrequencySeries:
         )
 
     @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 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."""
@@ -127,7 +133,9 @@ class FrequencySeries:
         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:
+    def noise_weighted_inner_product(
+        self, other: "FrequencySeries", psd: "FrequencySeries"
+    ) -> float:
         """
         Compute the noise-weighted inner product of two FrequencySeries.
 
@@ -144,9 +152,11 @@ class FrequencySeries:
             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)
+        return (4 * self.dt / self.ND) * xp.nansum(integrand)
 
-    def matched_filter_snr(self, other: "FrequencySeries", psd: "FrequencySeries") -> float:
+    def matched_filter_snr(
+        self, other: "FrequencySeries", psd: "FrequencySeries"
+    ) -> float:
         """
         Compute the signal-to-noise ratio (SNR) of a matched filter.
 
@@ -199,15 +209,15 @@ class FrequencySeries:
 
         # Create and return a TimeSeries object
         from .timeseries import TimeSeries
-        return TimeSeries(time_data, time)
 
+        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":
+        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.
 
@@ -216,9 +226,9 @@ class FrequencySeries:
         Wavelet
             The corresponding wavelet.
         """
-        from ..forward import from_freq_to_wavelet
-        return from_freq_to_wavelet(self, Nf=Nf, Nt=Nt, nx=nx)
+        from ..transforms.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."""
@@ -228,10 +238,8 @@ class FrequencySeries:
 
     def __copy__(self):
         return FrequencySeries(
-            xp.copy(self.data),
-            xp.copy(self.freq),
-            t0=self.t0
+            xp.copy(self.data), xp.copy(self.freq), t0=self.t0
         )
 
     def copy(self):
-        return self.__copy__()
+        return self.__copy__()

{pywavelet-0.0.5/src/pywavelet/transforms → pywavelet-0.1.1/src/pywavelet}/types/plotting.py

@@ -1,12 +1,11 @@
 import warnings
-from typing import Tuple, Optional, Union
-from scipy.signal import savgol_filter
-from scipy.interpolate import interp1d
+from typing import Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import LogNorm, TwoSlopeNorm
-from scipy.signal import spectrogram
+from scipy.interpolate import interp1d
+from scipy.signal import savgol_filter, spectrogram
 
 MIN_S = 60
 HOUR_S = 60 * MIN_S
@@ -53,8 +52,13 @@ def __get_smoothed_y(x, z, y_grid):
         # Interpolate to fill NaNs in y before smoothing
         nan_mask = ~np.isnan(y)
         if np.isnan(y).any():
-            interpolator = interp1d(x[nan_mask], y[nan_mask], kind='cubic', bounds_error=False,
-                                    fill_value="extrapolate")
+            interpolator = interp1d(
+                x[nan_mask],
+                y[nan_mask],
+                kind="cubic",
+                bounds_error=False,
+                fill_value="extrapolate",
+            )
             y = interpolator(x)  # Fill NaNs with interpolated values
 
         # Smooth the curve
@@ -64,8 +68,6 @@ def __get_smoothed_y(x, z, y_grid):
     return y
 
 
-
-
 def plot_wavelet_grid(
     wavelet_data: np.ndarray,
     time_grid: np.ndarray,
@@ -80,8 +82,8 @@ def plot_wavelet_grid(
     norm: Optional[Union[LogNorm, TwoSlopeNorm]] = None,
     cbar_label: Optional[str] = None,
     nan_color: Optional[str] = "black",
-    detailed_axes:bool = False,
-    show_gridinfo:bool = True,
+    detailed_axes: bool = False,
+    show_gridinfo: bool = True,
     trend_color: Optional[str] = None,
     whiten_by: Optional[np.ndarray] = None,
     **kwargs,
@@ -153,7 +155,9 @@ def plot_wavelet_grid(
 
     # Validate the dimensions
     if (Nf, Nt) != (len(freq_grid), len(time_grid)):
-        raise ValueError(f"Wavelet shape {Nf, Nt} does not match provided grids {(len(freq_grid), len(time_grid))}.")
+        raise ValueError(
+            f"Wavelet shape {Nf, Nt} does not match provided grids {(len(freq_grid), len(time_grid))}."
+        )
 
     # Prepare the data for plotting
     z = wavelet_data.copy()
@@ -162,14 +166,15 @@ def plot_wavelet_grid(
     if absolute:
         z = np.abs(z)
 
-
     # Determine normalization and colormap
     if norm is None:
         try:
             if np.all(np.isnan(z)):
                 raise ValueError("All wavelet data is NaN.")
             if zscale == "log":
-                norm = LogNorm(vmin=np.nanmin(z[z > 0]), vmax=np.nanmax(z[z<np.inf]))
+                norm = LogNorm(
+                    vmin=np.nanmin(z[z > 0]), vmax=np.nanmax(z[z < np.inf])
+                )
             elif not absolute:
                 vmin, vmax = np.nanmin(z), np.nanmax(z)
                 vcenter = 0.0
@@ -177,7 +182,9 @@ def plot_wavelet_grid(
             else:
                 norm = None  # Default linear scaling
         except Exception as e:
-            warnings.warn(f"Error in determining normalization: {e}. Using default linear scaling.")
+            warnings.warn(
+                f"Error in determining normalization: {e}. Using default linear scaling."
+            )
             norm = None
 
     if cmap is None:
@@ -195,7 +202,7 @@ def plot_wavelet_grid(
     im = ax.imshow(
         z,
         aspect="auto",
-        extent=[time_grid[0],time_grid[-1], freq_grid[0], freq_grid[-1]],
+        extent=[time_grid[0], time_grid[-1], freq_grid[0], freq_grid[-1]],
         origin="lower",
         cmap=cmap,
         norm=norm,
@@ -203,13 +210,25 @@ def plot_wavelet_grid(
         **kwargs,
     )
     if trend_color is not None:
-        plot_wavelet_trend(wavelet_data, time_grid, freq_grid, ax, color=trend_color, freq_range=freq_range, freq_scale=freq_scale)
+        plot_wavelet_trend(
+            wavelet_data,
+            time_grid,
+            freq_grid,
+            ax,
+            color=trend_color,
+            freq_range=freq_range,
+            freq_scale=freq_scale,
+        )
 
     # Add colorbar if requested
     if show_colorbar:
         cbar = fig.colorbar(im, ax=ax)
         if cbar_label is None:
-            cbar_label = "Absolute Wavelet Amplitude" if absolute else "Wavelet Amplitude"
+            cbar_label = (
+                "Absolute Wavelet Amplitude"
+                if absolute
+                else "Wavelet Amplitude"
+            )
         cbar.set_label(cbar_label)
 
     # Configure axes scales
@@ -239,14 +258,12 @@ def plot_wavelet_grid(
             bbox=dict(boxstyle="round", facecolor=None, alpha=0.2),
         )
 
-
     # Adjust layout
     fig.tight_layout()
 
     return fig, ax
 
 
-
 def plot_freqseries(
     data: np.ndarray,
     freq: np.ndarray,
@@ -277,9 +294,10 @@ def plot_periodogram(
     flow = np.min(np.abs(freq))
     ax.set_xlabel("Frequency [Hz]")
     ax.set_ylabel("Periodigram")
-    ax.set_xlim(left=flow, right=nyquist_frequency/2)
+    ax.set_xlim(left=flow, right=nyquist_frequency / 2)
     return ax.figure, ax
 
+
 def plot_timeseries(
     data: np.ndarray, time: np.ndarray, ax=None, **kwargs
 ) -> Tuple[plt.Figure, plt.Axes]:
@@ -314,7 +332,6 @@ def plot_spectrogram(
 
     _fmt_time_axis(t, ax)
 
-
     ax.set_ylabel("Frequency [Hz]")
     ax.set_ylim(top=fs / 2.0)
     cbar = plt.colorbar(cm, ax=ax)
@@ -322,20 +339,24 @@ def plot_spectrogram(
     return ax.figure, ax
 
 
-
 def _fmt_time_axis(t, axes, t0=None, tmax=None):
     if t[-1] > DAY_S:  # If time goes beyond a day
-        axes.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x / DAY_S:.1f}"))
+        axes.xaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x / DAY_S:.1f}")
+        )
         axes.set_xlabel("Time [days]")
     elif t[-1] > HOUR_S:  # If time goes beyond an hour
-        axes.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x / HOUR_S:.1f}"))
+        axes.xaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x / HOUR_S:.1f}")
+        )
         axes.set_xlabel("Time [hr]")
     elif t[-1] > MIN_S:  # If time goes beyond a minute
-        axes.xaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x / MIN_S:.1f}"))
+        axes.xaxis.set_major_formatter(
+            plt.FuncFormatter(lambda x, _: f"{x / MIN_S:.1f}")
+        )
         axes.set_xlabel("Time [min]")
     else:
         axes.set_xlabel("Time [s]")
     t0 = t[0] if t0 is None else t0
     tmax = t[-1] if tmax is None else tmax
     axes.set_xlim(t0, tmax)
-