Home | History | Annotate | Download | only in processing
      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