mirror of
https://gitlab.science.ru.nl/mthesis-edeboone/m-thesis-introduction.git
synced 2024-11-13 18:13:31 +01:00
114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
import numpy as np
|
|
from collections import namedtuple
|
|
|
|
from lib import direct_fourier_transform as dtft
|
|
|
|
import matplotlib.pyplot as plt # for debug plotting
|
|
|
|
passband = namedtuple("passband", ['low', 'high'], defaults=[0, np.inf])
|
|
|
|
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 bandpass_samples(samples, samplerate, band=passband()):
|
|
"""
|
|
Bandpass the samples with this passband.
|
|
This is a hard filter.
|
|
"""
|
|
fft, freqs = get_freq_spec(samples, samplerate)
|
|
|
|
fft[ ~ self.freq_mask(freqs) ] = 0
|
|
|
|
return np.fft.irfft(fft)
|
|
|
|
def bandpass_mask(freqs, band=passband()):
|
|
low_pass = abs(freqs) <= band[1]
|
|
high_pass = abs(freqs) >= band[0]
|
|
|
|
return low_pass & high_pass
|
|
|
|
def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True, debug_ax=False):
|
|
bins = 0
|
|
fft, freqs = get_freq_spec(samples, 1/samplerate)
|
|
bandmask = [False]*len(freqs)
|
|
|
|
if band[1] is None:
|
|
# Only a single frequency given
|
|
# use a DTFT for finding the power
|
|
time = np.arange(0, len(samples), 1/samplerate)
|
|
|
|
real, imag = dtft(band[0], time, samples)
|
|
power = np.sum(np.abs(real**2 + imag**2))
|
|
else:
|
|
bandmask = bandpass_mask(freqs, band=band)
|
|
|
|
if normalise_bandsize:
|
|
bins = np.count_nonzero(bandmask, axis=-1)
|
|
else:
|
|
bins = 1
|
|
|
|
bins = max(1, bins)
|
|
|
|
power = 1/bins * np.sum(np.abs(fft[bandmask])**2)
|
|
|
|
# Prepare plotting variables if an Axes is supplied
|
|
if debug_ax:
|
|
if any(bandmask):
|
|
min_f, max_f = min(freqs[bandmask]), max(freqs[bandmask])
|
|
else:
|
|
min_f, max_f = 0, 0
|
|
|
|
if band[1] is None:
|
|
min_f, max_f = band[0], band[0]
|
|
|
|
if debug_ax is True:
|
|
debug_ax = plt.gca()
|
|
|
|
l = debug_ax.plot(freqs, np.abs(fft), alpha=0.9)
|
|
|
|
amp = np.sqrt(power)
|
|
|
|
if min_f != max_f:
|
|
debug_ax.plot( [min_f, max_f], [amp, amp], alpha=0.7, color=l[0].get_color(), ls='dashed')
|
|
debug_ax.axvspan(min_f, max_f, color=l[0].get_color(), alpha=0.2)
|
|
else:
|
|
debug_ax.plot( min_f, amp, '4', alpha=0.7, color=l[0].get_color(), ms=10)
|
|
|
|
return power
|
|
|
|
def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None, debug_ax=False, mode='sine'):
|
|
if noise_band is None:
|
|
noise_band = signal_band
|
|
|
|
if noise is None:
|
|
noise = samples
|
|
|
|
if debug_ax is True:
|
|
debug_ax = plt.gca()
|
|
|
|
if mode == 'sine':
|
|
noise_power = bandpower(noise, samplerate, noise_band, debug_ax=debug_ax)
|
|
noise_amplitude = np.sqrt(noise_power)
|
|
|
|
signal_power = bandpower(samples, samplerate, signal_band, debug_ax=debug_ax)
|
|
signal_amplitude = np.sqrt(signal_power)
|
|
|
|
elif mode == 'pulse':
|
|
noise_amplitude = np.sqrt(np.mean(noise**2))
|
|
|
|
signal_amplitude = max(np.abs(samples))
|
|
|
|
if debug_ax:
|
|
l1 = debug_ax.plot(noise, alpha=0.5)
|
|
debug_ax.axhline(noise_amplitude, alpha=0.9, color=l1[0].get_color())
|
|
|
|
l2 = debug_ax.plot(samples, alpha=0.5)
|
|
debug_ax.axhline(signal_amplitude, alpha=0.9, color=l2[0].get_color())
|
|
else:
|
|
raise NotImplementedError("mode not in ['sine', 'pulse']")
|
|
|
|
return signal_amplitude/noise_amplitude
|