A demonstration of Frequency Estimation Based on SCDE.

scde_demo.py

# -*- coding: utf-8 -*-
"""A demonstration of Frequency Estimation Based on SCDE.

Copyright (C) 2025 by Akira TAMAMORI

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import numpy as np
import numpy.polynomial.polynomial as poly
import numpy.typing as npt
from scipy.signal import butter, convolve, filtfilt


def apply_bandpass_filter(
    signal: npt.NDArray[np.float64],
    fs: float,
    lowcut: int,
    highcut: int,
    order: int = 4,
) -> npt.NDArray[np.float64]:
    """Apply a Butterworth bandpass filter to the signal.

    Args:
        signal (np.ndarray): Input discrete-time signal.
        fs (float): Sampling frequency in Hz.
        lowcut (int): Lower cutoff frequency.
        highcut (int): Higher cutoff frequency.
        order (int): Order of bandpass filter coeff.

    Returns:
        filtered_signal (np.ndarray): Bandpass-filtered signal.
    """
    nyquist = 0.5 * fs
    low = lowcut / nyquist
    high = highcut / nyquist
    filter_coef = butter(order, [low, high], btype="band")
    assert (
        isinstance(filter_coef, tuple)
        and len(filter_coef) == 2
        and isinstance(filter_coef[0], np.ndarray)
        and isinstance(filter_coef[1], np.ndarray)
    )
    filtered_signal: npt.NDArray[np.float64] = filtfilt(
        filter_coef[0], filter_coef[1], signal
    )
    return filtered_signal


def get_kth_even_derivative(
    s: npt.NDArray[np.float64], k_order: int, dt: float
) -> tuple[npt.NDArray[np.float64], int]:
    """Calculates the k-th order even derivative of a signal using centered differences.

    Args:
        s (np.ndarray): Input signal.
        k_order (int): The order of the derivative (must be an even number).
        dt (float): Sampling period.

    Returns:
        tuple: (derived_signal, pad_length_per_side)
               derived_signal: The calculated derivative.
               pad_length_per_side: Number of points trimmed from each
                                    side of the original signal.
    """
    if k_order % 2 != 0:
        raise ValueError("Only even order derivatives are supported.")

    # Generate the centered difference kernel for the specified even order
    kernel = np.array([1.0])  # Use float to ensure float output
    for _ in range(k_order // 2):
        kernel = np.convolve(kernel, np.array([1.0, -2.0, 1.0]), mode="full")

    # convolve with mode='valid' shortens the signal by k_order
    # points in total (k_order/2 from each side)
    derived_s = convolve(s, kernel, mode="valid") / (dt**k_order)
    return derived_s, k_order // 2


# Helper to calculate the sum of products for covariance matrix elements (Sm,n)
def calc_covmat_mn(
    m: int,
    n: int,
    signal_original: npt.NDArray[np.float64],
    dt: float,
    global_max_pad: int,
) -> np.float64:
    """Calculates the sum of products of (2m)-th and (2n)-th derivatives.

    Corresponds to Sm,n in the paper.
    The k-th even derivative d^(2k)x/dt^(2k) is approximated as:
    1/(dt^2k) * Sum_{j=-k}^{k} C_{2k, k-j} x[n+j]

    Args:
        m (int): The order of the derivative.
        n (int): The order of the derivative
        signal_original (np.darray): Original signal.
        dt (float): Sampling period.
        global_max_pad (int): Maximum of padding space.

    Returns:
        covmat_mn (np.ndarray): The (m, n) element of covariance matrix.
    """
    # Get the derivatives. [0] for the array, [1] for pad_len for each deriv.
    deriv_x_m_full, pad_m = get_kth_even_derivative(signal_original, 2 * m, dt)
    deriv_x_n_full, pad_n = get_kth_even_derivative(signal_original, 2 * n, dt)

    # Trim derivatives to match the length of the highest derivative (max_pad_len)
    # This ensures all derivatives correspond to the same central time segment.
    trimmed_deriv = {}
    start_pos = global_max_pad - pad_m
    end_pos = len(deriv_x_m_full) - (global_max_pad - pad_m)
    trimmed_deriv["m"] = deriv_x_m_full[start_pos:end_pos]

    start_pos = global_max_pad - pad_n
    end_pos = len(deriv_x_n_full) - (global_max_pad - pad_n)
    trimmed_deriv["n"] = deriv_x_n_full[start_pos:end_pos]

    # Ensure they have the same length after trimming for element-wise multiplication
    min_len = min(len(trimmed_deriv["m"]), len(trimmed_deriv["n"]))
    covmat_mn: np.float64 = np.sum(
        trimmed_deriv["m"][:min_len] * trimmed_deriv["n"][:min_len]
    )
    return covmat_mn


def solve_poly(
    alphas: npt.NDArray[np.float64], n_sinusoids: int
) -> npt.NDArray[np.float64]:
    """Solve polynomial equation for the squared angular frequency.

    Args:
        alphas (np.ndarray): Coefficients of differential equation of SCDE.
        n_sinusoids (int): Number of Sinusoids.

    Returns:
        omegas_squared (np.darray): Estimated squared angular frequency.
    """
    # The polynomial is Ω^3 - α1Ω^2 + α2Ω - α3 = 0 (based on paper's Eq. 13 and 18)
    # polyroots expects coefficients in ascending order of power:
    # [coeff_0, coeff_1, ..., coeff_n]
    # So, for n=3, the coefficients are [-alpha3, alpha2, -alpha1, 1]
    coeffs_list = []
    for i in reversed(range(n_sinusoids)):
        coeffs_list.append((-1) ** (i + 1) * alphas[i])
    coeffs_list.append(1)
    polynomial_coeffs = np.array(coeffs_list, dtype=np.float64)
    omegas_squared: npt.NDArray[np.float64] = poly.polyroots(polynomial_coeffs)
    return omegas_squared


def estimate_frequencies_scde(
    signal: npt.NDArray[np.float64], fs: float, n_sinusoids: int = 3
) -> npt.NDArray[np.float64]:
    """Estimates the frequencies of three sinusoids by SCDE.

    This function implements the estimation method of frequencies using the Sinusoidal
    Constraint Differential Equation (SCDE) method, as described in the following paper:

    Kenta Yamada, Yoshiki Masuyama, Yukoh Wakabayashi, and Nobutaka Ono
    "Simultaneous Frequency Estimation for Three or More Sinusoids Based on Sinusoidal
    Constraint Differential Equation," Proceedings of 2022 APSIPA ASC, pp. 975-978.
    https://ieeexplore.ieee.org/document/9980228

    Args:
        signal (np.ndarray): Input discrete-time signal.
        fs (float): Sampling frequency in Hz.
        n_sinusoids (int): Number of Sinusoids.

    Returns:
        np.ndarray: An array of estimated frequencies in Hz, sorted in ascending order.
                    Returns an empty array if no valid frequencies are found.
    """
    dt = 1.0 / fs

    # Determine the maximum padding required (for the highest derivative)
    # This ensures all derivatives are trimmed to the same length and time segment.
    # For 6th order derivative (k_order=6), pad_len is 3
    max_pad_len = n_sinusoids

    # Construction of the Covariance Matrix 'B' and Right-Hand Side Vector 'V'
    # For n=3 sinusoids, the system of equations for alphas (alpha1, alpha2, alpha3)
    # corresponds to equation (11) in the paper.
    # The notation S_a,b in the paper's matrix elements corresponds to the sum of
    # (2a)-th derivative * (2b)-th derivative.

    # Construct the coefficient matrix B (from paper's Eq. 11)
    b_mat = np.zeros((n_sinusoids, n_sinusoids), dtype=np.float64)
    for m in range(n_sinusoids):
        for n in range(n_sinusoids):
            b_mat[m, n] = calc_covmat_mn(
                n_sinusoids - 1 - n, n_sinusoids - 1 - m, signal, dt, max_pad_len
            )

    # Calculate individual elements of the V vector (right-hand side)
    rhs = np.zeros(n_sinusoids, dtype=np.float64)
    for i in range(n_sinusoids):
        rhs[i] = calc_covmat_mn(
            n_sinusoids, n_sinusoids - 1 - i, signal, dt, max_pad_len
        )

    # Solving the system of linear equations
    # e.g., B * [alpha1, alpha2, alpha3]^T = -V
    try:
        alphas = np.linalg.solve(b_mat, -rhs).astype(np.float64)
    except np.linalg.LinAlgError:
        print("Warning: Failed to solve the linear system (e.g., singular matrix).")
        return np.array([])

    # Solving the polynomial equation for roots
    omegas_squared_complex = solve_poly(alphas, n_sinusoids)

    estimated_frequencies_hz = []
    for omega_sq in omegas_squared_complex:
        # Ω (omega_sq) should ideally be real and positive,
        # as it represents ω^2 (squared angular frequency).

        # Small imaginary parts or negative real parts might be due to
        # numerical errors or noise.
        if np.isreal(omega_sq) and omega_sq.real >= 0:
            # Convert angular frequency (rad/s) to Hz
            frequency_hz = np.sqrt(omega_sq.real) / (2 * np.pi)
            estimated_frequencies_hz.append(frequency_hz)

    return np.array(sorted(estimated_frequencies_hz))


def main() -> None:
    """Perform demonstration."""
    fs = 44100  # Sampling frequency (Hz)
    duration = 0.04  # Signal duration (seconds) - 40ms
    t = np.linspace(0, duration, int(fs * duration), endpoint=False)

    # True frequencies of the sinusoids
    f_true = np.array([440.0, 460.0, 480.0])  # Example from the paper

    # Generate the test signal (sum of three sinusoids)
    amplitude = 1.0
    signal = {}
    signal["pure"] = amplitude * np.cos(2 * np.pi * f_true[0] * t)
    for i in range(1, len(f_true)):
        signal["pure"] += np.cos(2 * np.pi * f_true[i] * t)

    print("--- Test without noise ---")
    estimated_freqs = estimate_frequencies_scde(signal["pure"], fs, len(f_true))
    print(f"True frequencies: {f_true} Hz")
    print(f"Estimated frequencies: {estimated_freqs} Hz")
    if len(estimated_freqs) == len(f_true):
        print(f"Estimation errors: {estimated_freqs - f_true} Hz")
    else:
        print("Number of estimated frequencies does not match the true number.")

    print("\n--- Test with band-limited noise ---")
    # Add noise to the signal (approx. 10dB SNR)
    # Calculate noise power for desired SNR
    signal_power = np.var(signal["pure"])
    snr_db = 10
    noise_power = signal_power / (10 ** (snr_db / 10))
    noise = np.random.normal(0, np.sqrt(noise_power), len(signal["pure"]))
    signal["noisy"] = signal["pure"] + apply_bandpass_filter(noise, fs, 430, 490)
    estimated_freqs = estimate_frequencies_scde(signal["noisy"], fs, len(f_true))
    print(f"True frequencies: {f_true} Hz")
    print(f"Estimated frequencies with noise: {estimated_freqs} Hz")
    if len(estimated_freqs) == len(f_true):
        print(f"Estimation errors: {estimated_freqs - f_true} Hz")
    else:
        print("Number of estimated frequencies does not match the true number.")

    print("\n--- Test with very short signal duration (2ms) ---")
    duration_short = 0.002  # 2ms
    t_short = np.linspace(0, duration_short, int(fs * duration_short), endpoint=False)
    signal["short"] = amplitude * np.cos(2 * np.pi * f_true[0] * t_short)
    for i in range(1, len(f_true)):
        signal["short"] += np.cos(2 * np.pi * f_true[i] * t_short)

    estimated_freqs = estimate_frequencies_scde(signal["short"], fs, len(f_true))
    print(f"True frequencies: {f_true} Hz")
    print(f"Estimated frequencies for 2ms signal: {estimated_freqs} Hz")
    if len(estimated_freqs) == len(f_true):
        print(f"Estimation errors: {estimated_freqs - f_true} Hz")
    else:
        print("Number of estimated frequencies does not match the true number.")


if __name__ == "__main__":
    main()