#!/usr/bin/env python3

# code to demonstrate usage of the fast Fourier Transform
# notation generally follows the lecture notes
# you will have to get from the DFT to the power spectrum by yourself

import numpy as np
import matplotlib.pyplot as plt

# as a demonstration, set up an oscillation with frequency 1/100 
# we do this in terms of n rather than t, 
# this is equivalent to using a timestep of 1
nmax = 1000
n = np.arange(0,nmax,1)
x = np.sin(2*np.pi*n/100)

# calculate Discrete Fourier Transform
dft = np.fft.fft(x)
freq = n/nmax
plt.figure()
plt.plot(freq,dft)
plt.xlabel('Frequency')
plt.ylabel('x_k')

plt.show()