# Generated by pandoc-plot 1.9.1
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2019)

import numpy as np
import matplotlib.pyplot as plt

t = np.linspace(-1, 1, 1024)
f = np.exp(1j * 2 * 2 * np.pi * t) + np.exp(1j * 2 * np.pi * 5 * t)

# For technical reasons, the frequency components of the Fourier transform
# are not arranged as one would expect. In order for the plot to look OK, we
# need to use `np.fft.fftshift`
fhat = np.fft.fftshift(np.fft.fft(f))
w = np.fft.fftshift(np.fft.fftfreq(f.shape[0], d=t[1] - t[0]))

fig, (ax_t, ax_w) = plt.subplots(2, 1, figsize=(8, 6))

ax_t.plot(t, np.real(f), color="k")
ax_t.set_xlim([t.min(), t.max()])
ax_t.set_xlabel("Time [s]")

ax_w.plot(w, np.abs(fhat)**2, ".", color="k")
ax_w.set_xlim([-2, 9])
ax_w.set_ylabel("Spectral power [a.u.]")
ax_w.set_xlabel("Angular frequency [rad/s]")

for freq in [2, 5]:
    ax_w.axvline(x=freq, linestyle='--', linewidth=1, color='k')
    ax_w.annotate(f"$\omega={freq}$", xy=(1.05*freq, 0.95*np.square(np.abs(fhat)).max()), xycoords='data')