Testing our DFT

Testing our DFT#

A good thing to check is that if we compute the DFT and then the inverse DFT of the result, do we get the original data back (up to roundoff error)?

First let’s copy over our dft() function and write the inverse transform

import numpy as np
import matplotlib.pyplot as plt
def dft(f_n):
    """perform a discrete Fourier transform"""
    
    N = len(f_n)
    n = np.arange(N)
    
    f_k = np.zeros((N), dtype=np.complex128)

    for k in range(N):
        f_k[k] = np.sum(f_n * np.exp(-2.0 * np.pi * 1j * n * k / N))
    return f_k
def idft(F_k):
    """perform a discrete Fourier transform"""
    
    N = len(F_k)
    k = np.arange(N)
    
    f_n = np.zeros((N), dtype=np.float64)

    for n in range(N):
        f_n[n] = (1/N) * np.sum(F_k * np.exp(2.0 * np.pi * 1j * n * k / N)).real
    return f_n

Let’s create some functions that we will transform

def f_single_sine(x):
    nu_0 = 0.2
    return np.sin(2.0 * np.pi * nu_0 * x)

def f_single_sine_phase(x):
    nu_0 = 0.2
    return np.sin(2.0 * np.pi * nu_0 * x + np.pi / 4)

def f_single_cosine(x):
    nu_0 = 0.2
    return np.cos(2.0 * np.pi * nu_0 * x)

def f_two_freq(x):
    nu_0 = 0.2
    nu_1 = 0.5
    return 0.5 * (np.sin(2.0 * np.pi * nu_0 * x) + np.sin(2.0 * np.pi * nu_1 * x))

Now let’s make a function that plots the original function, the DFT, the power spectrum, and the inverse transform of the DFT

def plot_dft(f, N=64, xmax=50):
    
    fig, ax = plt.subplots(nrows=4)
    
    x = np.linspace(0.0, xmax, N, endpoint=False)
    
    # get the transform   
    f_n = f(x)
    F_k_raw = dft(f_n)
    
    # normalize and cut off the duplicate data
    # since we are discarding 1/2 the data, the normalization if N/2
    F_k = F_k_raw[0:N//2] * 2 / N
    
    # compute the frequencies
    nu_k = np.arange(len(F_k)) / xmax
    
    # plot the original function
    ax[0].plot(x, f_n)
    ax[0].set_xlabel("$x$")
    ax[0].set_ylabel("$f(x)$")
    ax[0].set_xmargin(0.0)
    
    # plot the transform
    ax[1].plot(nu_k, F_k.real, label=r"$\mathrm{Re}(F_k)$")
    ax[1].plot(nu_k, F_k.imag, label=r"$\mathrm{Im}(F_k)$")    
    ax[1].set_xlabel(r"$\nu$")
    ax[1].set_ylabel("$F_k$")
    ax[1].set_xmargin(0.0)
    ax[1].legend(fontsize="small", ncol=2, frameon=False)
    
    # plot the power spectrum
    ax[2].plot(nu_k, np.abs(F_k)**2)
    ax[2].set_xlabel(r"$\nu$")
    ax[2].set_ylabel("$|F_k|^2$")
    ax[2].set_xmargin(0.0)
    
    # take the inverse transform -- note we need to use the unnormalized data
    f_n_new = idft(F_k_raw)
    ax[3].plot(x, f_n_new)
    ax[3].set_xlabel("$x$")
    ax[3].set_ylabel("$f(x)$")
    ax[3].set_xmargin(0.0)
    
    fig.tight_layout()
    
    return fig

Now let’s look at the transform of a single mode sine

fig = plot_dft(f_single_sine, N=256)
../_images/53e2082597d6ed2db278e058d7175d3b3b12d521e08b440883567e29730e7f23.png

We get what we expect:

  • The transform only has non-zero weights in the imaginary part

  • All of the power is at the single frequency that we constructed the sine with

  • The inverse transform of the transformed data gives us back our original function

Important:

The Fourier transform treats the data as if it is periodic (since the sine and cosine functions are). The input data is assumed to be equally spaced and periodic. The means that the first and last point need to be distinct.

This is why we construct the x data with endpoint=False—if we have data at endpoint, then the start and end refer to the exact same data sample, but the Fourier transform will treat them as if they are space \(\Delta x\) apart, and we will pick up an extra signal.

Try leaving off the endpoint=False and see what happens.

Now let’s look at cosine—it is the opposite of sine

fig = plot_dft(f_single_cosine, N=256)
../_images/8ea7ca10020e05cbb4231169f8c523216d8bb03b8307e66e36da395a8bacaf07.png

A \(\pi/4\) phase shift puts equal weight into the sine and cosine components:

fig = plot_dft(f_single_sine_phase, N=256)
../_images/387a633b3e4bb3804161045b63268b93d3511c73ce64b55b2b0d64121eaf0f88.png

The two frequency version shows power at the two frequencies, as expect:

fig = plot_dft(f_two_freq, N=256)
../_images/cac399a0c83705a2099572fb00f296d1064a9c418b41be540fea12aa8e70017d.png