m-thesis-introduction/fourier/mylib/fft.py

44 lines
1.1 KiB
Python

"""
Simple FFT stuff
"""
import numpy as np
import scipy.fftpack as ft
def get_freq_spec(val,dt):
"""From earsim/tools.py"""
fval = np.fft.fft(val)[:len(val)//2]
freq = np.fft.fftfreq(len(val),dt)[:len(val)//2]
return fval, freq
def ft_spectrum( signal, sample_rate=1, ftfunc=None, freqfunc=None, mask_bias=False, normalise_amplitude=False):
"""Return a FT of $signal$, with corresponding frequencies"""
if True:
return get_freq_spec(signal, 1/sample_rate)
n_samples = len(signal)
if ftfunc is None:
real_signal = np.isrealobj(signal)
if False and real_signal:
ftfunc = ft.rfft
freqfunc = ft.rfftfreq
else:
ftfunc = ft.fft
freqfunc = ft.fftfreq
if freqfunc is None:
freqfunc = ft.fftfreq
normalisation = 2/len(signal) if normalise_amplitude else 1
spectrum = normalisation * ftfunc(signal)
freqs = freqfunc(n_samples, 1/sample_rate)
if not mask_bias:
return spectrum, freqs
else:
return spectrum[1:], freqs[1:]