Mercurial > pygpibtoolkit / file revision

import numpy

def thd(spectrum, squared=True, db=True):
    """
    compute THD from a spectrum data, given by a simple 1D data
    vector.

    The result is the couple (fmax, thd), this latter being expressed
    in 'db' if db is True; if not, in percentage.
    
    """
    n = len(spectrum)
    # first, find the fundamental frequency
    fmax = spectrum.argmax()
    vmax = spectrum[fmax]

    # then, the compute the several max values for harmonics
    h_freqs = numpy.arange(0, n, fmax)
    h_vals = []
    h_freqs2 = [fmax]
    for i, f in enumerate(h_freqs[2:]):
        fm = spectrum[f-5:f+5].argmax()
        h_vals.append(spectrum[f-5+fm])
        h_freqs2.append((f-5.+fm)/(i+2))
    h_vals = numpy.array(h_vals)
    h_freqs2 = numpy.array(h_freqs2)
    if not squared:
        h_vals = h_vals**2
        vmax = vmax**2
    thd = h_vals.sum()/vmax
    if db:
        thd = 10*numpy.log10(thd)
    else:
        thd = 100.0*numpy.sqrt(thd)
    return h_freqs2.mean(), thd

def thd_n(spectrum, squared=True, db=True):
    """
    compute THD+N from a spectrum data, given by a simple 1D data
    vector.

    The result is the thdN, expressed in 'db' if db is True; if not,
    in percentage.
    
    """
    if not squared:
        spectrum = spectrum**2
    fmax = spectrum.argmax()
    w = int(fmax/10) # argh
    vmax = spectrum[fmax-w:fmax+w].sum()
    thd = (spectrum.sum()-vmax)/(vmax)
    if db:
        thd = 10*numpy.log10(thd)
    else:
        thd = 100.0*numpy.sqrt(thd)
    return thd