THD calculation

thd_mat.py

import math

import numpy as np
import scipy as sci
from numpy.linalg.linalg import solve
from scipy.signal import correlate, fftconvolve
import matplotlib.pyplot as plt
import cmath

def parabolic(f, x):
    xv = 1/2. * (f[x-1] - f[x+1]) / (f[x-1] - 2 * f[x] + f[x+1]) + x
    yv = f[x] - 1/4. * (f[x-1] - f[x+1]) * (xv - x)
    return (xv, yv)

def indexes(y, thres=0.3, min_dist=1, thres_abs=False):
    """Peak detection routine borrowed from
    https://bitbucket.org/lucashnegri/peakutils/src/master/peakutils/peak.py
    """
    if isinstance(y, np.ndarray) and np.issubdtype(y.dtype, np.unsignedinteger):
        raise ValueError("y must be signed")

    if not thres_abs:
        thres = thres * (np.max(y) - np.min(y)) + np.min(y)

    min_dist = int(min_dist)

    # compute first order difference
    dy = np.diff(y)

    # propagate left and right values successively to fill all plateau pixels (0-value)
    zeros, = np.where(dy == 0)

    # check if the signal is totally flat
    if len(zeros) == len(y) - 1:
        return np.array([])

    if len(zeros):
        # compute first order difference of zero indexes
        zeros_diff = np.diff(zeros)
        # check when zeros are not chained together
        zeros_diff_not_one, = np.add(np.where(zeros_diff != 1), 1)
        # make an array of the chained zero indexes
        zero_plateaus = np.split(zeros, zeros_diff_not_one)

        # fix if leftmost value in dy is zero
        if zero_plateaus[0][0] == 0:
            dy[zero_plateaus[0]] = dy[zero_plateaus[0][-1] + 1]
            zero_plateaus.pop(0)

        # fix if rightmost value of dy is zero
        if len(zero_plateaus) and zero_plateaus[-1][-1] == len(dy) - 1:
            dy[zero_plateaus[-1]] = dy[zero_plateaus[-1][0] - 1]
            zero_plateaus.pop(-1)

        # for each chain of zero indexes
        for plateau in zero_plateaus:
            median = np.median(plateau)
            # set leftmost values to leftmost non zero values
            dy[plateau[plateau < median]] = dy[plateau[0] - 1]
            # set rightmost and middle values to rightmost non zero values
            dy[plateau[plateau >= median]] = dy[plateau[-1] + 1]

    # find the peaks by using the first order difference
    peaks = np.where(
        (np.hstack([dy, 0.0]) < 0.0)
        & (np.hstack([0.0, dy]) > 0.0)
        & (np.greater(y, thres))
    )[0]

    # handle multiple peaks, respecting the minimum distance
    if peaks.size > 1 and min_dist > 1:
        highest = peaks[np.argsort(y[peaks])][::-1]
        rem = np.ones(y.size, dtype=bool)
        rem[peaks] = False

        for peak in highest:
            if not rem[peak]:
                sl = slice(max(0, peak - min_dist), peak + min_dist + 1)
                rem[sl] = True
                rem[peak] = False

        peaks = np.arange(y.size)[~rem]

    return peaks

def find_period(signal):
    signal -= np.mean(signal)  # Remove DC offset
    corr = fftconvolve(signal, signal[::-1], mode='full')
    corr = corr[len(corr)//2:]

    # Find the first peak on the left
    i_peak = indexes(corr, thres=0.8, min_dist=5)[0]
    i_interp = parabolic(corr, i_peak)[0]

    #print(i_interp)

    #plt.plot(corr)
    #plt.plot(signal*10)
    #plt.show()

    return i_interp

def thd_mat(signal, fs):
    n_samples = len(signal)
    period = int(np.round(find_period(signal)))
    omega = (1 / period) * fs * 2 * np.pi       # Fundamental frequency in rad/s.
    delta_t = 1 / fs                            # Time between samples in s.

    print(omega, delta_t)


    # Initialize the coefficient matrix C
    c = []
    size = period
    window = signal[:size]
    window -= np.mean(window)
    for i in range(0, size):
        row = []
        for j in range(0, size):
            row.append(cmath.exp(1j*i*omega*j*delta_t))
        c.append(row)

    # Solve for harmonic coefficients
    a = solve(c, window)
    a = a[0:size//2]
    #print(np.abs(a))
#    plt.bar(range(0, size), np.abs(a))
#    plt.show()

    harmonics_power = np.sum(abs(a[2:])**2)
    fundamental_power = abs(a[1])**2
    thd = math.sqrt(harmonics_power/fundamental_power)
    return thd




samples = 1024
fs = 44100
f = samples * (1000 / fs)
winsize = 60
signal = np.sin(2*np.pi*f*np.linspace(0, 1, samples)) + 0.1*np.sin(2*np.pi*f*2*np.linspace(0, 1, samples))
period = find_period(signal)
print(thd_mat(signal, fs))