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