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

pywavelet/transforms/types/plotting.py

@@ -0,0 +1,341 @@
+import warnings
+from typing import Tuple, Optional, Union
+from scipy.signal import savgol_filter
+from scipy.interpolate import interp1d
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import LogNorm, TwoSlopeNorm
+from scipy.signal import spectrogram
+
+MIN_S = 60
+HOUR_S = 60 * MIN_S
+DAY_S = 24 * HOUR_S
+
+
+def plot_wavelet_trend(
+    wavelet_data: np.ndarray,
+    time_grid: np.ndarray,
+    freq_grid: np.ndarray,
+    ax: Optional[plt.Axes] = None,
+    freq_scale: str = "linear",
+    freq_range: Optional[Tuple[float, float]] = None,
+    color: str = "black",
+):
+    x = time_grid
+    y = __get_smoothed_y(x, np.abs(wavelet_data), freq_grid)
+    if ax == None:
+        fig, ax = plt.subplots()
+    ax.plot(x, y, color=color)
+
+    # Configure axes scales
+    ax.set_yscale(freq_scale)
+    _fmt_time_axis(time_grid, ax)
+    ax.set_ylabel("Frequency [Hz]")
+
+    # Set frequency range if specified
+    freq_range = freq_range or (freq_grid[0], freq_grid[-1])
+    ax.set_ylim(freq_range)
+
+
+def __get_smoothed_y(x, z, y_grid):
+    Nf, Nt = z.shape
+    y = np.zeros(Nt)
+    dy = np.diff(y_grid)[0]
+    for i in range(Nt):
+        # if all values are nan, set to nan
+        if np.all(np.isnan(z[:, i])):
+            y[i] = np.nan
+        else:
+            y[i] = y_grid[np.nanargmax(z[:, i])]
+
+    if not np.isnan(y).all():
+        # 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")
+            y = interpolator(x)  # Fill NaNs with interpolated values
+
+        # Smooth the curve
+        window_length = min(51, len(y) - 1 if len(y) % 2 == 0 else len(y))
+        y = savgol_filter(y, window_length, 3)
+        y[~nan_mask] = np.nan
+    return y
+
+
+
+
+def plot_wavelet_grid(
+    wavelet_data: np.ndarray,
+    time_grid: np.ndarray,
+    freq_grid: np.ndarray,
+    ax: Optional[plt.Axes] = None,
+    zscale: str = "linear",
+    freq_scale: str = "linear",
+    absolute: bool = False,
+    freq_range: Optional[Tuple[float, float]] = None,
+    show_colorbar: bool = True,
+    cmap: Optional[str] = None,
+    norm: Optional[Union[LogNorm, TwoSlopeNorm]] = None,
+    cbar_label: Optional[str] = None,
+    nan_color: Optional[str] = "black",
+    detailed_axes:bool = False,
+    show_gridinfo:bool = True,
+    trend_color: Optional[str] = None,
+    whiten_by: Optional[np.ndarray] = None,
+    **kwargs,
+) -> Tuple[plt.Figure, plt.Axes]:
+    """
+    Plot a 2D grid of wavelet coefficients.
+
+    Parameters
+    ----------
+    wavelet_data : np.ndarray
+        A 2D array containing the wavelet coefficients with shape (Nf, Nt),
+        where Nf is the number of frequency bins and Nt is the number of time bins.
+
+    time_grid : np.ndarray, optional
+        1D array of time values corresponding to the time bins. If None, uses np.arange(Nt).
+
+    freq_grid : np.ndarray, optional
+        1D array of frequency values corresponding to the frequency bins. If None, uses np.arange(Nf).
+
+    ax : plt.Axes, optional
+        Matplotlib Axes object to plot on. If None, creates a new figure and axes.
+
+    zscale : str, optional
+        Scale for the color mapping. Options are 'linear' or 'log'. Default is 'linear'.
+
+    freq_scale : str, optional
+        Scale for the frequency axis. Options are 'linear' or 'log'. Default is 'linear'.
+
+    absolute : bool, optional
+        If True, plots the absolute value of the wavelet coefficients. Default is False.
+
+    freq_range : tuple of float, optional
+        Tuple specifying the (min, max) frequency range to display. If None, displays the full range.
+
+    show_colorbar : bool, optional
+        If True, displays a colorbar next to the plot. Default is True.
+
+    cmap : str, optional
+        Colormap to use for the plot. If None, uses 'viridis' for absolute values or 'bwr' for signed values.
+
+    norm : matplotlib.colors.Normalize, optional
+        Normalization instance to scale data values. If None, a suitable normalization is chosen based on `zscale`.
+
+    cbar_label : str, optional
+        Label for the colorbar. If None, a default label is used based on the `absolute` parameter.
+
+    nan_color : str, optional
+        Color to use for NaN values. Default is 'black'.
+
+    trend_color : bool, optional
+        Color to use for the trend line. Not shown if None.
+
+    **kwargs
+        Additional keyword arguments passed to `ax.imshow()`.
+
+    Returns
+    -------
+    Tuple[plt.Figure, plt.Axes]
+        The figure and axes objects of the plot.
+
+    Raises
+    ------
+    ValueError
+        If the dimensions of `wavelet_data` do not match the lengths of `freq_grid` and `time_grid`.
+    """
+
+    # Determine the dimensions of the data
+    Nf, Nt = wavelet_data.shape
+
+    # 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))}.")
+
+    # Prepare the data for plotting
+    z = wavelet_data.copy()
+    if whiten_by is not None:
+        z = z / whiten_by
+    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]))
+            elif not absolute:
+                vmin, vmax = np.nanmin(z), np.nanmax(z)
+                vcenter = 0.0
+                norm = TwoSlopeNorm(vmin=vmin, vcenter=vcenter, vmax=vmax)
+            else:
+                norm = None  # Default linear scaling
+        except Exception as e:
+            warnings.warn(f"Error in determining normalization: {e}. Using default linear scaling.")
+            norm = None
+
+    if cmap is None:
+        cmap = "viridis" if absolute else "bwr"
+        cmap = plt.get_cmap(cmap)
+        cmap.set_bad(color=nan_color)
+
+    # Set up the plot
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.get_figure()
+
+    # Plot the data
+    im = ax.imshow(
+        z,
+        aspect="auto",
+        extent=[time_grid[0],time_grid[-1], freq_grid[0], freq_grid[-1]],
+        origin="lower",
+        cmap=cmap,
+        norm=norm,
+        interpolation="nearest",
+        **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)
+
+    # 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.set_label(cbar_label)
+
+    # Configure axes scales
+    ax.set_yscale(freq_scale)
+    _fmt_time_axis(time_grid, ax)
+    ax.set_ylabel("Frequency [Hz]")
+
+    # Set frequency range if specified
+    freq_range = freq_range or (freq_grid[0], freq_grid[-1])
+    ax.set_ylim(freq_range)
+
+    if detailed_axes:
+        ax.set_xlabel(r"Time Bins [$\Delta T$=" + f"{1 / Nt:.4f}s, Nt={Nt}]")
+        ax.set_ylabel(r"Freq Bins [$\Delta F$=" + f"{1 / Nf:.4f}Hz, Nf={Nf}]")
+
+    label = kwargs.get("label", "")
+    NfNt_label = f"{Nf}x{Nt}" if show_gridinfo else ""
+    txt = f"{label}\n{NfNt_label}" if label else NfNt_label
+    if txt:
+        ax.text(
+            0.05,
+            0.95,
+            txt,
+            transform=ax.transAxes,
+            fontsize=14,
+            verticalalignment="top",
+            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,
+    nyquist_frequency: float,
+    ax=None,
+    **kwargs,
+):
+    if ax == None:
+        fig, ax = plt.subplots()
+    ax.plot(freq, data, **kwargs)
+    ax.set_xlabel("Frequency Bin [Hz]")
+    ax.set_ylabel("Amplitude")
+    ax.set_xlim(0, nyquist_frequency)
+    return ax.figure, ax
+
+
+def plot_periodogram(
+    data: np.ndarray,
+    freq: np.ndarray,
+    nyquist_frequency: float,
+    ax=None,
+    **kwargs,
+):
+    if ax == None:
+        fig, ax = plt.subplots()
+
+    ax.loglog(freq, np.abs(data) ** 2, **kwargs)
+    flow = np.min(np.abs(freq))
+    ax.set_xlabel("Frequency [Hz]")
+    ax.set_ylabel("Periodigram")
+    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]:
+    """Custom method."""
+    if ax == None:
+        fig, ax = plt.subplots()
+    ax.plot(time, data, **kwargs)
+
+    ax.set_ylabel("Amplitude")
+    ax.set_xlim(left=time[0], right=time[-1])
+
+    _fmt_time_axis(time, ax)
+
+    return ax.figure, ax
+
+
+def plot_spectrogram(
+    timeseries_data: np.ndarray,
+    fs: float,
+    ax=None,
+    spec_kwargs={},
+    plot_kwargs={},
+) -> Tuple[plt.Figure, plt.Axes]:
+    f, t, Sxx = spectrogram(timeseries_data, fs=fs, **spec_kwargs)
+    if ax == None:
+        fig, ax = plt.subplots()
+
+    if "cmap" not in plot_kwargs:
+        plot_kwargs["cmap"] = "Reds"
+
+    cm = ax.pcolormesh(t, f, Sxx, shading="nearest", **plot_kwargs)
+
+    _fmt_time_axis(t, ax)
+
+
+    ax.set_ylabel("Frequency [Hz]")
+    ax.set_ylim(top=fs / 2.0)
+    cbar = plt.colorbar(cm, ax=ax)
+    cbar.set_label("Spectrogram Amplitude")
+    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.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.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.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)
+

pywavelet/transforms/types/timeseries.py

@@ -0,0 +1,280 @@
+import matplotlib.pyplot as plt
+from typing import Tuple, Optional, Union
+from scipy.signal.windows import tukey
+from scipy.signal import butter, sosfiltfilt
+
+from ...logger import logger
+from .common import is_documented_by, xp, rfft, rfftfreq, fmt_timerange, fmt_time, fmt_pow2
+from .plotting import plot_timeseries, plot_spectrogram
+
+__all__ = ["TimeSeries"]
+
+
+class TimeSeries:
+    """
+    A class to represent a time series, with methods for plotting and converting
+    the series to a frequency-domain representation.
+
+    Attributes
+    ----------
+    data : xp.ndarray
+        Time domain data.
+    time : xp.ndarray
+        Array of corresponding time points.
+    """
+
+    def __init__(self, data: xp.ndarray, time: xp.ndarray):
+        """
+        Initialize the TimeSeries with data and time arrays.
+
+        Parameters
+        ----------
+        data : xp.ndarray
+            Time domain data.
+        time : xp.ndarray
+            Array of corresponding time points. Must be the same length as `data`.
+
+        Raises
+        ------
+        ValueError
+            If `data` and `time` do not have the same length.
+        """
+        if len(data) != len(time):
+            raise ValueError("data and time must have the same length")
+        self.data = data
+        self.time = time
+
+    @is_documented_by(plot_timeseries)
+    def plot(self, ax=None, **kwargs) -> Tuple[plt.Figure, plt.Axes]:
+        return plot_timeseries(self.data, self.time, ax=ax, **kwargs)
+
+    @is_documented_by(plot_spectrogram)
+    def plot_spectrogram(
+            self, ax=None, spec_kwargs={}, plot_kwargs={}
+    ) -> Tuple[plt.Figure, plt.Axes]:
+        return plot_spectrogram(
+            self.data,
+            self.fs,
+            ax=ax,
+            spec_kwargs=spec_kwargs,
+            plot_kwargs=plot_kwargs,
+        )
+
+    def __len__(self):
+        """Return the number of data points in the time series."""
+        return len(self.data)
+
+    def __getitem__(self, item):
+        """Return the data point at the specified index."""
+        return self.data[item]
+
+    @property
+    def sample_rate(self) -> float:
+        """
+        Return the sample rate (fs).
+
+        The sample rate is the inverse of the time resolution (Δt).
+        """
+        return float(xp.round(1.0 / self.dt, decimals=14))
+
+    @property
+    def fs(self) -> float:
+        """Return the sample rate (fs)."""
+        return self.sample_rate
+
+    @property
+    def duration(self) -> float:
+        """Return the duration of the time series in seconds."""
+        return len(self) / self.fs
+
+    @property
+    def dt(self) -> float:
+        """Return the time resolution (Δt)."""
+        return float(self.time[1] - self.time[0])
+
+    @property
+    def nyquist_frequency(self) -> float:
+        """Return the Nyquist frequency (fs/2)."""
+        return self.fs / 2
+
+    @property
+    def t0(self) -> float:
+        """Return the initial time point in the series."""
+        return float(self.time[0])
+
+    @property
+    def tend(self) -> float:
+        """Return the final time point in the series."""
+        return float(self.time[-1]) + self.dt
+
+    @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 series."""
+        return len(self)
+
+    def __repr__(self) -> str:
+        """Return a string representation of the TimeSeries."""
+        trange = fmt_timerange((self.t0, self.tend))
+        T = " ".join(fmt_time(self.duration, units=True))
+        n = fmt_pow2(len(self))
+        return f"TimeSeries(n={n}, trange={trange}, T={T}, fs={self.fs:.2f} Hz)"
+
+    def to_frequencyseries(self) -> 'FrequencySeries':
+        """
+        Convert the time series to a frequency series using the one-sided FFT.
+
+        Returns
+        -------
+        FrequencySeries
+            The frequency-domain representation of the time series.
+        """
+        freq = rfftfreq(len(self), d=self.dt)
+        data = rfft(self.data)
+
+        from .frequencyseries import FrequencySeries  # Avoid circular import
+        return FrequencySeries(data, freq, t0=self.t0)
+
+    def to_wavelet(
+            self,
+            Nf: Union[int, None] = None,
+            Nt: Union[int, None] = None,
+            nx: Optional[float] = 4.0,
+    ) -> 'Wavelet':
+        """
+        Convert the time series to a wavelet representation.
+
+        Parameters
+        ----------
+        Nf : int
+            Number of frequency bins for the wavelet transform.
+        Nt : int
+            Number of time bins for the wavelet transform.
+        nx : float, optional
+            Number of standard deviations for the `phi_vec`, controlling the
+            width of the wavelets. Default is 4.0.
+
+        Returns
+        -------
+        Wavelet
+            The wavelet-domain representation of the time series.
+        """
+        hf = self.to_frequencyseries()
+        return hf.to_wavelet(Nf, Nt, nx=nx)
+
+
+    def __add__(self, other: 'TimeSeries') -> 'TimeSeries':
+        """Add two TimeSeries objects together."""
+        if self.shape != other.shape:
+            raise ValueError("TimeSeries objects must have the same shape to add them together")
+        return TimeSeries(self.data + other.data, self.time)
+
+    def __sub__(self, other: 'TimeSeries') -> 'TimeSeries':
+        """Subtract one TimeSeries object from another."""
+        if self.shape != other.shape:
+            raise ValueError("TimeSeries objects must have the same shape to subtract them")
+        return TimeSeries(self.data - other.data, self.time)
+
+    def __eq__(self, other: 'TimeSeries') -> bool:
+        """Check if two TimeSeries objects are equal."""
+        shape_same = self.shape == other.shape
+        range_same = self.t0 == other.t0 and self.tend == other.tend
+        time_same =  xp.allclose(self.time, other.time)
+        data_same = xp.allclose(self.data, other.data)
+        return shape_same and range_same and data_same and time_same
+
+    def __mul__(self, other: float) -> 'TimeSeries':
+        """Multiply a TimeSeries object by a scalar."""
+        return TimeSeries(self.data * other, self.time)
+
+    def zero_pad_to_power_of_2(self, tukey_window_alpha:float=0.0)->'TimeSeries':
+        """Zero pad the time series to make the length a power of two (useful to speed up FFTs, O(NlogN) versus O(N^2)).
+
+        Parameters
+        ----------
+        tukey_window_alpha : float, optional
+            Alpha parameter for the Tukey window. Default is 0.0.
+            (prevents spectral leakage when padding the data)
+
+        Returns
+        -------
+        TimeSeries
+            A new TimeSeries object with the data zero-padded to a power of two.
+        """
+        N, dt, t0 = self.ND, self.dt, self.t0
+        pow_2 = xp.ceil(xp.log2(N))
+        n_pad = int((2 ** pow_2) - N)
+        new_N = N + n_pad
+        if n_pad > 0:
+            logger.warning(
+                f"Padding the data to a power of two. "
+                f"{N:,} (2**{xp.log2(N):.2f}) -> {new_N:,} (2**{pow_2}). "
+            )
+        window = tukey(N, alpha=tukey_window_alpha)
+        data = self.data * window
+        data = xp.pad(data, (0, n_pad), "constant")
+        time = xp.arange(0, len(data) * dt, dt) + t0
+        return TimeSeries(data, time)
+
+    def highpass_filter(self, fmin: float, tukey_window_alpha:float=0.0, bandpass_order: int = 4) -> 'TimeSeries':
+        """
+        Filter the time series with a highpass bandpass filter.
+
+        (we use sosfiltfilt instead of filtfilt for numerical stability)
+
+        Note: filtfilt should be used if phase accuracy (zero-phase filtering) is critical for your analysis
+        and if the filter order is low to moderate.
+
+
+        Parameters
+        ----------
+        fmin : float
+            Minimum frequency to pass through the filter.
+        bandpass_order : int, optional
+            Order of the bandpass filter. Default is 4.
+
+        Returns
+        -------
+        TimeSeries
+            A new TimeSeries object with the highpass filter applied.
+        """
+
+        if fmin <= 0 or fmin > self.nyquist_frequency:
+            raise ValueError(f"Invalid fmin value: {fmin}. Must be in the range [0, {self.nyquist_frequency}]")
+
+        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)
+        return TimeSeries(data, self.time)
+
+    def __copy__(self):
+        return TimeSeries(
+            self.data.copy(),
+            self.time.copy()
+        )
+
+    def copy(self):
+        return self.__copy__()
+
+    def __getitem__(self, key)->"TimeSeries":
+        if isinstance(key, slice):
+            # Handle slicing
+            return self.__handle_slice(key)
+        else:
+            # Handle regular indexing
+            return TimeSeries(self.data[key], self.time[key])
+
+    def __handle_slice(self, slice_obj)->"TimeSeries":
+        return TimeSeries(
+            self.data[slice_obj],
+            self.time[slice_obj]
+        )
+
+    @classmethod
+    def _EMPTY(cls, ND:int, dt:float)->"TimeSeries":
+        return cls(xp.zeros(ND), xp.arange(0, ND*dt, dt))