1#!/usr/bin/python
2
3# Copyright (C) 2012 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import numpy as np
18import scipy as sp
19import scipy.fftpack as fft
20import scipy.linalg as la
21import math
22
23def calc_thd(data, signalFrequency, samplingRate, frequencyMargin):
24    # only care about magnitude
25    fftData = abs(fft.fft(data * np.hanning(len(data))))
26    fftData[0] = 0 # ignore DC
27    fftLen = len(fftData)/2
28    baseI = fftLen * signalFrequency * 2 / samplingRate
29    iMargain = baseI * frequencyMargin
30    baseSignalLoc = baseI - iMargain / 2 + \
31        np.argmax(fftData[baseI - iMargain /2: baseI + iMargain/2])
32    peakLoc = np.argmax(fftData[:fftLen])
33    if peakLoc != baseSignalLoc:
34        print "**ERROR Wrong peak signal", peakLoc, baseSignalLoc
35        return 1.0
36    print baseI, baseSignalLoc
37    P0 = math.pow(la.norm(fftData[baseSignalLoc - iMargain/2: baseSignalLoc + iMargain/2]), 2)
38    i = baseSignalLoc * 2
39    Pothers = 0.0
40    while i < fftLen:
41        Pothers += math.pow(la.norm(fftData[i - iMargain/2: i + iMargain/2]), 2)
42        i += baseSignalLoc
43    print "P0", P0, "Pothers", Pothers
44
45    return Pothers / P0
46
47# test code
48if __name__=="__main__":
49    samplingRate = 44100
50    durationInSec = 10
51    signalFrequency = 1000
52    samples = float(samplingRate) * float(durationInSec)
53    index = np.linspace(0.0, samples, num=samples, endpoint=False)
54    time = index / samplingRate
55    multiplier = 2.0 * np.pi * signalFrequency / float(samplingRate)
56    data = np.sin(index * multiplier)
57    thd = calc_thd(data, signalFrequency, samplingRate, 0.02)
58    print "THD", thd
59
60