pygpibtoolkit/HP3562A/mathtools.py

# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
"""Copyright (c) 2007-2018 David Douard (Paris, FRANCE).
http://www.logilab.org/project/pygpibtoolkit -- mailto:david.douard@sdfa3.org

"""

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