Discrete Fourier Transform

my_dft.py

#!/usr/bin/python
# vim: set fileencoding=utf-8 :
# -*- coding: utf-8 -*-
########################################################################
# Copyright (C) 2012 by Carlos Veiga Rodrigues. All rights reserved.
# author: Carlos Veiga Rodrigues <[email protected]>
# version: 1.0, date: January 2012
#
# This program can be redistribuited and modified
# under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation,
# either version 3 of the License or any later version.
# This program is distributed WITHOUT ANY WARRANTY,
# without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.
# For more details consult the GNU Lesser General Public License
# at http://www.gnu.org/licenses/lgpl.html.
########################################################################

import numpy as np

def dft (U, NDFT=None, t=None, DT=None):
        if None==NDFT:
                NDFT = U.size
        if None==t and None==DT:
                DT, FS = 1
        if None==t:
                t = np.linspace(0, U.size-1, U.size)*DT
        if None==DT:
                DT = np.mean(np.diff(t))
                FS = 1.0/DT
        ## TAKE THE DFT
        udft = np.fft.fft(u, n=NDFT)
        ## NORMALIZE DFT TO HAVE DFT[0] = sum(z)*DT == AREA OF DISCRETE SIGNAL
        udft*= DT
        ## INDEXES OF POSITIVE COEFFICIENTS
        if 0==(NDFT % 2):
                ii = np.linspace(0, NDFT, NDFT/2+1)/float(2)
        else:
                ii = np.linspace(0, NDFT-1, (NDFT-1)/2+1)/float(2)
        ii = np.array(ii, dtype=int)
        ## CONSTRUCT FREQUENCY AXIS
        f = ii/float(NDFT*DT)
        udft = udft[ii]
        ## POWER SPECTRAL DENSITY
        upsd = (udft*np.conj(udft)).real
        #upsd = np.abs(udft)**2
        ## PSD NORMALIZATION:
        ## IF DFT IS NOT NORMALIZED
        #upsd*= DT/float(U.size)
        ## IF DFT IS NORMALIZED BY udft*= DT
        upsd/= float(DT*U.size)
        ## AMPLITUDE SPECTRAL DENSITY
        uasd = np.abs(udft)
        #uasd = np.sqrt(( udft*np.conj(udft)).real )
        ## ASD NORMALIZATION:
        ## IF DFT IS NOT NORMALIZED (2 SIDES, CORRECT)
        #uasd/= float(U.size)
        ## IF DFT IS NOT NORMALIZED (1 SIDE, RETURN EXACT AMP)
        #uasd[0]/= float(U.size)
        #uasd[1:]*= 2/float(U.size)
        ## IF DFT IS NORMALIZED BY udft*= DT (2 SIDES, CORRECT)
        #uasd/= float(DT*U.size)
        ## IF DFT IS NORMALIZED BY udft*= DT (1 SIDE, RETURN EXACT AMP)
        uasd[0]/= float(DT*U.size)
        uasd[1:]*= 2/float(DT*U.size)
        return f, udft, uasd, upsd

########################################################################
## TEST FUNCTION
########################################################################
if __name__ == '__main__':
        import sys
        ## TEST WITH SINUSOIDS
        N=201
        t=np.linspace(0., 100., N)
        DT=np.mean(np.diff(t))
        FS=1./DT
        frq = [0.2, 0.05]
        amp = [2., 5.]
        u0 = amp[0]*np.sin(2*np.pi*t*frq[0])
        u1 = amp[1]*np.sin(2*np.pi*t*frq[1])
        u = u0 + u1
        ## EXECUTE FUNCTION dft
        f, Udft, Uasd, Upsd = dft (u, DT=DT)
        ## PLOT
        import matplotlib as mpl
        import matplotlib.pyplot as plt
        mpl.rc('axes', titlesize="x-large", labelsize="large")
        mpl.rc(('xtick','ytick'), labelsize="medium")
        mpl.rc('legend', fontsize="large")
        mpl.rc(('xtick.major','xtick.minor','ytick.major','ytick.minor'),pad=6)
        fig=plt.figure(figsize=(8,8))
        fig.subplots_adjust(top=0.95, bottom=0.08, hspace=0.3,\
                left = 0.12, right = 0.92, wspace=0.2)
        #
        ax1=fig.add_subplot(311)
        ax1.plot(t, u0, c='LightGreen', ls='-',\
                label=r"$u_1 = 2 \, \sin(2 \pi 0.2 t)$")
        ax1.plot(t, u1, c='LightBlue', ls='-',\
                label=r"$u_2 = 5 \, \sin(2 \pi 0.05 t)$")
        ax1.plot(t, u, c='DarkRed', ls='-', marker='.', lw=1.2,\
                label=r"$u(t) = u_1 + u_2$")
        ax1.set_ylabel(r"$u(t)$")
        ax1.set_xlabel(r"$t \,,\;(s)$")
        ax1.legend(loc=1)
        #
        ax2=fig.add_subplot(312)
        ax2.plot(f, Uasd, c='DarkGreen', ls='-')
        ax2.set_ylabel(r"$A_u(\omega)$")
        #ax2.set_xlabel(r"Frequency $(1/s)$")
        ax2.axvline(frq[0], ls='--', lw=0.5, c='k')
        ax2.axvline(frq[1], ls='--', lw=0.5, c='k')
        plt.setp(ax2.get_xticklabels(), visible=False)
        #
        ax3=fig.add_subplot(313)
        ax3.plot(f, Upsd, c='DarkBlue', ls='-')
        ax3.set_ylabel(r"$S_u(\omega)$")
        ax3.set_xlabel(r"$\mathrm{Frequency} \,,\;(1/s)$")
        ax3.axvline(frq[0], ls='--', lw=0.5, c='k')
        ax3.axvline(frq[1], ls='--', lw=0.5, c='k')
        #
        #plt.savefig(figname, format='pdf')
        #plt.savefig("sines_spectra.png", format='png', dpi=110)
        plt.show()