Monday, March 11, 2019

SciPy - 4 (Fourier transforms)

Fourier transforms are one of the universal tools in computational physics, which appear over and over again in di erent contexts. SciPy provides functions for accessing the classic FFTPACK library from NetLib, which is an e cient and well tested FFT library written in FORTRAN. The SciPy API has a few additional convenience functions, but overall the API is closely related to the original FORTRAN library.

To use the fftpack module in a python program, include it using:

from numpy.fft import fftfreq
from scipy.fftpack import *

To demonstrate how to do a fast Fourier transform with SciPy, let's look at the FFT of the solution to the damped oscillator from the previous post:

from scipy . integrate import odeint
import numpy as N
from matplotlib import pyplot as plt
from numpy.fft import fftfreq
from scipy.fftpack import *

def dy(y, t, zeta, w0):
    """
    The right-hand side of the damped oscillator ODE
    """
    x, p = y[0], y[1]
    dx = p
    dp = -2 * zeta * w0 * p - w0**2 * x
    return [dx, dp]

# initial state:
y0 = [1.0, 0.0]

# time coodinate to solve the ODE for
t = N.linspace(0, 10, 1000)

w0 = 2*N.pi*1.0

# solve the ODE problem for three different values of the damping ratio
y1 = odeint(dy, y0, t, args=(0.0, w0)) # undamped
y2 = odeint(dy, y0, t, args=(0.2, w0)) # under damped
y3 = odeint(dy, y0, t, args=(1.0, w0)) # critial damping
y4 = odeint(dy, y0, t, args=(5.0, w0)) # over damped


N = len(t)
dt = t[1]-t[0]
# calculate the fast fourier transform
# y2 is the solution to the under-damped oscillator from the previous section
F = fft(y2[:,0])
# calculate the frequencies for the components in F
w = fftfreq(N, dt)

fig, ax = plt.subplots(figsize=(9,3))
ax.plot(w, abs(F));
plt.show()


The output of the program is shown below:


As the signal is real, the spectrum is symmetric as shown in the output. So we only need to plot the part that corresponds to the positive frequencies. To extract that part of the w and F we can use some of the indexing tricks for NumPy arrays as shown in the code below:

indices = where(w > 0)
w_pos = w[indices]
F_pos = F[indices]


Then we can plot the spectrum of the extracted part as shown below:

fig, ax = plt.subplots(figsize=(9,3))
ax.plot(w_pos, abs(F_pos))
ax.set_xlim(0, 5);
plt.show()


Add the above lines of code in the program and our output becomes:



In the above figure we now see a peak in the spectrum that is centered around 1, which is the frequency we used in the damped oscillator example.

Now let us make a program in which we create a signal as a superposition of a 50 Hz and 70 Hz sine wave (with a slight phase shift between them). We then Fourier transform the signal and plot the absolute value of the (complex) discrete Fourier transform coefficients against frequency, and expect to see peaks at 50Hz and 70Hz. See the program below:

import scipy
import matplotlib . pyplot as plt


pi = scipy .pi
signal_length = 0.5 #[ seconds ]
sample_rate =500 # sampling rate [Hz]
dt = 1./ sample_rate # time between two samples [s]
df = 1/ signal_length # frequency between points in frequency domain [Hz]
t= scipy . arange (0, signal_length ,dt) #the time vector
n_t=len(t) # length of time vector

# create signal
y= scipy .sin (2* pi *50* t)+ scipy .sin (2* pi *70* t+pi /4)

# compute fourier transform
f= scipy .fft(y)

# work out meaningful frequencies in fourier transform
freqs =df* scipy . arange (0 ,( n_t -1)/2. , dtype ='d') #d= double precision float
n_freq = len ( freqs )

# plot input data y against time
plt. subplot (2 ,1 ,1)
plt. plot (t,y, label ='input data ')
plt. xlabel ('time [s]')
plt. ylabel ('signal ')

# plot frequency spectrum
plt. subplot (2 ,1 ,2)
plt. plot (freqs ,abs(f[0: n_freq ]),
label ='abs( fourier transform )')
plt. xlabel ('frequency [Hz]')
plt. ylabel ('abs(DFT( signal )) ')


plt. show () #and display plot on screen


The resulting plots are shown below. The lower plot shows the discrete Fourier transform computed
from the data shown in the upper plot.



Here I am ending this post about Fourier transforms. Next post will discuss about linear algebra module. Till we meet next keep learning and practicing Python as Python is easy to learn!
Share:

0 comments:

Post a Comment