Source code for neurodsp.spectral.utils

"""Utility function for neurodsp.spectral."""

import numpy as np

# Alias a function that has moved, for backwards compatibility
from neurodsp.sim.utils import rotate_spectrum as rotate_powerlaw

###################################################################################################
###################################################################################################

[docs]def trim_spectrum(freqs, power_spectra, f_range): """Extract a frequency range of interest from power spectra. Parameters ---------- freqs : 1d array Frequency values for the power spectrum. power_spectra : 1d or 2d array Power spectral density values. If 2d, should be as [n_spectra, n_freqs]. f_range : list of [float, float] Frequency range to restrict to, as [f_low, f_high]. Returns ------- freqs_ext : 1d array Extracted frequency values for the power spectrum. power_spectra_ext : 1d or 2d array Extracted power spectral density values. Notes ----- This function extracts frequency ranges >= f_low and <= f_high. It does not round to below or above f_low and f_high, respectively. Examples -------- Trim the power spectrum of a simulated time series: >>> from neurodsp.sim import sim_combined >>> from neurodsp.spectral import compute_spectrum >>> sig = sim_combined(n_seconds=10, fs=500, ... components={'sim_powerlaw': {}, 'sim_oscillation' : {'freq': 10}}) >>> freqs, spectrum = compute_spectrum(sig, fs=500) >>> freqs_ext, spec_ext = trim_spectrum(freqs, spectrum, [1, 30]) """ # Create mask to index only requested frequencies f_mask = np.logical_and(freqs >= f_range[0], freqs <= f_range[1]) # Restrict freqs & spectrum to requested range # The if/else is to cover both 1d or 2d arrays freqs_ext = freqs[f_mask] power_spectra_ext = power_spectra[f_mask] if power_spectra.ndim == 1 \ else power_spectra[:, f_mask] return freqs_ext, power_spectra_ext
[docs]def trim_spectrogram(freqs, times, spg, f_range=None, t_range=None): """Extract a frequency or time range of interest from a spectrogram. Parameters ---------- freqs : 1d array Frequency values for the spectrogram. times : 1d array Time values for the spectrogram. spg : 2d array Spectrogram, or time frequency representation of a signal. Formatted as [n_freqs, n_time_windows]. f_range : list of [float, float] Frequency range to restrict to, as [f_low, f_high]. t_range : list of [float, float] Time range to restrict to, as [t_low, t_high]. Returns ------- freqs_ext : 1d array Extracted frequency values for the power spectrum. times_ext : 1d array Extracted segment time values spg_ext : 2d array Extracted spectrogram values. Notes ----- This function extracts frequency ranges >= f_low and <= f_high, and time ranges >= t_low and <= t_high. It does not round to below or above f_low and f_high, or t_low and t_high, respectively. Examples -------- Trim the spectrogram of a simulated time series: >>> from neurodsp.sim import sim_combined >>> from neurodsp.timefrequency import compute_wavelet_transform >>> from neurodsp.utils.data import create_times, create_freqs >>> fs = 500 >>> n_seconds = 10 >>> times = create_times(n_seconds, fs) >>> sig = sim_combined(n_seconds, fs, ... components={'sim_powerlaw': {}, 'sim_oscillation' : {'freq': 10}}) >>> freqs = create_freqs(1, 15) >>> mwt = compute_wavelet_transform(sig, fs, freqs) >>> spg = abs(mwt)**2 >>> freqs_ext, times_ext, spg_ext = trim_spectrogram(freqs, times, spg, ... f_range=[8, 12], t_range=[0, 5]) """ # Initialize spg_ext, to define for case in which neither f_range nor t_range is defined spg_ext = spg # Restrict frequency range of the spectrogram if f_range is not None: f_mask = np.logical_and(freqs >= f_range[0], freqs <= f_range[1]) freqs_ext = freqs[f_mask] spg_ext = spg_ext[f_mask, :] else: freqs_ext = freqs # Restrict time range of the spectrogram if t_range is not None: times_mask = np.logical_and(times >= t_range[0], times <= t_range[1]) times_ext = times[times_mask] spg_ext = spg_ext[:, times_mask] else: times_ext = times return freqs_ext, times_ext, spg_ext