Source code for vitalDSP.transforms.chroma_stft

"""
Signal Transforms Module for Physiological Signal Processing

This module provides comprehensive capabilities for physiological
signal processing including ECG, PPG, EEG, and other vital signs.

Author: vitalDSP Team
Date: 2025-01-27
Version: 1.0.0

Key Features:
- Object-oriented design with comprehensive classes
- Multiple processing methods and functions
- NumPy integration for numerical computations
- SciPy integration for advanced signal processing

Examples:
---------
Basic usage:
    >>> import numpy as np
    >>> from vitalDSP.transforms.chroma_stft import ChromaStft
    >>> signal = np.random.randn(1000)
    >>> processor = ChromaStft(signal)
    >>> result = processor.process()
    >>> print(f'Processing result: {result}')
"""

import numpy as np
from scipy.signal import get_window


[docs] class ChromaSTFT: """ A class to compute the Chroma Short-Time Fourier Transform (Chroma STFT) to analyze harmonic content in audio signals. The Chroma STFT is useful for identifying the harmonic structure of a signal by projecting its frequency content onto a small number of pitch classes (chroma bins). Attributes ---------- signal : numpy.ndarray The input audio signal. sample_rate : int The sample rate of the signal. n_chroma : int The number of chroma bins (usually 12 for the 12 pitch classes). n_fft : int The FFT size, determining the frequency resolution of the STFT. hop_length : int The number of samples between successive frames. Methods ------- compute_chroma_stft() Computes the Chroma STFT of the signal. Example Usage ------------- >>> signal = np.sin(2 * np.pi * 440 * np.linspace(0, 2, 16000)) # 2 seconds of A4 note (440Hz) >>> chroma_stft = ChromaSTFT(signal, sample_rate=16000, n_chroma=12, n_fft=2048, hop_length=512) >>> chroma_stft_result = chroma_stft.compute_chroma_stft() >>> print(chroma_stft_result.shape) # Output: (12, num_frames) """ def __init__( self, signal, sample_rate=16000, n_chroma=12, n_fft=2048, hop_length=512 ): """ Initialize the ChromaSTFT class with the signal and parameters. Parameters ---------- signal : numpy.ndarray The input audio signal to be analyzed. sample_rate : int, optional The sample rate of the signal (default is 16000 Hz). n_chroma : int, optional The number of chroma bins (default is 12). n_fft : int, optional The FFT size, determining the frequency resolution (default is 2048). hop_length : int, optional The number of samples between successive frames (default is 512). """ self.signal = signal self.sample_rate = sample_rate self.n_chroma = n_chroma self.n_fft = n_fft self.hop_length = hop_length def _compute_stft(self): """ Compute the Short-Time Fourier Transform (STFT) of the signal. Returns ------- stft_matrix : numpy.ndarray The STFT of the signal, where each column represents a time frame and each row a frequency bin. """ if len(self.signal) < self.n_fft: raise ValueError( "The length of the signal is shorter than the FFT size (n_fft)." ) window = get_window("hann", self.n_fft) num_frames = 1 + (len(self.signal) - self.n_fft) // self.hop_length # Check if there are enough frames to compute STFT if num_frames <= 0: raise ValueError( "The signal is too short for the given FFT size and hop length." ) stft_matrix = np.empty((self.n_fft // 2 + 1, num_frames), dtype=np.complex64) for i in range(num_frames): start = i * self.hop_length frame = self.signal[start : start + self.n_fft] * window stft_matrix[:, i] = np.fft.rfft(frame) return np.abs(stft_matrix) def _create_chroma_filter(self): """ Create a simulated chroma filter to map frequency bins to chroma bins. Returns ------- chroma_filter : numpy.ndarray A simulated chroma filter mapping frequency bins to chroma bins. """ chroma_filter = np.random.rand(self.n_chroma, self.n_fft // 2 + 1) return chroma_filter
[docs] def compute_chroma_stft(self): """ Compute the Chroma Short-Time Fourier Transform (Chroma STFT) of the signal. This function applies the STFT to the signal and then projects the result onto chroma bins. Returns ------- chroma_stft : numpy.ndarray The Chroma STFT of the signal. Each row represents a chroma bin, and each column represents a time frame. Examples -------- >>> signal = np.sin(2 * np.pi * 440 * np.linspace(0, 2, 16000)) # 2 seconds of A4 note (440Hz) >>> chroma_stft = ChromaSTFT(signal, sample_rate=16000, n_chroma=12, n_fft=2048, hop_length=512) >>> chroma_stft_result = chroma_stft.compute_chroma_stft() >>> print(chroma_stft_result.shape) # Output: (12, num_frames) """ stft_matrix = self._compute_stft() # Apply the chroma filter to the STFT result chroma_filter = self._create_chroma_filter() chroma_stft = np.dot(chroma_filter, stft_matrix) return chroma_stft