Created
February 3, 2012 19:50
-
-
Save cvr/1732018 to your computer and use it in GitHub Desktop.
Discrete Fourier Transform
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/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 <cvrodrigues@gmail.com> | |
| # 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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment