Square wave FFT

square_fft.py

#!/usr/bin/env python3

# build a square signal from sum of sin

import matplotlib.pyplot as plt
import numpy as np
from scipy import signal

# number of samples
Ns = 2000

# second between 2 samples
Ts = 1e-4  # 100 us

# N samples, every sample is time value (in s)
# 0 -> t_max with Ns items
t_max = Ns * Ts
t_samples = np.linspace(0.0, t_max, Ns)

# build square signal at 50 Hz
f = 50.0
w = 2 * np.pi * f

# by harmonic add
# y_samples = np.zeros(Ns)
#
# # add 2 sine waves
# for n in range(2):
#     y_samples += 1/(2*n + 1) * np.sin((2*n + 1) * w * t_samples)

# by scipy signal (pure square)
y_samples = np.zeros(Ns)
y_samples += signal.square(w * t_samples, duty=0.5)

# simulate "hole" in signal
# y_samples[0:200] = -1
# y_samples[600:800] = -1
# y_samples[1200:1400] = -1
# y_samples[1800:2000] = -1

#  fft of signal
# compute FFT
y_freq = np.fft.fft(y_samples)
xf = np.fft.fftfreq(y_samples.size, d=Ts)[:Ns // 2]
# yf between 0.0 and 1.0 for every xf step
yf = 1.0 / (Ns / 2.0) * np.abs(y_freq[:Ns // 2])

# compute mean of signal
# print(round(np.mean(y_samples), 2))
peak_ids = signal.argrelmax(yf, order=8)
print(xf[peak_ids])

# matplotlib display
fig, axes = plt.subplots(nrows=2)

# plot time signal:
axes[0].set_title('Signal')
axes[0].plot(t_samples, y_samples)
axes[0].grid()
axes[0].set_xlabel('Time (s)')
axes[0].set_ylabel('Amplitude')

# plot different spectrum types:
axes[1].set_title('Magnitude Spectrum')
axes[1].plot(xf, yf)
axes[1].scatter(xf[peak_ids], yf[peak_ids], color='r')
axes[1].grid()
axes[1].set_xlabel('Frequency (Hz)')
axes[1].set_ylabel('Magnitude')

fig.tight_layout()
plt.show()