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 17 import numpy as np 18 import scipy as sp 19 import scipy.fftpack as fft 20 import scipy.linalg as la 21 import math 22 23 def 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 48 if __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