fourier_polynomial.py

import numpy as np
import matplotlib.pyplot as plt


def fourier_polynomial(f: 'function', n: int, r: int, M: int) -> np.array:
    """ Computes the Fourier polynomial of degree n.
    
    Args:
        f: Any Python function that returns a float
        n: Degree of the Fourier polynomial
        r: r equidistant nodes that the Fourier polynomial is evaluated on
        M: number of (equidistant) samples that we want to have of f
    
    Returns:
        A tuple containing the 2*pi-grid (depending on r) and the polynomial
        evaluated on it.
    """
    
    # vectorize the function at values 2*pi*j/M:
    grid = np.array([2*np.pi*j / M for j in range(M)])
    y = np.vectorize(f)
    
    # compute the (scaled) Fourier transform of our function values
    c = 1/M * np.fft.fft(y(grid))
    
    # find number of zeros needed to fill the ct (c tilde) vector
    nz = r - c.shape[0]
    
    # fill the ct vector
    ct = np.r_[c[0:n], np.zeros(nz), c[n:]]
    
    # compute the Fourier transform of the ct values
    q = np.fft.fft(ct)
    p0 = np.r_[q[0], [q[r-j] for j in range(1, r)]]

    # return the x and y values of p0 
    return (np.linspace(0, 2*np.pi, r), p0)


def example_function(x):
    # this is the example function f(x) from the exercise sheet
    if (x >= 0 and x < np.pi / 2) or (x >= np.pi and x < 1.5 * np.pi):
        return 1.0
    else:
        return 0.0    


# Testing
M_values = [16, 32, 64, 128]
r = 256
n = 16

p = [fourier_polynomial(example_function, n, r, M) for M in M_values]

plt.figure()
for k, poly in enumerate(p):
    plt.subplot(4,1,(k+1))
    plt.title("Approximation for M=%s" % (M_values[k]))
    plt.plot(poly[0], poly[1])
    plt.plot(np.linspace(0, 2*np.pi, 256), np.vectorize(example_function)(np.linspace(0, 2*np.pi, 256)))
pylab.savefig('plots.pdf', bbox_inches='tight')
plt.show()