Code for amazing 3B1B lecture (https://youtu.be/spUNpyF58BY) about Fourier Transform intuition.

fourier_transform.py

from scipy.integrate import simpson
import matplotlib.pyplot as plt
import numpy as np


def almost_fourier_transform(g: np.ndarray, t: np.ndarray, f: float):
    # approximate center of mass as a mean
    t1 = np.min(t)
    t2 = np.max(t)
    scaling_factor = (1 / (t2 - t1))
    g_complex = signal_in_complex_plane(g, t, f)
    return scaling_factor * np.mean(g_complex)

def fourier_transform(g: np.ndarray, t: np.ndarray, f: float):
    # computes center of mass using numeric integration
    t1 = np.min(t)
    t2 = np.max(t)
    scaling_factor = (1 / (t2 - t1))
    g_complex = signal_in_complex_plane(g, t, f)
    return scaling_factor * simpson(g_complex, t)

def signal_in_complex_plane(g: np.ndarray, t: np.ndarray, f: float):
    return g * np.e ** (-2j * np.pi * f * t)

fig, [ax1, ax2, ax3] = plt.subplots(3, 1)
fig.set_figheight(15)
fig.set_figwidth(15)
# function g to be transformed
f_orig = 3 # Hz
t = np.linspace(0, 4, 1001)
g = np.cos(2*np.pi*t*f_orig) + np.cos(2*np.pi * t * 5) +np.cos(2*np.pi * t * 7)
ax1.plot(t, g)
ax1.set_title(f"Signal ({f_orig} Hz)")
ax1.set_xlabel("Time")
ax1.set_ylabel("Intensity")

# plot unit circle
t_unit = np.linspace(0,2*np.pi,101)
ax2.plot(np.cos(t_unit), np.sin(t_unit), label="Unit Circle")
ax2.set_title("Signal in Complex Plane")
ax2.set_xlabel("Re")
ax2.set_ylabel("Im")

# Wind the signal in complex plane, using particular frequency
f_wind = 2
g_complex = signal_in_complex_plane(g, t, f_wind)
real = [z.real for z in g_complex]
imag = [z.imag for z in g_complex]
ax2.scatter(real, imag, label=f"Signal in Complex Plane ({f_wind} Hz)", c="red", s=1)
center_of_mass = fourier_transform(g, t, f_wind)
ax2.scatter(center_of_mass.real, center_of_mass.imag, label="Center of Mass", c="black")
ax2.grid(True)
ax2.set_aspect('equal', 'datalim')
ax2.legend()

# Plot center of mass for different frequencies of the shape in complex plane
frequencies = np.linspace(0, 10, 1001)
centers_of_mass = []
for freq in frequencies:
    g_hat = fourier_transform(g, t, freq).real
    centers_of_mass.append(g_hat)
ax3.plot(frequencies, centers_of_mass, label="Frequencies")
ax3.set_title("Fourier Transform of Signal")
ax3.set_xlabel("Frequency")
ax3.set_ylabel("Center of Mass")
ax3.grid(True)
ax3.legend()

plt.show()