pysim/dft2.py

75 lines
1.8 KiB
Python
Raw Permalink Normal View History

2024-04-28 23:02:14 +02:00
import numpy as np
import matplotlib.pyplot as plt
def DFT(x):
"""
Compute the Discrete Fourier Transform (DFT) of the input signal x.
Parameters:
x (array-like): Input signal
Returns:
X (array-like): DFT of the input signal
"""
N = len(x)
n = np.arange(N)
k = n.reshape((N, 1))
e = np.exp(-2j * np.pi * k * n / N)
X = np.dot(e, x)
return X
def DFTI(X):
"""
Compute the Inverse Discrete Fourier Transform (IDFT) of the input signal X.
Parameters:
X (array-like): Input signal
Returns:
x (array-like): IDFT of the input signal
"""
N = len(X)
n = np.arange(N)
k = n.reshape((N, 1))
e = np.exp(2j * np.pi * k * n / N)
x = np.dot(e, X) / N
return x
# Define a sine function as the input signal
t = np.linspace(0, 1, 1000) # Time array
frequency = 5 # Frequency of the sine wave
x = np.sin(2 * np.pi * frequency * t) + np.sin(np.pi * frequency * t)
# Compute the DFT of the sine wave
X = DFT(x)
# Plot the input signal
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.plot(t, x)
plt.title('Input Signal')
plt.xlabel('Time')
plt.ylabel('Amplitude')
# Plot the magnitude spectrum (DFT)
freq = np.fft.fftfreq(len(x), d=t[1]-t[0])
plt.subplot(1, 3, 2)
plt.stem(freq, np.abs(X))
plt.title('Magnitude Spectrum (DFT)')
plt.xlabel('Frequency')
plt.ylabel('Magnitude')
plt.xlim(-10, 10) # Limit the x-axis for better visualization
# Compute the IDFT of the DFT signal
x_reconstructed = DFTI(X)
# Plot the reconstructed signal
plt.subplot(1, 3, 3)
plt.plot(t, x_reconstructed.real, label='Reconstructed Signal')
plt.plot(t, x, label='Original Signal')
plt.title('Reconstructed Signal vs Original Signal')
plt.xlabel('Time')
plt.ylabel('Amplitude')
plt.legend()
plt.tight_layout()
plt.show()