Skip to content

Instantly share code, notes, and snippets.

@GaryLee
Last active October 11, 2020 15:45
Show Gist options
  • Save GaryLee/d5180e12d2eaa5dae163b314f0eeed4d to your computer and use it in GitHub Desktop.
Save GaryLee/d5180e12d2eaa5dae163b314f0eeed4d to your computer and use it in GitHub Desktop.
Do DFT and plot with Bokeh.
from bokeh.plotting import figure, output_file, vplot, show
import numpy as np
import random
n = 1024
a = 1.0
sample_period = 0.001
t = np.linspace(0, n * sample_period, n)
dt = t[1] - t[0]
data = np.abs(np.sin(t * 2 * np.pi) * a)
def fft(data, sample_period, power=False, use_db=True):
dt = sample_period
sp = np.fft.rfft(data)
if power:
spectrum = (np.abs(sp) * 2 * dt) ** 2
else:
spectrum = np.abs(sp) * 2 * dt
if use_db:
max_input = np.max(data)
if power:
spectrum = 20 * np.log10(spectrum / max_input)
else:
spectrum = 10 * np.log10(spectrum / max_input)
n = len(data)
freqs = np.fft.fftfreq(n, sample_period)
# Ignore the negative part of frequency. It's because of symmetry of FFT.
idx = np.argsort(freqs)
idx = filter(lambda i: freqs[i] > 0, idx)
return freqs[idx], spectrum[idx].real
freqs, spectrum = fft(data, dt, use_db=False)
freqs_db, spectrum_db = fft(data, dt, use_db=True)
output_file("lines.html")
p1 = figure(width=800, height=300, title="data")
p1.line(t, data, legend="data", line_width=1, color="red")
p2 = figure(width=800, height=300, title="FFT", x_axis_label="Frequency(Hz)", y_axis_label="Amplitude")
p2.line(freqs, spectrum, legend="data", line_width=2, color="blue")
p3 = figure(width=800, height=300, title="FFT(dB)", x_axis_label="Frequency(Hz)", y_axis_label="Amplitude(dB)")
p3.line(freqs_db, spectrum_db, legend="data", line_width=2, color="blue")
p = vplot(p1, p2, p3)
# Show the results.
show(p)
print(spectrum.real[10])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment